Encoder-Decoder Architecture › Anatomy of a Decoder Block with Cross-Attention
Language AI HandbookAll chapters Encoder-Decoder Architecture › Anatomy of a Decoder Block with Cross-AttentionOne section · Sections in this chapter4 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Understanding the decoder block requires recognizing that it must accomplish two distinct but interrelated tasks: it must generate coherent, grammatical output (like any language model), while also staying faithful to the source sequence (which distinguishes translation from free generation). The architecture addresses these requirements through a carefully ordered sequence of three sublayers.
Sublayer 1: Masked Self-Attention. The decoder first attends to its own previous outputs. When generating the fourth target word, it considers the first three words it has already produced. This builds a representation of the generation context so far. The attention is masked (causal) to prevent the decoder from seeing future tokens, which is essential for autoregressive generation where we produce one token at a time.
Sublayer 2: Cross-Attention. Armed with context about what it has generated, the decoder now queries the encoder. This is where source information enters. The decoder's refined representation is queries, and the encoder's output provides keys and values. The cross-attention sublayer answers: "Given what I've generated so far, what source information is relevant for the next token?"
Sublayer 3: Feed-Forward Network. Finally, a position-wise FFN applies nonlinear transformation to each position independently. This adds representational capacity and allows the model to process the combined information from self-attention and cross-attention.
The ordering matters: self-attention establishes generation context, cross-attention enriches it with source information, and the FFN transforms the result. Each sublayer includes residual connections and layer normalization. This ensures stable gradient flow through deep stacks.
def rms_norm(x, gamma, eps=1e-6):
"""RMS Layer Normalization."""
rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + eps)
return gamma * (x / rms)
def gelu(x):
"""Gaussian Error Linear Unit activation."""
return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))
def create_causal_mask(seq_len):
"""Create causal mask for decoder self-attention."""
mask = np.triu(np.ones((seq_len, seq_len)) * float("-inf"), k=1)
return mask
class MultiHeadAttention:
"""Multi-head attention supporting both self and cross attention."""
def __init__(self, d_model, n_heads):
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
scale = np.sqrt(2.0 / (d_model + self.d_k))
self.W_q = np.random.randn(d_model, d_model) * scale
self.W_k = np.random.randn(d_model, d_model) * scale
self.W_v = np.random.randn(d_model, d_model) * scale
self.W_o = np.random.randn(d_model, d_model) * scale
def __call__(self, query_input, key_value_input=None, mask=None):
"""
Apply multi-head attention.
For self-attention: key_value_input is None, uses query_input for K, V
For cross-attention: key_value_input provides K, V from encoder
"""
if key_value_input is None:
key_value_input = query_input
q_len = query_input.shape[0]
kv_len = key_value_input.shape[0]
# Project to Q, K, V
Q = query_input @ self.W_q
K = key_value_input @ self.W_k
V = key_value_input @ self.W_v
# Reshape for multi-head
Q = Q.reshape(q_len, self.n_heads, self.d_k).transpose(1, 0, 2)
K = K.reshape(kv_len, self.n_heads, self.d_k).transpose(1, 0, 2)
V = V.reshape(kv_len, self.n_heads, self.d_k).transpose(1, 0, 2)
# Compute attention scores
scores = Q @ K.transpose(0, 2, 1) / np.sqrt(self.d_k)
if mask is not None:
scores = scores + mask
weights = softmax(scores, axis=-1)
attended = weights @ V
# Reshape and project
attended = attended.transpose(1, 0, 2).reshape(q_len, self.d_model)
output = attended @ self.W_o
return output, weights
class FeedForward:
"""Position-wise feed-forward network."""
def __init__(self, d_model, d_ff):
scale1 = np.sqrt(2.0 / (d_model + d_ff))
scale2 = np.sqrt(2.0 / (d_ff + d_model))
self.W1 = np.random.randn(d_model, d_ff) * scale1
self.b1 = np.zeros(d_ff)
self.W2 = np.random.randn(d_ff, d_model) * scale2
self.b2 = np.zeros(d_model)
def __call__(self, x):
hidden = gelu(x @ self.W1 + self.b1)
return hidden @ self.W2 + self.b2
class EncoderDecoderBlock:
"""
Decoder block with cross-attention for encoder-decoder architecture.
Contains three sublayers:
1. Masked self-attention (causal)
2. Cross-attention to encoder
3. Feed-forward network
"""
def __init__(self, d_model, n_heads, d_ff):
self.d_model = d_model
# Three attention/FFN sublayers
self.self_attn = MultiHeadAttention(d_model, n_heads)
self.cross_attn = MultiHeadAttention(d_model, n_heads)
self.ffn = FeedForward(d_model, d_ff)
# Three layer norms (pre-norm configuration)
self.norm1_gamma = np.ones(d_model)
self.norm2_gamma = np.ones(d_model)
self.norm3_gamma = np.ones(d_model)
def __call__(self, x, encoder_output, self_attn_mask=None):
"""
Forward pass through decoder block.
Args:
x: Decoder input of shape (tgt_len, d_model)
encoder_output: Encoder output of shape (src_len, d_model)
self_attn_mask: Causal mask for self-attention
Returns:
Output of shape (tgt_len, d_model)
"""
# Sublayer 1: Masked self-attention
normed = rms_norm(x, self.norm1_gamma)
self_attn_out, _ = self.self_attn(normed, mask=self_attn_mask)
x = x + self_attn_out
# Sublayer 2: Cross-attention to encoder
normed = rms_norm(x, self.norm2_gamma)
cross_attn_out, cross_weights = self.cross_attn(
normed, key_value_input=encoder_output
)
x = x + cross_attn_out
# Sublayer 3: Feed-forward
normed = rms_norm(x, self.norm3_gamma)
ffn_out = self.ffn(normed)
x = x + ffn_out
return x, cross_weightsEncoder-Decoder Block Test -------------------------------------------------- Encoder output (source): (6, 256) Decoder input (target): (4, 256) Decoder output: (4, 256) Cross-attention weights: (8, 4, 6) (n_heads, tgt_len, src_len) Each target position attends to all source positions. Cross-attention weight sums (should be 1.0): [1. 1. 1. 1.]
The decoder block successfully processes target embeddings while attending to encoder output. The cross-attention weights have shape (n_heads, tgt_len, src_len), showing that each of the 8 attention heads learns independent patterns for how target positions query source positions. Weight sums of 1.0 across each target position confirm proper softmax normalization. The output maintains the same shape as the input (tgt_len, d_model), ready to be passed to subsequent decoder layers.
Each attention head learns distinct patterns for how target positions attend to source positions. Some heads may specialize in local alignment while others capture longer-range dependencies.