In-Context Learning as Nonparametric Conditional Probability Estimation: Risk Bounds and Optimality

2025-08-12

AI Summary

This paper investigates conditional probability estimation in multiclass classification via in-context learning (ICL), evaluating performance via expected excess risk—i.e., the average truncated KL divergence. We derive a novel KL-based oracle inequality and, for the first time in a nonparametric setting, establish that ICL estimators achieve the minimax optimal convergence rate for conditional probability estimation. Moreover, we show that under controlled empirical covering entropy, both Transformers and even standard MLPs attain this optimality. By unifying the analysis of the log-likelihood function classes of Transformers and MLPs, we derive tight upper and lower bounds on the expected excess risk. Our theoretical results rigorously confirm that ICL achieves statistically optimal efficiency under appropriate design conditions.

Technology Category

• Machine Learning: Information Theory
• Natural Language Processing: Learning & Optimization for NLP
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

This paper investigates the expected excess risk of In-Context Learning (ICL) for multiclass classification. We model each task as a sequence of labeled prompt samples and a query input, where a pre-trained model estimates the conditional class probabilities of the query. The expected excess risk is defined as the average truncated Kullback-Leibler (KL) divergence between the predicted and ground-truth conditional class distributions, averaged over a specified family of tasks. We establish a new oracle inequality for the expected excess risk based on KL divergence in multiclass classification. This allows us to derive tight upper and lower bounds for the expected excess risk in transformer-based models, demonstrating that the ICL estimator achieves the minimax optimal rate - up to a logarithmic factor - for conditional probability estimation. From a technical standpoint, our results introduce a novel method for controlling generalization error using the uniform empirical covering entropy of the log-likelihood function class. Furthermore, we show that multilayer perceptrons (MLPs) can also perform ICL and achieve this optimal rate under specific assumptions, suggesting that transformers may not be the exclusive architecture capable of effective ICL.

Problem

• Analyzing excess risk in In-Context Learning for multiclass classification
• Establishing risk bounds for conditional probability estimation in transformers
• Demonstrating optimal ICL performance achievable by non-transformer architectures

Innovation

• KL divergence bounds for ICL risk
• Transformer models achieve minimax optimality
• MLPs perform ICL under specific conditions