Architecture Hyperparameters › Depth vs Width Trade-offs

Michael BrenndoerferUpdated June 18, 20257 min read
Language AI HandbookAll chapters
Part 16 of 67Transformer Architectures
Chapter 116Architecture Hyperparameters › Depth vs Width Trade-offs
Reading section 3 of 10 · Depth vs Width Trade-offs · 7 minSection 3 of 10 · 7 min

Book contents

Section 3 of 10 · 7 min read

Architecture Hyperparameters › Depth vs Width Trade-offsOne section · Sections in this chapter3 / 10 ↓
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

Given a fixed parameter budget, should you build a deep, narrow model or a shallow, wide one? This fundamental question shapes every transformer design decision. Think of it like allocating floors and square footage in a building: you can build a tall skyscraper with modest floor plans, or a sprawling low-rise complex. Both might have similar total space, but they serve different purposes.

Depth, the number of transformer layers, determines how many sequential processing stages information passes through. Each layer can refine and combine representations from the previous layer, then abstract them. This sequential refinement matters for tasks requiring multi-step reasoning: understanding that "the trophy wouldn't fit in the suitcase because it was too big" requires first parsing the sentence, then resolving "it" to either trophy or suitcase, then applying physical reasoning. Deeper models have more opportunities for such compositional processing.

Width, primarily governed by dmodeld_{\text{model}}, determines how rich each representation can be. A wider model can encode more detail at each layer: more semantic dimensions, finer-grained syntactic features, more factual associations. Think of width as the vocabulary of concepts available at each processing stage.

Research has revealed important patterns in how these two dimensions contribute to capability. Deeper models tend to develop hierarchical representations, with early layers capturing surface patterns (word shapes, common phrases) and later layers capturing abstractions (sentiment, intent, reasoning chains). Wider models have more capacity per layer but may not develop the same hierarchical structure, instead learning richer but flatter representations.

Depth for Abstraction, Width for Capacity

Empirical studies suggest that depth is more important for tasks requiring multi-step reasoning or hierarchical understanding. Width matters more for tasks requiring rich representations at a single level of abstraction. Most practical tasks benefit from a balance of both.

Understanding this trade-off requires knowing how each dimension affects parameter count. Let's derive the relationship from first principles. For a simplified transformer (ignoring biases and layer norms for clarity), the main parameter contributors are:

  • Embedding layer: V×dmodelV \times d_{\text{model}} parameters
  • Attention (per layer): 4×dmodel24 \times d_{\text{model}}^2 parameters (QKV and output projections)
  • FFN (per layer): 2×dmodel×dff2 \times d_{\text{model}} \times d_{\text{ff}} parameters (two weight matrices)
  • Output projection: Often tied to embeddings, so 0 additional parameters

Summing these components gives us the total parameter count for a transformer model:

Total≈V⋅dmodel+L⋅(4⋅dmodel2+2⋅dmodel⋅dff)\text{Total} \approx V \cdot d_{\text{model}} + L \cdot (4 \cdot d_{\text{model}}^2 + 2 \cdot d_{\text{model}} \cdot d_{\text{ff}})

where:

  • VV: vocabulary size (number of unique tokens)
  • dmodeld_{\text{model}}: model dimension (hidden size)
  • LL: number of transformer layers
  • dffd_{\text{ff}}: feed-forward hidden dimension
  • 4⋅dmodel24 \cdot d_{\text{model}}^2: attention parameters per layer (four projection matrices: Q, K, V, O)
  • 2⋅dmodel⋅dff2 \cdot d_{\text{model}} \cdot d_{\text{ff}}: FFN parameters per layer (two weight matrices: up-projection and down-projection)

For large models, the embedding term V⋅dmodelV \cdot d_{\text{model}} becomes negligible compared to the layer terms. With the standard expansion ratio dff=4×dmodeld_{\text{ff}} = 4 \times d_{\text{model}}, we can simplify by substituting this relationship into the per-layer FFN term:

2⋅dmodel⋅dff=2⋅dmodel⋅(4⋅dmodel)=8⋅dmodel22 \cdot d_{\text{model}} \cdot d_{\text{ff}} = 2 \cdot d_{\text{model}} \cdot (4 \cdot d_{\text{model}}) = 8 \cdot d_{\text{model}}^2

Adding the attention term gives us the simplified total:

Total≈L⋅(4⋅dmodel2+8⋅dmodel2)=12⋅L⋅dmodel2\text{Total} \approx L \cdot (4 \cdot d_{\text{model}}^2 + 8 \cdot d_{\text{model}}^2) = 12 \cdot L \cdot d_{\text{model}}^2

This formula reveals a key insight about scaling behavior: parameters scale linearly with depth LL but quadratically with width dmodeld_{\text{model}}. Doubling the depth doubles the parameters; doubling the width quadruples them.

In[6]:
Code
def count_parameters(V, L, d_model, d_ff, include_output=False):
    """
    Count approximate transformer parameters.

    Args:
        V: Vocabulary size
        L: Number of layers
        d_model: Model dimension
        d_ff: Feed-forward dimension
        include_output: Whether to include separate output projection (vs tied embeddings)

    Returns:
        Dictionary with parameter breakdown
    """
    params = {}

    # Embedding layer
    params["embedding"] = V * d_model

    # Per layer: attention (Q, K, V, O projections) and FFN
    params["attention_per_layer"] = 4 * d_model * d_model
    params["ffn_per_layer"] = 2 * d_model * d_ff
    params["layer_norm_per_layer"] = 2 * d_model  # Two norms per layer

    params["per_layer_total"] = (
        params["attention_per_layer"]
        + params["ffn_per_layer"]
        + params["layer_norm_per_layer"]
    )

    params["all_layers"] = L * params["per_layer_total"]

    # Final layer norm
    params["final_norm"] = d_model

    # Output projection (often tied with embedding)
    params["output"] = V * d_model if include_output else 0

    params["total"] = (
        params["embedding"]
        + params["all_layers"]
        + params["final_norm"]
        + params["output"]
    )

    return params


# Compare depth vs width trade-off
# Configuration A: Deep and narrow
deep_narrow = count_parameters(V=32000, L=48, d_model=2048, d_ff=8192)

# Configuration B: Shallow and wide (similar total params)
shallow_wide = count_parameters(V=32000, L=12, d_model=4096, d_ff=16384)
In[7]:
Code
def format_params(n):
    """Format parameter count with appropriate suffix (K, M, B)."""
    if n >= 1e9:
        return f"{n / 1e9:.2f}B"
    elif n >= 1e6:
        return f"{n / 1e6:.2f}M"
    else:
        return f"{n / 1e3:.2f}K"
Out[8]:
Console
Depth vs Width Trade-off Comparison
==================================================

Configuration A: Deep and Narrow
  Layers: 48, d_model: 2048, d_ff: 8192
  Total parameters: 2.48B
  - Embedding: 65.54M
  - All layers: 2.42B

Configuration B: Shallow and Wide
  Layers: 12, d_model: 4096, d_ff: 16384
  Total parameters: 2.55B
  - Embedding: 131.07M
  - All layers: 2.42B

Ratio of layer parameters (Wide/Deep): 1.00x

Although both configurations have similar total parameters, their characteristics differ significantly. The deep model processes information through 48 stages of refinement. The wide model has only 12 stages but can represent much richer information at each stage. Research by Tay et al. (2021) and others suggests that for language modeling, depth tends to be more important, but the optimal ratio depends on the specific task and compute constraints.

Out[9]:
Visualization
Line plot showing linear growth of parameters with depth for different model widths.
Parameter count as a function of depth and width. Each curve represents a fixed width (d_model), showing how parameters grow linearly with depth. Doubling width shifts the curve up by 4x due to quadratic scaling.

The visualization shows why width is so costly: moving from dmodel=2048d_{\text{model}} = 2048 to dmodel=4096d_{\text{model}} = 4096 quadruples the parameters at every depth level. In contrast, adding layers adds parameters linearly. This explains why many successful architectures (like GPT-3 and LLaMA) achieve large parameter counts primarily through depth, keeping width moderate.

The Scaling Laws Perspective

The depth-width trade-off isn't just theoretical. Large-scale empirical studies have quantified how each dimension affects model capability as a function of compute budget. Scaling laws research, including the Chinchilla paper by Hoffmann et al. (2022), provides rigorous guidance on optimal model shapes.

The Chinchilla findings were striking: most large language models were undertrained. Given a fixed compute budget, models achieved better performance with fewer parameters and more training data than the prevailing GPT-3 style of maximizing model size. But within the model size constraint, how should parameters be distributed between depth and width?

Studies have found that:

  • Depth should scale roughly as the square root of the parameter count
  • Width should scale similarly to depth
  • The ratio of depth to width (in terms of dmodeld_{\text{model}}) tends to stay fairly constant across scales

Modern frontier models follow these principles. LLaMA-7B has 32 layers and dmodel=4096d_{\text{model}} = 4096, giving a ratio of about 128. LLaMA-70B has 80 layers and dmodel=8192d_{\text{model}} = 8192, maintaining a similar ratio of about 100. This consistency suggests that the field has converged on roughly optimal proportions.