Decoder Architecture › Anatomy of a Decoder Block
Language AI HandbookAll chapters Decoder Architecture › Anatomy of a Decoder BlockOne section · Sections in this chapter5 / 16 ↓
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
Now that we understand causal masking, let's see how it fits into a complete decoder block. A decoder block is structurally similar to an encoder block, with one critical difference: its self-attention layer applies causal masking. But the block is more than just masked attention; it's a carefully designed composition of components that work together to transform token representations.
The Two-Sublayer Structure
Each decoder block consists of two sublayers that serve complementary purposes:
-
Causal Multi-Head Self-Attention: This sublayer enables cross-position communication. Each token gathers information from previous tokens (and itself), building a representation that incorporates relevant context. The causal mask ensures this communication respects temporal order.
-
Feed-Forward Network (FFN): This sublayer provides position-wise transformation. The same two-layer MLP is applied independently to each position, adding non-linear representational capacity. Unlike attention, the FFN doesn't mix information across positions; it processes each token's representation in isolation.
These sublayers address different aspects of representation learning. Attention handles where to look in the sequence. The FFN handles how to transform what was gathered. Together, they enable the block to both aggregate context and compute complex functions of that context.
Supporting Components
Two additional components wrap each sublayer to enable deep networks:
-
Residual Connections: Skip connections around each sublayer create additive shortcuts. Instead of computing , the block computes . This means the sublayer only needs to learn the delta, or refinement, to add to the representation. Residuals also provide direct gradient paths that bypass potentially problematic nonlinearities.
-
Layer Normalization: Normalizes activations to stabilize training. Modern decoders use Pre-LN architecture, applying normalization before each sublayer rather than after. Pre-LN produces more stable gradients at initialization and enables training without careful learning rate warmup.
The data flow through a Pre-LN decoder block follows this pattern:
- Normalize → Attention → Add residual
- Normalize → FFN → Add residual
Let's implement each component and assemble them into a complete decoder block.
Utility Functions
First, we need a few utility functions. The softmax function converts raw scores to probabilities. RMS normalization scales activations to have unit root-mean-square. GELU provides a smooth nonlinearity for the feed-forward network:
def softmax(x, axis=-1):
"""Numerically stable softmax."""
x_max = np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x - x_max)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
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)))
The smooth transition of GELU can improve gradient flow during training compared to ReLU's sharp corner at zero.
Causal Multi-Head Attention
The attention layer is the heart of the decoder block. It projects input tokens into QKV representations, then computes attention with causal masking. Multiple heads allow the model to attend to different aspects of the context simultaneously:
class CausalMultiHeadAttention:
"""
Multi-head attention with causal masking.
Each head learns different aspects of token relationships,
but all heads respect the causal constraint.
"""
def __init__(self, d_model, n_heads):
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
# Initialize projection matrices
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, x):
"""
Apply causal multi-head attention.
Args:
x: Input tensor of shape (seq_len, d_model)
Returns:
Output tensor of shape (seq_len, d_model)
"""
seq_len = x.shape[0]
# Project to queries, keys, values
Q = x @ self.W_q # (seq_len, d_model)
K = x @ self.W_k
V = x @ self.W_v
# Reshape for multi-head: (seq_len, n_heads, d_k)
Q = Q.reshape(seq_len, self.n_heads, self.d_k)
K = K.reshape(seq_len, self.n_heads, self.d_k)
V = V.reshape(seq_len, self.n_heads, self.d_k)
# Transpose to (n_heads, seq_len, d_k) for batch matmul
Q = Q.transpose(1, 0, 2)
K = K.transpose(1, 0, 2)
V = V.transpose(1, 0, 2)
# Compute attention scores: (n_heads, seq_len, seq_len)
scores = Q @ K.transpose(0, 2, 1) / np.sqrt(self.d_k)
# Apply causal mask
causal_mask = create_causal_mask(seq_len)
scores = scores + causal_mask # Broadcasting adds mask to each head
# Softmax to get attention weights
attn_weights = softmax(scores, axis=-1)
# Apply attention to values: (n_heads, seq_len, d_k)
attn_output = attn_weights @ V
# Transpose back and reshape: (seq_len, d_model)
attn_output = attn_output.transpose(1, 0, 2).reshape(
seq_len, self.d_model
)
# Output projection
output = attn_output @ self.W_o
return output
def num_parameters(self):
"""Count total parameters in this layer."""
return 4 * self.d_model * self.d_modelThe attention layer computes the same scaled dot-product attention as an encoder, but the causal mask ensures the upper triangle of attention weights becomes zero. Each position's output is a weighted sum of value vectors from itself and all previous positions.
Position-Wise Feed-Forward Network
After attention aggregates context, the FFN transforms each position independently. It expands the representation to a larger hidden dimension, applies a nonlinearity, then projects back to the model dimension. This expansion-contraction pattern adds significant representational capacity:
class FeedForward:
"""
Position-wise feed-forward network.
Applies the same two-layer MLP independently to each position.
"""
def __init__(self, d_model, d_ff):
self.d_model = d_model
self.d_ff = d_ff
# Initialize weights
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):
"""
Apply feed-forward transformation.
Args:
x: Input of shape (seq_len, d_model)
Returns:
Output of shape (seq_len, d_model)
"""
# First linear + GELU activation
hidden = gelu(x @ self.W1 + self.b1)
# Second linear
output = hidden @ self.W2 + self.b2
return output
def num_parameters(self):
return (
self.d_model * self.d_ff
+ self.d_ff
+ self.d_ff * self.d_model
+ self.d_model
)Assembling the Decoder Block
Now we combine these pieces into a complete decoder block. The Pre-LN architecture applies normalization before each sublayer, which produces more stable gradients and enables training without warmup:
class DecoderBlock:
"""
A single transformer decoder block.
Uses Pre-LN architecture: normalization before each sublayer.
Causal masking is applied in the self-attention.
"""
def __init__(self, d_model, n_heads, d_ff):
self.d_model = d_model
# Sublayers
self.attention = CausalMultiHeadAttention(d_model, n_heads)
self.ffn = FeedForward(d_model, d_ff)
# Layer norm parameters (for RMS norm, just scale)
self.norm1_gamma = np.ones(d_model)
self.norm2_gamma = np.ones(d_model)
def __call__(self, x):
"""
Process input through the decoder block.
Args:
x: Input of shape (seq_len, d_model)
Returns:
Output of shape (seq_len, d_model)
"""
# Self-attention with residual
normed = rms_norm(x, self.norm1_gamma)
attn_out = self.attention(normed)
x = x + attn_out
# FFN with residual
normed = rms_norm(x, self.norm2_gamma)
ffn_out = self.ffn(normed)
x = x + ffn_out
return x
def num_parameters(self):
"""Count parameters in this block."""
attn_params = self.attention.num_parameters()
ffn_params = self.ffn.num_parameters()
norm_params = 2 * self.d_model # Two RMS norms
return attn_params + ffn_params + norm_paramsTesting the Decoder Block
Let's instantiate a decoder block and verify it works correctly. We'll use typical hyperparameters: 256-dimensional model, 8 attention heads (32 dimensions each), and a 4x FFN expansion:
Decoder Block Configuration ======================================== Model dimension (d_model): 256 Number of heads: 8 Head dimension: 32 FFN hidden dim (d_ff): 1024 Parameters: 788,224 Input shape: (10, 256) Output shape: (10, 256)
The block contains roughly 790K parameters, dominated by the feed-forward network (which uses a 4x expansion ratio: ). The input and output shapes match, allowing blocks to be stacked without dimension changes. This "residual-friendly" design is what enables deep networks: each block can refine the representation without fundamentally altering its structure.
The decoder block transforms each position independently while allowing controlled information flow from past positions through causal attention. The FFN adds non-linear transformation capacity without any cross-position interaction.