Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Transformer Block Assembly › A Worked Example: Token by Token

Michael BrenndoerferUpdated June 13, 20257 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 110Transformer Block Assembly › A Worked Example: Token by Token

Reading section 6 of 16 · A Worked Example: Token by Token · 7 minSection 6 of 16 · 7 min

Section 6 of 16 · 7 min read

Transformer Block Assembly › A Worked Example: Token by TokenOne section · Sections in this chapter6 / 16 ↓

Sections in this chapter16

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

To truly understand how data flows through a transformer block, let's trace a small concrete example with actual numbers. We'll use a tiny configuration so you can follow each computation:

$d_{\text{model}} = 4$ : each token is represented by just 4 numbers
$n_{\text{heads}} = 2$ : we split attention into 2 parallel heads
$d_{\text{ff}} = 8$ : the FFN's hidden layer has 8 dimensions
$n = 3$ : our sequence contains 3 tokens

In[14]:

Code

# Tiny configuration for hand-traceable computation
d_model_tiny = 4
n_heads_tiny = 2
d_ff_tiny = 8
seq_len_tiny = 3


# Create a tiny transformer block with standard FFN for simplicity
class TinyTransformerBlock:
    """Minimal transformer block for demonstration."""

    def __init__(self):
        self.d_model = d_model_tiny

        # RMSNorm weights (all ones for simplicity)
        self.norm1_weight = np.ones(d_model_tiny)
        self.norm2_weight = np.ones(d_model_tiny)

        # Attention weights (small random values)
        scale = 0.5
        self.W_q = np.random.randn(d_model_tiny, d_model_tiny) * scale
        self.W_k = np.random.randn(d_model_tiny, d_model_tiny) * scale
        self.W_v = np.random.randn(d_model_tiny, d_model_tiny) * scale
        self.W_o = np.random.randn(d_model_tiny, d_model_tiny) * scale

        # FFN weights
        self.W1 = np.random.randn(d_model_tiny, d_ff_tiny) * scale
        self.W2 = np.random.randn(d_ff_tiny, d_model_tiny) * scale

    def rmsnorm(self, x, weight):
        """Apply RMSNorm."""
        rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + 1e-6)
        return weight * (x / rms)

    def attention(self, x):
        """Simplified single-head attention for demonstration."""
        Q = x @ self.W_q
        K = x @ self.W_k
        V = x @ self.W_v

        # Attention scores
        scale = np.sqrt(self.d_model)
        scores = (Q @ K.T) / scale
        attn_weights = softmax(scores, axis=-1)

        # Apply attention
        context = attn_weights @ V
        output = context @ self.W_o

        return output, attn_weights

    def ffn(self, x):
        """Standard FFN with GELU."""
        hidden = gelu(x @ self.W1)
        output = hidden @ self.W2
        return output, hidden

    def forward_with_intermediates(self, x):
        """Forward pass returning all intermediate values."""
        intermediates = {"input": x.copy()}

        # First sublayer: attention
        x_norm1 = self.rmsnorm(x, self.norm1_weight)
        intermediates["after_norm1"] = x_norm1.copy()

        attn_out, attn_weights = self.attention(x_norm1)
        intermediates["attention_output"] = attn_out.copy()
        intermediates["attention_weights"] = attn_weights.copy()

        h = x + attn_out  # Residual
        intermediates["after_residual1"] = h.copy()

        # Second sublayer: FFN
        h_norm = self.rmsnorm(h, self.norm2_weight)
        intermediates["after_norm2"] = h_norm.copy()

        ffn_out, ffn_hidden = self.ffn(h_norm)
        intermediates["ffn_hidden"] = ffn_hidden.copy()
        intermediates["ffn_output"] = ffn_out.copy()

        y = h + ffn_out  # Residual
        intermediates["output"] = y.copy()

        return y, intermediates


# Create tiny block and sample input
tiny_block = TinyTransformerBlock()
x_tiny = np.array(
    [[1.0, 0.5, -0.3, 0.8], [-0.2, 0.7, 0.4, -0.1], [0.6, -0.4, 0.2, 0.5]]
)

# Forward pass with intermediates
output_tiny, intermediates = tiny_block.forward_with_intermediates(x_tiny)

Out[15]:

Console

Transformer Block Walkthrough
============================================================

Input (3 tokens, 4 dimensions):
[[ 1.   0.5 -0.3  0.8]
 [-0.2  0.7  0.4 -0.1]
 [ 0.6 -0.4  0.2  0.5]]

--- ATTENTION SUBLAYER ---

After RMSNorm (normalized input):
[[ 1.421  0.711 -0.426  1.137]
 [-0.478  1.673  0.956 -0.239]
 [ 1.333 -0.889  0.444  1.111]]

Attention weights (which tokens attend to which):
[[0.057 0.828 0.114]
 [0.323 0.326 0.35 ]
 [0.101 0.77  0.129]]

Attention output:
[[-0.501  0.291  0.083 -0.598]
 [-0.255  0.105 -0.12  -0.139]
 [-0.467  0.271  0.056 -0.548]]

After residual connection (input + attention output):
[[ 0.499  0.791 -0.217  0.202]
 [-0.455  0.805  0.28  -0.239]
 [ 0.133 -0.129  0.256 -0.048]]

--- FFN SUBLAYER ---

After RMSNorm:
[[ 1.017  1.613 -0.442  0.413]
 [-0.915  1.617  0.563 -0.48 ]
 [ 0.831 -0.808  1.601 -0.303]]

FFN hidden activations (8 dimensions, showing first 4):
[[ 0.217 -0.144  0.387  0.409]
 [ 1.806  1.054  0.741 -0.005]
 [-0.061 -0.11   0.814  0.56 ]]

FFN output:
[[ 1.075 -1.042  0.461 -0.019]
 [ 1.119  0.147 -0.608 -0.293]
 [ 0.307  0.067 -0.009  0.341]]

--- FINAL OUTPUT ---

Block output (intermediate + FFN output):
[[ 1.574 -0.251  0.244  0.184]
 [ 0.664  0.953 -0.328 -0.532]
 [ 0.439 -0.062  0.246  0.292]]

Let's visualize how each token's representation evolves through the block stages:

Out[16]:

Visualization

Line plot showing 3 tokens' L2 norms changing through 6 processing stages of the transformer block. — Evolution of token representations through the transformer block. Each line traces one token's representation magnitude as it passes through normalization, attention, residual addition, and FFN stages. The residual connections visibly preserve and augment the representations.

Understanding the Attention WeightsLink Copied

The attention weights matrix shows how each token attends to others in the sequence. Let's visualize this pattern:

Out[17]:

Visualization

Heatmap showing 3x3 attention weight matrix with varying attention patterns per token. — Attention pattern from our worked example. Each row shows how one token distributes its attention across all tokens. Token 1 attends broadly, while Tokens 2 and 3 focus more on specific positions.

Each row in the attention matrix sums to 1.0 (due to softmax normalization). If we denote the attention weight from token $i$ to token $j$ as $\alpha_{ij}$ , then:

\sum_{j=1}^{n} \alpha_{ij} = 1 \quad \text{for all } i

where $n$ is the sequence length. This constraint ensures that each token's output is a valid weighted average of the value vectors from all positions. No attention weight can be negative, and the weights form a probability distribution over the sequence.

Tracing the Residual ContributionLink Copied

One key insight from residual connections is that the output preserves information from the input. Let's quantify this:

In[18]:

Code

# Decompose the output into residual and sublayer contributions
input_contribution = x_tiny  # Original input (passes through both residuals)
attn_contribution = intermediates["attention_output"]
ffn_contribution = intermediates["ffn_output"]

# The output is: input + attention_output + ffn_output
reconstructed = input_contribution + attn_contribution + ffn_contribution

Out[19]:

Console

Output Decomposition
==================================================

Original input contribution:
  Magnitude: 1.8682

Attention sublayer contribution:
  Magnitude: 1.1856

FFN sublayer contribution:
  Magnitude: 2.0974

Reconstruction check (should match output):
  Match: True

Relative contributions:
  Input: 36.3%
  Attention: 23.0%
  FFN: 40.7%

Out[20]:

Visualization

Horizontal stacked bar chart showing input, attention, and FFN contributions for 3 tokens. — Decomposition of transformer block output into its component contributions. The stacked bar chart shows how the final representation for each token combines the original input with additive contributions from attention and FFN sublayers.

The decomposition reveals how each component contributes to the final representation. Mathematically, the output $\mathbf{y}$ can be written as:

\mathbf{y} = \mathbf{x} + \Delta_{\text{attn}} + \Delta_{\text{ffn}}

where:

$\mathbf{x}$ : the original input, which passes through both residual connections unchanged
$\Delta_{\text{attn}} = \text{Attention}(\text{Norm}(\mathbf{x}))$ : the attention sublayer's contribution
$\Delta_{\text{ffn}} = \text{FFN}(\text{Norm}(\mathbf{h}))$ : the FFN sublayer's contribution

The original input remains a large part of the output. This shows how residual connections preserve information while allowing sublayers to add refinements. This additive structure also explains why gradients flow easily backward through transformer blocks: they can bypass the nonlinear sublayers entirely via the residual paths.