PaperlayerView exact revision on arXiv

Source-linked AI summary

Surgical Fine-Tuning Improves Adaptation to Distribution Shifts

Yoonho Lee, Annie S. Chen, Fahim Tajwar, Ananya Kumar, Huaxiu Yao, Percy Liang, Chelsea Finn

arXiv:2210.11466v3cs.LGcs.AI

TL;DR

Distribution shifts can make pretrained neural networks brittle, motivating adaptation with small labeled target datasets while preserving useful source information. The paper introduces surgical fine-tuning, evaluates different layer subsets across seven tasks and three shift types, and proves an idealized two-layer result. Selective tuning can outperform full fine-tuning, with the best subset depending on the shift and first-layer tuning improving CIFAR-10-C by almost 3% on average.

  • Problem

    Pretrained neural networks can be brittle under distribution shifts, while transfer learning must adapt to small labeled target datasets without discarding useful source information.

  • Method

    The paper studies surgical fine-tuning, adapting only a small contiguous subset of layers and comparing layer choices across seven tasks, three shift types, and an idealized two-layer theory.

  • Results

    Selective tuning can outperform full fine-tuning, with the best layer subset varying by shift type; first-block tuning improves unseen CIFAR-10-C performance by almost 3% on average.

  • Takeaways & Limitations

    Layer selection should reflect the type of distribution shift, because different shifts require adapting different parts of the network.

  • Takeaways & Limitations

    Parameter freezing is most relevant with limited target data; when target data is plentiful, choosing a bigger tuned subset can be beneficial, and automatic selection does not outperform manual layer choice.

Abstract

from arXiv · show

A common approach to transfer learning under distribution shift is to fine-tune the last few layers of a pre-trained model, preserving learned features while also adapting to the new task. This paper shows that in such settings, selectively fine-tuning a subset of layers (which we term surgical fine-tuning) matches or outperforms commonly used fine-tuning approaches. Moreover, the type of distribution shift influences which subset is more effective to tune: for example, for image corruptions, fine-tuning only the first few layers works best. We validate our findings systematically across seven real-world data tasks spanning three types of distribution shifts. Theoretically, we prove that for two-layer neural networks in an idealized setting, first-layer tuning can outperform fine-tuning all layers. Intuitively, fine-tuning more parameters on a small target dataset can cause information learned during pre-training to be forgotten, and the relevant information depends on the type of shift.

1 INTRODUCTION

The paper studies surgical fine-tuning, which adapts only selected layers while freezing the rest, and finds that the best layers depend on the distribution shift. Across seven real-world tasks and theoretically in two-layer networks, selective tuning can outperform full fine-tuning, including nearly 3% average improvement for CIFAR-10-C.

  • Empirical findings: The best layer subset varies by shift type: early layers suit input shifts, later blocks feature shifts, and the last layer output shifts.Figure 1 illustrates these patterns with CIFAR-C, Entity-30, and CelebA.
  • Empirical findings: Fine-tuning the first block outperforms full fine-tuning by almost 3% on average for unseen CIFAR-10-C corruptions.This result comes from adapting a CIFAR-10-pretrained model with a small labeled corrupted-image dataset.
  • Contribution: Surgical fine-tuning updates a small contiguous subset of pretrained layers while freezing the remaining parameters.The paper studies this approach across input-level, feature-level, and output-level distribution shifts.
  • Evaluation: The findings are validated across seven real-world distribution-shift scenarios and three shift categories.The paper combines empirical evaluation with an idealized two-layer neural-network analysis.
  • Theory: The theory shows that first-layer tuning can suit input perturbations, whereas last-layer tuning can suit label perturbations.It also constructs a setting where first-layer tuning provably outperforms updating all parameters by preserving source-distribution information.

2 SURGICAL FINE-TUNING: FREEZING PARAMETERS DURING ADAPTATION

The paper studies adapting a source-trained model to a smaller target dataset by freezing most parameters and tuning a selected contiguous layer subset. Across distribution shifts, surgical fine-tuning can outperform full fine-tuning, with the best layers depending on which part of the data-generating process changes.

  • Problem setting: The transfer-learning setting uses a large source dataset and a relatively small target dataset from different distributions.The objective is to retain broad source-domain performance while making minor adjustments for the target distribution.
  • Method: Surgical fine-tuning updates only a subset of layers while freezing the remaining parameters at their pre-trained values.The framework includes tuning all layers, the last layer, the last few layers, or earlier layers such as the first layer.
  • Scope: The benefit of freezing parameters is tied to limited target data, so tuning a larger subset can become beneficial when target data is plentiful.Full fine-tuning may update layers unnecessarily and underperform on held-out target data because of overfitting.
  • Experiments: Across nine real-world distribution shifts, the best layer block depends on shift type: first for input-level, middle for feature-level, and last for output-level shifts.The experiments include corrupted images, changed class subpopulations, altered labels or spurious correlations, and natural shifts across hospitals or regions.
  • Results: On every evaluated domain, tuning one network block outperforms tuning all parameters on the target distribution.For natural shifts, the earliest embedding layer works best for Camelyon17, while later attention blocks work best for FMoW.
  • Results: For image corruptions, early-layer adaptation remains effective with limited or online target data, while full fine-tuning can deteriorate as more online data is added.On CIFAR-C, first-block tuning can match or outperform full tuning across varying data amounts; in online adaptation, distortion of pre-trained features explains the deterioration of full fine-tuning.

3 ANALYSIS OF SURGICAL FINE-TUNING

The analysis shows that the most effective layer to tune depends on the distribution shift, and that selective tuning can outperform full fine-tuning when target data are limited. Synthetic experiments and theory support this layer-specific advantage.

  • 3.1 Layer Choice and Expressivity: Input perturbations favor first-layer tuning, which can achieve zero target loss when the target input is an invertibly transformed source input.The first layer can undo the perturbation, whereas changing only the last layer may be insufficient.
  • 3.1 Layer Choice and Expressivity: Label perturbations favor last-layer tuning because the inputs remain unchanged while the outputs are transformed.Tuning only the first layer may lack sufficient expressivity after information is destroyed by the ReLU activation.
  • 3.2 Can Surgical Fine-Tuning Outperform Full Fine-Tuning?: First-layer tuning can provably outperform full fine-tuning on small target datasets, even for two-layer linear networks.The theorem gives a setting where first-layer tuning reaches 0 loss while full fine-tuning retains higher non-zero loss throughout training.
  • 3.2 Can Surgical Fine-Tuning Outperform Full Fine-Tuning?: Fine-tuning unnecessary parameters can overfit the small target dataset and forget information learned during pre-training.The analysis attributes this risk to jointly changing the feature extractor and head when only a few new input directions must be learned.
  • 3.3 Surgical Fine-Tuning on Synthetic Distribution Shifts: Synthetic shifts created by adding noise to parameter blocks show that tuning the responsible subset outperforms tuning other subsets or all layers.The experiment uses a pretrained ResNet-26 on CIFAR-10 and separately perturbs three layer blocks or the last layer.

4 AUTOMATICALLY SELECTING WHICH LAYERS TO TUNE

The paper evaluates automatic criteria for selecting layers for surgical fine-tuning across seven real-data domains. Auto-RGN is the strongest one-run method, while Cross-Val performs best overall but requires a run for each block.

  • 4.1 Methods: The study compares automatic layer-selection criteria with full fine-tuning and regularization methods across seven real-data domains and synthetic shifts.The evaluated criteria include Relative Gradient Norm and Signal-to-Noise Ratio, alongside Cross-Val and existing baselines.
  • 4.1 Methods: Auto-RGN selects layers using relative gradient norms and converts normalized per-tensor values into per-tensor learning rates.It requires no additional hyperparameters over tuning all layers and uses larger gradients to identify potentially useful layers.
  • 4.1 Methods: Auto-SNR normalizes gradient signal-to-noise ratios and freezes layers whose SNR falls below a tuned threshold.The criterion is intended to identify noisy gradients that may distort the function learned during pre-training.
  • 4.2 Results on Real World Datasets: Cross-Val performs best overall but requires a separate surgical fine-tuning run for each network block.It selects the best block using a held-out target-distribution validation set.
  • 4.2 Results on Real World Datasets: Auto-RGN matches or outperforms full fine-tuning across all domains and is the most competitive method requiring only one fine-tuning run.It does not quite match Cross-Val, while Auto-SNR performs worse on most shifts and gradual unfreezing is not consistently better than full fine-tuning.
  • 4.3 Synthetic Distribution Shifts: On synthetic shifts, Auto-RGN assigns higher weights to layers in the block responsible for the shift and outperforms full fine-tuning.Its selected layers match the performance pattern of tuning only the noisy parameter subset.

5 RELATED WORK

Related work covers parameter freezing, transfer-learning methods that preserve pretrained information, and robustness to distribution shifts. These lines of work motivate selective adaptation and frame the paper’s contribution.

  • Parameter Freezing: Prior work uses parameter freezing in domain adaptation, early stopping, generative models, and gradient-based meta-learning.Prompt tuning is also cited as an effective approach for fast adaptation of large language models.
  • Transfer Learning: Transfer-learning research studies adapting pretrained features to target distributions while regularizing fine-tuning to preserve information from pre-training.The cited methods include approaches based on transfer learning and fine-tuning regularization.
  • Distribution Shifts: Distribution-shift research commonly frames robustness as zero-shot generalization from source data to a target distribution.The paper instead studies adaptation using a small amount of labeled target data.

6 DISCUSSION

Surgical fine-tuning adapts only a small contiguous parameter subset, and the most effective subset depends on the distribution shift. Automatic layer-selection criteria help address the extra cost of manually choosing layers, but do not yet match the best surgical approach.

  • 6 DISCUSSION: Surgical fine-tuning adapts only a small contiguous subset of network parameters during target-distribution fine-tuning.It is equivalently defined as freezing all but a few layers.
  • 6 DISCUSSION: The effective layer subset depends on the distribution shift, with earlier layers outperforming full or later-layer fine-tuning for image-corruption shifts.This pattern is stated for input-level shifts such as image corruption.
  • 6 DISCUSSION: Automatic layer-selection criteria are analyzed because manually choosing frozen layers requires more fine-tuning runs than tuning all layers.The paper studies two criteria for automatic selection.
  • 6 DISCUSSION: Auto-RGN consistently improves over full fine-tuning but does not match the best surgical fine-tuning approach.The paper identifies more effective automatic-selection criteria as a direction for future work.

A.1 PROOFS FOR SECTION 3.1

For two-layer networks, the theoretically preferred layer depends on the shift type: first-layer tuning can handle input transformations, whereas last-layer tuning can handle label transformations. In each case, tuning the opposite layer may be insufficient.

  • A.1 PROOFS FOR SECTION 3.1: For invertible input transformations with unchanged labels, a first-layer transformation can achieve zero target loss.The construction uses B = ˆBsrcA−1.
  • A.1 PROOFS FOR SECTION 3.1: For the same input-shift setting, there are transformations where no last-layer choice achieves zero target loss.The counterexample sets A = −I and uses ReLU activations to make the transformed representation zero.
  • A.1 PROOFS FOR SECTION 3.1: For unchanged inputs with labels scaled by t, a last-layer scaling can achieve zero target loss.The construction sets v = tˆvsrc.
  • A.1 PROOFS FOR SECTION 3.1: For some label transformations, changing only the first layer cannot achieve zero target loss.With t = −1 and positive source parameters, the fixed head cannot express the negative target labels.

A.2 PROOF OF THEOREM 1 IN SECTION 3.2

The proof compares first-layer tuning with full fine-tuning in a two-layer model under a target distribution that mixes source and orthogonal subspaces. Under the stated assumptions, first-layer tuning reaches zero target loss, while full fine-tuning retains positive loss.

  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: The proof establishes Theorem 1 by combining a zero-loss result for first-layer tuning with a positive-loss result for full fine-tuning.Two lemmas provide the respective loss bounds.
  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: The proof assumes that the conditional label distribution is shared by source and target and imposes a non-degeneracy condition on the optimal model.These assumptions support the construction and the contradiction used in the full-fine-tuning argument.
  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: The source distribution has density on a dsrc-dimensional subspace, with dsrc less than the input dimension d.The orthogonal distribution has density on a dorth-dimensional subspace orthogonal to the source support.
  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: The target distribution is an equal mixture of the source distribution and an orthogonal distribution.Sampling selects the source or orthogonal distribution with probability 0.5 each.
  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: With probability at least 1−δ, first-layer tuning reaches zero target loss at convergence when n > 10dorth log 2/δ.The argument uses examples spanning the orthogonal subspace and correctness on the source and orthogonal supports.
  • A.2 PROOF OF THEOREM 1 IN SECTION 3.2: When the representation dimension is one and dsrc > n, full fine-tuning has positive target loss at every time t.Because the training examples do not span the source distribution, the proof shows that zero target loss cannot persist.

A.3 ADDITIONAL DATASET DETAILS

The additional dataset details cover image-corruption, spurious-correlation, medical-hospital, and geographic satellite-image shifts. Target fine-tuning uses small labeled subsets whose sizes and balancing schemes vary by dataset.

  • A.3 ADDITIONAL DATASET DETAILS: Waterbirds and CelebA evaluate spurious correlations using balanced target subsets and 400 target images for fine-tuning.Each target set is evenly split across four bird-background or hair-color-gender groups.
  • A.3 ADDITIONAL DATASET DETAILS: Camelyon17 measures a hospital-related shift across medical images collected from five hospitals.Differences in devices and data processing produce the natural distribution shift.
  • A.3 ADDITIONAL DATASET DETAILS: FMoW evaluates 62-way building or land-use classification under a geographic shift to the Africa test subset.The source contains satellite images from five geographic regions, while the target subset has roughly 2,500 images.

A.4 ADDITIONAL DETAILS FOR SUPERVISED TRANSFER LEARNING EXPERIMENTS

Additional supervised-transfer experiments report implementation settings and show that surgical fine-tuning remains effective across architectures and evaluation metrics. On ImageNet-C, mCE yields conclusions similar to average accuracy, while a CLIP ViT-B/16 experiment on Living-17 finds a middle block strongest.

  • Experimental settings: All supervised-transfer experiments use early stopping based on the best accuracy on a held-out labeled-target validation subset.The remaining settings vary by dataset, including training duration, optimizer, learning-rate grid, and weight decay.
  • Dataset-specific settings: The Camelyon17 setup uses a CLIP ViT-B/16 pretrained on CLIP datasets before target fine-tuning.The supplied passage begins the Camelyon17 procedure but does not include all subsequent details.
  • Vision transformer experiment: A CLIP ViT-B/16 experiment on Living-17 finds that tuning one middle transformer block outperforms full fine-tuning.Living-17 represents a feature-level shift in this experiment.
  • Results: mCE and average accuracy produce similar conclusions on ImageNet-C, with Cross-Val and Auto-RGN performing best.mCE weights target-distribution error by corruption difficulty.

A.5 MORE LARGE VISION TRANSFORMER EXPERIMENTS

Additional large-vision-transformer experiments evaluate surgical fine-tuning with AdamW on Camelyon17 and FMoW. Surgical tuning still improves over no fine-tuning, but its relative advantage over full fine-tuning is smaller than with SGD, and full fine-tuning is unstable for ViT-B/16 on FMoW.

  • Experimental setup: The AdamW experiments fine-tune CLIP ViT-B/16 or ViT-L/14 models on Camelyon17 and FMoW target data.Camelyon17 uses three pretraining epochs, while FMoW uses five; both use cosine annealing and batch size 32.
  • Evaluation: Table 8 reports OOD-set accuracies for surgically fine-tuning different parameters, including standard deviations from three random seeds.The table compares parameter choices across the two WILDS datasets.
  • Results: Surgical fine-tuning with AdamW improves over no fine-tuning, but its relative advantage over full fine-tuning is smaller than with SGD.The experiments use CLIP ViT-B/16 and ViT-L/14 models on two WILDS datasets.
  • Results: Full fine-tuning of a ViT-B/16 network is unstable for FMoW target distributions.Later-layer tuning appears to improve performance without the reported instability.

A.6 COMPLETE UNSUPERVISED ADAPTATION RESULTS

Complete unsupervised-adaptation experiments apply MEMO with parameter freezing to corruption datasets under episodic and online settings. The results are organized by layer-freezing choices and representative corruptions, with ImageNet-C evaluated on a computationally selected subset.

  • Method: MEMO adapts one test image at a time by minimizing marginal entropy over standard augmentations before predicting its label.The experiments evaluate both episodic and online variants.
  • Method: Episodic MEMO resets the network to pretrained weights after each prediction, whereas online MEMO retains updates across test images.These settings distinguish whether adaptation changes accumulate over the test stream.
  • Related adaptation method: TENT parameter-freezing experiments were not pursued because TENT updates only batch-normalization modules while MEMO updates all parameters.The reported mismatch prevented the expected results in the authors’ experiments.
  • Evaluation caveat: The ImageNet-C evaluation uses 2000 test images rather than the full 50,000-image test set because of computational constraints.The subset contains the first two test images per class across 1000 classes, and no-adaptation baselines are included for comparison.
  • Results: Tables 9 and 10 report MEMO with parameter freezing on CIFAR-10-C and ImageNet-C at severity level 5 using representative corruptions.The ImageNet-C table covers 14 representative corruptions, while the CIFAR-10-C table covers 10.
Loading 2210.11466v3…