Transformer Parameter Calculation

Introduction

In our last lesson, we traced the complete data flow through the encoder and decoder blocks, clarifying how information is processed, transformed, and communicated between the different parts of the model. You now have a solid mental map of the Transformer's operational sequence for both training and inference.

Today, we shift from the process to the substance. Our learning outcome is to calculate the number of parameters in a transformer model given its configuration (layers, heads, dimensions). Understanding where a model's parameters come from is a fundamental engineering skill. It allows you to:

• Estimate the memory footprint for training and deployment.
• Approximate computational requirements.
• Verify your own implementations against published model architectures.
• Develop an intuition for how architectural choices impact model scale.

By the end of this lesson, you will be able to deconstruct a model like BERT or GPT into its constituent parts and account for every parameter.

1. Where Do the Parameters Live?

A Transformer's parameters, or weights, are the learnable numbers that are adjusted during training. They aren't spread uniformly throughout the architecture. They are concentrated in specific components. Let's start with the "bookends" of the model.

For a great intuitive overview of how parameters contribute to the model's total size, using GPT-3 as a running example, have a look at the following video.

As the video highlights, we have two major sources of parameters outside the main Transformer blocks:

• Token Embedding Layer (W_E): This is a large lookup table with a row for every token in the vocabulary and d_model columns.

  • Shape: (V, d_model)
  • Parameters: V * d_model

• Unembedding Layer / Language Model Head (W_U): This layer projects the final d_model-dimensional output vector back to the vocabulary space to produce token probabilities.

  • Shape: (d_model, V)
  • Parameters: d_model * V

A common practice called weight tying sets W_U = W_E^T. This not only saves a significant number of parameters but has also been shown to improve performance. For the rest of this lesson, we'll assume weight tying is used, meaning the unembedding layer adds no new parameters.

While these embedding parameters are numerous, in large models, they are dwarfed by the parameters inside the stack of L Transformer layers.

2. Deconstructing a Transformer Block

Let's dissect a single Transformer block and count the parameters in each of its sub-layers. For this, we'll use a few key hyperparameters:

• d_model: The model's main hidden dimension.
• n_heads: The number of attention heads.
• d_head: The dimension of each attention head (d_head = d_model / n_heads).
• d_ff: The hidden dimension of the feed-forward network.

The following article provides a fantastic, detailed breakdown with clear formulas. We will use it as our primary reference for the exact calculations.

Let's summarize the key formulas from the article (ignoring biases for a moment to build intuition).

A. Multi-Head Attention (MHA)

The MHA block has four primary weight matrices:

1. Query projection (W_Q): Projects the input X to the query vectors Q.
2. Key projection (W_K): Projects X to K.
3. Value projection (W_V): Projects X to V.
4. Output projection (W_O): Projects the concatenated head outputs back to d_model.

Each of these is a linear transformation from d_model to d_model. Therefore, each weight matrix has a shape of (d_model, d_model).

• Parameters (MHA) ≈ 4 * d_model * d_model = 4 * d_model^2

B. Position-wise Feed-Forward Network (FFN)

The FFN consists of two linear layers:

1. An "up-projection" from d_model to d_ff.
2. A "down-projection" from d_ff back to d_model.

• Parameters (FFN) ≈ (d_model * d_ff) + (d_ff * d_model) = 2 * d_model * d_ff

A very common architectural choice is to set d_ff = 4 * d_model. Substituting this in, the FFN parameter count becomes:

• Parameters (FFN) ≈ 2 * d_model * (4 * d_model) = 8 * d_model^2

Notice that the FFN contains roughly twice as many parameters as the MHA block.

C. Layer Normalization

Each LayerNorm layer has two learnable parameter vectors: gamma (scale) and beta (shift). Each vector has a size equal to the dimension it's normalizing, which is d_model.

• Parameters (LayerNorm) = d_model (for gamma) + d_model (for beta) = 2 * d_model
  This is a small number compared to the MHA and FFN matrices, but we include it for completeness.

3. Assembling the Total Count

Now, let's put it all together for a standard decoder-only model (like GPT). Each block contains one MHA layer, one FFN, and two LayerNorm layers.

Parameters per Transformer Block:

For large d_model, the 4 * d_model term is negligible. So, a good rule of thumb is:

Total Model Parameters:
Now, we can write the formula for the entire model, which has L layers.

This simplified formula reveals a crucial insight: parameters scale linearly with the number of layers (L), but quadratically with the model width (d_model). This is why doubling the width of a model is far more expensive than doubling its depth.

The following article provides a fantastic, comprehensive guide to these architectural hyperparameters and includes a Python function that implements a detailed parameter count.

4. A Practical Example: BERT-base

Let's apply these formulas to a well-known model, BERT-base (treating it as a decoder-only for this calculation).

• V = 30,522
• L = 12
• d_model = 768
• d_ff = 3072 (4 * 768)

Let's calculate the approximate number of parameters.

1. Embedding Parameters:
   V * d_model = 30522 * 768 = 23,441,664 (~23.4M)

2. Parameters per Block:
   12 * d_model^2 = 12 * 768^2 = 12 * 589,824 = 7,077,888 (~7.1M)

3. Total Layer Parameters:
   L * (Params per Block) = 12 * 7,077,888 = 84,934,656 (~84.9M)

4. Total Model Parameters (Approximate):
   Embeddings + Layers = 23,441,664 + 84,934,656 = 108,376,320 (~108.4M)

The official BERT-base has ~110M parameters. Our approximation is very close! The small difference comes from biases, positional encoding parameters (if learned), and exact LayerNorm counts, all of which are handled by the full_parameter_count function from the resource.

Conclusion

Today you've acquired a critical skill for any aspiring LLM engineer: the ability to reason quantitatively about a model's architecture.

Key Takeaways:

• A Transformer's parameters are primarily in the token embedding matrix and the stack of Transformer blocks.
• Within each block, the Feed-Forward Network (FFN) contains about twice as many parameters as the Multi-Head Attention (MHA) mechanism (8 * d_model^2 vs. 4 * d_model^2).
• The total number of parameters can be approximated with the formula Total ≈ V*d_model + L * 12 * d_model^2.
• This shows that model size scales linearly with depth (L) but quadratically with width (d_model), making width a much more "expensive" hyperparameter to increase.

What's Next?

This lesson concludes Module 2. You now have a robust understanding of the original Transformer's components, data flow, and parameterization.

We are now ready to move into Module 3: Attention Mechanisms and Architectural Variants. In the next lesson, we will dive into a crucial detail that enables autoregressive models like GPT to work: implementing causal masking to enable autoregressive language modeling. This will shift our focus from the static architecture to the dynamic mechanisms that control information flow for specific tasks.