Train Encoder-Decoder Seq2Seq Step

Implement train_seq2seq_step: run an encoder-decoder model with supplied target inputs, calculate all parameter gradients by hand, and take one SGD step.

Teacher forcing means feeding the known target prefix into the decoder during training. For target IDs [1, 4, 2, 3], the decoder receives [1, 4, 2] and must predict [4, 2, 3]. It does not feed back its own predictions. The causal mask still matters: a position must not read the later target input it is predicting.

Inputs and return value

• src_ids: (N, S), integer-valued source IDs in [0, V_src), delivered as floats.
• tgt_ids: (N, T+1), integer-valued target IDs in [0, V_tgt), delivered as floats.
• src_emb: (V_src, d); tgt_emb: (V_tgt, d), independent embedding tables.
• enc_pos_embed: (S, d); dec_pos_embed: (T, d), position embeddings.
• enc_blocks: (E, 6, d, d); dec_blocks: (L, 12, d, d), block weights.
• w_head: (d, V_tgt), an independent vocabulary projection.
• num_heads: a positive integer dividing d.
• lr: a finite, nonnegative learning rate.

N, S, T, d and the vocabulary sizes are positive. E and L may be zero. Source and target lengths and vocabularies may differ. Parameters are finite and share a floating-point dtype and device; IDs are on that device. Shapes agree. Cast IDs to integers and leave every input unchanged.

Return one flat tensor, concatenating the row-major flattened updated parameters: [src_emb, tgt_emb, enc_pos_embed, dec_pos_embed, enc_blocks, dec_blocks, w_head]. Preserve parameter dtype and device. Return weights, not gradients or loss. With lr=0, the result contains the original parameters in that order.

Forward pass

Set decoder_inputs = tgt_ids[:, :-1] and loss targets = tgt_ids[:, 1:]. All positions participate; there is no padding mask, EOS stopping or special handling for token zero. The final supplied target token is a label only.

Shared operations

Split each projected vector into consecutive heads of width d/num_heads. For each head, divide q @ k.T by sqrt(d/num_heads), apply stable softmax over key positions, and use those weights to combine value vectors. Concatenate head outputs in their original order before applying the output projection. Keep batch examples separate.

LN normalizes each token across its d features using population variance: (token - mean) / sqrt(variance + 1e-5). It has no learned scale or bias. Every sublayer uses post-normalization: add the residual, then normalize.

The feed-forward width is d. Use the elementwise tanh approximation: GELU(u) = 0.5u(1 + tanh(sqrt(2/pi)(u + 0.044715u**3))). There are no linear biases, dropout, extra embedding scaling or additional final normalization.

1. Encode the source

Start with enc = src_emb[src_ids] + enc_pos_embed, broadcasting positions over the batch. Each encoder block contains [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)

Every source position can attend to every source position. Keep the final enc as the fixed context for all decoder blocks. With E=0, the context is simply the source token-plus-position embeddings.

2. Decode the supplied target inputs

Start with dec = tgt_emb[decoder_inputs] + dec_pos_embed. Each decoder block has:

• Slots 0 through 3: self-attention matrices [w_q_self, w_k_self, w_v_self, w_o_self].
• Slots 4 through 7: cross-attention matrices [w_q_cross, w_k_cross, w_v_cross, w_o_cross].
• Slots 8 and 9: feed-forward matrices [w_mlp1, w_mlp2].
• Slots 10 and 11: unused. Ignore their values, even if they are nonzero.

Apply these three sublayers in order, using the updated dec at each step:

q, k, v = dec @ w_q_self, dec @ w_k_self, dec @ w_v_self
dec = LN(dec + causal_attention(q, k, v) @ w_o_self)
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)

In target self-attention, query position t may read key positions 0 through t, including itself. Exclude later positions before softmax, using negative infinity rather than a finite score penalty. For T=3, the allowed keys are [0], [0,1], then [0,1,2].

Cross-attention uses decoder queries and source keys/values. Its score shape per head is (T, S), and every query can read all S source positions. Do not reuse the target causal mask here, even when source and target lengths match.

Loss and manual backward pass

Project the final decoder state with w_head. Use mean cross-entropy against the shifted targets over all NT decoder positions. The logit gradient is (softmax(logits) - one_hot(targets)) / (NT). Use stable softmax, with no extra averaging when forming parameter gradients. The loss value is not returned.

Differentiate w_head, then visit decoder blocks in reverse order: feed-forward, cross-attention, then causal self-attention. Each post-LN sublayer first backpropagates through its own normalization, then through both the residual skip and the sublayer. Save each normalization's output and standard deviation.

Cross-attention produces two input gradients: one through decoder queries, and one through encoder keys and values. Sum the encoder-output contributions from every decoder block before differentiating the encoder stack. In self-attention, add the query, key, value and residual contributions to the same input gradient.

Accumulate embedding gradients for every occurrence of each lookup ID. Repeated IDs add to the same row. Sum position-embedding gradients across examples. The last target token is not an input lookup. Unused decoder slots 10 and 11 have zero gradients and remain unchanged, including when initially nonzero.

With E=0, cross-attention gradients flow directly to the source embedding and position tables. With L=0, the output has no source dependency: source embeddings, source positions and encoder blocks remain unchanged. The target embeddings, target positions and output head still train.

Update

Compute all gradients from the original weights, then update each parameter as parameter - lr * gradient. There is no clipping, weight decay, momentum or optimizer state. Build the forward operations and derivatives from basic tensor operations in both frameworks; do not use autograd, backward(), jax.grad, packaged transformer layers or optimizer helpers.