Linear layers in code — Linear algebra for LLMs

What nn.Linear computes

\displaystyle y = xW^\top + b

nn.Linear(in, out) stores $W$ as $(\text{out} \times \text{in})$ and a bias of length $\text{out}$ . Input $(\dots, \text{in})$ , output $(\dots, \text{out})$ : any leading axes (batch, tokens) come along for free.

Write

y = xW^\top + b

,

W\!: (\text{out}, \text{in})

,

b\!: (\text{out},)

.

A small one by hand

+

=

In 2, out 3. $xW^\top = [1, 2, 3]$ , plus the bias $[0, 0, 1]$ .

Each output is a dot product with a row of W

\quad y_3 = x \cdot W_{3,:} + b_3 = [1, 2] \cdot [1, 1] + 1 = 4

Row $j$ of the stored weight belongs to output feature $j$ . That’s why PyTorch stores $(\text{out}, \text{in})$ : each row is one output’s “detector”, contiguous in memory.

Write

y_j = x \cdot W_{j,:} + b_j

: one row of

W

per output.

Count the parameters

\displaystyle \text{params} = \text{in} \times \text{out} + \text{out}

\displaystyle 4096 \times 14336 = 58{,}720{,}256 \quad (\text{no bias})

Llama’s linear layers have no biases, so a $4096 \to 14336$ projection is just its weight: almost 59 million numbers, one of three in each MLP.

✏️ Your turn

GPT-2's MLP input

\displaystyle \texttt{nn.Linear(768, 3072)}

How many parameters, bias included?

Work it out on paper, then go to the next card to check.

Answer

\displaystyle 768 \times 3072 + 3072 = 2{,}362{,}368

In code

lin = nn.Linear(4096, 14336, bias=False)
lin.weight.shape          # torch.Size([14336, 4096])
y = lin(x)                # x: (B, T, 4096) → y: (B, T, 14336)

In a transformer

When you load a checkpoint, mlp.up_proj.weight with shape [14336, 4096] reads as (out, in): this layer maps 4,096 features to 14,336.

📌 Recap

Linear layers in one card

$y = xW^\top + b$ , with $W$ stored $(\text{out}, \text{in})$ .
Output $j$ is $x \cdot W_{j,:} + b_j$ .
Parameters: $\text{in} \cdot \text{out} + \text{out}$ (Llama: no bias).

Write

The three lines above.