PaperlayerView exact revision on arXiv

Source-linked AI summary

Two-Timescale Fine-tuning Provably Learns New Features for Two-Layer ReLU Networks

Etienne Boursier, Nicolas Flammarion

arXiv:2609.34667v1stat.MLcs.LG

TL;DR

Theoretical understanding of fine-tuning from structured pre-trained weights remains limited, particularly for learning features absent from the pre-trained representation. The paper analyzes two-timescale fine-tuning in a Gaussian multi-index setting and shows recovery of a new feature from O(d) samples, independent of the number of pre-trained features.

  • Problem

    Theoretical analyses have largely focused on reusing pre-trained features, leaving limited understanding of how fine-tuning learns new task-specific features while preserving useful existing ones.

    Source: §1 · p. 1

  • Method

    The paper analyzes two-timescale training of a two-layer ReLU network in a setting with m pre-trained teacher features and one feature to learn during fine-tuning.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    O(d) fine-tuning samples recover a new feature independently of the number of pre-trained features, whereas learning from scratch requires sample complexity that grows with the number of features.

    Source: §1 · p. 2, §3.2 · p. 8

  • Takeaways & Limitations

    In the analyzed setting, pre-training can induce an implicit bias that lets fine-tuning learn an absent feature while preserving pre-trained features.

    Source: §1 · pp. 1–2

    The guarantee requires d ≳ m^2 so that the new neuron is initialized in the target weight's attraction basin; rigorously extending the argument to multiple new neurons remains open.

    Source: §4 · p. 8

Abstract

from arXiv · show

Fine-tuning pre-trained models on specialized tasks with scarce data is central to modern deep learning. Despite its empirical success, theoretical understanding of fine-tuning remains limited. We introduce a Gaussian multi-index setting to study fine-tuning from pre-trained weights, where the teacher network has $m+1$ features, $m$ of which are learned during pre-training and one of which must be learned during fine-tuning. For two-layer ReLU networks, we show that two-timescale training, i.e., updating the outer weights infinitely faster than the hidden ones, learns the new task-specific feature while preserving the pre-trained ones in the model representation. Moreover, only $\mathcal{O}(d)$ fine-tuning samples are required for this recovery, independently of the number of pre-trained features. In contrast, with random initialization, the same number of samples is insufficient to recover the target parameters. Our results therefore demonstrate that pre-training can induce an implicit bias with a clear statistical advantage over random initialization, enabling feature learning from scarce fine-tuning data.

1 Introduction

The paper asks whether pre-training changes what limited fine-tuning data can teach, beyond making optimization easier. It analyzes how two-timescale fine-tuning can learn an absent feature while preserving pre-trained features, and compares this with learning from scratch.

1 Introduction

  • The paper addresses how fine-tuning can learn a task-specific feature absent from the pre-trained representation while retaining useful pre-trained features.Earlier analyses emphasized preservation and reuse, whereas empirical evidence also indicates that fine-tuning can modify representations to learn new features.Source: §1 · p. 1
  • Two-timescale fine-tuning recovers a new feature from O(d) samples, independent of the number of pre-trained features, while learning from scratch requires complexity that grows with feature count.The analysis characterizes this as a statistical advantage over single-task learning.Source: §1 · p. 2

1.1 Related work

  • Prior theoretical analyses largely reduce fine-tuning to adapting pre-trained representations and do not capture rich feature learning; this work studies that learning in two-layer ReLU networks.The setting requires reusing pre-trained features and learning a genuinely new one.Source: §1.1 · p. 2

1.2 Outline.

  • The paper studies limiting two-timescale dynamics, proves recovery guarantees, discusses their relation to prior work, and illustrates the theory with a toy numerical experiment.The analysis covers population-level recovery and recovery from n ≳ d samples.Source: §1.2 · p. 3

2 Setting and limit dynamics

The setting models fine-tuning when a teacher has m + 1 features, m already present in pre-training and one absent. In the vanishing-initialization limit, two-timescale dynamics let the absent feature evolve while pre-trained features remain fixed.

2 Setting and limit dynamics

  • The teacher has m + 1 features, with m recovered during pre-training and the remaining feature absent and required to be learned during fine-tuning.The absent hidden weight is initialized at small scale α in a random unit direction, idealizing finite-time pre-training on the first m features.Source: §2 · p. 3

2.1 Two-timescale dynamics

  • As the output-to-hidden learning-rate ratio tends to infinity, output weights are fitted instantaneously while hidden weights evolve more slowly.The analysis models gradient-based optimization with gradient flow, the continuous-time behavior in the small-step-size limit; uniqueness of the output-weight minimizer is ensured along the trajectory under the stated conditioning assumptions.Source: §2.1 · pp. 3–4

2.2 Separated neuron dynamics

  • ReLU homogeneity keeps each neuron's norm constant, so the dynamics can be tracked through normalized neurons.The normalized-neuron evolution rate scales as the inverse square of its initialization norm.Source: §2.2 · p. 4
  • Because the new neuron's initialization scale is α, it evolves infinitely faster as α approaches zero; in the limit, it alone evolves while pre-trained neurons remain fixed.The rescaled trajectories converge on compact time intervals toward solutions of the limiting differential inclusion, provided they remain in a set satisfying the conditioning assumption.Source: §2.2 · pp. 4–5

3 Two-timescale small initialization dynamics

In the small-initialization limit, two-timescale dynamics keep pre-trained feature directions fixed while the new feature evolves; its output weights continue adapting. Population dynamics explain how the new feature converges, and the empirical analysis establishes recovery with O(d) samples, independent of m.

3 Two-timescale small initialization dynamics

  • In the α → 0 limit, pre-trained feature directions remain fixed while only the new feature evolves; the pre-trained output weights still adapt.Those output weights initially compensate for the missing feature and reach their optimal value only at convergence.Source: §3 · p. 5
  • Pre-trained output weights transiently compensate for the missing feature, so the dynamics are not equivalent to freezing the pre-trained neurons and training only the new one.This transient also explains why Linear Probe then Fine-Tune does not achieve zero test loss in this setting.Source: §3 · p. 5

3.1 Warm up: gradient flow on population loss

  • The new feature’s component along its target grows faster than its component along previously learned directions, keeping the latter small when initialization is suitable.This mechanism supports continued application of the proof conditions along the trajectory.Source: §3.1 · p. 6
  • Under orthogonal teacher features, population gradient flow converges linearly to the optimal new feature after an initialization delay of order m.Initially, its update is only of order 1/m, so it takes order m time to leave the slow-evolution phase.Source: §3.1 · pp. 6–7
  • The convergence guarantee requires suitable initialization because the dynamics have spurious local minima and are not globally attracted to the target.The stated dimensional requirement is d ≳ m^2; the initialization must lie outside the spurious minima’s attraction basins.Source: §3.1 · p. 7

3.2 Gradient flow on empirical loss

  • Empirical-loss analysis transfers the population-dynamics argument using concentration of empirical quantities around their population counterparts.The proof also requires tighter concentration near the new feature’s initialization.Source: §3.2 · p. 7
  • O(d) fine-tuning samples suffice to recover the new parameter and achieve zero test loss, independently of the m pre-trained features.The theorem’s stated conditions include d ≥ C m^2 ln(1/δ) and n ≥ C(d + ln(1/δ)).Source: §3.2 · pp. 7–8

4 Discussion and limitations

The discussion contrasts fine-tuning strategies that preserve pre-trained features with those that adapt them, and identifies assumptions and dimensional conditions that bound the analysis. The idealized initialization is motivated by pre-training and complemented by experiments using an actually pre-trained model.

Comparison with other fine-tuning methods

  • Full fine-tuning can distort pre-trained features, linear probing cannot learn tasks needing new features, and two-timescale fine-tuning offers an intermediate strategy.LP-FT delays distortion but does not prevent it; L2-SP preserves the representation but may hinder needed adaptation of pre-trained weights.Source: §4 · p. 8
  • Two-timescale fine-tuning uses smaller inner-layer than outer-layer learning rates, allowing new features to be learned while the pre-trained representation remains essentially unchanged.In the vanishing-void-feature limit, the small feature magnitude is compensated by large outer weights, amplifying its gradients so it evolves faster than pre-trained features.Source: §4 · p. 8

Requirement d ≳ m^2

  • The guarantee requires d ≳ m^2 so a single randomly initialized new neuron falls within the target feature's attraction basin with high probability.Lower dimensions lack this high-probability guarantee for one neuron; training multiple new neurons might help, but rigorous analysis is more difficult.Source: §4 · p. 8

Pre-trained initialization of the weights

  • The idealized initialization places the first m neurons at optimal pre-trained weights and gives the extra neuron a small norm and random direction.The small norm is motivated by suppression of unused features, including through weight decay; experiments from an actually pre-trained model show consistent dynamics.Source: §4 · p. 9

Orthogonality of features

  • Orthogonality is an analytical convenience that can be relaxed to allow unequal feature norms and nonzero correlations, subject to a weaker angular condition.That condition ensures the target feature's signal dominates those of pre-trained features along the relevant dynamics.Source: §4 · p. 9

5 Experiments

Experiments compare fine-tuning methods on a Gaussian-feature teacher setup. Two-timescale fine-tuning and L2-SP generalize best while preserving pre-trained features, with a potential cost in training time; the paper also reports its use of generative AI tools.

5 Experiments

  • The experiments use Gaussian teacher features and compare fine-tuning algorithms with m = 20, d = 500, and n = 4000.The first m network weights are initialized at the corresponding optimal weights, while the additional weight has a random direction and small norm.Source: §5 · p. 9
  • LP-FT initially performs well but later deteriorates, while full fine-tuning distorts pre-trained features from the start; linear probing is limited by its expressivity.L2-SP performs well when only the hidden layer is regularized, whereas regularizing all parameters substantially degrades performance.Source: §5 · p. 9
  • Two-timescale fine-tuning and L2-SP achieve the best generalization while preserving pre-trained features, but may need more epochs; nearly all methods eventually reach near-zero training loss.The Figure 1a comparison shows that the test-performance differences are attributed to the methods' different implicit regularization, rather than their ability to minimize the training objective.Source: §5 · p. 9, Figure 1 · p. 10
  • Full fine-tuning and LP-FT cause comparable feature distortion, whereas two-timescale fine-tuning causes almost none and L2-SP also keeps distortion low.The two-timescale result is consistent with the limiting dynamics applying at finite initialization scales and learning-rate ratios.Source: §5 · p. 10

AI use statement

  • The authors report using generative AI to polish writing, code, and literature references and to explore mathematical arguments; human authors developed and verified all proofs.Source: AI use statement · p. 10

A.1 Experimental details

The experimental appendix describes Gaussian teacher and initialization choices, training settings, and a hidden-layer-only L2-SP implementation. Random initialization reaches nearly zero training loss but has much larger test loss.

A.1 Experimental details

  • The experimental setup draws teacher features and outer weights independently from standard Gaussians, initializes the first m weights at ground truth, and gives the extra feature covariance 10^-6 I_d.The extra output weight also has variance 10^-6, while the first m output weights are initialized as unit-variance Gaussians.Source: §A.1 · p. 13
  • Training uses SGD with batches of 200; two-timescale learning rates are 10^-2 for the output layer and 10^-5 for the hidden layer.L2-SP regularizes only hidden parameters, a choice reported to substantially improve its performance.Source: §A.1 · p. 13
  • With random initialization, the model reaches nearly zero training loss but much larger test loss, illustrating a statistical rather than purely optimization-based benefit of pre-training.When new features start near zero, hidden-only L2-SP can move them in essentially any direction at little regularization cost while scaling their output weights.Source: §A.1 · p. 14

A.2 Cosine similarities of learned features

The cosine-similarity analysis shows how fine-tuning methods acquire the new feature and affect pre-trained features. Similar findings hold with an overparameterized student.

A.2 Cosine similarities of learned features

  • Two-timescale fine-tuning and L2-SP perfectly learn the new feature, while some other algorithms may not.The comparison is shown after 5000 training epochs.Source: Figure 2 · p. 14, §A.2 · p. 14
  • LP-FT learns the new feature only at a late stage; at epoch 160, when its test loss is minimal, it has not learned it correctly.The accompanying discussion attributes later performance degradation to distortion of pre-trained features, even when cosine similarity is approximately 99.7%.Source: §A.2 · pp. 14–15

A.3 Overparameterized student

  • The empirical observations are very similar with a 100-neuron overparameterized student and the same 4000 fine-tuning samples.The first 20 neurons are initialized from ground-truth weights; the other 80 are initialized as small-variance Gaussian weights.Source: §A.3 · p. 15

A.4 Fine-tuning from real pre-training

Experiments from actual pre-trained weights broadly match the paper’s main insights, though good performance requires more samples and training epochs than in the idealized setting. A proof also shows that the hidden-weight norms remain constant over time.

A.4 Fine-tuning from real pre-training

  • The fine-tuning task uses 8000 samples and adds one new feature to 10 randomly selected features from a 20-feature pre-training teacher.The setup uses dimension d = 500; the authors describe it as more challenging because it uses actual pre-trained weights, an overparameterized model, and some pre-trained features irrelevant to the task.Source: §A.4 · pp. 15–16
  • Two-timescale fine-tuning and L2-SP recover the new feature with little pre-trained-feature distortion; full fine-tuning and LP-FT suffer distortion.Two-timescale appears to perform better than L2-SP in this setting, while LP-FT’s distortion is less pronounced than full fine-tuning’s.Source: Figure 4 · p. 16, §A.4 · p. 16
  • Good performance from actual pre-trained weights requires a larger sample size and more epochs than in the idealized setting.The authors state that the theoretical setting does not directly translate to this practical setting.Source: §A.4 · p. 16
  • The experiments support the idealized initialization as a model of practical fine-tuning for the pre-training setups considered.Source: §A.4 · p. 16

B.1 Proof of Lemma 1

  • The proof establishes that each hidden-weight norm remains constant over time.It does so by computing the time derivative of the squared norm.Source: §B.1 · p. 17

B.2 Proof of Proposition 1

The proof establishes a limit trajectory for the rescaled dynamics and shows that it solves the limiting differential inclusion. Its population-dynamics analysis also characterizes the outer-weight minimizer and the loss gradient under the stated assumptions.

B.2 Proof of Proposition 1

  • Uniform boundedness and Lipschitz continuity yield a uniformly convergent subsequence, whose limit solves the differential inclusion in Equation (4).The proof uses Arzelà–Ascoli for subsequence convergence, then establishes that the limit satisfies the inclusion.Source: §B.2 · pp. 17–19
  • Under Assumption 1, the outer-weight minimizer is unique, and the map W ↦ a_n(W) is Lipschitz on T.Invertibility of the feature Gram matrix gives the unique minimizer; boundedness and Lipschitz continuity on T yield the Lipschitz property.Source: §B.2 · p. 17

C.1 Notations and preliminaries

  • The preliminaries define ϕ_v(x) = σ(vᵀx), the Gram matrix of pretrained features, and a projection orthogonal to their span.The feature functions are considered in L2(N(0, I_d)); the projection is defined when the Gram matrix is invertible.Source: §C.1 · p. 19

C.2 Population dynamics

  • When D(w_{m+1}) > 0, the population outer-weight minimizer is unique and its new-feature coefficient is H(w_{m+1}).The result is stated for w_{m+1} on the unit sphere under Assumption 2.Source: §C.2 · p. 19
  • The population-loss subdifferential is expressed using Ψ(w) = 2∇H(w) − H(w)D(w)∇D(w).The derivation relates the gradients on the sphere to the Riemannian gradients of the corresponding extensions.Source: §C.2 · pp. 20–21
  • For D̃(w_{m+1}) > 0, the population loss is minimized at a unique value of a_{m+1}.The excerpt identifies the condition for uniqueness but does not give the minimizer's full formula here.Source: §C.2 · p. 20

C.3 Key inequalities

This section develops inequalities for the trajectory proof, including a population lower bound and empirical-loss results with concentration of empirical quantities toward their population versions. Because uniform concentration does not provide the desired bound near $w^\star$, it uses different bounding techniques there and also presents concentration bounds developed by Soltanolkotabi (2017a).

C.3 Key inequalities

  • For sufficiently large m, Lemma 4 gives H(w) ≥ c m for w ∈ T(c).Source: §C.3 · p. 22

C.4 Proof of Theorem 1

  • The proof shows that the fine-tuned feature trajectory wm+1(t) remains in T(c) for all t ≥ 0.Source: §C.4 · p. 27

D.2 Empirical loss: expression and gradient

  • Under Assumption 3, the empirical outer-weight optimizer is unique when Dn(wm+1) > 0, and its new-feature coefficient equals Hn(wm+1).Source: §D.2 · pp. 28–29

D.3 Concentration of key quantities

  • Proposition 2 establishes high-probability concentration of key empirical quantities toward their population versions over T(c), under its sample-size condition.Source: §D.3 · p. 29
  • Near w⋆, uniform concentration bounds fail to provide the desired bound, so the analysis instead uses different bounding techniques.The passage says a quantity involving m+1 vanishes, but does not identify it fully in the supplied excerpt.Source: §D.3 · p. 30
  • The analysis also uses concentration bounds developed by Soltanolkotabi (2017a) for a quantity involving m+1.Source: §D.3 · p. 30

D.3.1 Uniform concentration bounds

This section develops concentration controls for relevant directions, including unit vectors satisfying an alignment condition with w⋆, and uses them to establish local convergence near the target direction.

D.3.1 Uniform concentration bounds

  • Lemma 7 gives simultaneous uniform concentration of intermediate quantities, and Lemma 8 establishes invertibility of F_n^⊤F_n with a uniform bound for w.Both lemmas assume Assumption 2; Lemma 8 also assumes the items of Lemma 7 hold.Source: §D.3.1 · pp. 30–31

2. The proof of this item follows the exact same lines, since

  • The empirical Jacobian-projection quantity approximates one-half of ∇eD(w) within Cε, first at differentiable points and then for all w.The extension to all w follows by convexity and continuity.Source: §2 · p. 33

D.3.2 Concentration near w⋆

  • For n ≥ C(d + ln(1/δ)), Lemma 9 gives simultaneous high-probability bounds, including λ_min(M) > 0.09.The stated sample-size condition is n ≥ C(d + ln(1/δ)); the eigenvalue bound is a positive universal constant.Source: §D.3.2 · pp. 33–34
  • The cited passage states that the first inequality of Lemma 10 is also a consequence of Equation (38), when taking η to a specified value.Source: §D.3.2 · p. 36
  • Near w⋆, concentration bounds control both terms in a decomposition, with the first dominating the second to ensure local convergence.The decomposition is useful because Corollary 4 controls both terms near w⋆.Source: §D.3.2 · p. 38
  • For sufficiently small universal ε and η0, the text states a conclusion for unit vectors w satisfying a condition involving w⋆, but the conclusion is omitted.The passage ends before specifying the alignment condition or the resulting bound.Source: §D.3.2 · p. 39
  • The proof invokes Equation (45) and then uses the decomposition in Equation (44) to reach its conclusion.The excerpt does not include either equation or specify the conclusion.Source: §D.3.2 · p. 39

D.3.3 Proof of Proposition 2

The proof of Proposition 2 develops uniform concentration bounds over a restricted region, while a separate initialization argument uses pointwise concentration. The latter supports a lower bound on the empirical quantity H_n at the random initialization.

D.3.3 Proof of Proposition 2

  • Proposition 2 states that for ε∈(0,1), under Assumption 2 with m≥m0 and n≥C ε^-2 ln(1/ε)^2(d+ln(1/δ)), Assumption 3 holds with probability at least 1−δ, with the statement ranging over every w∈T(c).Source: §D.3.3 · p. 39

D.3.4 Concentration at initialization

  • Pointwise concentration at initialization yields tighter bounds than uniform concentration over w and v, which requires a larger sample complexity.The proof uses concentration at a random initialization chosen independently of the training data.Source: §D.3.4 · p. 40
  • Corollary 5 states that H_n(w˘) has a positive lower bound with high probability for an initialization drawn independently of the training data.The excerpt omits the full sample-size condition and the exact lower-bound expression.Source: §D.3.4 · p. 41

D.4 Proof of Theorem 2

Theorem 2’s proof confines the new-feature trajectory to T(c) and concludes its first-part claim. An auxiliary lemma extends finite-time trajectory containment to augmented sets under Assumption 1.

D.4 Proof of Theorem 2

  • Under the stated sample-size regime, the proof combines high-probability events from Corollary 5 and Proposition 2, then follows the population-case argument.One stated event is Hn(w◦m+1(0)) ≥ c m.Source: §D.4 · p. 42
  • The proof concludes the first part of Theorem 2 after establishing the trajectory remains in T(c) for all t ≥ 0.The proof also bounds the reparametrized trajectory’s projected norm under suitable choices of m0 and ε.Source: §C.4 · p. 27, §D.4 · pp. 42–43
  • Invertibility and positivity imply full rank throughout T(c), allowing compactness and continuity to establish Assumption 1 for T.The proof uses Lemma 8 for the invertibility and positivity conditions.Source: §D.4 · p. 43

E Additional lemmas

  • If solutions of Equation (4) stay in a set satisfying Assumption 1 on [0,T], sufficiently large k keeps the rescaled trajectory inside an augmented set also satisfying Assumption 1.The lemma’s contradiction argument establishes lim infk τk > T.Source: E Additional lemmas · p. 43

E.1 Useful concentration bounds

This section develops concentration bounds for function classes under Gaussian inputs. The bounds give sample-size conditions for uniform concentration, with a simpler requirement when the set class is a singleton.

E.1 Useful concentration bounds

  • The concentration proof uses chaining and VC-subgraph arguments to control empirical processes, then a covering net to extend bounds over vector choices.The cited lemma bounds a class formed by products of h1 and h2 restricted to sets in A.Source: §E.1 · pp. 44–45
  • n ≥ C ε^-2 ln(1/ε)^2 (VC(A) + p1 + p2 + ln(1/δ)) suffices for the stated uniform concentration bound with probability at least 1 − δ.The bound assumes a VC class A and sub-Gaussian h1 and h2 under X ∼ N(0, I_d).Source: §E.1 · pp. 44–45
  • When A = {R^d}, the covering argument requires only n ≥ C ε^-2 (p1 + p2 + ln(1/δ)).In this case the function class reduces to a single function, so the proof uses Bernstein’s inequality before covering.Source: §E.1 · p. 46
  • The authors note that the ln(1/ε)^2 factor in Lemma 15’s sample complexity could likely be improved or removed, but is sufficient when ε is constant-order.They attribute this possible tightening to the available empirical-process theorem.Source: §E.1 · p. 46
Loading 2609.34667v1…