When is Fine-Tuning Scalable? Certified Recovery from Noisy Data and Linear-Cost Fine-Tuning on Structured Task Shifts

Anastasis Kratsios; Tin Sum Cheng; Bum Jun Kim; Daniel Roy

stat.MLAug 31, 2025

When is Fine-Tuning Scalable? Certified Recovery from Noisy Data and Linear-Cost Fine-Tuning on Structured Task Shifts

Authors: Anastasis Kratsios, Tin Sum Cheng, Bum Jun Kim, Daniel Roy

Organizations: Department of Mathematics, McMaster University, Canada · Department of Mathematics and Computer Science, University of Basel, Switzerland · Department of Statistical Sciences, University of Toronto, Canada · Vector Institute, Canada

Abstract

Transformers often transfer across related tasks with little fine-tuning; we study this question in noisy nonparametric regression for continuous maps $f,g:[0,1]^d\to\mathbb{R}^D$ , representing the source and shifted tasks, using a single-query, single-head transformer $\mathcal{T}(\cdot)=V\operatorname{Softmax}(-α(\cdot,K))$ with $\ell^{\infty}$ -based alignment scores $α$ . Rather than assume that gradient-based training reaches an empirical risk minimizer, we give a model-specific two-stage procedure certifying source-task recovery and out-of-distribution (OOD) adaptation. With high probability, the resulting attention head recovers $f$ to arbitrary accuracy $\varepsilon>0$ in every fixed $L^p([0,1]^d)$ norm, $1\le p<\infty$ , from noisy observations. For Lipschitz targets, its rate is minimax-optimal up to logarithmic factors without assuming oracle access to an asymptotic ERM. If $D=1$ and $f$ is Pfaffian than the values matrix factors into $AC$ with $A$ scaling at-most linearly $\mathcal{O}(\varepsilon^{-1})$ in the reciprocal approximation error $\varepsilon>0$ . Importantly, our training procedure encodes all our identified (potentially) shared geometric structure from $f$ only into the values matrix factor $C$ and the keys matrix $K$ allowing for rapid access, and all $f$ -specific structure only into $A$ . Our main result concerns the OOD phase; where we identify geometric structures, which if shared by the source $f$ and OOD $g$ task; then only only the values matrix factor $A$ needs to be fine-tuned to guarantee an $L^p$ -recovery of $g$ . Critically, fine-tuning only $A$ is a convex problem which involving an at-most linear number of parameters $\mathcal{O}(\varepsilon^{-1})$ ; dramatically cutting down on the worst-case scaling of $\mathcal{O}(\varepsilon^{-\mathbf{d}})$ unavoidable in the source-training phase.

Figures & tables

Figure 1: Shared symmetries across tasks. Each panel depicts a different task on $[0,1]^{2}$ at a scale of $\tfrac{1}{q}=\tfrac{1}{8}$ . In each panel, the same groups of sub-squares share a common greyscale value, while those greyscale values change across panels. This illustrates the full-compatibility case of the fixed-grid class-value adaptation mechanism: the class assignment is shared across tasks, while only the corresponding class values change; here $\kappa=4$ . The general oracle theorem also allows a supplied compatible family of cells and controls the number of target-dependent class-value entries.

Figure 2: Examples of symmetries in Pfaffian functions. The left panel shows a familiar sinusoidal pattern, for which clustering by output value occurs as expected. The right panel illustrates the quantitative conclusion of Theorems 5 and 6 : under their stated non-degeneracy conditions (e.g.,

\|\nabla f(x)\|>0

for each

x

) and the Pfaffian format assumptions and the scale

\eta_{q}=cq^{-r}

, the largest approximate value class has cardinality

\Theta(q^{d-r})

. In both panels, each curve segment is coloured according to its assigned cluster. The cluster assignment depends both on the cell containing the segment and on the value of the curve at the midpoint of the corresponding cube.

Panel of Figure 3 — Figure 3: Interpretation of the keys ( $K$ ) and query ( $x$ ).

Panel of Figure 4 — Figure 4: Interpretation of the factors $A$ and $C$ of the values matrix $V$ in ( 2 ).

Panel of Figure 7 — Figure 7: Step 1 - Pre-Processing: Constructing a Noiseless Virtual Dataset via Algorithm 1 . The cells are learned in a data-driven manner and need not lie on a grid.

Panel of Figure 10 — Figure 10: Step 2 - Source: Domain Learning via Algorithm 2 .

Figure 11: Dimension Ablation. As the problem dimension

d

increases, the recovery performance of the student MLP deteriorates when it is trained either on the raw noisy samples (A) or on the hybrid samples (C). In contrast, the student MLP trained on the data processed by Algorithm 1 (B) exhibits substantially more reliable recovery. The error measured against the noisy responses

Y

, rather than the latent signal

f(X)

, deteriorates at a comparable rate as

d

increases, although the hybrid method (C) performs slightly better.

Panel of Figure 14 — Figure 14: Effect of noise level ( $\sigma$ ) and sample size ( $N$ ). Across both ablations, the student MLP trained on the dataset produced by Algorithm 1 (B) consistently achieves the best $f$ -recovery. The advantage is most pronounced when the noise level is small and when the sample size is limited. In these regimes, Algorithm 1 improves the signal-to-noise ratio by averaging nearby observations while retaining enough local geometric information for the student to learn the teacher reliably. As $\sigma$ grows, the intrinsic difficulty of the recovery problem increases for all methods. As $N$ becomes large, the direct learner already has enough data to average out noise on its own, so the additional gain from preprocessing becomes smaller.

Figure 15: Dimension (

d

) Ablation. Algorithm

2

closely matches Adam on the denoised data and attains the lowest mean

f

-recovery error at every tested dimension. Adam trained directly on the raw sample generally recovers

f

less accurately, with the clearest separation occurring at the intermediate dimensions

d=5

and

d=10

. At the larger dimensions, the three recovery errors become closer, although Algorithm

2

remains marginally best. By contrast, the test error against the noisy responses remains nearly unchanged across the two denoised-data methods and largely plateaus as the dimension increases. Thus, Algorithm

2

reproduces the performance of the optimized denoised model, while the gap to Adam (Raw) reflects the additional recovery benefit provided by Algorithm 1 .

Figure 16: Sample Size (

N

) Ablation. Algorithm

2

tracks Adam on the denoised data throughout the tested sample-size range and achieves a slightly lower mean

f

-recovery error in every experiment. Both denoised-data methods substantially outperform Adam (Raw), particularly for small and moderate sample sizes, while the gap narrows as

N

becomes large. Within these experiments, Algorithm

2

remains a close surrogate for Adam on the denoised data, while purification improves recovery relative to direct training on the raw sample.

Figure 17: Number of Symmetry Classes (

\kappa

) Ablation. Fine-tuning only the shared class values matches or improves upon updating all cell values. The advantage is clearest when

\kappa\ll q^{d},

while the two procedures agree when

\kappa=q^{d}.

This reflects the predicted gain from reusing a compressed symmetry representation.

Figure 18: Grid Scale (

q

) Ablation. Both errors increase as the grid becomes finer at fixed sample size, but the class-value update remains uniformly better. The same trend reflects the growing separation between a fixed

\mathcal{O}(\kappa D)

update and an unrestricted update over

q^{d}

cells.

Fine-Tuning

Fine-tuning adapts pretrained models to new tasks, domains, and behaviors by updating some or all of their parameters.

See all 250 papers

Sep 29, 2026

Theory on Attention Dynamics for Out-of-Distribution In-Context Learning

Transformers have demonstrated remarkable in-context learning (ICL) capabilities, enabling them to perform new tasks without additional fine-tuning. However, their performance often deteriorates when encountering out-of-distribution (OOD) inputs that deviate from the training distribution, and the underlying theory remains poorly understood. To fill this gap, we characterize the OOD error under the input distribution shift through the interplay between the dynamics of the so-called

α

-type and

β

-type attention weights, which represent the transformer's confidence in identifying the correct and incorrect features, respectively. Our results indicate that the OOD error for each feature depends on all pairwise interactions between the training features and OOD features, and under certain cases the transformer performs no better than random guessing. To improve the OOD generalization performance, we next investigate the impact of model finetuning with the OOD data, and particularly, characterize the model forgetting performance on the source domain. Interestingly, the performance on the source domain may not always degrade after finetuning, which highly depends on the nature of the feature shift: finetuning on OOD domain keeps enhancing the confidence of identifying correct features from the original distribution, while the interference from other incorrect features may either increase or decrease. Extensive experiments on both synthetic and real data are conducted to corroborate the theoretical insights.

Junze Deng, Daouda Sow, Sen Lin +1

Ohio State University Ohio, USA · University of Houston Houston, USA

In-Context Learning OOD Generalization Transformer Attention Attention Mechanisms

Ask PDF

Jul 18, 2026

First-Order Predictable but Pairwise Fragile: Local Task Adaptation in Trained Transformers

Task arithmetic, sequential fine-tuning, activation steering, and first-order random search all operate through relatively small perturbations around an already trained checkpoint, and they rely on different local approximations: individual perturbations should be first-order predictable, task updates should compose with controlled interference, useful tangent structure should be stable and possible to estimate, and weight edits should have counterparts in representation space. We measure 8 such properties with the same harness around a multitask LoRA operating point, on 9 transformers (82M-7B), with a prospectively registered property list, thresholds, and test split. We find a shared one-direction validity window up to the tested scale

10^{-2}

, but no universal radius for pairwise composition or update ordering. Along individual directions, changes of the probe loss remain first-order predictable throughout the grid: a perturbation's effect on the loss is essentially its projection onto the gradient, which is also what makes local random search work. Pairwise structure, however, proves to be far more fragile: on over a third of the measured (model, task pair) combinations, two-update order sensitivity sets in strictly inside that window; task-gradient subspaces rotate within tens of steps; additivity under our fixed activation probe fails at full task-vector scale on several models, including both held-out 7B models; and no model median passes the registered global mean-vector weight-to-steering correspondence bar. For two sequential task-gradient steps, the leading order-dependent term is the Lie bracket

H_B\textbf{g}_A-H_A\textbf{g}_B

; its normalized prediction

c(η)=ηκ+O(η^2)

tracks the measured defect at median ratio 1.002, while the onset scale

η^\dagger\approx0.10/κ

spans three orders of magnitude across models and task pairs.

Irina Piontkovskaia, Sergey Nikolenko

1DAIMLD, Moscow, Russia · 2St. Petersburg Department of the Steklov Institute of Mathematics, St. Petersburg, Russia · 3St. Petersburg State University

Fine-Tuning Neural Network Optimization Task Arithmetic Activation Steering

Ask PDF

Sep 13, 2026

Retrieval-Guided Fine-Tuning as Noisy Estimation: Risk bounds and Architectural Analysis

Retrieval-Guided Fine-Tuning (RAG-FT) incorporates retrieved data directly into the training objective, but the statistical consequences of noisy retrieval during training remain theoretically undercharacterized. We study this question by modeling RAG-FT as an estimation problem in a multi-task linear regression framework, using an OLS proxy for single-layer linear self-attention to obtain finite-sample risk bounds. Under homoscedastic retrieval noise, we show that retrieval failure decays exponentially with task separation relative to noise, and derive explicit finite-sample conditions under which RAG-FT achieves lower risk than both target-only and full-corpus training. We then introduce a Distance-Proportional Noise (DPN) model, in which retrieval quality degrades with rank, and compare two estimators under the same retrieval process: the OLS proxy and the literal, uniform-weight forward pass of linear self-attention. We prove that the attention estimator's bias diverges as

Θ(n^{2q})

even under exact retrieval, while OLS risk remains

Θ(d/n)

for every noise exponent

q>0

. These results locate the instability not in noisy retrieval itself, but in the fixed, unweighted aggregation of the literal LSA forward pass, which reweighting by reliability empirically removes. We validate the predicted rate separation through direct simulation of the DPN model.

Bhargav Lad, Yifan Hao

Department of Mathematics Indian Institute of Technology Delhi New Delhi, India · Department of Computer Science University of Illinois Urbana-Champaign Urbana, IL 61801

Fine-Tuning RAG Robustness Linear Attention

Ask PDF

Abstract

Fine-Tuning

Explore similar work

Theory on Attention Dynamics for Out-of-Distribution In-Context Learning

First-Order Predictable but Pairwise Fragile: Local Task Adaptation in Trained Transformers

Retrieval-Guided Fine-Tuning as Noisy Estimation: Risk bounds and Architectural Analysis