Seq2Seq 與 Attention：翻譯為什麼需要回頭看？

p(y_1, …, y_S | x)
  = ∏ᵢ p(y_i | y_1, …, y_(i−1), x)

e_ij = vᵀ tanh(W_q s_(i−1) + W_k h_j)
α_ij = softmax over valid source positions j of e_ij
c_i  = Σ_j α_ij h_j

α = softmax([1, 0, 1])
  = [e/(2e+1), 1/(2e+1), e/(2e+1)]
  ≈ [0.422319, 0.155362, 0.422319]

c = 0.422319 × [1, 0]
  + 0.155362 × [0, 2]
  + 0.422319 × [2, 1]
  + 0 × [100, 100]
  ≈ [1.266956, 0.733044]

"""Deterministic CPU/float64 Seq2Seq attention mechanics; no model training."""
from __future__ import annotations

import json
import math
import platform

import torch


def masked_softmax(scores: torch.Tensor, valid_lens: torch.Tensor) -> torch.Tensor:
    """Softmax over valid keys; reject any query with no valid key.

    Args:
        scores: (batch, queries, keys) unnormalized attention scores.
        valid_lens: (batch,) or (batch, queries), number of valid leading keys.
    """
    if scores.ndim != 3:
        raise ValueError("scores must have shape (batch, queries, keys)")
    batch, queries, keys = scores.shape
    valid_lens = torch.as_tensor(valid_lens, dtype=torch.long, device=scores.device)
    if valid_lens.ndim == 1 and valid_lens.shape[0] == batch:
        flat_lens = valid_lens.repeat_interleave(queries)
    elif valid_lens.ndim == 2 and tuple(valid_lens.shape) == (batch, queries):
        flat_lens = valid_lens.reshape(-1)
    else:
        raise ValueError("valid_lens must have shape (batch,) or (batch, queries)")
    if torch.any(flat_lens < 1):
        raise ValueError("each query must have at least one valid key; empty key sets are rejected")
    if torch.any(flat_lens > keys):
        raise ValueError("valid key count cannot exceed the key dimension")

    positions = torch.arange(keys, dtype=torch.long, device=scores.device)
    keep = positions.unsqueeze(0) < flat_lens.unsqueeze(1)
    flat_scores = scores.reshape(batch * queries, keys)
    masked_scores = flat_scores.masked_fill(~keep, float("-inf"))
    return torch.softmax(masked_scores, dim=-1).reshape(batch, queries, keys)


def rounded(values: torch.Tensor, digits: int = 12):
    if values.ndim == 0:
        return round(float(values.item()), digits)
    if values.ndim == 1:
        return [round(float(value), digits) for value in values.tolist()]
    return [rounded(row, digits) for row in values]


def main() -> None:
    # One fixed query, three real source positions, and one padded position.
    # Key depth is one, so scaled dot product (divide by sqrt(d_k)) equals q @ k.T.
    queries = torch.tensor([[[1.0]]], dtype=torch.float64, device="cpu")
    keys = torch.tensor([[[1.0], [0.0], [1.0], [2.0]]], dtype=torch.float64, device="cpu")
    values = torch.tensor(
        [[[1.0, 0.0], [0.0, 2.0], [2.0, 1.0], [100.0, 100.0]]],
        dtype=torch.float64,
        device="cpu",
    )
    valid_lens = torch.tensor([3], dtype=torch.long, device="cpu")
    key_depth = queries.shape[-1]
    scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(key_depth)
    weights = masked_softmax(scores, valid_lens)
    context = torch.bmm(weights, values)

    # Independent scalar calculation using the same fixed inputs.
    raw_scores = [1.0, 0.0, 1.0]
    exp_scores = [math.exp(score - max(raw_scores)) for score in raw_scores]
    py_weights = [value / sum(exp_scores) for value in exp_scores]
    py_context = [
        sum(py_weights[i] * values[0, i, dim].item() for i in range(3))
        for dim in range(values.shape[-1])
    ]
    assert max(abs(weights[0, 0, i].item() - py_weights[i]) for i in range(3)) < 1e-14
    assert weights[0, 0, 3].item() == 0.0
    assert abs(float(weights.sum().item()) - 1.0) < 1e-14
    assert max(abs(context[0, 0, i].item() - py_context[i]) for i in range(2)) < 1e-14

    try:
        masked_softmax(scores, torch.tensor([0], dtype=torch.long, device="cpu"))
        empty_key_result = "unexpectedly accepted"
        raise AssertionError("empty-key row must be rejected")
    except ValueError as exc:
        empty_key_result = str(exc)

    # Deliberate counterexample: finite -1e6 sentinels become equal logits
    # when every position is masked, so softmax returns a uniform distribution.
    # This is not a valid empty-key attention result and is not used above.
    all_masked_finite = torch.full((1, 1, 3), -1e6, dtype=torch.float64, device="cpu")
    finite_sentinel_bad_result = torch.softmax(all_masked_finite, dim=-1)

    # Fixed illustrative logits: not outputs of a trained model.
    vocab = ["ils", "nous", "regardent", "sommes", "<eos>"]
    toy_logits = {
        "<bos>": [0.2, 1.0, -1.0, -2.0, -3.0],
        "ils": [-2.0, -1.0, 1.2, 0.0, -3.0],
        "regardent": [-3.0, -3.0, -2.0, -2.0, 1.3],
        "nous": [-3.0, -2.0, -3.0, 1.1, -1.0],
        "sommes": [-3.0, -3.0, -3.0, -2.0, 1.2],
    }
    target = ["ils", "regardent", "<eos>"]
    teacher_inputs = ["<bos>"] + target[:-1]
    teacher_forced = []
    for step, (decoder_input, label) in enumerate(zip(teacher_inputs, target)):
        predicted = vocab[int(torch.argmax(torch.tensor(toy_logits[decoder_input], dtype=torch.float64)).item())]
        teacher_forced.append({
            "step": step,
            "decoder_input_from_gold_prefix": decoder_input,
            "target_label": label,
            "argmax_from_fixed_logits": predicted,
        })

    autoregressive = []
    decoder_input = "<bos>"
    stopped_on_eos = False
    max_steps = 6
    for step in range(max_steps):
        predicted = vocab[int(torch.argmax(torch.tensor(toy_logits[decoder_input], dtype=torch.float64)).item())]
        autoregressive.append({"step": step, "decoder_input_from_previous_prediction": decoder_input,
                               "argmax_from_fixed_logits": predicted})
        if predicted == "<eos>":
            stopped_on_eos = True
            break
        decoder_input = predicted
    assert stopped_on_eos

    result = {
        "scope": "Hand-set one-query Q/K/V and fixed decoder-logit lookup only; no encoder, learned parameters, corpus, training, translation benchmark, or claim of trained output.",
        "environment": {"python": platform.python_version(), "torch": torch.__version__,
                        "device": str(queries.device), "dtype": str(queries.dtype)},
        "attention": {
            "shapes": {"queries": list(queries.shape), "keys": list(keys.shape),
                       "values": list(values.shape), "scores": list(scores.shape),
                       "weights": list(weights.shape), "context": list(context.shape)},
            "query": rounded(queries[0, 0]),
            "keys_including_pad": rounded(keys[0, :, 0]),
            "values_including_pad": rounded(values[0]),
            "valid_key_count": int(valid_lens[0].item()),
            "scores_scaled_dot_product": rounded(scores[0, 0]),
            "weights_after_valid_key_mask": rounded(weights[0, 0]),
            "weight_sum": round(float(weights.sum().item()), 12),
            "weighted_context": rounded(context[0, 0]),
            "pad_value_excluded": True,
            "independent_python_scalar_reference": {"valid_weights": [round(x, 12) for x in py_weights],
                                                     "context": [round(x, 12) for x in py_context],
                                                     "max_torch_difference_lt_1e-14": True},
        },
        "empty_key_boundary": {
            "valid_len_zero": "rejected",
            "error": empty_key_result,
            "finite_minus_1e6_all_masked_counterexample": rounded(finite_sentinel_bad_result[0, 0]),
            "counterexample_note": "D2L's finite sentinel makes every all-masked logit equal; softmax is uniform, not a valid empty-key attention distribution. This prototype explicitly rejects the row.",
        },
        "teacher_forcing_vs_autoregressive": {
            "gold_target_including_eos": target,
            "teacher_forcing_input_shift": teacher_inputs,
            "vocabulary_order": vocab,
            "fixed_logits_by_previous_token": toy_logits,
            "teacher_forced_steps": teacher_forced,
            "autoregressive_steps": autoregressive,
            "autoregressive_stopped_on_eos": stopped_on_eos,
            "max_steps_guard": max_steps,
            "note": "Logits are hand-written solely to demonstrate shifted gold-prefix inputs versus previous-prediction inputs. They are not from a trained Seq2Seq model. EOS ends this toy loop; PAD masking of target loss is outside this prototype.",
        },
        "checks": "PASS: scaled QK scores, mask excludes PAD key, valid weights sum to one, weighted context matches independent scalar calculation, empty valid-key row rejected, fixed-logit autoregressive loop terminates at EOS.",
    }
    print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))


if __name__ == "__main__":
    main()