Transformer Block Forward Pass

hardFREETransformer ComponentsLLM 1: Foundations

Implement a single transformer block forward pass.

A transformer block consists of:

Multi-head self-attention with residual connection and layer norm: x = LayerNorm(x + Attention(x))
Feed-forward network with residual connection and layer norm: x = LayerNorm(x + FFN(x))

For simplicity:

Single-head attention (h=1), no projection weights: $Q = K = V = x$ , and $\text{Attention}(x) = \text{softmax}(x x^T / \sqrt{d})\, x$ with the softmax taken over each row
FFN: ReLU(x @ W1) @ W2, where $W_1$ is $d \times d_{ff}$ and $W_2$ is $d_{ff} \times d$ . Call np.random.seed(42), then draw $W_1$ first and $W_2$ second from a standard normal distribution, each multiplied by 0.1
LayerNorm: (x - mean) / sqrt(var + 1e-5) over each row, no affine params
d_ff = 4 * d

Input: X, a NumPy array of shape (n, d) (sequence length n, model dimension d).

Output: The output matrix after one transformer block, shape (n, d). The tests print each value rounded to 4 decimal places.

Example

Input

2 2
1 0
0 1

Output

[1.0000 -1.0000]
[-1.0000 1.0000]

Reasoning

Input. $X$ holds two tokens with $d = 2$ features: $[1, 0]$ and $[0, 1]$ .
Attention. With $Q = K = V = X$ , the scores are $XX^T/\sqrt{2} = I/\sqrt{2}$ : each token scores $0.7071$ on itself and $0$ on the other. Row-wise softmax gives weights $[0.6698,\ 0.3302]$ and $[0.3302,\ 0.6698]$ ; since $X = I$ , these weights are also the attention output.
Residual + LayerNorm 1. $X + \text{Attn}$ has rows $[1.6698,\ 0.3302]$ and $[0.3302,\ 1.6698]$ . Each has mean $1$ and variance $0.4486$ , so normalising gives $H \approx$ rows $[1, -1]$ and $[-1, 1]$ .
FFN. With $W_1$ ( $2 \times 8$ ) and $W_2$ ( $8 \times 2$ ) scaled by $0.1$ , $\text{ReLU}(HW_1)W_2$ is tiny: rows $[-0.0393,\ -0.0065]$ and $[-0.0088,\ -0.0091]$ .
Residual + LayerNorm 2. The sums $[0.9607,\ -1.0065]$ and $[-1.0088,\ 0.9909]$ normalise to $\pm 0.99999$ , which rounds to [1.0000 -1.0000] and [-1.0000 1.0000].
Notice that with $d = 2$ , LayerNorm sends any row with two different entries to $\pm 1$ (minus a sliver from $\epsilon$ ), so the output just records which feature is larger.

Constraints

1 <= n <= 5, 2 <= d <= 8
Single head attention (no projections)
d_ff = 4 * d
Layer norm with ε=1e-5, no affine
np.random.seed(42) before generating FFN weights
Round to 4 decimal places

Learn this in

Deep LearningCV Study Plan Introduction to LLMsLLM Study Plan Transformer ArchitectureLLM Study Plan CNNs & Sequence ModelsML Study Plan

More from LLM 1: Foundations

Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium Causal Attention Maskmedium

Concept primer

Background Knowledge

The Transformer architecture, introduced in the paper "Attention is All You Need" by Vaswani et al., revolutionized the field of Natural Language Processing (NLP). It relies heavily on self-attention mechanisms, which allow the model to attend to different parts of the input sequence simultaneously and weigh their importance. In the context of this problem, we're dealing with a single Transformer Block, which is a fundamental component of the Transformer architecture. This block consists of two main sub-layers: Multi-head Self-Attention and a Feed-Forward Network (FFN), each followed by a residual connection and Layer Normalization.

The Multi-head Self-Attention mechanism allows the model to jointly attend to information from different representation subspaces at different positions. However, for simplicity, this problem reduces it to single-head attention, eliminating the need for projection weights. The Feed-Forward Network (FFN), on the other hand, applies a transformation to each position separately and identically, which includes two linear layers with a ReLU activation function in between. Layer Normalization is used to normalize the inputs to each sub-layer, which helps in stabilizing the learning process and reducing the dependency on the choice of hyperparameters.

Understanding the mathematical formulations of these components is crucial. For instance, self-attention can be represented as $Attention(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d}})V$ , where $Q$ , $K$ , and $V$ are the query, key, and value matrices, respectively, and $d$ is the dimensionality of the input. Layer Normalization can be formulated as $\frac{x - \text{mean}(x)}{\sqrt{\text{var}(x) + \epsilon}}$ , where $\epsilon$ is a small constant for numerical stability.

Algorithm/Approach

The general approach to solving this problem involves implementing the two main components of a Transformer Block: Multi-head Self-Attention (simplified to single-head for this problem) and the Feed-Forward Network (FFN), along with Layer Normalization and residual connections. The algorithm pattern typically follows the sequence: input preparation, self-attention calculation, feed-forward network application, and finally, applying layer normalization and residual connections after each sub-layer.

Step-by-Step Strategy

Read Input: Read the sequence length $n$ and model dimension $d$ , followed by the input matrix $X$ .
Self-Attention Calculation: Compute the attention output using the simplified single-head self-attention mechanism.
Apply Layer Norm and Residual Connection: Apply layer normalization to the sum of the input and the attention output, followed by a residual connection.
Feed-Forward Network (FFN): Apply the FFN to the output from the previous step, which involves two linear transformations with a ReLU activation in between.
Apply Layer Norm and Residual Connection Again: Apply layer normalization to the sum of the output from the FFN and its input, followed by another residual connection.
Output: The final output after one transformer block, rounded to 4 decimal places.

Common Pitfalls

Incorrect implementation of the self-attention mechanism or the FFN.
Forgetting to apply layer normalization and residual connections appropriately.
Not using the correct scale for the FFN weights or not setting the random seed as specified.

Time & Space Complexity

The time complexity of a Transformer Block is dominated by the self-attention mechanism, which is $O(n^2 \cdot d)$ , where $n$ is the sequence length and $d$ is the model dimension. The space complexity is $O(n \cdot d)$ , mainly due to the storage of the input and output matrices. Note that these complexities assume a simplified version of the Transformer Block as described in the problem statement.

solution.py

Python 3.13

Test results