Weight-Space Learning for Certifiable Few-shot Transfer Learning

In contemporary deep learning, a prevalent and effective workflow for solving low-data problems is adapting powerful pre-trained foundation models (FMs) to new tasks via parameter-efficient fine-tuning (PEFT). However, while empirically effective, the resulting solutions lack generalisation guarantees to certify their accuracy - which may be required for ethical or legal reasons prior to deployment in high-importance applications. In this paper we develop a novel transfer learning approach that is designed to facilitate non-vacuous learning theoretic generalisation guarantees for downstream tasks, even in the low-shot regime. Specifically, we first use upstream tasks to train a {\em distribution over PEFT parameters}. We then learn the downstream task by a {\em sample-and-evaluate} procedure -- sampling plausible PEFTs from the trained diffusion model and selecting the one with the highest likelihood on the downstream data. Crucially, this confines our model hypothesis to a {\em finite} set of PEFT samples. In contrast to the typical continuous hypothesis spaces of neural network weights, this facilitates tighter risk certificates. We instantiate our bound and show non-trivial generalization guarantees compared to existing learning approaches which lead to vacuous bounds in the low-shot regime.

Modern AI models can be quickly adapted to new tasks using just a handful of examples. For instance, personalising a language model for an individual user, or teaching an image classifier to recognise a new type of bird. They usually work well in practice, but in high-stakes settings like medicine or law, "usually works" is not enough: we need a mathematical guarantee that the model will be accurate on new, unseen data. For modern AI adapted from only a few examples, no existing method can provide such a guarantee. Instead of fine-tuning the model in the usual way, we train a generator — similar in spirit to the image generators behind tools like Stable Diffusion, but it produces small AI customisations rather than pictures. To adapt to a new task, we sample a finite menu of candidates from this generator and pick the one that best fits the few available examples. Choosing from a finite menu, rather than infinitely many possibilities, is what makes mathematical guarantees possible. Our method delivers the first meaningful accuracy guarantees in this few-shot setting, on both language and vision tasks, while matching the practical accuracy of standard methods that offer no such guarantees.