Technical Summary: A Theory of How Pretraining Shapes Inductive Bias in Fine-Tuning

Problem Statement

Modern deep learning relies heavily on the pretraining–fine-tuning (PT+FT) pipeline, where models learn broad features on large general-purpose datasets and subsequently adapt them to specific target tasks with limited data. While practical success is well-documented, the theoretical mechanisms governing how initialization choices influence feature learning, reuse, and refinement across these two stages remain elusive. Specifically, it is unclear how weight initialization shifts the training dynamics between "lazy" (fixed-feature/kernel) and "rich" (feature-learning) regimes during fine-tuning, and how these shifts interact with task statistics (e.g., feature overlap and sparsity) to determine generalization performance.

Methodology

The authors develop an end-to-end analytical theory using diagonal linear networks, a tractable model known to exhibit both lazy and rich learning behaviors depending on initialization scales.

1. Theoretical Framework:

   • Architecture: A one-hidden-layer diagonal linear network $f_{w,v}(x) = \beta(w,v)^T x$, where $\beta$ is parameterized by positive and negative pathways ($w_+, w_-, v_+, v_-$) to allow initialization at $\beta=0$ without saddle points.
   • Training Protocol: The model undergoes gradient flow training in two stages: pretraining (PT) on task $X_{PT}$ and fine-tuning (FT) on task $X_{FT}$. Both stages train to zero training error.
   • Initialization Parameters: The theory focuses on three key parameters:
     • $c_{PT}$: Absolute scale of pretraining weights.
     • $\lambda_{PT}$: Relative scale between the first and second layers during pretraining.
     • $\gamma_{FT}$: Re-initialization scale of the readout weights before fine-tuning.
   • Data Generative Model: A teacher-student setup where ground-truth vectors $\beta^*_{PT}$ and $\beta^*_{FT}$ are sampled from a joint distribution. The model explicitly accounts for sparsity (active dimensions) and overlap (shared active dimensions between PT and FT tasks) using parameters $\rho_{PT}$, $\rho^{\text{shared}}_{FT}$, and $\rho^{\text{new}}_{FT}$.

2. Analytical Tools:

   • Implicit Bias Derivation: The authors derive the exact implicit regularization (inductive bias) induced by the PT+FT process. They show that the gradient flow solution converges to the minimizer of a specific potential function $Q_k(\beta_{FT})$ subject to data constraints.
   • Replica Theory: To characterize generalization error in the high-dimensional limit ($D \to \infty$), the authors employ the replica method. This reduces the high-dimensional estimation problem to a scalar denoising problem governed by fixed-point equations for effective noise and regularization parameters.

Key Contributions

1. Analytical Characterization of PT+FT Bias: The paper derives exact expressions for the generalization error as a function of initialization parameters ($c_{PT}, \lambda_{PT}, \gamma_{FT}$), task parameters, and data scale. This provides the first integrated characterization of implicit bias in a PT+FT setup.
2. Identification of Four Learning Regimes: The theory reveals that initialization choices place the network into four distinct regimes, distinguished by two metrics: $\ell$-order (measuring sparsity bias, ranging from $\ell_1$ to $\ell_2$) and Pretraining Dependence (PD) (measuring reliance on pretrained features).
   • Regime I (Rich, Pretraining-Independent): $\ell$-order $\approx 1$, PD $\approx 0$. The network learns new features with a sparse bias, ignoring pretrained features.
   • Regime II (Lazy, Pretraining-Dependent): $\ell$-order $\approx 2$, PD $\approx -1$. The network reuses pretrained features with minimal adaptation (weighted $\ell_2$ norm).
   • Regime III (Lazy, Pretraining-Independent): $\ell$-order $\approx 2$, PD $\approx 0$. The network behaves like a standard kernel machine, neither reusing features nor benefiting from sparsity.
   • Regime IV (Rich, Pretraining-Dependent): An intermediate regime where the network simultaneously benefits from sparsity ($\ell$-order $< 2$) and feature reuse (PD $< 0$).
3. Universal Trade-off: The authors identify a fundamental trade-off: it is impossible to simultaneously maximize pretraining dependence (PD = -1) and sparsity bias ($\ell$-order = 1) in the limit. However, intermediate initialization scales allow the network to operate in Regime IV, balancing these properties.
4. Role of Relative Scale: The paper highlights that the relative scale of initialization across layers ($\lambda_{PT}$) is a critical lever for controlling the inductive bias. Specifically, a smaller initialization scale in earlier layers (relative to later layers) enables the network to both reuse and refine features.

Results

• Theoretical Predictions:

  • Task Overlap: When PT and FT tasks share dimensions, Regime II (Lazy, Pretraining-Dependent) is optimal. When they share no dimensions, Regime I (Rich, Pretraining-Independent) is optimal.
  • Intermediate Regime: For tasks where FT features are a subset of PT features, Regime IV (Rich, Pretraining-Dependent) yields superior generalization. This regime is accessed by tuning $\lambda_{PT}$ to intermediate values.
  • Imperfect Pretraining: When pretraining is imperfect ($\alpha_{PT} < 1$), the relative scale $\lambda_{PT}$ becomes even more critical. Negative $\lambda_{PT}$ allows the model to refine imperfectly learned structures, improving performance even in full-overlap scenarios where it would otherwise hurt.

• Empirical Validation:

  • Diagonal Networks: Simulations on diagonal linear networks confirm the theoretical generalization curves across varying initialization scales and task overlaps.
  • ResNets on Vision Tasks: Experiments on ResNet-18 trained on CIFAR-100 (pretraining) and fine-tuned on subsets of classes or SVHN demonstrate that:
    • Decreasing the relative scale of early layers (negative $\lambda_{PT}$ analog) improves fine-tuning generalization.
    • Decreasing the absolute scale ($c_{PT}$) degrades performance.
    • Decreasing the readout re-initialization scale ($\gamma_{FT}$) improves performance for intermediate sample sizes.
  • Representation Analysis: Measurements of participation ratio (PR) and effective number of shared dimensions (ENSD) confirm that the successful interventions induce a "pretraining-dependent rich" signature: the representation compresses (lower PR) while maintaining alignment with the pretrained state (high ENSD).
  • Transformers: Experiments on Transformers performing modular arithmetic tasks show that smaller balancedness (analogous to negative relative scale) improves generalization, particularly on tasks involving novel prime factors.

Significance and Claims

The paper claims to provide an end-to-end analytical understanding of how data and network initialization interact to shape fine-tuning generalization. Its primary significance lies in:

1. Moving Beyond Extremes: Prior work often focused on extreme cases (purely lazy or purely rich). This theory demonstrates that intermediate regimes are often optimal and depend critically on the specific task statistics (overlap and sparsity).
2. Actionable Levers: It identifies the relative scale of initialization across layers as a principled, actionable mechanism to control the balance between reusing pretrained features and inducing continued feature learning.
3. Bridging Theory and Practice: By validating theoretical predictions on deep, nonlinear architectures (ResNets, Transformers) and diverse tasks (vision, modular arithmetic), the work suggests that the insights derived from diagonal linear networks are relevant to practical deep learning scenarios.

The authors maintain modesty, noting that their theoretical analysis is based on a simplified model (diagonal linear networks) and that practical networks involve nonlinearities, residual connections, and batch normalization. They acknowledge that future work is needed to carefully study these parameters in more complex settings, but assert that their results highlight a crucial role for initialization scales in enabling continued feature learning during fine-tuning.