[Paper Review] What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization

Yufeng Zhang, Fengzhuo Zhang+2 more|arXiv (Cornell University)|May 30, 2023

Topic Modeling4 citations

TL;DR

This paper frames in-context learning (ICL) in large language models as Bayesian model averaging (BMA) implicitly implemented via the transformer attention mechanism, establishing an O(1/T) regret bound for ICL performance and proving that pretraining error decays exponentially with depth and sublinearly with dataset size. It provides a unified theoretical framework linking attention, pretraining, and generalization in ICL.

ABSTRACT

In this paper, we conduct a comprehensive study of In-Context Learning (ICL) by addressing several open questions: (a) What type of ICL estimator is learned by large language models? (b) What is a proper performance metric for ICL and what is the error rate? (c) How does the transformer architecture enable ICL? To answer these questions, we adopt a Bayesian view and formulate ICL as a problem of predicting the response corresponding to the current covariate, given a number of examples drawn from a latent variable model. To answer (a), we show that, without updating the neural network parameters, ICL implicitly implements the Bayesian model averaging algorithm, which is proven to be approximately parameterized by the attention mechanism. For (b), we analyze the ICL performance from an online learning perspective and establish a $\mathcal{O}(1/T)$ regret bound for perfectly pretrained ICL, where $T$ is the number of examples in the prompt. To answer (c), we show that, in addition to encoding Bayesian model averaging via attention, the transformer architecture also enables a fine-grained statistical analysis of pretraining under realistic assumptions. In particular, we prove that the error of pretrained model is bounded by a sum of an approximation error and a generalization error, where the former decays to zero exponentially as the depth grows, and the latter decays to zero sublinearly with the number of tokens in the pretraining dataset. Our results provide a unified understanding of the transformer and its ICL ability with bounds on ICL regret, approximation, and generalization, which deepens our knowledge of these essential aspects of modern language models.

Motivation & Objective

• To understand what type of estimator ICL implements in large language models.
• To define a proper performance metric for ICL and quantify its error rate.
• To explain how the transformer architecture enables ICL through both prompting and pretraining.
• To provide a unified theoretical analysis of ICL, pretraining, and generalization using Bayesian and online learning frameworks.

Proposed method

• Formulates ICL as a Bayesian inference problem where examples are drawn from a latent variable model parameterized by a hidden concept z*.
• Shows that ICL without parameter updates implements Bayesian model averaging (BMA), with the attention mechanism approximately encoding BMA.
• Introduces an ICL regret metric based on online learning, proving an O(1/T) regret bound for perfectly pretrained models.
• Applies the PAC-Bayes framework to analyze pretraining error, decomposing it into approximation and generalization error components.
• Proves that approximation error decays exponentially with transformer depth and generalization error decays sublinearly with pretraining data size.
• Uses concentration inequalities and matrix norm bounds to analyze attention mechanism stability and generalization.

Experimental results

Research questions

• RQ1 What type of ICL estimator is implicitly learned by large language models without parameter updates?
• RQ2 What is a suitable performance metric for ICL, and what is its theoretical error rate?
• RQ3 How does the transformer architecture enable in-context learning during both pretraining and prompting?
• RQ4 How do approximation and generalization errors in pretraining affect downstream ICL performance?

Key findings

• ICL without parameter updates implicitly implements Bayesian model averaging (BMA), with the attention mechanism approximately encoding this process.
• The ICL regret after T examples in the prompt is bounded by O(1/T), up to the statistical error of the pretrained model.
• Approximation error in the pretrained model decays exponentially fast with increasing transformer depth.
• Generalization error decays sublinearly with the number of tokens in the pretraining dataset.
• The standard softmax attention mechanism approximately encodes BMA, especially as prompt length increases.
• The theoretical framework unifies ICL, pretraining, and generalization, providing bounds on all three components.