📘

Transformer Block Forward Pass

HardTransformer Components 📚 LLM 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
FFN: ReLU(x @ W1) @ W2 with np.random.seed(42), scale 0.1
LayerNorm: (x - mean) / sqrt(var + 1e-5), no affine params
d_ff = 4 * d

Input:

Line 1: n d (sequence length, model dimension)
Next n lines: input matrix X

Output: Output matrix after one transformer block, rounded to 4 decimal places.

Example:

Input:

2 2
1 0
0 1

Output:

[ 0.9998 -0.9998]
[-0.9998  0.9998]

Reasoning:

First, we calculate the self-attention: since it's single-head attention with no projection weights, Attention(x) simplifies to $x \cdot x^T$ normalized by the square root of the sequence length ( $\sqrt{2}$ ). For the input matrix $X = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , Attention(x) becomes $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ .
Then, we apply the layer norm to the sum of x and Attention(x): $x = \frac{x + Attention(x) - mean}{\sqrt{var + 1e-5}}$ . Since $x + Attention(x) = \begin{bmatrix} 2 & 0 \\ 0 & 2 \end{bmatrix}$ , $mean = 1$ , and $var = 0$ , $x$ becomes $\begin{bmatrix} 2 - 1 \\ 2 - 1 \end{bmatrix} / \sqrt{0 + 1e-5} \approx \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , but due to the specifics of the layer norm calculation, it actually results in $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ .
Next, the feed-forward network (FFN) is applied: $FFN(x) = ReLU(x \cdot W1) \cdot W2$ . Given $d_{ff} = 4 \cdot d = 8$ , $W1$ is a $2 \times 8$ matrix and $W2$ is an $8 \times 2$ matrix, both initialized with $\np.random.seed(42)$ and scale $0.1$ . After computation, $FFN(x)$ is calculated and then added to $x$ .
Finally, another layer norm is applied to the sum of $x$ and $FFN(x)$ , resulting in the output matrix after one transformer block, which after computation and rounding to 4 decimal places gives the output: $\begin{bmatrix} 0.9998 & -0.9998 \ -0.9998 & 0.9998 \end{b

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

More from LLM 1: Foundations

Feed-Forward NetworkEasy Layer NormalizationEasy KV-Cache AttentionHard Multi-Head AttentionMedium Causal Attention MaskMedium

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

0/0

Run code to see test results.