Context-Adaptive Inference: A Unified Statistical and Foundation-Model View
Authors: Yue Yao, Caleb N. Ellington, Jingyun Jia, Baiheng Chen, Dong Liu, Rikhil Rao, Jiaqi Wang, Samuel Wales-McGrath, +4 more
Organizations: Department of Statistics, University of Wisconsin-Madison · Computational Biology Department, Carnegie Mellon University · Department of Computer Science, Yale University · Department of Computer Science, University of Wisconsin - Madison · Paul G. Allen School of Computer Science & Engineering, University of Washington · Department of Computer Science and Engineering, The Ohio State University · Department of Computer Sciences, University of Wisconsin-Madison · Machine Learning Department, Carnegie Mellon University; Mohamed bin Zayed University of Artificial Intelligence
Modern predictive systems are expected to adapt their behavior to the specific situation they are facing. A clinical model should not treat every patient the same; a retrieval-augmented model should change its answer when given different evidence; a mixture-of-experts model should route different inputs to different experts. We call this capability context-adaptive inference: before predicting, the system uses information about the current context to specialize its parameters or computation for that instance. This article provides a unified view of context-adaptive inference across three traditions that are usually treated separately: (i) explicit adaptation in statistics (e.g. varying-coefficient models, local regression, hierarchical sharing), (ii) rapid task-specific adaptation in meta-learning and transfer, and (iii) implicit adaptation in large foundation models via prompting, retrieval, and expert routing. We formalize these approaches under a common objective: to map context c to adapted parameters θ(c), then to predict via f(x;θ(c)). Under squared loss, linear prediction heads, and fixed features, we prove that explicit parameter adaptation and implicit routing are mathematically equivalent to kernel ridge regression on joint features of inputs and context. Building on this bridge, we propose practical design principles and evaluation metrics including adaptation-efficiency, routing stability, and context-specific robustness to guide when to specialize, how to constrain that specialization, and how to audit context-adaptive models in deployment. Finally, we identify open problems in identifiability, robustness under distribution shift, and efficient large-scale adaptation, outlining design principles for methods that are scalable, reliable, and transparent in real-world settings.
Bayesian predictive inference provides a principled framework for uncertainty quantification, data efficiency, and robust generalization. However, exact inference is often intractable, and scalable approximations may remain computationally expensive or require restrictive modeling assumptions that degrade predictive performance. Prior-Data Fitted and in-context models have recently emerged as an amortized alternative by learning to map datasets directly to predictive distributions, but existing approaches are tightly coupled to the support of the training prior and lack explicit mechanisms for adapting to new priors at test time, resulting in limited robustness under distribution shift. We introduce a multi-task in-context learning framework for amortized hierarchical Bayesian predictive inference that explicitly represents prior information as a prefix of in-context datasets. A transformer trained on sequences of prior and target tasks learns to adapt its predictions across families of priors. On a suite of evaluations with increasing difficulty, including out-of-meta-distribution priors and priors with high-dimensional latent structures, our method matches oracle Bayesian predictors while being orders of magnitude faster. We further demonstrate its practical relevance on a real-world spatiotemporal temperature prediction benchmark. Code is available at https://github.com/martianmartina/multi-task-bayesian-icl/.
In-context learning (ICL) has attracted increasing attention for enabling models to perform new tasks using only a few ``input--output'' prompt examples. However, existing approaches suffer from \textbf{shallow task adaptation}, where prompts are primarily used as contextual cues to implicitly infer task intent through semantic representations, while the underlying computational process remains unchanged. This limitation restricts task-specific adaptation and compromises inference interpretability. We argue that prompts should not only condition feature representations but also dynamically regulate the model's computation pathways. To this end, we propose \textbf{PromptPath}, an adaptive ICL framework that enables computation-level adaptation through prompt-conditioned dynamic pathways. Specifically, PromptPath introduces a prompt-driven routing mechanism to selectively activate and compose lightweight low-rank experts, forming task-specific computational pathways tailored to different prompts. By integrating prompt information directly into the inference process, PromptPath dynamically reconfigures model computation to enhance task specialization and interpretability. Extensive experiments on 3D point cloud and 2D visual recognition benchmarks demonstrate that PromptPath consistently outperforms state-of-the-art ICL baselines while exhibiting strong cross-domain and cross-task generalization.
Adapting pretrained models typically involves a trade-off between the high training costs of backpropagation and the heavy inference overhead of memory-based or in-context learning. We propose FAAST, a forward-only associative adaptation method that analytically compiles labeled examples into fast weights in a single pass. By eliminating memory or context dependence, FAAST achieves constant-time inference and decouples task adaptation from pretrained representation. Across image classification and language modeling benchmarks, FAAST matches or exceeds backprop-based adaptation while reducing adaptation time by over 90% and is competitive to memory/context-based adaptation while saving memory usage by up to 95%. These results demonstrate FAAST as a highly efficient, scalable solution for supervised task adaptation, particularly for resource-constrained models. We release the code and models at https://github.com/baoguangsheng/faast.