Google

In a Transformer architecture, the Feed-Forward Network (FFN) sub-layer is typically implemented as a two-layer network. The standard mathematical formulation for this FFN is:

$$ \mathrm{FFN}(\mathbf{h}) = \sigma(\mathbf{h} \mathbf{W}_h + \mathbf{b}_h) \mathbf{W}_f + \mathbf{b}_f $$

Here, $$\mathbf{h}$$ is the input vector. The network's parameters consist of:
*   $$\mathbf{W}_h \in \mathbb{R}^{d \times d_h}$$ and $$\mathbf{b}_h \in \mathbb{R}^{d_h}$$: The weight matrix and bias vector for the initial linear transformation.
*   $$\mathbf{W}_f \in \mathbb{R}^{d_h \times d}$$ and $$\mathbf{b}_f \in \mathbb{R}^{d}$$: The weight matrix and bias vector for the subsequent linear transformation.

The dimension $$d$$ represents the input and output size, whereas $$d_h$$ indicates the hidden layer's size. The function $$\sigma(\cdot)$$ is the non-linear activation function utilized in the hidden layer, with the Rectified Linear Unit (ReLU) being a widespread choice.

Feed-Forward Network (FFN) Formula and Component Dimensions in Transformers

The Rectified Linear Unit (ReLU) is a common choice for the activation function $$\sigma(\cdot)$$ within the hidden layers of neural networks. It is defined to output the positive portion of its argument. When applied to an input vector $$\mathbf{h}$$, the ReLU function is given by the formula: $$\sigma_{\mathrm{relu}}(\mathbf{h}) = \mathrm{max}(0, \mathbf{h})$$.

ReLU (Rectified Linear Unit)

In the practical implementation of Large Language Models (LLMs), increasing the hidden size parameter, denoted as $$d_h$$, is generally beneficial for performance. However, deploying and training models with a very large hidden size introduces significant computational challenges. Because of these constraints, the careful design and selection of the activation function play a relatively more critical role in the effectiveness of such wide Feed-Forward Networks (FFNs).

Importance of Activation Function Design in Wide FFNs

In a standard two-layer feed-forward network (FFN) within a Transformer, an input vector `h` has a dimension of `d = 512`. The network's hidden layer has a dimension of `d_h = 2048`. The FFN is defined by the operation: `Output = σ(h * W_h + b_h) * W_f + b_f`, where `σ` is a non-linear activation function. What must be the dimensions of the weight matrix `W_f` for the output vector to have the same dimension as the input vector `h`?

A Transformer's Feed-Forward Network (FFN) takes an input vector of dimension `d = 768` and processes it through a hidden layer of dimension `d_h = 3072` before producing an output vector of the same dimension as the input (`d = 768`). The intermediate vector, after the first linear transformation and activation function, correctly has a dimension of 3072. If a dimension mismatch error occurs during the second linear transformation, what are the required dimensions for the second weight matrix (`W_f`) to resolve this error? Explain your reasoning based on the rules of matrix multiplication.

Troubleshooting FFN Dimension Mismatch

A standard Feed-Forward Network (FFN) in a Transformer model processes an input vector `h` of dimension `d` using the formula: `FFN(h) = σ(h * W_h + b_h) * W_f + b_f`. The intermediate hidden layer has a dimension `d_h`. Match each component from the formula to its correct description.

You’re debugging a Transformer block in an interna...

You are reviewing a teammate’s implementation of a...

You’re implementing a single Transformer block in ...

You are writing a one-page implementation spec for a new internal Transformer-block API that must be unambiguous enough for two different teams (training + inference) to implement independently and still produce identical tensor shapes and computation order.

Constraints:
- The block input is H ∈ R^{m×d} (m = sequence length, d = model width).
- The block contains exactly two sub-layers in this order: (1) multi-head self-attention, (2) a 2-layer position-wise FFN.
- Multi-head attention uses n_head heads with per-head dimension d_k such that concatenation returns to width d.
- The FFN must expand to hidden width d_h and return to width d using the standard formula FFN(h)=σ(hW_h+b_h)W_f+b_f.
- You must choose either a pre-norm or post-norm scheme and specify precisely where LayerNorm is applied relative to F(·) and the residual addition for BOTH sub-layers.

Create the spec by writing:
1) A step-by-step computation graph (as numbered equations) for the full block from input H to output H_out, including residual connections and LayerNorm placement.
2) The required matrix dimensions for W_q, W_k, W_v, the output projection W_o, and the FFN matrices W_h and W_f (use d, d_h, n_head, d_k; you may assume d = n_head·d_k).

Your answer must be internally consistent: every addition must be shape-compatible, and your norm placement must match the scheme you chose.

Design a Transformer Block Spec for a New Internal LLM Library (Shapes + Norm Placement)

You are reviewing a teammate’s refactor of a Transformer block used in an internal LLM service. The block should contain (1) a multi-head self-attention sub-layer and (2) a two-layer position-wise FFN, each with a residual connection and layer normalization. The model uses sequence length m and model width d, and the attention module uses h heads.

During integration testing, you see two symptoms: (A) training becomes unstable when stacking many blocks (loss spikes and occasional divergence), and (B) a shape error appears in the FFN path when d=512 and the hidden width is set to d_h=2048.

Write a technical diagnosis that:
1) Identifies the most likely normalization/residual ordering mistake that would explain symptom (A), explicitly contrasting pre-norm vs post-norm computation for a sub-layer and why one tends to be more stable in very deep stacks.
2) Explains, using correct tensor/matrix dimensions, how multi-head self-attention can take an input H ∈ R^{m×d} and still return an output in R^{m×d} (include what happens across heads conceptually, and why concatenation must be followed by a projection).
3) Pinpoints the FFN dimension mismatch that could cause symptom (B) and states the required shapes of W_h, b_h, W_f, and b_f so that the FFN maps an input of width d back to width d.

Your answer should connect these three parts into a single coherent explanation of how a correct Transformer block maintains consistent dimensions while using residuals and layer normalization, and how an incorrect norm placement can interact with deep stacking to destabilize training.

Diagnosing a Transformer Block Refactor: Attention/FFN Shapes and Norm Placement

You are reviewing two candidate implementations of a Transformer block for an internal LLM that must be scaled from 12 to 96 layers without changing the model dimension d. Each block has (1) a multi-head self-attention sub-layer that maps an input H ∈ R^{m×d} to an output in R^{m×d} by running multiple attention heads in parallel, concatenating their outputs, and applying a final linear projection, and (2) a position-wise two-layer FFN applied independently to each token: FFN(h)=σ(hW_h+b_h)W_f+b_f with W_h ∈ R^{d×d_h} and W_f ∈ R^{d_h×d}. Both designs use residual connections and layer normalization (LN), where LN normalizes each token’s d features using that token’s mean and standard deviation and then applies learnable gain and bias.

Design A (post-norm) uses, for each sub-layer: y = LN(x + F(x)).
Design B (pre-norm as defined here) uses, for each sub-layer: y = LN(F(x)) + x.

In early training runs at 96 layers, Design A frequently diverges (loss becomes NaN) while Design B trains but shows slightly slower early loss reduction.

Write an engineering recommendation memo (as an essay) that: (a) argues which design you would choose for the 96-layer model and why, explicitly linking your reasoning to how LN placement interacts with residual connections across many stacked blocks; (b) demonstrates that you understand the required tensor shapes through the attention and FFN sub-layers (i.e., why both F(x) terms can be added to x and why the FFN must use W_h and W_f with the given dimensions); and (c) explains one plausible tradeoff your choice introduces for model behavior or optimization (e.g., gradient flow, representational scaling, or sensitivity to initialization), grounded in the two formulas above rather than generic statements.

Choosing Pre-Norm vs Post-Norm for a Deep Transformer: Stability, Shapes, and Sub-layer Semantics

You inherit a production LLM codebase where a teammate made a “minor cleanup” to the Transformer block. After the change, training becomes unstable (loss spikes and occasional NaNs) only when scaling from 12 to 48 layers; the 12-layer model still trains. The teammate claims they only (a) moved LayerNorm, and (b) refactored the attention and FFN code for readability.

Assume the model uses token representations H ∈ R^{m×d}. The self-attention sub-layer is multi-head self-attention: for each head j, Q^{[j]} = H W_j^q, K^{[j]} = H W_j^k, V^{[j]} = H W_j^v; each head output is computed via scaled dot-product attention, head outputs are concatenated, then projected back to dimension d. The FFN is two linear layers with a nonlinearity: FFN(h) = σ(h W_h + b_h) W_f + b_f, where W_h ∈ R^{d×d_h} and W_f ∈ R^{d_h×d}. LayerNorm normalizes each token’s d features and has learnable gain/bias.

Write an engineering memo that (1) proposes the two most plausible implementation mistakes that could simultaneously explain “works at 12 layers but diverges at 48 layers” and are consistent with the teammate’s description, and (2) for each mistake, explains the mechanism of failure by explicitly connecting (i) residual + LayerNorm placement (pre-norm vs post-norm), (ii) how multi-head attention and FFN preserve/return to dimension d, and (iii) why depth amplifies the issue. Conclude with a concrete, minimal patch (in words or pseudocode) that would fix each mistake and a quick sanity-check you would run to confirm the fix.

Root-Cause Analysis of Training Instability After a “Minor” Transformer Block Change

You are reviewing a teammate’s Transformer block implementation for an internal LLM service. The model uses hidden size d=1024, sequence length m=256, number of attention heads h=16 (so per-head dimension is 64), and FFN hidden size d_h=4096 with ReLU. The teammate reports that training becomes unstable (loss spikes and occasional NaNs) after a refactor that was intended to be “behavior-preserving.” They provide the following pseudocode for one block:

Input: H (shape m×d)
1) A = MultiHeadSelfAttention(LNorm(H))
2) H1 = LNorm(H + A)
3) F = FFN(LNorm(H1))
4) H2 = H1 + LNorm(F)
Output: H2

Assume MultiHeadSelfAttention follows the standard pattern: for each head j, Q[j]=X Wq[j], K[j]=X Wk[j], V[j]=X Wv[j], scaled dot-product attention is computed per head, head outputs are concatenated, then projected back to d.

Case task: Identify whether this block is consistently implementing a pre-norm scheme, a post-norm scheme, or an inconsistent mixture, and explain (a) the most likely stability-related issue caused by the mixture in terms of residual-path “cleanliness” and normalization placement, and (b) one concrete corrected block formula (in equations or pseudocode) that makes the normalization placement consistent while keeping all tensor dimensions valid for both the attention sub-layer and the FFN sub-layer.

Production Bug Triage: Transformer Block Norm Placement vs Attention/FFN Interface Contracts

You are reviewing a teammate’s pull request that “converts a Transformer block to pre-norm for better stability” in an internal LLM used for document triage. The model uses representation size d=512, FFN hidden size d_h=2048, sequence length m=128, and multi-head self-attention with n_head=8 (so each head uses d_head=64). The PR includes the following pseudocode for one block (self-attention sub-layer then FFN sub-layer):

1) a = LN(x)
2) attn_out = MultiHeadSelfAttention(a)   # returns shape (m, 512)
3) y = LN(attn_out + x)
4) f = LN(y)
5) ffn_out = ReLU(f * W_h + b_h) * W_f + b_f
6) out = LN(ffn_out + y)

The author claims this is “pre-norm” because LN is applied before each function. During training, you still see instability and slower convergence than expected.

As the reviewer, identify whether this block is actually pre-norm, post-norm, or a hybrid, and explain (a) the minimal change(s) needed to make it a true pre-norm block for both sub-layers, and (b) the required dimensions of W_h and W_f so that the FFN preserves the (m, d) interface expected by the residual connections. Your answer must explicitly reference how residual connections, layer normalization placement, and the attention/FFN output shapes interact in this block.

Post-Norm vs Pre-Norm Migration: Verifying Tensor Shapes and Correct Sub-layer Wiring

You are reviewing a production incident in an internal LLM service. A teammate made an “optimization” to a Transformer block to reduce compute and simplify code. After the change, training no longer diverges, but model quality drops noticeably (worse long-context retrieval and weaker instruction following) while throughput improves. You are given the following implementation notes for one block (sequence length m, model width d):

- The block has two sub-layers in order: (1) multi-head self-attention, (2) a 2-layer FFN.
- The teammate changed the FFN hidden size from d_h = 4d to d_h = d.
- They also changed normalization placement from post-norm to pre-norm, but their code now does:
  1) y = x + Attention(LNorm(x))
  2) z = y + FFN(LNorm(y))
  3) output = LNorm(z)
- They claim this is “equivalent but faster” because the FFN is narrower and “extra LNorm at the end keeps things stable.”

As the reviewer, identify the most likely *two* root causes of the quality regression that follow from the interaction of (a) multi-head self-attention’s role, (b) the FFN’s dimensionality/structure, and (c) layer normalization placement (pre-norm vs post-norm). Then propose one concrete code-level correction (in words or pseudocode) that would address the regression while keeping training stable, and justify why it helps.

Learn Before

Related