Layers and parameter counts

The last lesson ended with one layer, Y = X @ W + b, run on a whole batch at once. A network stacks several of them. Two questions follow, and this lesson answers both.

First, what does a stack compute, and why does every network put a nonlinearity between its layers? Without one, a hundred layers turn out to be no better than one, and you will see exactly why. Second, how big is a network: how many numbers must it store, and how much arithmetic does one example cost? The same counting that sizes a small digit classifier also explains a phrase you hear every day, "a 7B model".

#A layer is a matrix product plus a bias

Every network in the rest of the course is built from one piece, so it is worth writing it down carefully.

A dense layer (also called fully connected) in the code convention is

with $\mathbf{X}$ of shape $(\text{batch},\allowbreak d_{\text{in}})$, $\mathbf{W}$ of shape $(d_{\text{in}},\allowbreak d_{\text{out}})$, and the bias $\mathbf{b}$ added to every row by broadcasting. Every output feature reads every input feature, which is where the name "dense" comes from. A two-layer network puts a nonlinearity $\phi$ between two of them, applied to every entry separately:

A common choice is the ReLU, $\phi(z) = \max(0,\allowbreak z)$: it keeps positive numbers and turns negative ones into 0. Follow the shapes: $\mathbf{X}$ is $(\text{batch},\allowbreak d_{\text{in}})$ and $\mathbf{W}_1$ is $(d_{\text{in}},\allowbreak h)$, so the hidden layer $\mathbf{H}$ is $(\text{batch},\allowbreak h)$; with $\mathbf{W}_2$ of shape $(h,\allowbreak d_{\text{out}})$ the output is $(\text{batch},\allowbreak d_{\text{out}})$. The hidden width $h$ is a free choice; it is the inner size that has to match on both sides.

A worked example with one input, so every number is visible. Take a network with 2 inputs, 2 hidden units and 1 output:

and the input row $\mathbf{x} = (0,\allowbreak 2)$.

1. Layer 1 product. Entry $j$ of $\mathbf{x}\mathbf{W}_1$ is $\mathbf{x}$ dotted with column $j$ of $\mathbf{W}_1$: $\big((0)(1) + (2)(-1),\allowbreak \ (0)(2) + (2)(1)\big) = (-2,\allowbreak 2)$.
2. Add the bias. $(-2 + 0,\allowbreak \ 2 + 1) = (-2,\allowbreak 3)$.
3. ReLU. $(\max(0,\allowbreak -2),\allowbreak \ \max(0,\allowbreak 3)) = (0,\allowbreak 3)$. The negative value became 0.
4. Layer 2. $(0)(2) + (3)(1) = 3$, plus $b_2 = -1$, gives the output $2$.

Why the nonlinearity? Leave out step 3 and see what happens. Layer 2 then acts on $(-2,\allowbreak 3)$ directly: $(-2)(2) + (3)(1) - 1 = -4 + 3 - 1 = -2$. Now try a single dense layer with weight $\mathbf{W}_1\mathbf{W}_2$ and bias $\mathbf{b}_1\mathbf{W}_2 + \mathbf{b}_2$ (the algebra below says where these come from). The weight is $(2 \times 2)(2 \times 1) = 2 \times 1$, with entries $(1)(2) + (2)(1) = 4$ and $(-1)(2) + (1)(1) = -1$. The bias is $(0)(2) + (1)(1) - 1 = 0$. That one layer gives $\mathbf{x}\mathbf{W} + b = (0)(4) + (2)(-1) + 0 = -2$: exactly what the two layers gave without the ReLU. With the ReLU the answer was 2 instead, because the ReLU zeroed the negative hidden value.

Here is the same collapse in general, for any input:

That is a single dense layer, with weight $\mathbf{W}_1\mathbf{W}_2$ and bias $\mathbf{b}_1\mathbf{W}_2 + \mathbf{b}_2$. Stacking a hundred layers without nonlinearities still gives one layer, so depth buys nothing. Neurons and layers picks up exactly here: the nonlinearity is what lets each layer bend space so that the next one can do something new.

Now run a small network both ways yourself, with the ReLU and without it, and collapse it.

#Counting parameters and work

The numbers a network stores decide how much memory it needs, and, as you will see, nearly the same count decides how much arithmetic it does. So it pays to count them.

A dense layer from $d_{\text{in}}$ to $d_{\text{out}}$ has $d_{\text{in}} \times d_{\text{out}}$ weights and $d_{\text{out}}$ biases. Together these are its parameters, the values training adjusts. A forward pass over a batch costs $\text{batch} \times d_{\text{in}} \times d_{\text{out}}$ multiply-adds for the matrix product, by the $mnp$ rule from Matrix multiplication with $m = \text{batch}$, $n = d_{\text{in}}$ and $p = d_{\text{out}}$; the bias adds only $\text{batch} \times d_{\text{out}}$ additions on top.

The shape widget from the last lesson counts both for you, per layer and in total. Here it holds a tiny network, 4 features in, a hidden layer of 3 and 2 outputs, run on a batch of 2 examples. The table under the formula writes each layer's count out, such as $4 \cdot 3 + 3 = 15$ for the parameters of layer 1.

A worked example at a realistic size: a network for small digit images with $28 \times 28 = 784$ pixels, one hidden layer of 128 units and 10 outputs, one per digit.

• Layer 1: $784 \times 128 = 100{,}352$ weights plus 128 biases, $100{,}480$ parameters.
• Layer 2: $128 \times 10 = 1{,}280$ weights plus 10 biases, $1{,}290$ parameters.
• Total: $100{,}480 + 1{,}290 = 101{,}770$ parameters.

The forward pass for one example costs $100{,}352 + 1{,}280 = 101{,}632$ multiply-adds, one per weight, and a batch of 64 costs $64 \times 101{,}632 = 6{,}504{,}448$.

Now count a slightly deeper network yourself.

The dense layer in the next exercise checks its shapes before it computes anything, a habit that pays off for the rest of the course. In Backpropagation you will derive the gradient of a layer's weights entry by entry, and shapes will guide and check the result: the gradient of W must have the same shape as W, and only one arrangement of the available matrices produces it. A recurrent network (Sequences and state) adds a time axis to every array, and an LSTM (The forward and backward pass) stacks the weights of four gates into one matrix and splits the result apart again. Each of those steps is a shape to get right.

Now turn both sections of this lesson into code: a dense layer that checks its shapes, a two-layer forward pass, a parameter counter and the collapse.

A 7B model, counted

The same arithmetic scales all the way up. Here is a published design, the first LLaMA model from Meta (2023) at its 7B size: width 4,096, 32 layers and a vocabulary of 32,000 tokens. Each layer has four $4096 \times 4096$ attention matrices (Attention explains what they do) and an MLP made of three matrices with $4096 \times 11008$ entries each, a gated variant of the two-matrix MLP above. There are no biases. Each layer also has two small vectors of normalization weights, 4,096 numbers each, and there is one more such vector after the last layer.

Count it one piece at a time. A $4096 \times 4096$ matrix holds $4096^2 = 16{,}777{,}216$ numbers and a $4096 \times 11008$ matrix holds $45{,}088{,}768$, so:

The third row adds the first two, $67{,}108{,}864 + 135{,}266{,}304 = 202{,}375{,}168$ per layer, and multiplies by 32. The normalization row counts $2 \times 32 + 1 = 65$ vectors. The token embedding is a table with one row of 4,096 numbers for each of the 32,000 tokens, and the output layer is a matrix that turns the final 4,096-number vector into one score per token. The total is about 6.7 billion parameters, called "7B". Almost all of them sit in weight matrices, and the MLPs alone hold $32 \times 135{,}266{,}304 \approx 4.3$ billion, nearly two thirds of the total.

The code below redoes the count, converts it to memory at 16 bits (2 bytes) per parameter, and counts the multiply-adds for one token. Expect a total of 6,738,415,616, about 13.5 GB, and about 6.61 billion multiply-adds.

To recap: a dense layer is a matrix product plus a bias; without a nonlinearity between them, layers collapse into one; a layer stores $d_{\text{in}}d_{\text{out}} + d_{\text{out}}$ numbers whatever the batch; and a forward pass costs about one multiply-add per weight per example, which is why a parameter count tells you both memory and compute. Finish with mixed reps from this lesson and the last one: shapes, transposes, broadcasting and counts, with new numbers each time.