Back to Deep Learning
AI/ML · Deep Learning

Seq2Seq with Attention

Translating a sentence with a plain encoder-decoder means reading the whole source, compressing it into one fixed vector, and then generating from that vector alone. Everything the decoder will ever know about the source has to fit there — and it does not, which is why the approach degraded sharply on long sentences and why reversing the source words used to help. Attention removes the compression. Keep every encoder state, and at each output step compute a score between what the decoder is currently doing and each source position, turn those scores into weights that sum to one, and read a weighted average. The decoder is no longer remembering the source; it is looking at it, and choosing where to look.

SupervisedHybridSequence modelingGeneration

Classified hybrid because the recurrence and the attention are separate mechanisms doing separate jobs: a recurrent encoder and decoder carry order, and an attention layer bridges them. The transformer keeps the second and discards the first, which is why this entry sits between them.

Also known as: Encoder-decoder, Bahdanau attention, Luong attention, Neural machine translation, Soft alignment

Intuition

What it actually does

Translating a sentence with a plain encoder-decoder means reading the whole source, compressing it into one fixed vector, and then generating from that vector alone. Everything the decoder will ever know about the source has to fit there — and it does not, which is why the approach degraded sharply on long sentences and why reversing the source words used to help. Attention removes the compression. Keep every encoder state, and at each output step compute a score between what the decoder is currently doing and each source position, turn those scores into weights that sum to one, and read a weighted average. The decoder is no longer remembering the source; it is looking at it, and choosing where to look.

Assumptions it makes
  • Input and output are sequences whose alignment is learnable but not necessarily monotonic — which is what distinguishes this from the CTC formulation
  • Paired examples exist: this is supervised transduction and needs source-target pairs, not just target sequences
  • The source is short enough that attention over every position is affordable, since the cost is the product of the two lengths
  • Teacher forcing during training approximates the inference regime well enough, which is exactly the assumption exposure bias violates
  • A single attention distribution per output step suffices, which multi-head attention later relaxed
Objective

Loss function

Maximize the probability of the target sequence one token at a time, conditioning on the source and on what has been generated so far. The second expression is where the architecture lives: the context is a convex combination of every encoder state, with weights chosen per output step. That single change converts a fixed-capacity memory into a content-addressed lookup, and it is why performance stops degrading with source length. The conditioning on y-less-than-u also names the method’s characteristic problem — during training those tokens are the true ones and at inference they are the model’s own, so errors compound in a regime the model was never shown.

encoder state at source position t — all of them are kept, which is the whole change
attention weight: how much output step u reads from source position t, summing to one over t
the context vector — a fresh weighted average of the source, recomputed at every output step
the previously generated tokens; during training these are the true ones, which is teacher forcing
Optimization

How it is minimized

Teacher-forced maximum likelihood by backpropagation through both recurrences; beam search at inference

the compatibility score between decoder state and source position — additive here, a dot product in the multiplicative variant
depends only on the encoder, so it is computed once per source and reused at every output step
Bahdanau’s additive scoring: a small network, more expressive and slower than a dot product
normalizes the scores into a distribution over source positions — which is why the weights read as a soft alignment

Training is teacher-forced maximum likelihood: feed the true previous token, predict the next, and backpropagate through the decoder, the attention and the encoder together. Nothing about the optimizer is unusual; three details are. The additive score has a term that depends only on the encoder, so projecting all encoder states once per source rather than once per output step removes a factor of U from the attention cost — the single most consequential implementation detail here. Padded source positions must be masked before the softmax, or the model attends to padding and the weights become nonsense in a way that trains without complaint. And decoding is a search problem the loss says nothing about: greedy decoding takes the best token at each step and is myopic, beam search keeps several hypotheses and is a heuristic with no optimality guarantee, and without length normalization it systematically prefers short outputs because every additional token multiplies in another probability below one.

Hyperparameters
  • attention type

    Additive (Bahdanau) or multiplicative (Luong). Additive is more expressive; multiplicative is a matrix product and far faster, which is what the transformer adopted (multiplicative with scaling once dimensions are large)

  • beam width

    How many hypotheses to keep at decode time. Returns flatten quickly and very wide beams can score worse, which is a genuinely counterintuitive result (4 to 10 for translation)

  • length penalty

    Normalization of the sequence score by length. Not optional: without it the search prefers short outputs by construction (alpha 0.6 to 1.0 in the standard formulation)

  • teacher forcing ratio

    How often the true previous token is fed rather than the model’s own. Scheduled sampling anneals it to reduce exposure bias

  • encoder direction

    Bidirectional encoders are standard, since a source position benefits from both sides; the decoder must stay causal

  • coverage penalty

    Discourages attending repeatedly to the same source position, which is the mechanism behind repeated and dropped output

Convergence & cost

Trains reliably given the usual recurrent care — clipping, and gating in both recurrences. The failure modes are about generation rather than fitting. Exposure bias: teacher forcing means the model never sees its own mistakes during training, so at inference an early error puts it in a state the training distribution never contained and errors compound. Repetition and omission: nothing constrains the attention to cover the source once, so a decoder can attend to the same span repeatedly and produce loops, or skip a span entirely — coverage penalties address the symptom rather than the cause. Length bias: the sequence score is a product of probabilities, so longer is always less likely and the search must be corrected for it. And the decoder remains sequential, so generation cannot be parallelized even though the encoder can.

Encoder: O(T · h^2) and sequential in T. Attention: O(T · U · h) — the product of source and target length, which is where the cost is and which is quadratic when the two are comparable. Decoder: O(U · h^2), sequential in U. Beam search multiplies the decode cost by the beam width. Memory holds all T encoder states rather than one summary vector, which is the trade that buys the lookup and is negligible relative to the activations.

Applications

How, where, and why it is used

Time-Series Forecasting

Viable
How

Encode the history with a recurrent encoder, then decode the horizon step by step with attention back over the encoder states — so the forecast for a distant horizon can look directly at the relevant part of the history rather than at a summary of it. Known-future covariates such as holidays or promotions enter the decoder at each step, which is a genuine structural advantage over a model that only sees the past.

Where
  • Multi-horizon forecasting where different horizons depend on different parts of the history, such as a weekly and an annual pattern
  • Series with known future covariates, which the decoder can consume at each step
  • The direct ancestor of attention-based forecasters, where the attention weights over the history are shown to analysts as an explanation
  • Global models over many series, where the encoder-decoder structure shares one set of weights across all of them
Why

The attention lets a long history stay accessible instead of being compressed, which is the same argument as in translation and applies just as well to a seasonal series. The interpretable weights are a real secondary benefit in a forecasting context, where the question "what is this based on" is asked more often than in most domains. Against it: the decoder is sequential, so multi-step generation is slow; errors compound across the horizon exactly as they do in text; and on a single ordinary business series it loses to exponential smoothing, because none of this apparatus is identifiable from a few hundred observations.

Featurization
  • Scale per series and encode calendar structure explicitly, as for any neural forecaster trained across series
  • Feed known-future covariates into the decoder rather than the encoder — that is what the split is for, and it is frequently wired wrongly
  • Mask padded history positions before the attention softmax, or the model attends to padding and trains without complaint
  • Predict all horizons directly rather than recursively where possible, which avoids compounding the model’s own errors
Evaluation

Rolling-origin backtesting with MASE against seasonal-naive and against exponential smoothing. Inspect the attention weights as a diagnostic rather than as an explanation — they show where the model looked, which is not the same as why the forecast is what it is.

Pitfalls
  • Reading attention weights as causal attribution, which they are not — they are one layer’s routing, and a well-known caution in the interpretability literature
  • Recursive decoding compounding errors across a long horizon, which teacher-forced training never exhibits
  • Leaking future covariates into the encoder, which is a subtle and complete invalidation of the backtest

Anomaly Detection

Not applicable

The architecture maps one sequence to another under supervision from paired examples; it has no notion of normality or of outlyingness, and a reconstruction variant of it is the autoencoder rather than this entry.

Optimization

Adapted
How

Decoding is a search problem the training loss says nothing about, and it is the clearest example in this category of a combinatorial search bolted to a learned model. The space is every possible output sequence, which is exponential; greedy decoding is the myopic heuristic; beam search keeps a bounded frontier and is still a heuristic with no optimality guarantee. Attention itself is a soft assignment — a relaxation of choosing one source position into a distribution over all of them, which is what makes it differentiable.

Where
  • Beam search as a bounded-width heuristic over an exponential sequence space, with no guarantee of finding the highest-probability output
  • Length normalization as a correction for a scoring function that is biased by construction
  • Attention as a differentiable relaxation of a hard alignment, which is the same move that makes many discrete decisions trainable
Why

It is worth studying because the model and the search are separable and both matter: a better model with greedy decoding frequently loses to a worse model with a good beam, and the loss the model was trained on does not measure decoding quality at all. The length-bias result is also instructive — a sequence score that is a product of probabilities is monotonically decreasing in length, so the search prefers short outputs for reasons that have nothing to do with the data, and the fix is an explicit correction rather than a better model.

Featurization
  • Normalize the sequence score by length, or the search systematically truncates
  • Add a coverage term when repetition or omission is observed, since nothing in the objective constrains the attention to cover the source once
Evaluation

Compare greedy, beam and exhaustive decoding on short instances where exhaustive is feasible — the gap between beam and exhaustive is a direct measurement of what the heuristic costs, and it is usually larger than expected.

Pitfalls
  • Assuming a wider beam is monotonically better, when very wide beams frequently score worse — a genuinely counterintuitive and well-replicated result
  • Comparing models under different decoding settings, which confounds the model with the search
  • Omitting length normalization and then attributing the short outputs to the model
Also applied in

Natural Language

Primary fit
How

The formulation neural machine translation was built on. A bidirectional recurrent encoder reads the source; a decoder generates the target one token at a time, computing an attention distribution over source positions at every step and consuming the resulting context vector alongside its own state. The attention weights form a soft alignment between source and target that can be read directly.

Where
  • Neural machine translation, where this architecture replaced phrase-based statistical systems
  • Abstractive summarization, where the decoder must select from a long source rather than compress it
  • Speech recognition with attention-based decoders, as the alternative to the CTC formulation
  • Any sequence-to-sequence transduction where input and output lengths differ and are not monotonically aligned
Why

It solved the specific problem that limited encoder-decoder translation: a fixed context vector could not hold a long sentence, and quality fell off with length. Attention removed that ceiling and, as a side effect, produced alignments a linguist could inspect — which is rare. It has been superseded by the transformer, which kept the attention and discarded the recurrence to gain parallelism, so this entry is now the explanation of where attention came from rather than a deployment choice. What survives operationally is the encoder-decoder structure itself, which the transformer inherited unchanged.

Featurization
  • Mask padded source positions before the softmax, or attention mass leaks onto padding and training proceeds without any error
  • Bucket by length before batching, since attention cost is the product of source and target lengths and a mixed batch wastes most of it
  • Share embeddings between the decoder input and the output projection where the vocabularies overlap, which is a large parameter saving
  • Use subword units rather than words, which bounds the vocabulary and removes the unknown-token problem
Evaluation

BLEU or a learned metric on a held-out set, always reporting the decoding configuration alongside — a model compared at different beam widths is not a comparison of models. Inspect alignments qualitatively on known-hard cases such as reordering and long-distance agreement.

Pitfalls
  • Reporting scores without the beam width and length penalty, which makes the numbers incomparable
  • Exposure bias at generation, where an early mistake puts the decoder in a state training never contained
  • Repetition loops from unconstrained attention, which coverage penalties treat symptomatically
  • Reading attention as explanation rather than as routing, a caution the interpretability literature has made repeatedly
Deployment

What it costs to run

Training

Two recurrences plus an attention that costs the product of source and target lengths. Both recurrences are sequential, so the encoder parallelizes across the batch and not across time, and training wall-clock is dominated by that rather than by the attention arithmetic.

Inference

Generation is sequential in output length and multiplied by the beam width, so a beam of eight is roughly eight times the decode cost. The encoder runs once per input and its states are cached for the whole decode, which is the standard and necessary optimization.

Retraining

Periodic and full, warm-started. Vocabulary changes force more than a refit — a new subword model reindexes every embedding, so vocabulary and weights are versioned together or not at all.

Drift & monitoring
  • Track output length distribution against the reference: systematic shortening is the length-bias signature and points at the decoder configuration rather than the model
  • Monitor repetition rate, since attention loops produce degenerate output that aggregate quality metrics smooth over
  • Watch the share of attention mass falling on padding or on the end-of-sequence position, which indicates a masking bug
  • Compare teacher-forced and free-running loss — a widening gap is exposure bias rather than a fitting problem
Production gotchas
  • Source padding must be masked before the softmax, not after — masking afterwards leaves the weights normalized over padding and the context vector is wrong while training proceeds cleanly
  • Length normalization must match between evaluation and deployment, or the deployed model produces systematically different lengths from the one that was benchmarked
  • Encoder states must be cached across the decode; recomputing them per output step multiplies the cost by the target length and is an easy accidental regression
  • Beam search state — hypotheses, scores, finished flags — is intricate, and an off-by-one in the finished handling produces plausible output that is systematically worse
  • The end-of-sequence token must be handled explicitly in the beam, or hypotheses that have finished keep accumulating score and crowd out live ones
Trade-offs

Pros and cons, in context

Strengths
  • Removes the fixed-size bottleneck entirely

    Quality stops degrading with source length, which was the ceiling on encoder-decoder models before it. The single change that made neural translation work, and the reason the mechanism survived into the transformer unchanged.

  • Attention weights are inspectable as a soft alignment

    A genuinely rare property — the model produces something a domain expert can read. Worth being careful with: it shows where the layer looked, which is routing rather than explanation.

  • Handles unequal and non-monotonic input-output lengths

    Reordering, insertion and deletion are all representable, which is what translation requires and what a monotonic formulation like CTC cannot express.

  • The encoder-decoder structure generalizes far beyond translation

    Summarization, speech, forecasting and code generation all use it unchanged. Learning it once pays across the category, which is why it sits here rather than in an applied domain.

Limitations
  • The decoder is sequential, so generation cannot be parallelized

    The encoder parallelizes across the batch and the decode does not, which caps throughput. The transformer removed the recurrence and kept the attention for exactly this reason.

  • Attention cost is the product of source and target lengths

    Quadratic when the two are comparable, which bounds usable sequence length. The same cost the transformer inherited and that a long line of efficient-attention work exists to reduce.

  • Exposure bias between training and generation

    Teacher forcing never shows the model its own mistakes, so an early error at inference lands it outside the training distribution and compounds. Scheduled sampling reduces it and does not remove it.

  • Decoding quality depends on a search the loss never optimized

    Beam width and length penalty change output substantially and are not part of training. A model comparison at different decoding settings is not a model comparison at all, and this is frequently overlooked.

Implementation

Intuitive, idiomatic, optimized

The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.

1. Intuitive

The math, transcribed literally. No dependencies.

"""Additive attention and greedy decoding - transcribed.

The decoder no longer receives a single summary of the source. At every output
step it scores itself against each encoder state, normalizes those scores into
weights, and reads a weighted average. That is the whole mechanism.
"""

import math


def additive_score(decoder_state, encoder_state, W_s, W_h, v):
    """e_ut = v^T tanh(W_s s_{u-1} + W_h h_t)

    Bahdanau's additive form: a small network rather than a dot product. More
    expressive, and slower - which is why the multiplicative variant is what
    the transformer adopted.
    """
    hidden = len(v)
    total = 0.0

    for j in range(hidden):
        projected = 0.0
        for k in range(len(decoder_state)):
            projected += W_s[k][j] * decoder_state[k]
        for k in range(len(encoder_state)):
            projected += W_h[k][j] * encoder_state[k]
        total += v[j] * math.tanh(projected)

    return total


def attention(decoder_state, encoder_states, W_s, W_h, v):
    """Scores, softmax, weighted average. One context vector per output step."""
    scores = [
        additive_score(decoder_state, state, W_s, W_h, v) for state in encoder_states
    ]

    # Softmax with the max subtracted: exp of a large score overflows, and the
    # naive form silently produces NaN.
    largest = max(scores)
    exponentials = [math.exp(score - largest) for score in scores]
    partition = sum(exponentials)
    weights = [value / partition for value in exponentials]

    # The context: a convex combination of EVERY encoder state, recomputed
    # fresh at each output step. Nothing is compressed, which is the whole
    # difference from a fixed context vector.
    hidden = len(encoder_states[0])
    context = [0.0] * hidden
    for t, state in enumerate(encoder_states):
        for j in range(hidden):
            context[j] += weights[t] * state[j]

    return context, weights


def greedy_decode(encoder_states, start_token, max_length, step_fn, W_s, W_h, v):
    """Take the highest-probability token at every step.

    Myopic by construction: a token that looks best now may lead nowhere, and
    nothing here can revise it. Beam search keeps several hypotheses precisely
    because of this.
    """
    decoder_state = [0.0] * len(encoder_states[0])
    token = start_token
    output = []
    alignments = []

    for _ in range(max_length):
        context, weights = attention(decoder_state, encoder_states, W_s, W_h, v)
        alignments.append(weights)

        # step_fn advances the decoder recurrence and returns log-probabilities
        # over the vocabulary.
        decoder_state, log_probabilities = step_fn(decoder_state, token, context)

        best = 0
        for index in range(1, len(log_probabilities)):
            if log_probabilities[index] > log_probabilities[best]:
                best = index

        if best == 0:      # end-of-sequence
            break

        output.append(best)
        token = best

    return output, alignments

O(U * T * h^2) for the attention alone, since the additive score reprojects every encoder state at every output step.

2. IdiomaticNumPy

How this language actually wants it written.

"""Attention - typed, batched, masked, with the encoder projection cached."""

from dataclasses import dataclass
from typing import Literal

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]

ScoreType = Literal["additive", "multiplicative"]


@dataclass
class AttentionLayer:
    """Additive or multiplicative scoring over a cached encoder projection."""

    W_s: Matrix        # (decoder_hidden, attention_hidden)
    W_h: Matrix        # (encoder_hidden, attention_hidden)
    v: Vector          # (attention_hidden,)
    score_type: ScoreType = "additive"

    def project_encoder(self, encoder_states: Tensor) -> Tensor:
        """W_h @ h_t, computed ONCE per source rather than per output step.

        This term depends only on the encoder, so recomputing it at every
        decoder step - which the naive formulation does - repeats the same
        product U times. Hoisting it is the single most consequential
        implementation detail in additive attention.
        """
        if encoder_states.ndim != 3:
            raise ValueError(f"expected (batch, time, hidden), got {encoder_states.shape}")
        return encoder_states @ self.W_h

    def __call__(
        self,
        decoder_state: Matrix,
        encoder_states: Tensor,
        projected: Tensor,
        source_mask: NDArray,
    ) -> tuple[Matrix, Matrix]:
        """One output step for a whole batch. Returns (context, weights)."""
        if source_mask.shape != encoder_states.shape[:2]:
            raise ValueError(
                f"mask shape {source_mask.shape} does not match "
                f"{encoder_states.shape[:2]} source positions"
            )

        if self.score_type == "additive":
            # (batch, 1, attn) + (batch, time, attn) broadcasts to every
            # source position without materializing a per-step projection.
            combined = np.tanh(decoder_state @ self.W_s)[:, None, :] + projected
            scores = combined @ self.v
        else:
            # Luong: a plain dot product against the projected states. Far
            # cheaper, and what the transformer generalized.
            scores = np.einsum("bh,bth->bt", decoder_state @ self.W_s, projected)

        # MASK BEFORE THE SOFTMAX. Masking afterwards leaves the weights
        # normalized over padding, so the context vector is wrong while
        # training proceeds without any error - the classic silent bug here.
        scores = np.where(source_mask, scores, -np.inf)

        largest = np.max(scores, axis=1, keepdims=True)
        exponentials = np.exp(scores - largest)
        weights = exponentials / exponentials.sum(axis=1, keepdims=True)

        # The context: a convex combination of every encoder state, fresh at
        # each output step. Nothing is compressed.
        context = np.einsum("bt,bth->bh", weights, encoder_states)
        return context, weights


def beam_search(
    encoder_states: Tensor,
    projected: Tensor,
    source_mask: NDArray,
    step_fn,
    attention_layer: AttentionLayer,
    beam_width: int = 5,
    max_length: int = 100,
    length_penalty: float = 0.7,
    eos: int = 0,
) -> list[int]:
    """Bounded-width search over an exponential space.

    A heuristic with no optimality guarantee: the highest-probability sequence
    can fall off the beam at any step and is unrecoverable afterwards.
    """
    if beam_width < 1:
        raise ValueError(f"beam_width must be at least 1, got {beam_width}")

    hidden = encoder_states.shape[2]
    beams: list[tuple[float, list[int], Vector]] = [(0.0, [], np.zeros(hidden))]
    finished: list[tuple[float, list[int]]] = []

    for _ in range(max_length):
        candidates: list[tuple[float, list[int], Vector]] = []

        for score, tokens, state in beams:
            context, _ = attention_layer(
                state[None, :], encoder_states, projected, source_mask
            )
            next_state, log_probabilities = step_fn(state, tokens, context[0])

            top = np.argpartition(log_probabilities, -beam_width)[-beam_width:]
            for token in top:
                extended = score + float(log_probabilities[token])
                if token == eos:
                    # LENGTH NORMALIZATION. The sequence score is a product of
                    # probabilities, so it decreases monotonically with length
                    # and the search prefers short outputs by construction.
                    # Without this correction the model looks terse and the
                    # cause is the search rather than the model.
                    normalized = extended / ((len(tokens) + 1) ** length_penalty)
                    finished.append((normalized, tokens))
                else:
                    candidates.append((extended, [*tokens, int(token)], next_state))

        if not candidates:
            break
        candidates.sort(key=lambda item: item[0], reverse=True)
        beams = candidates[:beam_width]

    if not finished:
        return beams[0][1]
    return max(finished, key=lambda item: item[0])[1]

What changed: Two changes carry the value and neither is the batching. The encoder projection is hoisted out of the decoder loop — it depends only on the encoder, so the naive formulation repeats the same product once per output token, and caching it removes a factor of U from the attention cost. And the source mask is applied before the softmax rather than after, which is the silent bug in this architecture: masking afterwards leaves the weights normalized over padding, so the context vector is wrong and training proceeds without any error. Beam search is added with length normalization, because the sequence score is a product of probabilities and therefore biased toward short outputs by construction rather than by anything in the data.

Explicit type hints on every public signatureDataclasses or NamedTuples over ad-hoc dictsRaise specific exceptions, never bare exceptGuard clauses over nested conditionals

O(T * h * a) once per source for the projection, then O(U * T * a) for the scores — a factor of h removed from the inner loop.

3. OptimizedNumPy

Faster, and the justification for each change.

"""Attention - scores as one GEMM, beam state batched, buffers reused."""

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


class BatchedAttention:
    """Multiplicative scoring, and the beam treated as a batch.

    Two changes. The additive score is replaced by a dot product, which turns
    the whole attention into one matrix product per step - this is exactly the
    move the transformer made, and it trades a little expressiveness for a
    kernel the hardware is built for.

    And the beam stops being a Python list of hypotheses stepped one at a
    time. A beam of width B is B independent decoder states, so it becomes a
    batch dimension: one forward pass per step for the entire beam rather than
    B of them. On a beam of ten that is an order of magnitude.
    """

    def __init__(self, encoder_hidden: int, decoder_hidden: int, max_source: int,
                 beam_width: int) -> None:
        # Single projection matrix for multiplicative scoring: the additive
        # form needs two projections and a tanh, this needs one product.
        self._W = np.zeros((decoder_hidden, encoder_hidden), dtype=np.float32)

        # Buffers sized for the beam, allocated once and rewritten per step.
        self._scores = np.empty((beam_width, max_source), dtype=np.float32)
        self._weights = np.empty((beam_width, max_source), dtype=np.float32)
        self._context = np.empty((beam_width, encoder_hidden), dtype=np.float32)

    def step(
        self,
        decoder_states: Matrix,      # (beam, decoder_hidden)
        encoder_states: Matrix,      # (source, encoder_hidden) - one source
        source_mask: NDArray,        # (source,)
    ) -> tuple[Matrix, Matrix]:
        """Attention for the entire beam in one pass."""
        beam, source = decoder_states.shape[0], encoder_states.shape[0]
        scores = self._scores[:beam, :source]
        weights = self._weights[:beam, :source]

        # One GEMM for every (hypothesis, source position) pair. The
        # projection and the score fuse: (S @ W) @ H^T evaluated right to
        # left would build a (decoder, source) intermediate per hypothesis;
        # this ordering never materializes one.
        np.matmul(decoder_states @ self._W, encoder_states.T, out=scores)

        # MASK BEFORE THE SOFTMAX. Masking afterwards leaves the weights
        # normalized over padding, so the context is wrong while training
        # proceeds cleanly - the silent bug in this architecture.
        scores[:, ~source_mask] = -np.inf

        # Softmax written through the score buffer: no temporary per step,
        # and with a beam of ten over a long source that is the dominant
        # allocation otherwise.
        np.subtract(scores, scores.max(axis=1, keepdims=True), out=scores)
        np.exp(scores, out=scores)
        np.divide(scores, scores.sum(axis=1, keepdims=True), out=weights)

        context = self._context[:beam]
        np.matmul(weights, encoder_states, out=context)
        return context, weights

    def rank(
        self,
        beam_scores: Vector,          # (beam,)
        log_probabilities: Matrix,    # (beam, vocab)
        beam_width: int,
    ) -> tuple[NDArray, NDArray, Vector]:
        """Top-k over the flattened (beam, vocab) grid in one call.

        The naive beam loops over hypotheses, takes the top tokens of each,
        then sorts the union - three passes and a Python sort. Flattening the
        grid makes it one partition over beam * vocab, which is a single
        vectorized call regardless of beam width.
        """
        combined = (beam_scores[:, None] + log_probabilities).ravel()

        # argpartition, not argsort: the top-k is an O(n) selection and
        # ordering the rest is work whose result is discarded.
        top = np.argpartition(combined, -beam_width)[-beam_width:]
        top = top[np.argsort(combined[top])[::-1]]

        vocab = log_probabilities.shape[1]
        return top // vocab, top % vocab, combined[top]

What changed: Two changes, and the second is the larger one. Additive scoring becomes multiplicative, which turns the whole attention into one matrix product per step — exactly the move the transformer made, trading a little expressiveness for a kernel the hardware is built for. Then the beam stops being a list of hypotheses stepped one at a time: a beam of width B is B independent decoder states, so it becomes a batch dimension and a step costs one forward pass instead of B. Top-k selection over the flattened beam-by-vocabulary grid replaces a loop plus a sort with a single partition, and every per-step buffer is allocated once, since with a wide beam over a long source those temporaries dominate.

  • Vectorize to NumPy/BLAS (@, np.dot, einsum)

    Multiplicative scoring makes the attention one GEMM over every hypothesis-position pair, and the context is a second GEMM against the weights.

    Trade-off: Multiplicative scoring is less expressive than additive and needs scaling as dimensions grow, or the dot products saturate the softmax into near-one-hot weights that pass almost no gradient.

  • Eliminate Python-level loops over samples

    The beam becomes a batch dimension, so one decoder pass serves every hypothesis and the top-k is one partition over the flattened grid rather than a loop plus a sort.

    Trade-off: Hypotheses now advance in lockstep, so a finished one still occupies a slot and consumes compute until the whole beam terminates — wasted work the sequential version avoided.

  • Pre-allocate output arrays and use in-place operations

    Score, weight and context buffers are allocated once for the widest beam and longest source, and the softmax is written through the score buffer rather than materializing temporaries per step.

    Trade-off: Fixes the beam width and source length at construction, and the in-place softmax destroys the raw scores — which anything wanting to inspect the pre-softmax alignment must copy first.

O(B * T * h) per decode step for the whole beam, versus B separate passes. Illustrative, not a measured benchmark.