Implementing Linear Layers in PyTorch

This article starts from a working PyTorch class rather than from a definition. You will read the class once, predict how it behaves, and then check that prediction against the explanation. The goal is a mental model you can test against the code, so that when a LoRA layer produces an unexpected result you already know which line to examine first.

A LoRA layer, in the form used here, is a PyTorch module with two learnable matrices called A and B and a scalar setting called alpha. The module takes an input, passes it through A and then through B, scales the result by alpha, and returns it. Everything else in this article explains why those few lines are arranged the way they are.

Two details of the constructor matter most. The first is the pair of shapes: A is in_dim by rank, and B is rank by out_dim. Multiplying x by A and then by B only works because the inner dimension of the first product matches the outer dimension of the second, and that shared inner size is the rank. The second is that A and B receive different initial values, which the next sections examine closely.

The class is a compact expression of the forward computation alpha times x @ A @ B. The matrices A and B are the learned part of the adaptation, and alpha is the dial that controls how much their result is allowed to affect the output. Because the class is short, you can follow one input through it line by line, which is exactly what the next sections do.

The A matrix is filled with random values drawn using Kaiming uniform initialization, the same scheme PyTorch uses for Linear layers. Random values in A give the layer a non-trivial starting point, so gradients can flow into it during training. The B matrix, by contrast, is filled with zeros.

Combining these two choices gives a starting state in which the adaptation's output is exactly zero. Because the forward computation multiplies x by A and then by B, any zero in B makes the whole product zero, regardless of what A contains. The layer begins by adding nothing to the pretrained output and then learns how much to add. Initializing both matrices randomly would instead inject arbitrary noise into the model from the first step.

Follow one input through the forward method. The input x enters A, the result enters B, and that result is multiplied by alpha. The diagram below shows each step as one operation, and the nodes show the width of the value after each step.

The rank-sized intermediate value is the only place where the input is compressed. The input is first projected down to rank values and then projected back up to out_dim. Because the rank is usually much smaller than the input and output sizes, the adaptation stays small. The final shape always matches out_dim, so the LoRA output can be combined with the output of a layer of that width.

Rank and alpha are the two settings that most affect a LoRA layer, and they do different jobs. The rank sets the inner dimension shared by A and B, and that single number determines how many extra parameters LoRA adds to a layer. The source describes this as a balance: a larger rank gives the adaptation more capacity to change the layer, at the cost of more parameters, while a smaller rank keeps the layer cheap but limits what it can represent.

Alpha works differently. Rank decides how many parameters exist, while alpha decides how strongly the low-rank output is allowed to influence the result. In the class, alpha multiplies the whole adaptation result. Doubling alpha doubles the contribution of the LoRA path to the output while leaving A and B exactly as they are. A small alpha keeps the adaptation quiet, and a large alpha lets it dominate. Alpha is therefore a dial on the strength of the update, not a dial on its size in parameter terms.

A LoRA layer is not usually used alone. The typical goal is to add it to an existing Linear layer, so that weight updates are applied to the pretrained weights rather than replacing them. The original layer computes its output from the pretrained weights W. The LoRA layer computes a low-rank approximation of the weight update, which the source writes as delta W. The final output is the sum of the two results.

Both results are computed for the same input, and they meet at the summing step. Read this way, the adapted layer is the original layer plus a correction. The pretrained weights remain intact, and the learned part lives in A and B. When B is still zero, the correction term is zero and the adapted layer returns exactly the pretrained output.

Most LoRA bugs in this implementation come from a small number of misreadings. Each mistake below shows where the symptom appears and which part of the reasoning to check.