alphaFeedback

WorksMathematics

Nonparametric In-Context Learning and Inference via Measure-Indexed Kernel Ridge Regression

AI contributions
Not recorded for this paper
View details

Abstract

In-Context learning enables an estimator, trained across multiple tasks, to adapt to a new unseen task solely through its in-context samples, without any parameter refitting. Despite growing interest in the statistical foundations of in-context learning, two gaps still remain. First, existing convergence results do not characterize the distributional behavior of the in-context learning estimators. Second, classical nonparametric methods are not designed to handle the measure-valued inputs that in-context learning naturally produces. We address both gaps by developing a measure-indexed reproducing kernel Hilbert space framework, in which the adaptivity of the estimator is captured through the dependence of a mean functional on the empirical measure of the context. Within this framework, we propose a measure-indexed kernel ridge regression estimator computed from multiple independent tasks and establish its minimax-optimal convergence rate. We further prove that the proposed estimator is pointwise asymptotically normal, enabling the construction of confidence intervals for the mean and prediction intervals for the response on an unseen task. Notably, the estimator achieves valid uncertainty quantification without any refitting when adapting to a new task. As a special case, our framework recovers classical distributed nonparametric regression when all tasks share the same underlying function. The proposed framework is illustrated through a latent Gaussian process example.

Introduction

In-Context learning (ICL) refers to an emerging prediction paradigm in which an estimator, trained once across a collection of tasks, automatically adapts its predictions to an unseen task solely through the task’s in-context samples, without any refitting of model parameters ([6, 12]). This contrasts sharply with conventional statistical learning, where a separate model must be retrained for each new task. The adaptivity of ICL calls for a theoretical explanation, for which the Bayesian perspective provides a compelling account ([38]); the optimal predictor is the posterior predictive mean which is a fixed functional of the in-context

In-Context Learning in Measure-Indexed RKHS Framework

samples and the query that automatically adapts to any new task through the context alone. This perspective has motivated a growing body of work connecting ICL to classical statistical estimators. For linear function classes, a single layer of linear attention can implement ridge regression on the in-context samples [3, 35, 4]. For nonparametric regression over smooth function classes, Ching et al. [10] showed that transformers trained via meta-learning can attain minimax-optimal prediction rates, and Wakayama and Suzuki [37] established a rigorous connection between ICL and Bayesian inference for characterizing the optimal predictor and its generalization error.

Despite this rapid progress, two gaps remain. The first is concerned with uncertainty quantification. Existing convergence results establish that the ICL predictor converges to the truth, but say nothing about its distributional behavior, though the adaptivity of ICL is most valuable precisely when one needs to quantify uncertainty about a prediction made on a previously unseen task. Constructing confidence or prediction intervals in this setting requires asymptotic normality, which no existing result provides. The second gap is concerned with the geometry of the input domain. Directly using the raw context sample as input can lead to the “curse of dimensionality,” motivating a dimension-reduction strategy that transforms the input space into a more structured representation. The natural mathematical object representing an in-context sample is its empirical distribution [20]. The ICL predictor is therefore a function whose argument includes a probability measure as part of its input. Classical nonparametric methods, including the reproducing kernel Hilbert space (RKHS) based methods such as kernel ridge regression or related ones, are mostly built for Euclidean domains. They do not directly accommodate this measure-valued structure. Adapting them requires a nontrivial generalization of the standard RKHS framework.

In this paper, we propose a new RKHS framework, called measure-indexed RKHS (MI-RKHS), with measure-valued inputs for ICL that addresses both gaps simultaneously. The key observation is that ICL defines a fixed mean functional from a (context, query) pair to predictions; the adaptivity of the estimator across tasks is naturally captured by the dependence of this functional on the empirical distribution of the context. We model this function as an element in an MI-RKHS defined over the product of a space of probability measures and a feature space, which allows classical RKHS analysis to be extended to the measure-valued setting. For instance, we propose a measure-indexed kernel ridge regression (MI-KRR) estimator computed from multiple independent tasks and establish its convergence rate. We prove that the MI-KRR estimator is pointwise asymptotically normal, enabling the construction of confidence intervals for the mean and prediction intervals for the response on an unseen task. Notably, the estimator adapts to a new task without any refitting, yet still achieves asymptotic normality. The proposed MI-KRR method extends classical kernel ridge regression to the measure-valued setting, with both convergence rate and asymptotic normality established within this new framework. As a special case, when all tasks share the same fixed underlying function, our framework degenerates to the well-known distributed nonparametric regression, recovering known results in that literature such as optimal convergence rate [43, 29, 30, 21, 32]. More generally, our framework is conceptually related to meta-learning or transfer learning since both paradigms require the underlying tasks are different but still possess similarity.

The rest of this paper is organized as follows. In § 3, we introduce the MI-RKHS framework as well as a MI-KRR regression for estimating the mean functional. In § 4.2, we derive the convergence rate and asymptotic distribution for the proposed MI-KRR estimator. An example on latent Gaussian process model is provided to examine the proposed MI-RKHS framework. Proofs of the theorems and lemmas are deferred to § 5.

In-Context Learning in Measure-Indexed RKHS Framework

Let X⊆Rd\mathcal{X} \subseteq\mathbb{R}^{d} be a dd-dimensional feature space, and HX\mathcal{H}_{X} be an RKHS of functions on X\mathcal{X} endowed with inner product ⟨⋅,⋅⟩HX\langle\cdot,\cdot\rangle_{\mathcal{H}_{X}} and reproducing kernel R⟨X⟩:X×X→RR^{\langle X\rangle}:\mathcal{X}\times\mathcal{X}\to\mathbb{R}. Suppose that μ\mu is a known deterministic probability measure on X\mathcal{X}. We further assume that R⟨X⟩R^{\langle X\rangle} is bounded.

Assumption 1. CX≜sup⁡x∈XR⟨X⟩(x,x)<∞C_{X}\triangleq\sup_{x\in\mathcal{X}}R^{\langle X\rangle}(x,x)<\infty.

The Assumption 1 is standard in the nonparametric literature which is satisfied for commonly used kernels such as Sobolev kernel and Gaussian kernels [28, 40]. Let ff be a latent random function with distribution FF, where FF is supported on HX\mathcal{H}_{X}. Suppose X1,…,Xn+1X_{1},\ldots,X_{n+1} are i.i.d. feature variables from μ\mu and independent of ff. The response variables YiY_{i} satisfy

Yi=f(Xi)+εi,i=1,…,n+1,(1)Y_{i}=f(X_{i})+\varepsilon_{i},\qquad i=1,\ldots,n+1, \tag*{(1)}

where εi\varepsilon_i’s are i.i.d. with mean zero and finite variance, independent of f,X1,…,Xn+1f,X_{1},\ldots,X_{n+1}. We are interested in estimating the following conditional mean functional

M∗(Dn,x)=E[Yn+1∣Dn,Xn+1=x],x∈X,(2)M_{*}(D_{n},x)=\mathbb{E}[Y_{n+1}\mid D_{n},X_{n+1}=x],\qquad x\in\mathcal{X}, \tag*{(2)}

where Dn={Z1,…,Zn}D_{n}=\{Z_{1},\ldots,Z_{n}\} with Zi=(Xi,Yi)Z_{i}=(X_{i},Y_{i}). Estimation of M∗M_{*} is closely related to ICL. In classic ICL literature, DnD_{n} is regarded as the context data and Zn+1Z_{n+1} is the new query pair [14, 35, 41, 24, 16, 2, 34, 19, 37, 10]. The target is to estimate M∗M_{*} which will be useful to predict the next query output Yn+1Y_{n+1} based on DnD_{n} and Xn+1X_{n+1}. Directly estimating M∗M_{*} generally suffers from the “curse-of-dimensionality” since the input DnD_{n} is n(d+1)n(d+1)-dimensional which is high as nn diverges. A dimension-reduction strategy is usually needed to make this estimation problem well posed.

An observation is that, although the ambient dimension of DnD_{n} increases with nn, the estimation target M∗(Dn,x)M_{*}(D_{n},x) may admit a substantially simpler representation. In particular, under the conditional i.i.d. structure of DnD_{n}, M∗M_{*} is invariant to the ordering of the observations in DnD_{n} and hence depends on the DnD_{n} only through its empirical measure. This observation is formalized as follows.

Lemma 1. Under the conditional i.i.d. assumption on Z1,…,Zn+1Z_{1},\ldots,Z_{n+1} given ff, it holds that

E[Yn+1∣Dn,Xn+1]=E[Yn+1∣π^n,Xn+1]a.s.,\mathbb{E}[Y_{n+1}\mid D_{n},X_{n+1}]=\mathbb{E}[Y_{n+1}\mid\widehat{\pi}_{n},X_{n+1}]\quad\text{a.s.},

where π^n=n−1∑i=1nδZi\widehat{\pi}_{n}=n^{-1}\sum_{i=1}^{n}\delta_{Z_{i}} is the empirical measure of Z1,…,ZnZ_{1},\ldots,Z_{n}.

Lemma 1 shows that, for the purpose of predicting Yn+1Y_{n+1}, the ordered DnD_{n} can be replaced by its empirical measure π^n\widehat{\pi}_{n}. Thus, although DnD_{n} is a high-dimensional ordered object, M∗M_{*} is invariant to permutations of the observations and can be represented through the measure-valued input π^n\widehat{\pi}_{n} together with the query point xx. Accordingly, with a slight abuse of notation, we henceforth regard M∗M_{*} as a function of (π^n,x)(\widehat{\pi}_{n},x) and write

M∗(π^n,x)≡M∗(Dn,x).M_{*}(\widehat{\pi}_{n},x)\equiv M_{*}(D_{n},x).

Therefore, M∗M_{*} is a function defined on the domain P^(Z)×X\widehat{\mathcal{P}}(\mathcal{Z})\times\mathcal{X}, where

P^(Z)={n−1∑i=1nδzi:z1,…,zn∈Z,n≥1}.\widehat{\mathcal{P}}(\mathcal{Z})=\left\{n^{-1}\sum_{i=1}^{n}\delta_{z_i}:z_1,\ldots,z_n\in\mathcal{Z},n\geq1\right\}.

is the space of empirical measures and Z=X×R\mathcal{Z} = \mathcal{X} \times\mathbb{R} is the range of feature-response pairs. This allows us to characterize the regularity of M∗M_* with respect to π^n\widehat{\pi}_n, and hence, reduce the complexity of M∗M_*. This measure-valued representation is related to the recent ICL framework of Liu et al. [20], where empirical measures are used as a permutation-invariant representation of the context data to facilitate a unified analysis across different context lengths and transformer architectures. In contrast, Lemma 1 shows that the use of π^n\widehat{\pi}_n provides an alternative representation of the conditional mean functional M∗M_* through information-preserving reduction of DnD_n, which will facilitate our statistical analysis.

A General MI-RKHS Framework

In this section, we establish a general MI-RKHS framework that characterizes the regularity of M∗M_*. Recall that ff is the unobservable latent function drawn from FF, and therefore f∈HXf \in\mathcal{H}_X almost surely. We further assume that ∥f∥HX\|f\|_{\mathcal{H}_X} has a finite second moment.

Assumption 2. E[∥f∥HX2]<∞\mathbb{E}[\|f\|_{\mathcal{H}_X}^2] < \infty.

Assumption 2 says that ff has a finite second moment in terms of the ∥⋅∥HX\|\cdot\|_{\mathcal{H}_X} norm. It ensures that the complexity of ff is well controlled in expectation. By (1), (2) and independence between Xn+1X_{n+1} and (f,Dn)(f, D_n), we have

M∗(π^n,x)=E[f(Xn+1)+ϵn+1∣π^n,Xn+1=x]=E[f∣π^n](x).(3)M_*(\widehat{\pi}_n, x) = \mathbb{E}[f(X_{n+1}) + \epsilon_{n+1} \mid\widehat{\pi}_n, X_{n+1} = x] = \mathbb{E}[f \mid\widehat{\pi}_n](x). \tag*{(3)}

Let H∗(π^n)=E[f∣π^n]H_*(\widehat{\pi}_n) = \mathbb{E}[f \mid\widehat{\pi}_n]. By Assumption 2, we have

E[∥H∗(π^n)∥HX2]≤E[∥f∥HX2]<∞,\mathbb{E}[\|H_*(\widehat{\pi}_n)\|_{\mathcal{H}_X}^2] \leq\mathbb{E}[\|f\|_{\mathcal{H}_X}^2] < \infty,

so H∗(π^n)∈HXH_*(\widehat{\pi}_n) \in\mathcal{H}_X almost surely, implying that H∗H_* is a mapping from P^(Z)\widehat{\mathcal{P}}(\mathcal{Z}) to HX\mathcal{H}_X. To characterize the regularity of H∗H_* over a richer domain, let

P(Z)={π:π is a probability measure on Z with ∫Zy2 dπ(x,y)<∞}.\mathcal{P}(\mathcal{Z}) = \left\{ \pi: \pi\text{ is a probability measure on } \mathcal{Z} \text{ with } \int_{\mathcal{Z}} y^2 \,\mathrm{d}\pi(x,y) < \infty\right\}.

References

  1. [1]Adams, R. A. and Fournier, J. J. (2003). Sobolev Spaces, volume 140. Elsevier.
  2. [2]Agarwal, R., Singh, A., Zhang, L., Bohnet, B., Rosias, L., Chan, S., Zhang, B., Anand, A., Abbas, Z., Nova, A., Co-Reyes, J., Chu, E., Behbahani, F., Faust, A., and Larochelle, H. (2024). Many- shot in-context learning. In Advances in Neural Information Processing Systems, volume 37, pages 76930–76966. Curran Associates, Inc.
  3. [3]Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023). What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations.arxiv.org/abs/2211.15661
  4. [4]Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, volume 36.arxiv.org/abs/2306.04637
  5. [5]Banerjee, M., Durot, C., and Sen, B. (2019). Divide and conquer in nonstandard problems and the super-efficiency phenomenon. The Annals of Statistics, 47(2):720 – 757.DOI
  6. [6]Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. (2020). Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877–1901. Curran Associates, Inc.arxiv.org/abs/2005.14165
  7. [7]Cai, T. T. and Pu, H. (2024). Transfer learning for nonparametric regression: Non-asymptotic minimax analysis and adaptive procedure. arXiv preprint arXiv:2401.12272.arxiv.org/abs/2401.12272
  8. [8]Cai, T. T. and Wei, H. (2021). Transfer learning for nonparametric classification: Minimax rate and adaptive classifier. The Annals of Statistics, 49(1):100 – 128.arxiv.org/abs/1906.02903
  9. [9]Carl, B. and Stephani, I. (1990). Entropy, compactness and the approximation of operators, volume 98 of Cambridge Tracts in Mathematics. Cambridge University Press, Cambridge.DOI
  10. [10]Ching, M., Popescu, I., Smith, N., Ma, T., Underwood, W. G., and Samworth, R. J. (2026). Efficient and minimax optimal in-context nonparametric regression with transformers. In Proceedings of the 43rd International Conference on Machine Learning, volume 306 of Proceedings of Machine Learning Research, pages 19581–19610. PMLR.
  11. [11]Conway, J. B. (1990). A Course in Functional Analysis, volume 96 of Graduate Texts in Mathematics. Springer, 2nd edition.
  12. [12]Dong, Q., Li, L., Dai, D., Zheng, C., Ma, J., Li, R., Xia, H., Xu, J., Wu, Z., Liu, T., Chang, B., Sun, X., Li, L., and Sui, Z. (2022). A survey on in-context learning. arXiv preprint arXiv:2301.00234.arxiv.org/abs/2301.00234
  13. [13]Evgeniou, T., Micchelli, C., and Pontil, M. (2005). Learning multiple tasks with kernel methods. Journal of Machine Learning Research, 6(Apr):615–637.DOI
  14. [14]Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems, volume 35, pages 30583–30598.
  15. [15]Gu, C. (2013). Smoothing Spline ANOVA Models, 2nd edition. Springer, New York.
  16. [16]Kim, J., Nakamaki, T., and Suzuki, T. (2024). Transformers are minimax optimal nonparametric in-context learners. In Advances in Neural Information Processing Systems (NeurIPS).arxiv.org/abs/2408.12186
  17. [17]Kosorok, M. R. (2008). Introduction to Empirical Processes and Semiparametric Inference. Springer.DOI
  18. [19]Liu, C., Tan, F., Xie, C., Zeng, Y., and Zhu, L. (2025). In-context learning as nonparametric conditional probability estimation: Risk bounds and optimality. arXiv preprint arXiv:2508.08673.arxiv.org/abs/2508.08673
  19. [20]Liu, C., Xie, C., Tan, F., Zeng, Y., and Zhu, L. (2026). A unified framework for in-context learning with causal and masked language models. arXiv preprint arXiv:2607.04081.arxiv.org/abs/2607.04081
  20. [21]Liu, M., Shang, Z., and Cheng, G. (2020). Nonparametric distributed learning under general designs. Electronic Journal of Statistics, 14:3070–3102.DOI
  21. [22]Micchelli, C. A. and Pontil, M. (2005). On learning vector-valued functions. Neural Computation, 17(1):177–204.pubmed.ncbi.nlm.nih.gov/15563752
  22. [23]Muandet, K., Fukumizu, K., and Scholkopf, B. (2017). Kernel mean embedding of distributions: A review and beyond. Foundations and Trends in Machine Learning, 10:1–141.
  23. [24]Oko, K., Song, Y., Suzuki, T., and Wu, D. (2024). Pretrained transformer efficiently learns low-dimensional target functions in-context. In Advances in Neural Information Processing Systems (NeurIPS).arxiv.org/abs/2411.02544
  24. [25]Pinelis, I. (1994). Optimum bounds for the distributions of martingales in banach spaces. The Annals of Probability, 22(4):1679–1706.arxiv.org/abs/1208.2200
  25. [26]Sai Li, T. T. C. and Li, H. (2023). Transfer learning in large-scale Gaussian graphical models with false discovery rate control. Journal of the American Statistical Association, 118(543):2171–2183.DOI
  26. [27]Sai Li, Linjun Zhang, T. T. C. and Li, H. (2024). Estimation and inference for high-dimensional generalized linear models with knowledge transfer. Journal of the American Statistical Association, 119(546):1274–1285.
  27. [28]Shang, Z. and Cheng, G. (2013). Local and global asymptotic inference in smoothing spline models. Annals of Statistics, 41:2608–2638.DOI
  28. [29]Shang, Z. and Cheng, G. (2017). Computational limits of a distributed algorithm for smoothing spline. Journal of Machine Learning Research, 18(108):1–37.
  29. [30]Shang, Z., Hao, B., and Cheng, G. (2019). Nonparametric Bayesian aggregation for massive data. Journal of Machine Learning Research, 20(140):1–81.
  30. [31]Shang, Z., Sang, P., and Jin, C. (2026). Bootstrap nonparametric inference under data integration. Stat, 15(4):e70175.arxiv.org/abs/2501.01610
  31. [32]Szabó, B. and van Zanten, H. (2020). Adaptive distributed methods under communication constraints. The Annals of Statistics, 48(4):2347 – 2380.arxiv.org/abs/1804.00864
  32. [33]Tian, Y. and Feng, Y. (2023). Transfer learning under high-dimensional generalized linear models. Journal of the American Statistical Association, 118(544):2684–2697.DOI
  33. [34]Vladymyrov, M., von Oswald, J., Sandler, M., and Ge, R. (2024). Linear transformers are versatile in-context learners. In Advances in Neural Information Processing Systems, volume 37, pages 48784–48809.arxiv.org/abs/2402.14180
  34. [35]von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151–35174. PMLR.arxiv.org/abs/2212.07677
  35. [36]Wahba, G. (1990). Spline Models for Observational Data. SIAM CBMS-NSF Regional Conference Series in Applied Mathematics, 59.
  36. [37]Wakayama, T. and Suzuki, T. (2026). In-context learning is provably Bayesian inference: A generalization theory for meta-learning. In Proceedings of the Forty-Third International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR.arxiv.org/abs/2510.10981
  37. [38]Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. (2022). An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations.arxiv.org/abs/2111.02080
  38. [40]Yang, Y., Shang, Z., and Cheng, G. (2020). Non-asymptotic analysis for nonparametric testing. In 33rd Annual Conference on Learning Theory, pages 1–47. ACM.
  39. [41]Zhang, R., Frei, S., and Bartlett, P. L. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25:1–55.arxiv.org/abs/2306.09927
  40. [42]Zhang, Y., Duchi, J., and Wainwright, M. (2013). Divide and conquer kernel ridge regression. In Shalev-Shwartz, S. and Steinwart, I., editors, Proceedings of the 26th Annual Conference on Learning Theory, volume 30 of Proceedings of Machine Learning Research, pages 592–617, Princeton, NJ, USA. PMLR.
  41. [43]Zhang, Y., Duchi, J., and Wainwright, M. (2015). Divide and conquer kernel ridge regression: A distributed algorithm with minimax optimal rates. Journal of Machine Learning Research, 16:3299–3340.
  42. [44]Álvarez, M. A., Rosasco, L., and Lawrence, N. D. (2012). Kernels for vector-valued functions: A review. Foundations and Trends in Machine Learning, 4(3):195–266.

Paper details

vlnc:W_fmpb6t8k
v1 · latestHistory (1)

published · Fresh — added to the Hub within the last 30 days

added ·

Usage

views · 8

downloads · 0

A view counts once per browser per day, a download each time the PDF is requested.

CC BY 4.0

Original source

  • Mathematics

Not recorded for this paper · View statement →

Looking for related papers…
Contents