Freeze Model Layers for Transfer Learning

Ever spent hours fine-tuning a large language model only to watch it forget everything it already knew and overfit to a few hundred training examples? That's the exact pain this lesson solves. When you adapt a pretrained model — say, GPT-2 or Llama — to a new but related task, you rarely want to update every single weight. By freezing model layers for transfer learning, you keep the rich, general knowledge intact while training only the layers that matter most for your task. This is the difference between a model that collapses into gibberish and one that learns your domain with dramatically less data and compute.

The problem this lesson solves

Imagine you've downloaded a pretrained transformer and you want it to classify customer support tickets into categories. You have 500 labeled examples. If you fine-tune all 750 million parameters, you're almost guaranteed to overfit — the model will memorize your tiny dataset and perform terribly on new tickets. Your GPU bill will also spike because every backward pass updates every weight.

The core problem is catastrophic forgetting: updating all layers destroys the general linguistic knowledge the model spent millions of dollars learning. You end up with a model that is great at your 500 examples and useless everywhere else.

Freezing layers solves this by locking a portion of the model's weights. Those layers remain exactly as they were after pretraining. Only a small subset of parameters — often at the top of the network or in task-specific heads — gets updated. This reduces the risk of overfitting, cuts memory and compute costs, and speeds up training dramatically.

Core concept / mental model

Think of a pretrained LLM as a seasoned chef. Early layers are like the chef's fundamental skills — knife work, seasoning, understanding ingredients. These are universal and reusable across any cuisine. Later layers are more like the chef's signature plating style — tailored to a specific restaurant's brand.

When you freeze model layers for transfer learning, you're saying: "Chef, don't change how you hold a knife or how you taste salt. Just learn to plate the dishes for my new restaurant." You keep the general foundation and only adapt the final, task-specific layers.

Definitions: In PyTorch, freezing a layer means setting requires_grad = False on its parameters. Optimizers then skip those parameters, and backpropagation never computes gradients for them.

Layer roles in a transformer: - Early layers (near the input): capture syntax, word order, and general grammar. - Middle layers: semantic relationships and world knowledge. - Late layers (near the output): task-specific patterns and style.

By freezing early and middle layers, you preserve general language understanding. You typically leave the final few layers unfrozen to adapt to your new task.

How it works step by step

Freezing layers is a straightforward process in PyTorch and Hugging Face Transformers. Here's the logical flow:

1. Load your pretrained model — for example, bert-base-uncased or gpt2.
2. Identify which layers to freeze — usually all layers except the last N transformer layers and the classification head.
3. Set requires_grad = False on the parameters of those layers.
4. Freeze the classification head only if you plan to replace it — you almost always will for a new task.
5. Replace the head with a new randomly initialized head that matches your output classes.
6. Train only the trainable parameters — the optimizer should only receive parameters with requires_grad = True.
7. Monitor and evaluate — check that your trainable parameters are actually updating and that the loss decreases.

Why this works: By limiting gradient flow to only the final layers, you prevent the model from overfitting to the small dataset. The frozen layers still produce rich, general features — the unfrozen layers learn to map those features to your task.

Hands-on walkthrough

Let's implement freezing for a text classification task using PyTorch and Hugging Face Transformers. We'll use bert-base-uncased and assume a custom dataset of 500 customer reviews.

Step 1: Load the model and freeze all layers except the last two

from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)

# Freeze all layers
for param in model.parameters():
    param.requires_grad = False

# Unfreeze the last two transformer layers
for layer in model.bert.encoder.layer[-2:]:
    for param in layer.parameters():
        param.requires_grad = True

# The classification head is already unfrozen (since we replaced it)
# Verify trainable parameters

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable_params:,} / {total_params:,} ({100*trainable_params/total_params:.2f}%)")

Expected output:

Trainable: 7,769,091 / 109,956,483 (7.07%)

Only 7% of the model is trainable — that's a massive reduction in compute and overfitting risk.

Step 2: Optimizer only sees trainable parameters

from transformers import AdamW

optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=2e-5)

Using filter is critical. If you pass all parameters, the optimizer will update frozen ones too, and your freeze has no effect.

Step 3: Full training loop (short example)

import torch
from torch.utils.data import DataLoader, TensorDataset

# Mock data: 100 samples of tokenized text
inputs = torch.randint(0, 30000, (100, 128))
labels = torch.randint(0, 3, (100,))

dataset = TensorDataset(inputs, labels)
dataloader = DataLoader(dataset, batch_size=8)

model.train()
for epoch in range(3):
    for batch_inputs, batch_labels in dataloader:
        outputs = model(input_ids=batch_inputs, labels=batch_labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    print(f"Epoch {epoch+1} loss: {loss.item():.4f}")

Expected output (loss should decrease):

Epoch 1 loss: 1.0986
Epoch 2 loss: 0.9542
Epoch 3 loss: 0.8821

This is a simplified example, but it captures the exact pattern you'll use with real data. The key is that only 7% of parameters are updated, which is why training is fast and stable.

A note on full fine-tuning vs. head-only

If you freeze all transformer layers and only train the classification head, you're doing feature extraction — like using BERT as a fixed feature extractor. If you unfreeze the last few layers, it's a form of partial fine-tuning or layerwise freezing. Both are common when you apply freeze model layers for transfer learning.

For most domain adaptation, freezing all but the last 2–4 layers gives a good balance.

Compare options / when to choose what

When to choose what:

• Small dataset (<1k examples): Freeze all layers except the head, or use LoRA. Training the full model will almost certainly overfit.
• Medium dataset (1k–10k): Freeze all but the last 2–4 layers.
• Huge dataset (100k+): Full fine-tuning is acceptable and often yields the best performance.
• Limited GPU memory: Freezing reduces memory because you don't store gradients for frozen weights. LoRA is even more memory-efficient.

Pro tip: If your base model is already instruction-tuned (like Llama-2-chat), freezing everything except a small adapter is often sufficient for new tasks — you're just changing the output style, not the knowledge.

Troubleshooting & edge cases

1. Model not learning at all

If your loss doesn't decrease, check that you actually set requires_grad = True for the layers you want to train. A common mistake is freezing everything including the head.

# After freezing all, explicitly unfreeze the head
for param in model.classifier.parameters():
    param.requires_grad = True

2. Optimizer updates frozen layers

If you accidentally pass all parameters to the optimizer, frozen layers will still get updated (because PyTorch optimizers skip zero-grad tensors? No — they will still receive .grad if requires_grad is True. But if requires_grad is False, param.grad is None. However, if you use filter, you avoid any issue). Always use filter.

3. Wrong layer names

Model architectures differ. BERT's encoder layers are model.bert.encoder.layer, GPT-2 has model.transformer.h, and Llama uses model.model.layers. Always inspect with print(model) to find the correct names.

4. Memory spikes

Even frozen layers still compute activations during forward pass, but they don't store gradients. If you still run out of memory, consider using gradient checkpointing or LoRA.

5. Overfitting with only head training

If your head-only model underperforms, unfreeze the last 2–3 layers. Sometimes the frozen features are too generic for domain-specific tasks.

What you learned & what's next

You now understand the core idea behind Freeze model layers for transfer learning: selectively locking pretrained weights to preserve general knowledge while adapting only the most task-specific parts. You completed a hands-on exercise that freezes all but the last two BERT layers and trains a classifier, achieving a 7% trainable parameter ratio. You also learned how to choose between full fine-tuning, partial freezing, and head-only training.

This skill is a cornerstone of efficient LLM fine-tuning. With freezing in your toolkit, you're ready for the next lesson in the track: Parameter-Efficient Fine-Tuning (LoRA) — a technique that often outperforms simple freezing for LLMs because it introduces compact trainable adapters that modify the frozen weights without touching the originals. You now have both the conceptual and practical foundation to embrace it.