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

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\mathrm{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 $\epsilon >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}({\epsilon }^{-1})$ in the reciprocal approximation error $\epsilon >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 \textit{at-most linear} number of parameters $\mathcal{O}({\epsilon }^{-1})$; dramatically cutting down on the worst-case scaling of $\mathcal{O}({\epsilon }^{-\mathbf{d}})$ unavoidable in the source-training phase.

arXiv comment28 Pages + Proofs