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
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→RD, representing the source and shifted tasks, using a single-query, single-head transformer T(⋅)=VSoftmax(−α(⋅,K)) with ℓ∞-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 ε>0 in every fixed Lp([0,1]d) norm, 1≤p<∞, 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 O(ε−1) in the reciprocal approximation error ε>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 Lp-recovery of g. Critically, fine-tuning only A is a convex problem which involving an at-most linear number of parameters O(ε−1); dramatically cutting down on the worst-case scaling of O(ε−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 q1=81 . 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 κ=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., ∥∇f(x)∥>0 for each x ) and the Pfaffian format assumptions and the scale ηq=cq−r , the largest approximate value class has cardinality Θ(qd−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.
Figure 3: Interpretation of the keys ( K ) and query ( x ).
Figure 4: Interpretation of the factors A and C of the values matrix V in ( 2 ).
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.
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.
Figure 14: Effect of noise level ( σ ) 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 σ 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 ( κ ) Ablation. Fine-tuning only the shared class values matches or improves upon updating all cell values. The advantage is clearest when κ≪qd, while the two procedures agree when κ=qd. 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 O(κD) update and an unrestricted update over qd cells.
1DAIMLD, Moscow, Russia · 2St. Petersburg Department of the Steklov Institute of Mathematics, St. Petersburg, Russia · 3St. Petersburg State University
Department of Mathematics Indian Institute of Technology Delhi New Delhi, India · Department of Computer Science University of Illinois Urbana-Champaign Urbana, IL 61801