Encoder-Decoder Beam Search

Implement seq2seq_beam_search for one source sequence. Keep a limited set of candidate target sequences, called beams, instead of committing to one token at every step.

A beam's score is the sum of the log probabilities of its generated tokens. Higher is better. Beam search can retain alternatives that greedy decoding would discard, but it is still an approximate search: pruning can remove the best eventual sequence. It does not guarantee better text or a global optimum.

Inputs and return value

• src_ids: (S,), integer-valued source IDs in [0, V_src), delivered as floats.
• src_emb: (V_src, d); tgt_emb: (V_tgt, d), token embeddings.
• enc_pos_embed: (S, d); dec_pos_embed: (P, d), position embeddings.
• enc_blocks: (E, 6, d, d); dec_blocks: (L, 12, d, d), transformer weights.
• w_head: (d, V_tgt), vocabulary projection.
• num_heads: a positive integer dividing d.
• start_token_id and eos_token_id: integers in [0, V_tgt).
• max_new_tokens: an integer limiting tokens appended after the start token.
• beam_width: an integer K with 1 <= K <= V_tgt.

S, d and the vocabulary sizes are positive. E and L may be zero. Let B = max(0, max_new_tokens); require P >= B. Parameters are finite and share a floating-point dtype and device; src_ids is on that device. Shapes agree. Leave all inputs unchanged.

Return a one-dimensional integer tensor on the input device, containing the highest-ranked retained sequence. Include the start token and any generated EOS. Its length is between 1 and 1+B. A nonpositive budget returns only the start token. Return IDs, not beam scores, padding or all candidate sequences.

Search rules

Encode the source once and keep that context fixed throughout the call. Initialize one unfinished beam: sequence [start_token_id], score 0. The start token has no probability contribution and does not consume the budget. Even when start_token_id equals EOS, this initial beam is unfinished.

For each of at most B rounds:

1. If every retained beam is finished, stop.
2. Carry each finished beam into the candidate pool unchanged. Do not decode it, append tokens or add to its score.
3. For each unfinished beam, run the full decoder on its current prefix. Apply stable log-softmax to the last position's vocabulary logits.
4. Choose its K highest-scoring next tokens, breaking exact ties by smaller token ID. Create one candidate per token by appending it and adding its log probability to the beam score. Mark it finished exactly when the appended token is EOS.
5. Rank the combined pool by total score, descending, and retain at most K candidates. For exactly equal scores, use lexicographic token-ID order: compare the first differing ID, smaller first; if one sequence is a prefix, the shorter comes first.

After the loop, return the first beam under the same ranking. It may be finished or unfinished: there is no final preference for EOS. Finished beams compete for the same K slots as unfinished ones. Do not remove them or keep them in a separate unlimited pool.

Use raw summed log probabilities, with no length normalization, length penalty, coverage bonus or sampling. Appending a token cannot increase a beam's score, which can favor shorter sequences. With K=1, these rules reduce to greedy decoding with the same smallest-ID tie rule and append-before-EOS check.

For example, suppose K=2 and one retained finished beam has score -0.8. If the next candidate scores are -1.0 and -1.2, retain the finished beam at -0.8 and the candidate at -1.0. Do not add another EOS or another probability to the finished beam.

Each decoder pass starts from the whole current prefix, using position rows from zero. There is no decoder key/value cache. The final appended token is returned without another pass, so B position rows suffice for 1+B output IDs.

Encoder and decoder architecture

Use the following post-normalization model. Cast source IDs to integers and start with enc = src_emb[src_ids] + enc_pos_embed. Each encoder block holds [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]. Apply blocks in order:

q, k, v = enc @ w_q, enc @ w_k, enc @ w_v
enc = LN(enc + attention(q, k, v) @ w_o)
enc = LN(enc + GELU(enc @ w_mlp1) @ w_mlp2)

Encoder attention can read every source position. Keep its final output fixed for the generation loop. With E=0, use source token-plus-position embeddings.

For each target prefix, start afresh with dec = tgt_emb[seq] + dec_pos_embed[:len(seq)]. Decoder slots 0 through 3 hold self-attention [w_q, w_k, w_v, w_o], slots 4 through 7 hold cross-attention [w_q_cross, w_k_cross, w_v_cross, w_o_cross], and slots 8 and 9 hold [w_mlp1, w_mlp2]. Ignore slots 10 and 11, even when nonzero. Each block applies:

q, k, v = dec @ w_q, dec @ w_k, dec @ w_v
dec = LN(dec + causal_attention(q, k, v) @ w_o)
q, k, v = dec @ w_q_cross, enc @ w_k_cross, enc @ w_v_cross
dec = LN(dec + attention(q, k, v) @ w_o_cross)
dec = LN(dec + GELU(dec @ w_mlp1) @ w_mlp2)

Target self-attention at position t reads positions 0 through t, including itself. Exclude later positions with negative infinity before softmax. Cross-attention reads every source position and has no causal mask. With L=0, project target token-plus-position embeddings directly with w_head.

Shared operations

Split projected features into consecutive heads of width d/num_heads. Divide q @ k.T by sqrt(d/num_heads), apply stable softmax over key positions, combine value vectors and concatenate heads in their original order before the output projection. No token ID, including zero, is treated as padding.

LN uses the feature mean and population variance of each token: (token - mean) / sqrt(variance + 1e-5), with no learned scale or bias. GELU(u) = 0.5u(1 + tanh(sqrt(2/pi)(u + 0.044715u**3))). The feed-forward width is d. There are no biases, dropout, extra embedding scaling or additional final normalization.