Replacing Linear Layers in GPTModel

LinearWithLoRA wraps an existing Linear layer and a LoRALayer together so it can replace the Linear layer in a network.

Programming

Swapping a Layer Without Breaking the Model

A GPTModel is built from many Linear layers, including those inside its attention and feed-forward modules. Earlier in this chapter you built a LoRALayer as a standalone component whose job is to produce a low-rank correction. The next step is to attach that correction to a layer that already exists in the network. This article follows one wrapper class, LinearWithLoRA, which holds an ordinary Linear layer and a LoRALayer side by side.

As you read the trace, keep one question in mind: at the very start, before any training has happened, does the wrapped layer behave differently from the plain Linear layer it replaces? Your answer tells you whether the wrapper is safe to drop into a pretrained model. By the end of the article you should be able to answer that question and justify it from the code's structure.

Two Components Inside One Wrapper

Definition

LinearWithLoRA wraps an existing Linear layer and a LoRALayer together so that it can replace the Linear layer in a network. It is a subclass of torch.nn.Module, which means it is an ordinary PyTorch module and can be placed wherever a Linear layer used to be.

The wrapper is designed for places such as the attention and feed-forward modules inside a GPTModel. Its constructor takes three inputs: the existing linear layer, a rank, and an alpha value. Two components end up inside the wrapper. The first is the original Linear layer, which carries the pretrained weights. The second is a LoRALayer, which produces the correction. The diagram below shows how these two pieces are related inside the wrapper.

What does LinearWithLoRA contain, and how are the original Linear layer and the LoRALayer related to each other inside it?

What the Constructor Does

The constructor performs two separate jobs. First, it stores the original linear layer as an attribute named self.linear. This matters because the pretrained weights are kept inside the wrapper rather than discarded. Second, it creates a LoRALayer whose input and output sizes are read directly from the original layer, using linear.in_features and linear.out_features. Because of this, the LoRA path always matches the shape of the layer it augments. The rank and alpha values are simply passed through to the LoRALayer.

python

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features,
            linear.out_features,
            rank,
            alpha,
        )

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features,
            linear.out_features,
            rank,
            alpha,
        )

Notes on the code

The wrapper is a regular torch.nn.Module subclass
self.linear keeps the pretrained layer, so its weights are not thrown away
The LoRA path takes its input and output sizes from the original layer
rank and alpha are passed through unchanged

Tracing the Forward Pass

The forward method is where the two components meet. The input x is sent through the original linear layer and also through the LoRALayer. The two results are then added together, and that sum is the output of the wrapper. Nothing else happens inside forward, so the wrapper's behavior is fully determined by these two paths.

When an input enters LinearWithLoRA, how does the data split into the original Linear path and the LoRA path, and where do the two outputs merge again?

Rule

Read this expression as a sum of two independent contributions. The term linear(x) is the result from the pretrained path, which uses the original weights. The term lora(x) is the correction produced by the LoRALayer. The output y has the same shape either way, because the LoRA path was built to match the original layer's dimensions. Both terms are computed from the same input x, and only the final addition combines them.

Worked example

Tracing one input through the wrapper

Suppose the original Linear layer produces linear(x) = [0.5, -1.0] for some input x, and the LoRALayer produces lora(x) = [0.2, 0.3] for the same x. What is the wrapper output y?

1
Send x down the original path
linear(x) = [0.5, -1.0]. This comes from the pretrained weights.
2
Send x down the LoRA path
lora(x) = [0.2, 0.3]. This is the low-rank correction.
3
Add the two results element by element
y = [0.5 + 0.2, -1.0 + 0.3] = [0.7, -0.7].
4
Check the shape
Both paths return two values, so the sum has the same shape as either path alone.

Result

y = [0.7, -0.7]. The wrapper output is the sum of the original output and the LoRA output.

Predicting the Starting Output

What do you think happens?

The wrapper was just built around a pretrained Linear layer, and no training step has run yet. What does the lora(x) term contribute at this moment?

It contributes a large random value that changes the output
It contributes exactly zero, so the output equals linear(x)
It contributes the same value as linear(x)

Reveal answer

Answer:

It contributes exactly zero, so the output equals linear(x).

The matrix B inside LoRALayer, stored as self.B, is initialized to zero. The product of A and B is therefore a zero matrix no matter what values A holds. Multiplying the input by a zero matrix gives zeros, so lora(x) is exactly zero at initialization.

The reason the prediction holds comes from how the LoRALayer's matrices are initialized. The matrix B is set to zero values. The product of A and B produces a zero matrix, and multiplying the input by a zero matrix yields zeros. The LoRA contribution at initialization is therefore exactly zero, regardless of the values in A. Notice that the argument depends only on B. Whatever A contains, the zero in B removes its effect at the start.

Why does the LoRA contribution start at zero, so that the wrapped layer's output equals the original layer's output before any training?

Why Zero Initialization Makes the Swap Safe

Adding a zero contribution does not change anything. The wrapper's output starts out identical to the original layer's output. This means that dropping the wrapper into a model does not alter the model's behavior until training begins to move the LoRA weights away from zero. The initialization is what makes the replacement safe at the moment of swapping.

Key point

Zero initialization of B means the LoRA contribution vanishes at the start. The replacement therefore does not change behavior until training moves the weights.

This is the answer to the question posed at the beginning. At the very start, the wrapped layer does not behave differently from the plain Linear layer it replaces. That property is what makes the wrapper safe to drop into a pretrained model.

Finding Where a Mismatch Starts

Suppose the wrapper's output does not match what you expected. The source does not list a catalogue of errors, but its design shows where to look. The forward method has only two paths and one join. The LoRA path is sized from linear.in_features and linear.out_features, so its output is built to match the original layer. Working through the forward path in order narrows down the likely location.

Which step in the forward path, from input to the summed output, is the most likely place for a mismatch between expected and actual results?

The most likely point of mismatch is the join where the two outputs are added. Both terms must have the same shape for the addition to represent a single correction to the original output. The LoRA path's sizes come from the original layer, so a mismatch would most likely come from that sizing step or from any code that changes the LoRA path after construction. The source describes the sum as the wrapper's only output step, so the sum is the place where any disagreement between the two paths would surface.

Common Misreadings of the Wrapper

Common mistake
Thinking the wrapper discards the pretrained layer
Why it fails
The constructor stores the original layer as self.linear, so the pretrained weights stay inside the wrapper.
Do this instead
Treat the wrapper as holding the original layer and adding a correction to it.
Common mistake
Believing the LoRA path starts with random output that changes the model
Why it fails
B is initialized to zero, so the product of A and B is zero and lora(x) is exactly zero at the start.
Do this instead
At initialization, the wrapped output equals linear(x).
Common mistake
Looking for extra computation inside forward
Why it fails
The source states that nothing else happens in forward beyond the two paths and their sum.
Do this instead
Trace only the two paths and the addition when checking behavior.

Practice Tracing the Wrapper

EASY

A LinearWithLoRA wraps a pretrained Linear layer that produces linear(x) = [2.0, 4.0] for some input x. The wrapper has just been built and no training has run. Without computing anything else, state the value of lora(x) and the wrapper output y. Then explain which matrix's initialization justifies your answer.

Hints

Ask which matrix inside LoRALayer is set to zero at initialization.
What does a zero matrix do when the input is multiplied by it?
Adding zero to a vector leaves that vector unchanged.

Your attempt (not saved)

MEDIUM

After a few training steps, the matrix B is no longer zero. Explain in two or three sentences why lora(x) is now generally nonzero, and why the wrapper output can now differ from linear(x).