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

[Submitted on 31 Aug 2025 (v1), last revised 26 Oct 2026 (this version, v2)]

View PDF HTML (experimental)

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]^dtomathbb{R}^D$, representing the source and shifted tasks, using a single-query, single-head transformer $mathcal{T}(cdot)=Voperatorname{Softmax}(-alpha(cdot,K))$ with $ell^{infty}$-based alignment scores $alpha$. Rather than assume that gradient-based training reaches an empirical risk minimizer, we give a model-specific two-stage procedure textbf{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, $1le p0$. 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 textit{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.

Submission history

From: Anastasis Kratsios [view email]
[v1] Sun, 31 Aug 2025 16:20:27 UTC (1,257 KB)
[v2] Mon, 26 Oct 2026 18:02:29 UTC (810 KB)