How Transfer Learning Works When Fine-Tuning Frozen vs Unfrozen Layers

Transfer learning adapts pre-trained models by either freezing the base architecture to train only the classification head (minimizing GPU memory and preventing catastrophic forgetting) or unfreezing all layers for full fine-tuning (maximizing accuracy when downstream tasks diverge from pre-training data).

The rohitg00/ai-engineering-from-scratch curriculum demonstrates practical transfer learning through a toggle mechanism that switches between frozen-body and full fine-tuning regimes. This approach allows practitioners to leverage massive pre-trained models while controlling computational costs and preserving general knowledge based on dataset size and domain similarity.

Freezing the Body: Head-Only Fine-Tuning

Freezing the model body involves setting requires_grad=False on all parameters except the final classification head, effectively treating the pre-trained layers as a fixed feature extractor.

Why Freeze the Base Layers?

The curriculum emphasizes that freezing preserves the broad linguistic or visual features learned during massive pre-training, which is critical when working with limited downstream data. As noted in the classifier fine-tuning documentation, "head-only training is fast, almost free in memory, and rarely overfits with this little data."

Memory and Computational Efficiency

When the body is frozen, the optimizer stores gradients only for head parameters. This reduces GPU memory requirements by roughly an order of magnitude and accelerates each training step, making it feasible to fine-tune large models on consumer hardware.

Unfrozen Full Fine-Tuning

Unfreezing the entire architecture allows gradients to flow through all layers, enabling the model to adapt its internal representations to the new domain.

Task-Specific Adaptation

Full fine-tuning reaches higher accuracy when the downstream domain drifts significantly from the pre-training corpus. The curriculum observes that "full fine-tuning… reaches higher accuracy when the downstream domain drifts from the pretraining corpus," as the model can adjust low-level features to match new task distributions.

Capacity for Large Datasets

When plentiful labeled data is available (>10,000 examples), updating every weight extracts more signal than a static body, leading to measurable gains in precision, recall, and F1 scores.

Critical Trade-offs: Memory, Speed, and Stability

Aspect Frozen Body (Head-Only) Unfrozen Body (Full Fine-Tuning)
GPU RAM Minimal (head parameters only) Requires full model + gradients (often >50 GB for 7B models)
Training Time Short (few epochs) Longer (more epochs, larger batches)
Catastrophic Forgetting Low risk—body stays unchanged High risk—body may lose pre-training knowledge if over-trained
Accuracy Ceiling Modest baseline Potentially higher on domain-shifted tasks

Implementation in the Curriculum

The rohitg00/ai-engineering-from-scratch repository explicitly implements these regimes through a unified training interface.

The freeze_body Toggle

In phases/19-capstone-projects/38-classifier-finetuning/docs/en.md, the curriculum provides a single train_classifier function that accepts a freeze_body boolean flag. When enabled, the implementation iterates through named parameters and sets requires_grad=False on all layers except the classifier, ensuring the optimizer "sees only the head and the body stays frozen."

Key Source Files

Practical Code Examples

The following PyTorch snippets mirror the pattern used in the Classifier Fine-tuning lesson, demonstrating both frozen and unfrozen approaches:

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "facebook/opt-350m"
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# ------------------------------------------------------------------

# 1️⃣  Frozen-body (head-only) fine-tuning

# ------------------------------------------------------------------

# Freeze all parameters except the classification head

for name, param in model.named_parameters():
    if "classifier" not in name:
        param.requires_grad = False

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

# Training loop (simplified)

for batch in dataloader:
    inputs = tokenizer(batch["text"], padding=True, truncation=True, return_tensors="pt")
    labels = batch["label"]
    outputs = model(**inputs, labels=labels)
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# ------------------------------------------------------------------

# 2️⃣  Unfrozen (full) fine-tuning

# ------------------------------------------------------------------

# Enable gradients for every parameter

for param in model.parameters():
    param.requires_grad = True

optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

# Same training loop as above but now the whole model is updated

When to Choose Each Regime

Scenario Recommended Approach
Few (< 1,000) labeled examples Freeze the body, train only the head
Large (> 10,000) domain-shifted dataset Unfreeze the body for full fine-tuning
Budget-constrained inference environments Freeze the body to minimize deployment costs
Need to retain zero-shot capabilities Freeze the body or partially unfreeze only the last block

Summary

  • Frozen-body fine-tuning sets requires_grad=False on base layers, optimizing only the classification head to save memory and prevent overfitting on small datasets.
  • Unfrozen fine-tuning enables gradients across all parameters, maximizing accuracy for domain-shifted tasks but requiring substantial GPU resources.
  • The rohitg00/ai-engineering-from-scratch curriculum implements these approaches via a freeze_body toggle in phases/19-capstone-projects/38-classifier-finetuning/docs/en.md.
  • Catastrophic forgetting poses minimal risk with frozen layers but becomes a significant concern when fully fine-tuning on specialized domains.
  • Choose head-only training for limited data and resource constraints; select full fine-tuning when abundant data and compute are available.

Frequently Asked Questions

What happens to GPU memory when freezing layers in transfer learning?

When you freeze the model body, the optimizer only tracks gradients for the unfrozen head parameters. This reduces GPU memory consumption by approximately an order of magnitude compared to full fine-tuning, which must store gradients for every parameter in the architecture (often exceeding 50 GB for 7B parameter models).

Can I unfreeze layers after initially training the head only?

Yes, this progressive unfreezing strategy is common in advanced implementations. You can first train the head with frozen layers to establish a baseline, then gradually unfreeze deeper layers (starting from the final transformer blocks) to allow domain adaptation while maintaining training stability. The curriculum's train_classifier function supports this by allowing dynamic toggling of the freeze_body flag between training phases.

What is catastrophic forgetting in the context of fine-tuning?

Catastrophic forgetting occurs when a model loses previously learned general knowledge during fine-tuning on a specific task. When layers remain frozen, the risk is eliminated because the pre-trained weights never change. However, with full fine-tuning (unfrozen layers), the model may overwrite its broad linguistic or visual capabilities if trained too aggressively on a narrow dataset, reducing its effectiveness on out-of-distribution examples.

Where is the freeze_body parameter implemented in the repository?

The freeze_body toggle is documented in phases/19-capstone-projects/38-classifier-finetuning/docs/en.md and implemented in the accompanying train_classifier function within phases/19-capstone-projects/38-classifier-finetuning/code/main.py. This parameter controls whether the training loop applies requires_grad=False to body parameters, effectively switching between head-only and full fine-tuning regimes.

Have a question about this repo?

These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:

Share the following with your agent to get started:
curl -s "https://instagit.com/install.md"

Works with
Claude Codex Cursor VS Code OpenClaw Any MCP Client

Maintain an open-source project? Get it listed too →