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
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 \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, $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 \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.
尚无 AI 解读，展示原文摘要