Decoder Architecture › Layer Stacking and Information Flow

Michael BrenndoerferUpdated June 19, 20254 min read
Language AI HandbookAll chapters
Part 16 of 67Transformer Architectures
Chapter 112Decoder Architecture › Layer Stacking and Information Flow
Reading section 7 of 16 · Layer Stacking and Information Flow · 4 minSection 7 of 16 · 4 min

Book contents

Section 7 of 16 · 4 min read

Decoder Architecture › Layer Stacking and Information FlowOne section · Sections in this chapter7 / 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

Decoder blocks are stacked sequentially. The output of block ll becomes the input to block l+1l + 1. Each layer refines the representations, building increasingly abstract features.

In[15]:
Code
def trace_layer_outputs(model, token_ids):
    """
    Trace how representations evolve through decoder layers.

    Returns the hidden states after each layer.
    """
    seq_len = len(token_ids)

    # Initial embeddings
    tok_emb = model.token_embedding[token_ids]
    pos_emb = model.position_embedding[:seq_len]
    x = tok_emb + pos_emb

    layer_outputs = [x.copy()]

    # Through each block
    for block in model.blocks:
        x = block(x)
        layer_outputs.append(x.copy())

    # After final norm
    x = rms_norm(x, model.final_norm_gamma)
    layer_outputs.append(x.copy())

    return layer_outputs


# Trace through our model
test_tokens = np.array([42, 100, 256, 789, 1000, 500, 333, 42])
layer_outputs = trace_layer_outputs(model, test_tokens)
Out[16]:
Console
Tracing 8 tokens through 6 layers

Layer output statistics (L2 norm of representations):
--------------------------------------------------
Embedding   : mean norm = 0.463, std = 0.028
Block 1     : mean norm = 24.058, std = 3.668
Block 2     : mean norm = 36.957, std = 4.267
Block 3     : mean norm = 44.214, std = 5.426
Block 4     : mean norm = 53.139, std = 4.717
Block 5     : mean norm = 60.062, std = 3.791
Block 6     : mean norm = 64.284, std = 2.553
Final Norm  : mean norm = 16.000, std = 0.000

The representation norms reveal how the network transforms token embeddings. Starting from the embedding layer with relatively small norms (around 0.3), the representations grow as they pass through successive blocks. This growth reflects the accumulation of information via residual connections: each block adds its contribution to the running representation. The final layer norm rescales the output. This ensures consistent magnitudes before the language model head computes logits.

Out[17]:
Visualization
Line plot showing L2 norm of representations for 8 tokens across embedding layer, 6 decoder blocks, and final normalization.
Representation norm growth through decoder layers. Each line represents one of the 8 input tokens. Norms generally increase through the decoder blocks as residual connections accumulate contributions, then are rescaled by the final layer norm.

As we go deeper, representation norms tend to grow or stabilize depending on initialization and normalization. The final layer norm ensures consistent scale before the language model head.

Out[18]:
Visualization
Heatmap showing cosine similarity between decoder layer outputs, with diagonal of 1.0 and off-diagonal values showing layer relationships.
Cosine similarity between layer representations for the final token position. Earlier layers show higher similarity; deeper layers diverge as they develop specialized features.

The similarity matrix reveals how representations evolve. Adjacent layers tend to be more similar (the residual connections preserve information), while distant layers may diverge significantly as the model builds higher-level abstractions.