Freezing a layer does not freeze its batch norm

"Freeze the backbone and train a new head" is one of the most common moves in transfer learning. In PyTorch it usually means one line that turns off gradients for the backbone's parameters, and the mental model that goes with it is simple: the frozen part no longer changes, so whatever it did before, it still does.

I wanted to check that mental model against what a batch-normalised network actually stores. A batch-norm layer, as Ioffe and Szegedy (2015)¹ defined it, holds two kinds of state. Its scale and shift, $\gamma$ and $\beta$, are parameters, and freezing stops the optimizer from touching them. Its running mean and variance are buffers, not parameters, and they are updated by the forward pass itself whenever the layer is in training mode:

No gradient is involved, so the gradient flag has no say. With PyTorch's default momentum of $m = 0.1$, the weight left on the original statistics after $t$ batches is $0.9^t$, which is below one in a thousand after 66 batches. A few minutes of training on new data replaces the statistics completely.

On paper, then, a "frozen" backbone trained in training mode is not frozen. It keeps its weights and re-measures its normalisation on whatever data passes through it. I wanted to see how much that matters.

The experiment

I trained a small convolutional network with three batch-normalised convolutions on task A: classify noisy copies of ten random 3×32×32 templates. Then I froze the backbone's weights, attached a second linear head, and trained only that head for 800 steps on task B: ten different templates, with the inputs shifted in brightness and stretched in contrast. Finally I checked the original head on task A again, on the same held-out images as before.

I did the freezing four ways, three seeds each:

1. gradients turned off and the model left in training mode, the one-line version;
2. the same, with the frozen backbone switched to evaluation mode so its batch norm uses stored statistics;
3. the whole model in evaluation mode;
4. the frozen weights shared with a second copy of the backbone that keeps its own batch-norm statistics for task B.

Task A scored 100% before the second stage in every run. Accuracies are means over three seeds.

Every row froze the weights exactly: the largest change in any weight was zero. In the first row the running statistics moved by up to 4 units, and task A dropped from 100% to 10.4%, which is chance for ten classes. The old head's outputs on the old images changed by up to 39 logits. Nothing that the optimizer could see had moved.

Why the old head broke

The old head was trained on features normalised with task A's statistics. After the second stage, the backbone normalises everything, including task A's images, with task B's statistics. Task B's inputs are brighter and higher in contrast, so in the first layer the stored mean and variance now describe a different distribution. Task A's activations come out shifted and squeezed. Each layer after that receives the wrong input distribution and normalises it with the wrong numbers too, so the error compounds through the network. The weights are untouched and the function is different.

Why it helped the new head

The first row also has the best task B accuracy, 43% against about 21–26% for the two rows that pinned the statistics. The drifting statistics adapted the backbone to task B's input distribution, renormalising its features for the new data for free. Li et al. (2016)² turned exactly this into a method, AdaBN: re-estimate a trained network's batch-norm statistics on the target domain and change nothing else. In the first row it happened by accident, and it took the old task down with it.

The two pinned rows keep task A intact but give the new head features normalised for the wrong distribution. A linear head can only undo part of that, which shows in its accuracy. The two pinned modes compute the same thing in this network, and the gap between them, 25.6% against 20.6%, is run-to-run noise.

Keeping both

The fourth row gets both results. The backbone's weights are shared and frozen, but task B has its own copy of the batch-norm buffers, and only that copy is in training mode during the second stage. Task A, run through the original buffers, is bit-identical: its logits changed by exactly zero. Task B, run through its own buffers, reaches 43%, the same as the adaptation that broke task A. Chang et al. (2019)³ use the same structure for domain adaptation, with domain-specific batch normalisation and all other parameters shared.

The cost is bookkeeping. At inference the model has to know which set of statistics an input belongs to, which is easy when the task is known and hard when it isn't.

The traps around eval()

Switching the frozen part to evaluation mode is the usual fix, and it has two traps of its own.

The first is that it doesn't stay pinned. Most training loops call model.train() at the start of every epoch, after validation. That call recursively sets every submodule back to training mode, including the frozen backbone, and the statistics start drifting again from the next batch. The pin has to be reapplied after every such call, for example by overriding the model's train method so that the frozen part always stays in evaluation mode.

The second is that switching the whole model to evaluation mode, the blunt version of the fix, changes more than batch norm. Dropout switches off. Some models also change what forward returns. Torchvision's detection models return a dictionary of losses in training mode and detections in evaluation mode, and Inception v3 returns its auxiliary classifier's output only in training mode. A training loop that relies on those outputs either breaks with an error or, worse, trains on without a loss term and reports nothing.

What I would check after any freeze

Three cheap checks would have caught every problem above:

• the largest change in every buffer as well as every parameter of the frozen part, since "frozen" has to mean both;
• the frozen model's outputs on a fixed batch of the old data, compared before and after, where any nonzero difference means the freeze leaked;
• the training mode of each submodule at the moment the loss is computed, not only when the loop starts.

Why running statistics fall behind even in an unfrozen network, when Adam keeps rotating the weights, is the subject of batch norm’s running statistics describe a network that has already moved. The safest way to add a task to a frozen network, with a branch that cannot touch the old outputs, is in adding an output without moving the old ones.

What I still don't know

My two tasks differ by a strong, simple shift in brightness and contrast, which is close to the worst case for shared statistics. With a gentler shift, drifting statistics might cost the old task little and help the new one less. I don't know where the crossover sits.

I also only tried linear heads on top of the frozen backbone. A deeper head, with its own normalisation, could absorb part of the shift itself and might narrow the gap between the pinned and the separate-statistics rows. And separate statistics assume the task is known at inference; when it isn't, the model would need some way to tell which statistics an input belongs to, and I haven't tried any.

References

1. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift — Sergey Ioffe and Christian Szegedy, ICML 2015
2. Revisiting Batch Normalization For Practical Domain Adaptation — Yanghao Li, Naiyan Wang, Jianping Shi, Jiaying Liu and Xiaodi Hou, arXiv 2016
3. Domain-Specific Batch Normalization for Unsupervised Domain Adaptation — Woong-Gi Chang, Tackgeun You, Seonguk Seo, Suha Kwak and Bohyung Han, CVPR 2019