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

TL;DR

This paper provides a theoretical analysis of in-context learning (ICL) for multiclass classification, establishing risk bounds and demonstrating that transformer-based models achieve near-optimal rates in estimating conditional probabilities.

Contribution

It introduces a new oracle inequality for ICL risk, derives tight bounds for transformers, and shows MLPs can also perform optimally in ICL under certain conditions.

Findings

Transformer-based ICL achieves minimax optimal rate (up to log factors).

A novel method for controlling generalization error via uniform empirical entropy.

MLPs can also perform ICL with optimal rates under suitable assumptions.

Abstract

This paper investigates the expected excess risk of in-context learning (ICL) for multiclass classification. We formalize each task as a sequence of labeled examples followed by a query input; a pretrained model then estimates the query's conditional class probabilities. The expected excess risk is defined as the average truncated Kullback-Leibler (KL) divergence between the predicted and true conditional class distributions over a specified family of tasks. We establish a new oracle inequality for this risk, based on KL divergence, in multiclass classification. This yields tight upper and lower bounds for transformer-based models, showing that the ICL estimator achieves the minimax optimal rate (up to logarithmic factors) for conditional probability estimation. From a technical standpoint, our results introduce a novel method for controlling generalization error via uniform empirical entropy. We further demonstrate that multilayer perceptrons (MLPs) can also perform ICL and attain the same optimal rate (up to logarithmic factors) under suitable assumptions, suggesting that effective ICL need not be exclusive to transformer architectures.