Skip to content

← Attention Variants step 11 of 14

Hard End-to-End

Encoder-Decoder Transformer Forward Pass

Implement encoder_decoder_forward: encode a source sequence, let a decoder read that encoded context, and return vocabulary scores at every supplied target position.

The encoder can use the entire source sequence. The decoder can use the entire encoded source but only its current and earlier target inputs. This function processes the supplied target sequence in one forward pass. It does not shift target IDs, choose new tokens or run a generation loop.

Inputs and output

  • src_ids: (N, S), integer-valued source IDs in [0, V_src), delivered as floats.
  • tgt_ids: (N, T), integer-valued target IDs in [0, V_tgt), delivered as floats.
  • src_emb: (V_src, d), source token embeddings.
  • tgt_emb: (V_tgt, d), target token embeddings.
  • enc_pos_embed: (S, d), source position embeddings.
  • dec_pos_embed: (T, d), target position embeddings.
  • enc_blocks: (E, 6, d, d), encoder weights.
  • dec_blocks: (L, 12, d, d), decoder weights.
  • w_head: (d, V_tgt), output projection.
  • num_heads: a positive integer dividing d.

N, S, T, V_src, V_tgt and d are positive. E and L may be zero. Source and target lengths and vocabulary sizes may differ. Parameters are finite and share a floating-point dtype and device; IDs are on that device. Shapes agree.

Return raw logits of shape (N, T, V_tgt), preserving parameter dtype and device. Leave all inputs unchanged. Cast IDs to integers before lookup. Use every ID as supplied, including zero: there are no padding masks or special-token rules.

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[tgt_ids] + 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.

3. Project to the target vocabulary

Return dec @ w_head. Do not apply a vocabulary softmax or argmax. With L=0, project target token-plus-position embeddings directly; the source then has no effect on the result.

The visualization below illustrates the attention patterns. The equations above define the exact normalization and weight layout for this exercise.

Loading visualization…