Source-linked AI summary
Minimax-optimal Scaling Laws in Ridge Regression via Transfer Learning
Yalda Shabanzadeh, Diyuan Wu, Marco Mondelli
TL;DR
The paper asks when transfer learning improves high-dimensional regression scaling and how the source model should be trained to obtain those gains. It analyzes ridge pretraining followed by target fitting regularized toward the source, and finds a regime where transfer attains the joint minimax exponent. In that regime the optimal source estimator is target-independent, and one-pass SGD fine-tuning achieves the same exponent.
Problem
The paper asks when pretraining improves downstream learning rates and how to train the source model for those gains, since source-task performance need not track downstream performance.
Source: §1 · pp. 1–2Method
Results
Takeaways & Limitations
Under the analyzed shared-regression-function and commuting power-law covariance setting, one source estimator can serve across target distributions while ridge or one-pass SGD fine-tuning attains the improved scaling exponent.
Source: §7 · p. 9The analysis is for regression; whether the improved scaling laws and source-estimator universality persist in classification remains an open direction.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Pretrained models are often reused across downstream tasks, motivating the question of whether this transfer learning procedure can provide a substantial performance improvement, namely, in terms of the scaling law. We address this question in the prototypical setting of high-dimensional regression, where (i) a pretrained \emph{source} model is trained via ridge regression, and (ii) a fine-tuned \emph{target} model is regularized via the $L_2$ distance from the source estimator. First, we establish a deterministic equivalent for the risk of this two-step procedure under both model and covariance shift. Next, by specializing the deterministic equivalent to a shared regression function and commuting power-law covariance spectra, we identify a regime where the resulting scaling law achieves the minimax-optimal exponent. Notably, in this regime, the optimal source regularization depends on source data only, indicating its universality across all targets. Furthermore, we show that the same optimal rate is achieved by performing the fine-tuning step via one-pass stochastic gradient descent. We support our theoretical analysis via experiments on both synthetic settings and language models.
1 INTRODUCTION
The paper studies when ridge-based transfer learning improves high-dimensional regression scaling and how source training should be chosen. It shows that optimally tuned transfer can attain the joint minimax rate in a specified regime, with a source estimator that is universal across targets and fine-tuning methods.
1 INTRODUCTION
- The paper is motivated by evidence that stronger source-task performance does not necessarily improve downstream results, alongside scaling laws that depend on model size and pretraining and fine-tuning data.The introduction notes that longer pretraining can improve source performance while worsening downstream performance, whereas stronger pretraining weight decay can improve downstream performance despite higher pretraining validation loss.Source: §1 · p. 1
- The paper studies a high-dimensional linear-regression setup allowing source and target tasks to differ in coefficients, covariate distributions, noise levels, and sample-size growth rates.It uses source ridge regression followed by target fitting with a squared ℓ2 penalty centered at the source estimator.Source: §1 · p. 2
- When the joint minimax rate is strictly faster than either dataset’s rate alone, optimally tuned transfer achieves that joint minimax rate.The result specializes to a common regressor and power-law covariance spectra.Source: §1 · p. 2
- The optimal source regularization depends only on source information, and the same source estimator remains optimal with one-pass SGD fine-tuning.The analysis therefore supports using one source estimator across various target distributions without target-dependent source tuning.Source: §1 · p. 2
- Experiments evaluate deterministic-equivalent predictions on synthetic data and PCA features from MNIST and CIFAR-10, and test whether similar transfer phenomena arise in language modeling.The experiments provide numerical evidence beyond the theoretical linear-model setting.Source: §1 · p. 2
2 RELATED WORK
Prior work studies empirical adaptation and scaling as well as theoretical transfer guarantees across several settings. This paper focuses on how covariance spectra and relative sample sizes determine scaling benefits in high-dimensional regression.
Empirical works
- Empirical studies examine adaptation methods and how transfer varies across tasks, while scaling analyses relate downstream performance to model size, data sizes, and adaptation procedure.Reported findings also connect distribution alignment and fine-tuning data size to downstream scaling, without consistent downstream gains from increasing pretraining data across metrics.Source: §2 · p. 2
Theoretical works
- Theoretical transfer research spans continual regression, classification, shared representations, and high-dimensional regression with heterogeneous tasks and covariances.This paper complements those lines by characterizing when ridge pretraining followed by ridge or SGD fine-tuning achieves minimax scaling and outperforms source-only and target-only learning.Source: §2 · p. 2
3 PROBLEM SETUP
The setup uses independent source and target training datasets with potentially distinct task parameters, covariances, and noise levels. The target estimator is fit using a ridge penalty centered at the source estimator.
3 PROBLEM SETUP
- The setup uses independent source and target datasets, with sample counts denoted n_sr and n_tr and covariance matrices Σ_sr and Σ_tr.The design distributions are centered and have the stated covariances.Source: §3 · pp. 2–3
- The design vectors are required to satisfy a concentration property, which the paper says covers examples such as independent sub-Gaussian entries.The concentration assumption is stated for both source and target designs.Source: §3 · p. 3
- Source and target tasks may differ in regression parameters, covariances, and noise levels, and target risk averages over label noise.The target model is obtained by fitting on target data with a ridge penalty centered at the source estimator.Source: §3 · p. 3
- The paper uses matrix notation for finite-dimensional problems and states that its arguments also apply to trace-class covariance operators on a separable Hilbert space.In that formulation, transposes are interpreted as adjoints.Source: §3 · p. 3
4 DETERMINISTIC EQUIVALENT FOR TARGET RISK
The section constructs a deterministic equivalent for the target risk using effective source and target regularization parameters. It also reports an additive approximation-error bound under the stated conditions.
4 DETERMINISTIC EQUIVALENT FOR TARGET RISK
- Effective regularization parameters τ_sr and τ_tr are defined as the unique solution of a system, alongside two scalar prefactors.The supplied passage identifies the quantities but does not specify their defining system or the prefactor formulas.Source: §4 · p. 3
- Theorem 1 establishes a deterministic equivalent for target risk under Assumption 1 and the stated boundedness condition.The theorem statement gives a probability guarantee and a risk expression in terms of the defined quantities.Source: §4 · p. 4
- The approximation error is additive and of order O(max(ε_sr, ε_tr)), reflecting asymmetric terms in the risk expression.The authors note that multiplicative error may be obtainable in structured settings such as the power-law model.Source: §4 · p. 4
5 SCALING LAWS FOR TRANSFER LEARNING
This section derives transfer-learning scaling laws under power-law covariance spectra and compares them with minimax rates using one or both datasets. It shows when ridge and one-pass SGD can attain the joint-data minimax exponent.
5.1 Power-law model
- With commuting power-law covariances and shared regression weights, the worst-case signal’s bandwise energy yields the same polynomial test-error exponent as the bias.The analysis assumes the signal lies outside the ellipsoid but reflects its worst case.Source: §5.1 · pp. 4–5
5.2 Minimax rates
- When δ > 1 and γc < γn < 1, using both datasets yields a strictly faster minimax rate than using either dataset alone.The joint-data result and a ridge estimator achieving it require simultaneous access to source and target data.Source: §5.2 · p. 5
5.3 Scaling laws for ridge estimator
- In the regime δ > 1 and γc < γn < 1, optimally tuned transfer ridge matches the joint-data minimax rate in its deterministic-equivalent risk.The authors expect the approximation to carry over to actual test risk with high probability and report numerical support.Source: §5.3 · p. 6
- The optimal source regularization depends only on source sample size and spectral decay, so one source estimator can be reused across targets without target-dependent source tuning.The result establishes existence, not uniqueness; other optimal ridge pairs may use target-dependent source regularization.Source: §5.3 · p. 6
5.4 Scaling laws for one-pass SGD
- Under δ > 1 and γc < γn < 1, one-pass SGD matches the exponent of optimally regularized ridge, which is minimax optimal, using the same rate-optimal source regularization.The variance dominates the test error in this regime, while signal-dependent bias is lower order, so ridge saturation does not alter the leading exponent.Source: §5.4 · p. 7
6 NUMERICAL RESULTS
The experiments compare transfer learning with target-only and source-only training across synthetic regression and language-model settings. Transfer improves the synthetic scaling law over either single dataset, while language-model gains vary with source-data size and target optimization.
Experiments on synthetic data
- In the synthetic setting, optimally tuned ridge transfer improves the target-risk scaling law over optimally tuned target-only and source-only estimators.The experiments also show that deterministic-equivalent predictions track empirical results and recover the corresponding theoretical decay rates.Source: §6 · p. 7, Figure 1 · p. 8
- For γc < γn < 1, transfer improves the risk decay rate over either dataset alone; the figure compares its exponent with the joint minimax rate.The gray band denotes the best achievable worst-case exponent using both datasets.Source: §6 · p. 7, Figure 2 · p. 8
Experiments on language models
- The language-model experiments vary source-data size and target sample size to examine whether fine-tuning changes the target-loss scaling rate.They use GPT-2 small pretrained on FineWeb-Edu and fine-tuned on Python.Source: §6 · p. 8
- Transfer fine-tuning shows a higher rate than target-only and source-only training, but its advantage depends on target optimization.With fixed source data, target-only training gradually closes the gap; with growing source data, the early rate advantage is smaller after additional target optimization.Source: §6 · p. 8
7 CONCLUSIONS
The paper analyzes ridge-regression transfer learning in high-dimensional regression and derives a deterministic equivalent for its risk. Under shared regression functions and power-law covariance spectra, transfer can attain the joint minimax exponent, with a source estimator that is universal across target distributions.
Conclusions
- The paper establishes a deterministic equivalent for the risk of ridge transfer under general conditions.The target estimator is regularized toward the source estimator.Source: §7 · p. 8
- With shared regression functions and power-law covariance spectra, transfer attains the joint minimax exponent whenever combining datasets is strictly better than either dataset alone.One-pass SGD fine-tuning achieves the same exponent, using an optimal source estimator that does not depend on the target task.Source: §7 · p. 9
Future directions
- The analysis leaves classification and settings with model size explicitly varying through random features as future directions.These extensions would examine whether the scaling gains and source-estimator universality persist, and how model capacity affects transfer.Source: §7 · p. 9
AI use statement
The authors describe their use of generative AI tools and retain responsibility for the final work.
AI use statement
- Generative AI tools were used to polish the writing.The authors also report using them to improve related-work coverage.Source: AI use statement · p. 10
- The tools assisted with straightforward code implementation and proof-related work.This included filling in proof details after key methods and arguments were identified and checking proofs for correctness.Source: AI use statement · p. 10
- The authors state that AI tools did not develop the core ideas or define the problem setting, and they take responsibility for the final content and artifacts.Source: AI use statement · p. 10
A PROOF OF THEOREM 1
The proof of Theorem 1 decomposes conditional excess risk into bias and two variance terms, then controls each using deterministic counterparts.
Proof setup
- The proof defines effective regularization parameters, matrices, and prefactors for both source and target stages.It also introduces source and target resolvents.Source: A PROOF OF THEOREM 1 · p. 13
Proof strategy
- The proof first decomposes risk exactly into bias and two variance terms, then controls each term by its deterministic counterpart.The decomposition is for conditional excess risk.Source: A PROOF OF THEOREM 1 · pp. 13, 16
Risk decomposition
- The risk is written as R = B + Vsr + Vtr, separating bias from source and target variance contributions.Source: A PROOF OF THEOREM 1 · p. 16
B PROOF OF THEOREM 3
The proof derives Theorem 3’s scaling law by determining effective regularization scales, then evaluating variance and bias contributions before combining them into total risk.
B PROOF OF THEOREM 3
- The proof proceeds from effective source and target regularization scales to spectral quantities in the bias and variance, then combines them to obtain total risk scaling.It uses power-law sum estimates for the final step.Source: B PROOF OF THEOREM 3 · pp. 17–19
- For power-law eigenvalues, the proof specializes generic effective-ridge scaling to the source and target problems and defines corresponding effective cutoffs.At either threshold, both the effective regularization parameter and its gap from sample size are of order the sample size.Source: B PROOF OF THEOREM 3 · pp. 17–19
- Fine-tuning variance depends only on the target spectrum, whereas pretraining variance and the second bias term depend on the source-target cross-trace.The proof therefore estimates that cross-trace separately before deriving the variance and bias scalings.Source: B PROOF OF THEOREM 3 · pp. 19–20
B.1 Technical lemmas
- The proof sets K = λ−1/α; for k ≤ K the denominator is ≍1 and that range contributes ≍P, while η < αγ ensures summability.Source: §B.1 · p. 23
C PROOF OF COROLLARY 1
This section proves Corollary 1 by establishing upper and matching lower bounds for the two-stage ridge scaling exponent, including admissible cutoffs, a maximal-source-scaling regime, and an SGD upper-bound argument. It also tests deterministic-equivalent risk predictions on image-derived features and Gaussian controls, examining how closely they track risk-curve shapes.
C PROOF OF COROLLARY 1
- The optimal exponent is determined by the maximum among the target-only, source-only, and transfer candidates, and the proposed cutoffs attain it.When the target-only candidate is maximal, target data alone is optimal; when the source-only candidate is maximal, source data alone is optimal.Source: C PROOF OF COROLLARY 1 · pp. 23, 26
C.1 Proof of Corollary 1
- For γ_n > γ_c = (1 − δ)/(1 + 2s), the source candidate exceeds the other exponents, and choosing u = γ_n, v = 1 − γ_n preserves the optimum.The proof shows the optimal cutoffs can be replaced by u = γ_n without changing the resulting exponent.Source: §C.1 · p. 27
D PROOF OF LEMMA 1
- The Lemma 1 proof conditions on the source estimator, applies the last-iterate SGD upper bound, evaluates deterministic equivalents, and derives a matching lower bound.The upper bound is from (Wu et al., 2022, Corollary 3.4).Source: D PROOF OF LEMMA 1 · p. 27
D.1 Upper bound
- Conditioning on the pretrained estimator makes its error independent of target data; the SGD upper bound then reduces to two terms depending on the source ridge estimator.The proof averages over source-label noise and computes deterministic equivalents for the resulting terms.Source: §D.1 · pp. 28–29
- The upper-bound proof controls the risk on a common source-design event and then takes expectations to complete the upper bound.The event is chosen so that the source-design bounds hold with complement probability at most CDm^-D.Source: §D.1 · p. 32
D.2 Lower bound
- The lower-bound proof isolates target-label noise and combines its bound with the upper bound to prove Lemma 1.The argument conditions on the source sample and target covariates, where the initialization error is fixed and the target noises are centered and independent.Source: §D.2 · p. 33
E.1 Deterministic equivalent predictions for empirical spectra
- On MNIST and CIFAR-10 PCA features, deterministic-equivalent predictions capture the risk curve’s main shape, including its peak near p = n and its shift with n.The Gaussian control follows the prediction more closely; image-feature deviations are larger, especially under weaker regularization.Source: §E.1 · p. 33, Figure 4 · p. 34
E.2 Details for experiments on language models
The language-model experiments use GPT-2 small with FineWeb-Edu as the source and Python as the target, comparing target-only training with source-initialized methods across nested target budgets.
Model and data
- GPT-2 small (123.6M parameters) is pretrained on FineWeb-Edu and evaluated on Python, with target validation and test data kept outside the source pool.Python training and held-out data are split by repository, and near-duplicate training files are removed.Source: §E.2 · p. 34
Setup
- Target budgets are nested within each seed, and γ = 0 keeps the source budget fixed while γ > 0 makes it grow with the target budget.Methods see the same target examples in the same order within each seed.Source: §E.2 · p. 34
Training
- Training uses AdamW with batches of 65,536 tokens, one source pass, and validation-selected target epochs, with the one-epoch branch reporting the one-pass result.Hyperparameters are tuned on seed 0 separately by target and method; experiments use seeds 0, 4, and 5 for paired comparisons.Source: §E.2 · pp. 34–35
Training procedures
- Target-only training starts from random initialization, while fine-tuning starts from the source checkpoint and FTSP pulls toward source parameters after each AdamW step.FTSP uses a decoupled update; the original L2-SP instead adds a penalty to the loss.Source: §E.2 · p. 35
Evaluation
- Test loss is average cross-entropy per token, and scaling is summarized from observed test loss because the irreducible loss floor is unknown.For windows containing more than two target budgets, one exponent is fit by least squares; ratios and differences are averaged across seeds with standard errors.Source: §E.2 · pp. 35–36