The linear algebra inside a language model

You have the operations; this is where they are actually used. A transformer layer is almost entirely matrix multiplication, and following the shapes through one layer — what size goes in, what comes out at each step — is the quickest way to stop a language model being a black box. It also makes the cost visible, which is what lets you see in a line of arithmetic why LoRA fine-tuning is so much cheaper than updating every weight.

Tokens become vectors#

An embedding table is a matrix with one row per vocabulary token. Looking up a token means picking its row. A sentence of 10 tokens becomes a 10 by d matrix.

It is worth seeing that a lookup and a matrix multiply are the same operation here, because papers write it both ways and beginners assume they are two different mechanisms. Multiplying a one-hot row — all zeros except a single 1 in position i — by the embedding table selects row i and nothing else, since every other row is multiplied by zero. Implementations index instead of multiplying because indexing is enormously cheaper, but the maths written in the paper is the multiply.

import numpy as np

table = np.array([[0.1, 0.2, 0.3],      # token 0
                  [1.0, 0.0, 0.0],      # token 1
                  [0.5, 0.5, 0.5]])     # token 2

one_hot = np.array([0.0, 1.0, 0.0])     # 'this is token 1'
print(one_hot @ table)                   # [1. 0. 0.]
print(table[1])                          # [1. 0. 0.]  the same thing, done cheaply

tokens = [2, 0, 1, 1]                    # a four-token sequence
print(table[tokens].shape)               # (4, 3)  four rows of width d

Attention is three matrix products#

X        [tokens, d]            the input
Q = X Wq [tokens, d_head]       what each token is looking for
K = X Wk [tokens, d_head]       what each token offers
V = X Wv [tokens, d_head]       what each token passes on

scores  = Q K^T / sqrt(d_head)  [tokens, tokens]   every token against every token
weights = softmax(scores)       each row sums to 1
out     = weights V             [tokens, d_head]   a weighted mix of values

The tokens by tokens score matrix is why attention cost grows with the square of the sequence length. Put numbers on it. At 8,192 tokens that matrix has 8,192 squared entries, which is 67 million numbers, or 134 megabytes in 16-bit precision — for one head, in one layer, for one sequence in the batch. Double the context and it quadruples. Every long-context technique you will read about exists to avoid holding that matrix.

Counting the parameters of a layer#

Model sizes stop being folklore once you can count them, and the count is four matrix shapes. Take d equal to 4,096. Attention holds four d by d matrices, for queries, keys, values and the output projection: four times 4,096 squared, which is 67.1 million. The feed-forward block holds matrices that widen to 4d and come back, and in the gated form used by most current models that is three matrices of d by 4d: three times 4,096 times 16,384, which is 201.3 million. So one layer is 268.4 million parameters, and sixteen d squared is a useful thing to remember as the shape of that number.

d, layers, vocab = 4096, 32, 128256

attention = 4 * d * d              # Q, K, V and the output projection
mlp = 3 * d * (4 * d)              # gated feed-forward: two up, one down
per_layer = attention + mlp

print(per_layer / 1e6)             # 268.435456 million per layer
print(layers * per_layer / 1e9)    # 8.589934592 billion in the layers
print(vocab * d / 1e6)             # 525.336576 million in the embedding table

That count is also the memory bill. At two bytes per parameter in 16-bit precision, 8.6 billion parameters is 17.2 gigabytes just to hold the weights, before activations, before the key-value cache, before the optimiser state that training adds on top. When someone says a model will not fit on a 24 gigabyte card, this multiplication is the whole of what they did.

Low rank: why LoRA is cheap#

A full weight update for a 4096 x 4096 matrix:   16,777,216 numbers

LoRA writes the update as B A, with B [4096 x 8] and A [8 x 4096]:
  4096 x 8 + 8 x 4096 = 65,536 numbers   (about 0.4% of the full update)

A product of a tall thin matrix and a short wide one has rank at most 8 here: it can only change the weights along 8 directions. The bet behind LoRA is that fine-tuning mostly needs a few directions, and in practice that often holds.

Scale that to the whole model with the count from above. Applying rank-8 adapters to all four attention matrices in all 32 layers is 32 times 4 times 65,536, which is 8.4 million trainable parameters against 8.6 billion in the model: under a tenth of one per cent. That is the difference between a fine-tune that needs a cluster and one that runs on a single card, and it comes entirely from the shape of two matrices.

rank = 8
lora_per_matrix = 2 * d * rank                    # 65,536
lora_total = layers * 4 * lora_per_matrix         # four attention matrices per layer

print(lora_per_matrix)                             # 65536
print(lora_total / 1e6)                            # 8.388608 million trainable
print(100 * lora_total / (layers * per_layer))     # 0.09765625 per cent

One caution on the arithmetic: fewer trainable parameters is not proportionally less compute. The forward pass still runs every one of the 8.6 billion weights, and the backward pass still has to propagate gradients through the frozen layers to reach the adapters. What LoRA removes is optimiser state and gradient storage for the frozen weights, which with Adam is the largest part of a training run's memory, and it makes each fine-tuned variant a few megabytes to store rather than a few gigabytes.