NNX Transformer Decoder Block

Read a decoder prefix and an encoder sequence

This is an encoder-decoder Transformer block. It receives decoder states x_q shaped (Tq, D) and encoder states x_kv shaped (Tk, D). The two sequences can have different lengths but share feature width D.

The block first mixes information within the decoder prefix, then reads the encoder sequence, then applies a token-wise feed-forward network. A decoder-only block typically omits the encoder cross-attention branch; the exact normalization and feed-forward choices vary by architecture.

Three sequential residual updates

Each branch normalizes its decoder input before transforming it. Write the intermediate values explicitly so the residual paths are clear:

h = ln1(x_q)
a = x_q + self_attention(h, h, causal=True)
b = a + cross_attention(ln2(a), x_kv, causal=False)
output = b + ff2(relu(ff1(ln3(b))))

Each residual adds the unnormalized input to that branch. The later branches consume the updated decoder states. The encoder states x_kv are used as supplied: this block does not apply a LayerNorm to them. All three norms have their own learned scale and bias and normalize across features separately at each decoder position.

Two attention modules with separate parameters

Both modules can use the same MHA class. Each owns four D-to-D linear projections, but self-attention and cross-attention must be distinct instances with different parameters.

Project Q from the first argument and K/V from the second. With H heads and Dh=D//H, Q has shape (H, Tq, Dh) and K/V have shape (H, Tk, Dh). Scores have shape (H, Tq, Tk). Divide dot products by sqrt(Dh), softmax over the key-position axis, multiply by V, then join heads and project back to D features.

In the self-attention call, both sequences are h and have equal length. Allow key position j only when j <= i for query position i. Include the diagonal. Replace forbidden scores with negative infinity before softmax, assuming finite allowed scores. Every row has at least one allowed position; an entirely masked row would have undefined softmax.

Cross-attention uses all supplied encoder positions here. This assumes the whole source sequence is available and all its positions are valid. Other tasks can require source padding or streaming masks; absence of a cross-attention mask is this exercise's contract, not a universal rule.

A small cross-attention example

One decoder query may attend to more encoder positions than its own index suggests. Here all content scores are zero, so each query averages three encoder values equally:

import jax
import jax.numpy as jnp

q = jnp.zeros((1, 2, 1))       # one head, two decoder queries
k = jnp.zeros((1, 3, 1))       # three encoder keys
v = jnp.array([[[2.0], [4.0], [9.0]]])
scores = q @ k.transpose(0, 2, 1)
weights = jax.nn.softmax(scores, axis=-1)
output = weights @ v
assert jnp.allclose(output, jnp.array([[[5.0], [5.0]]]))

Applying a lower-triangular mask to this cross-attention would instead yield 2 for the first query and 3 for the second. That would discard available source information. This example isolates attention; the full block also projects its inputs and applies three residual updates.

What causality guarantees here

For fixed encoder states, changing a future decoder row cannot affect an earlier output row. Causal self-attention enforces that restriction; feature-wise normalization, cross-attention queries and the FFN preserve it. This assumes x_kv itself does not contain information that should be hidden from the decoder. In next-token training, target shifting is a separate responsibility; the mask does not shift inputs or labels.

No positional embeddings, cache, dropout, padding masks, final LayerNorm or vocabulary projection are included. The output is a feature array, not token probabilities. If all three branch outputs are zero, the residual paths return x_q exactly.

The diagram's causal and cross-attention panels illustrate the two attention patterns. They do not show normalization or residual additions.