Skip to content

← Flax step 40 of 100

Hard Primitives

Transformer Decoder Block (Pre-LN)

Jump to the task

Read the decoder prefix and the encoder sequence

This block updates decoder representations in three stages:

  1. Causal self-attention reads the current and earlier decoder positions.
  2. Cross-attention reads the supplied encoder representations.
  3. A feed-forward network transforms each decoder position independently.

Each stage adds its result to the representation it received. This exercise uses pre-LN: normalize before the sublayer, then add the update. It is a particular encoder-decoder block, not a complete language model or a promise of stable training for every configuration.

Track three residuals and four normalizations

Name the intermediate decoder states x1 and x2:

x1  = x_q + SelfAttention(LN0(x_q), causal_mask)
x2  = x1  + CrossAttention(queries=LN1(x1), keys_and_values=LN2(x_kv))
out = x2  + Dense(D)(relu(Dense(d_ff)(LN3(x2))))

There are four separate LayerNorms in this reference, including one for the encoder input to cross-attention. Each has its own learned scale and offset. The two attention modules also have separate parameters.

The second residual adds to x1; the third adds to x2. Neither restarts from the original x_q. The encoder input x_kv supplies keys and values; it is not added to a decoder residual or updated by this block.

Apply the causal mask only to self-attention

For decoder length T_q, the self-attention mask has shape (T_q, T_q) and allows key position j exactly when j <= i, where i is the query position. The diagonal stays allowed so a token can attend to itself.

Cross-attention has a different key axis: encoder positions. This exercise treats every encoder position as valid and applies no cross-attention mask. A real padded encoder sequence may need a padding mask; the absence of a causal mask does not mean cross-attention can never use a mask.

With three decoder positions and two encoder positions, the permissions are:

import jax.numpy as jnp

self_mask = jnp.tril(jnp.ones((3, 3), dtype=bool))
cross_permissions = jnp.ones((3, 2), dtype=bool)

assert jnp.array_equal(self_mask, jnp.array([
    [True, False, False],
    [True, True, False],
    [True, True, True],
]))
assert self_mask[0, 0] and not self_mask[0, 1]
assert jnp.all(cross_permissions[0])  # Even the first query can read both keys.

cross_permissions illustrates the allowed pairs; you do not need to pass this all-true array to the module. With encoder inputs held fixed, later decoder inputs cannot affect earlier outputs in this block.

Follow the two sequence lengths

Let x_q have shape (T_q, D), x_kv have shape (T_kv, D_kv), H = num_heads and F = d_ff. Both attention modules use qkv_features=D, giving head width D // H.

Array Shape
Self-attention scores (H, T_q, T_q)
Cross-attention queries (T_q, H, D/H)
Cross-attention keys and values (T_kv, H, D/H) each
Cross-attention scores (H, T_q, T_kv)
Each attention output and decoder residual (T_q, D)
FFN hidden activation (T_q, F)
Final output (T_q, D)

The sequence lengths can differ. The input widths can differ too: separate projections map D and D_kv into the attention head space. Only D must be divisible by H. The default attention output width follows the query input, so both residual additions remain valid.

All four LayerNorms normalize the final feature axis within each row. The encoder LayerNorm uses D_kv features; the others use D.

Problem

Implement decoder_block_forward(seed, x_q, x_kv, num_heads, d_ff):

  1. Define a compact Linen module with integer num_heads and d_ff fields. Build the lower-triangular self-attention mask.
  2. Follow the three-stage computation above, creating a fresh unnamed nn.LayerNorm() for each LN0 through LN3, in that order.
  3. Use two unnamed nn.MultiHeadDotProductAttention modules, each with num_heads=H and qkv_features=D. Call the first on normalized decoder input with mask=causal_mask. Call the second with normalized decoder queries and normalized encoder input as its two positional arguments, without a mask.
  4. Build the FFN with unnamed nn.Dense(F), nn.relu and nn.Dense(D). Keep all layer defaults, including biases, LayerNorm epsilon 1e-6 and zero attention dropout. Do not add a final normalization.
  5. Initialize the module directly with jax.random.PRNGKey(int(seed)), x_q and x_kv. Apply those variables to the same inputs and return out.reshape(-1).

The default scopes are LayerNorm_0 through LayerNorm_3, MultiHeadDotProductAttention_0 and _1, and Dense_0 and Dense_1. Preserve those names and the call order to reproduce seeded parameters.

Both inputs are float32 arrays with positive dimensions. H and F are positive and D is divisible by H. Return a one-dimensional array of length T_q * D, with features varying fastest. Keep seed, num_heads and d_ff static if JIT-compiling the whole function.

This is a full-sequence forward pass, with no decoding cache or position encoding added here. The causal and cross-attention panels below illustrate its two attention stages. The bidirectional panel is a comparison.

Loading visualization…