Surgical Fine-Tuning Improves Adaptation to Distribution Shifts (2210.11466v3)

Abstract: 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.

• The paper shows that fine-tuning only selected neural network layers significantly improves performance under various distribution shifts.
• The methodology combines empirical evaluations and theoretical analysis, employing automated criteria like Auto-RGN and Auto-SNR for optimal layer selection.
• The research highlights that layer-specific tuning mitigates overfitting and enhances adaptation efficiency, especially in resource-constrained transfer learning scenarios.

Surgical Fine-Tuning under Distribution Shift

The paper under review explores advanced methods for fine-tuning neural networks in the context of transfer learning, specifically when there is a distribution shift between the source and target datasets. The core argument presented is that selectively fine-tuning a small subset of layers—termed "surgical fine-tuning"—can surpass traditional fine-tuning methods, particularly under various types of distribution shifts.

Key Areas of Investigation

1. Selective Layer Fine-Tuning: The authors identify that fine-tuning only a subset of layers instead of the entire network can yield better performance under distribution shifts.
2. Empirical Evaluation: They conduct thorough empirical evaluations across seven real-world tasks classified into three types: input-level shifts (e.g., image corruptions), feature-level shifts, and output-level shifts.
3. Theoretical Justifications: The paper includes theoretical analysis supporting the empirical findings, providing conditions under which surgical fine-tuning can outperform full fine-tuning in two-layer neural networks.

Empirical Findings

The empirical results are compelling, demonstrating significant performance improvements when applying surgical fine-tuning:

• Image Corruptions: For datasets like CIFAR-10-C, fine-tuning just the initial layers of the neural network yields better accuracy compared to tuning the entire network.
• Layer Specificity: The best-performing layers to fine-tune vary depending on the type of distribution shift:
  • Input-level Shifts: Earlier layers are more effective.
  • Feature-level Shifts: Middle layers provide optimal results.
  • Output-level Shifts: The final layers outperform others.

The research validates these claims systematically across multiple datasets and neural network architectures, reinforcing the robustness of surgical fine-tuning.

Theoretical Contributions

The paper corroborates its empirical results with theoretical analysis. One key theorem demonstrates that in settings with limited target data, fine-tuning just the first layer of a two-layer network can be more advantageous than full fine-tuning, which might lead to overfitting:

• First-Layer Tuning: Handles perturbations in input distributions by realigning the input space efficiently without the risk of overfitting inherent in full fine-tuning.
• Last-Layer Tuning: Addresses perturbations in label distributions.

These theoretical insights are important as they reveal the underlying mechanics of why certain layers are more suitable for fine-tuning under specific types of shifts.

Automated Layer Selection

To mitigate the computational burden of manually selecting which layers to fine-tune, the authors propose two automated criteria:

1. Relative Gradient Norm (Auto-RGN): Assigns tuning priority based on the gradient norm relative to the parameter norm, thus identifying layers that may need more adjustment.
2. Signal-to-Noise Ratio (Auto-SNR): Measures the noisiness of gradients across different inputs, with noisier gradients indicating less reliable information.

In their experiments, Auto-RGN consistently improves performance over full network fine-tuning, aligning it closely with the best manually selected surgical fine-tuning results.

Implications and Future Directions

The research carries several important implications for the field of AI, especially in real-world applications that involve dynamic or evolving data distributions.

• Practical Applications: Surgical fine-tuning can be especially beneficial in scenarios where computational resources are limited, and overfitting is a significant concern.
• Layer-wise Adaptation: Understanding layer-specific roles in adapting to distribution shifts opens new avenues for designing more resilient neural network architectures.

Future research could expand on developing more sophisticated automated techniques for layer selection and explore the relationship between layer functionality and different types of distribution shifts. Exploring these avenues could further optimize the adaptability and efficiency of neural networks in diverse and shifting data environments.