Toward Understanding In-context vs. In-weight Learning

Abstract: It has recently been demonstrated empirically that in-context learning emerges in transformers when certain distributional properties are present in the training data, but this ability can also diminish upon further training. We provide a new theoretical understanding of these phenomena by identifying simplified distributional properties that give rise to the emergence and eventual disappearance of in-context learning. We do so by first analyzing a simplified model that uses a gating mechanism to choose between an in-weight and an in-context predictor. Through a combination of a generalization error and regret analysis we identify conditions where in-context and in-weight learning emerge. These theoretical findings are then corroborated experimentally by comparing the behaviour of a full transformer on the simplified distributions to that of the stylized model, demonstrating aligned results. We then extend the study to a full LLM, showing how fine-tuning on various collections of natural language prompts can elicit similar in-context and in-weight learning behaviour.

• The paper introduces a novel bi-level model with a gating mechanism that selects between in-context and in-weight predictors.
• It presents theoretical error bounds and regret analysis, revealing how data distribution properties influence the emergence of ICL and IWL.
• Experimental validations on synthetic data and Omniglot show that factors like input noise, class imbalance, and context length critically affect learning behavior.

Toward Understanding In-context vs. In-weight Learning

The paper "Toward Understanding In-context vs. In-weight Learning" (2410.23042) investigates the distributional properties of data that lead to the emergence and subsequent disappearance of in-context learning (ICL) in transformers. It introduces a simplified model with a gating mechanism that selects between in-weight (IW) and in-context (IC) predictors, providing a theoretical framework supported by experiments on synthetic data, Omniglot, and a fine-tuned LLM.

Theoretical Model and Analysis

The paper presents a bi-level model that learns both an in-weight predictor $g$ and an in-context predictor $h$. A function $\alpha$ selects between these predictors based on the input $\tilde{x}$. The model is formalized within a tabular classification problem setting, assuming a finite input space and considering inputs with label noise. The IWL class $uses only the query, while the ICL class$ uses labels in the context to make a prediction, drawing inspiration from induction heads.

A generalization error and regret analysis reveals conditions for the emergence of ICL and IWL. Proposition 1 provides a generalization bound for the in-weight learner, showing that its test error converges to that of the optimal predictor at a rate of $O(1/\sqrt{N_x})$.

Figure 1: The theoretical error bounds of IC and IW predictors, illustrating how IC error increases with irrelevant contexts while IW error decreases with more samples.

Proposition 2 provides bounds on the error of the IC predictor, showing that it depends on the number of irrelevant labels in the context. Figure 1 illustrates these theoretical error bounds, showing how the lower bound of IC error increases as the number of irrelevant contexts, $k$, increases. The paper then presents a bi-level parameter update procedure (Algorithm 1) and proves that its regret is bounded by the sum of the regrets of learning $\alpha$ and $g$ (Proposition 3). Proposition 4 relates the online learning performance to the generalization error, showing that the average loss of $g(\cdot; w_t)$ exhibits similar behavior as the generalization-error guarantee.

Experimental Validation

The theoretical findings are supported by experiments on synthetic data and the Omniglot dataset.

Synthetic Classification

The synthetic classification task involves an imbalanced data distribution with high- and low-frequency classes. The experiments explore how different parameters, such as input noise, the probability of sampling high-frequency classes ($p_{high}$), and the number of relevant contexts, impact the emergence and transience of ICL. Results, such as those in Figure 2, show that ICL diminishes as $N$ increases, and IWL and ICL can emerge simultaneously.

Figure 2: Validation errors of IC predictor, IW predictor, and transformer on synthetic data, showing the influence of relevant/irrelevant context, and classes from the high-/low-frequency classes.

Varying distributional parameters such as $p_{high}$ and the number of low-frequency classes $|C_L|$ also affects ICL performance. The results indicate that ICL appears to be stronger with smaller $p_{high}$ and larger $p_{relevant}$.

Figure 3: 0-1 validation errors as a function of the dataset size $N$ on the synthetic data.

Increasing the context length $L$ prevents ICL from emerging, as shown in Figure 3.

Omniglot Dataset

Experiments on the Omniglot dataset corroborate the synthetic data findings. The model and data construction differ from previous studies, with a context length of $L = 2$ and a varying number of relevant contexts. Figure 4 shows trends similar to those observed in the synthetic data experiments.

Figure 4: Validation errors of IC predictor, IW predictor, and transformer on Omniglot data, showing the influence of input noise.

The transformer exhibits ICL for low-frequency classes, but IWL emerges more easily on common classes even with larger input noise.

Finetuning a Real LLM

To bridge the gap to practical applications, the paper demonstrates that fine-tuning a real LLM (Gemini Nano 1) to memorize specific data can reduce its ICL ability. The LLM is fine-tuned to memorize where certain people live, and the results indicate that the in-weight information can overwrite the in-context prediction present in the base model.

Conclusion

The paper provides a theoretical and empirical analysis of the conditions under which ICL emerges and becomes transient in transformers. The findings highlight the importance of data distributional properties and the interplay between in-weight and in-context learning. The results contribute to a better understanding of ICL in LLMs and may inform the design of training schedules. The work connects the phenomena to the learnability of the data, naturally affected by its distribution.