In-context learning and Occam's razor

We thank the reviewer for their time and valuable feedback. We are heartened to hear that the reviewer liked the perspective of our paper. We address the questions and concerns raised by the reviewer below.

Validity of theoretical claims relative to the tightness of the bound

We acknowledge the reviewer’s point regarding the tightness of the bound and believe that our initial discussions about it in the draft might have been distracting as this “approximate equality” is not necessary for our theoretical results. We have revised our draft to show that our quantity of interest (training loss + model complexity) is upper-bounded by the prequential code length, and minimizing this upper-bound is a valid proxy to minimizing the objective itself. Such an approach of minimizing a bound is in line with well-established methods like variational inference, where the negative ELBO is minimized which upper bounds negative log likelihood [1]. We kindly refer the reviewer to the updated Sections 2.2 and 2.3, which are now solely based on upper-bounds.

Prequential coding assumes that as the learner $T$ sees more data, it will generalize better on new data. Is this assumption valid? How is this quantitatively measured? Does it require assumptions on the distribution of the data? What is the sample complexity?

The reviewer brings up an interesting point. The assumption in question is valid as the learner is trained to minimize the prediction error on the next datapoint given prior context. This is similar in spirit to the maximum likelihood estimator, which generalizes better on new data as it seems more data. However, this might not hold true for us in certain OoD scenarios, for e.g. when the number of context points seen during inference are larger than those considered during training. This is a known problem for transformer models, called length generalization.

Our findings in Figure 2 precisely indicate that this assumption is valid when minimizing next-token prediction error. The monotonically decreasing curves imply that as $T$ sees more data through increasing context length (x-axis), its generalization error on new data decreases (y-axis).

Our setup requires very limited assumptions on the data distribution. Akin to maximum likelihood estimation, observing more data points $(x,y) \sim p$ leads to a better model in our setup as long as the underlying true data generating model $p$ is kept fixed, which is the standard setup in most machine learning approaches. In the extreme case where $y$ is independent of $x$, observing additional context would not lead to better generalization for either $T_\phi$ or any other learner like SGD.

Finally, Figure 2 also highlights that in low sample regimes prequential ICL outperforms train-risk ICL and SGD, demonstrating better sample complexity.

1) What is the decoding function used in the prequential coding algorithm?

The decoding function used in the prequential coding algorithm corresponds to the decoding algorithm used in arithmetic coding (a lossless entropy coding scheme), which reconstructs the original data given an optimally compressed code and a probability distribution. For our purposes, the point is that the algorithm is short to write (does not contribute to total program complexity) and allows an object $x$ to be compressed using only $-\log_2 p(x)$ bits and then recovered with 0 error using the arithmetic decoding function. Thank you for spotting this missing detail; we have updated the caption of Fig. 1a accordingly.

2) Is equation 4 on prequential code length missing the complexity of the learning algorithm $K(T)$?

The reviewer is correct that the prequential code length + the complexity of the learner together upper-bounds the joint complexity of the data and model:

$K(D|p_\theta) + K(p_\theta) = K(D, p_\theta) \le L_{preq}(D;T) + K(T)$

We argue that training the learner $T_\phi$ on meta-datasets that do not include the target dataset $D$ ensures that $K(T_\phi)$ is low (i.e. it cannot overfit to $D$). We point the reviewer to the last paragraph of Section 2.2 and the start of Section 2.3 for reference, with further details present in Appendix B.

3) What is there an approximation in equation (6)?

Equation (6) is an upper-bound that only becomes a tight approximation under certain circumstances. However, our theoretical results only require (6) to be an upper-bound, and therefore we have replaced all unnecessarily distracting uses of approximations with upper bounds in our revised manuscript.

Minor Suggestion

We agree with the reviewer’s observation that the first sentence in our abstract presents generalization as the sole objective of machine learning, rather than as one of its important objectives. We have amended this to say that generalization is a “central” goal in ML.

References

[1] Bishop, Christopher M. "Pattern recognition and machine learning." Springer google schola 2 (2006): 1122-1128.

We thank the reviewer for their time, feedback and for the positive assessment of our work. We address the key questions raised by the reviewer below, and hope that it addresses all the concerns raised.

1) Credit attribution to Delétang et al

In Delétang et al., the authors discuss the connection between (pre-trained) sequence models and compression algorithms through prequential coding, which is indeed an important result used in our theory, but not our main contribution. In fact, this is exactly why we explicitly discuss such related work in the section “sequence modeling and compression” to clarify that others have drawn connections between sequence models and compression. Our contribution, as outlined in that section, is that we frame the next-token prediction loss used to train sequence models in the meta-learning framework through in-context learning and draw its connections to minimizing prequential code length across a distribution of tasks. This yields learning algorithms that align with Occam’s razor. The meta-learning perspective is precisely what is absent from Delétang et al. (which does not aim to explain ICL, as we do), but it is one of the cornerstones of our theory.

2) Validity of theoretical claims relative to the tightness of the bound

We acknowledge the reviewer’s point regarding the approximate equality and the tightness of the bound and believe that our initial discussions about it in the draft might have been distracting as this “approximate equality” is not necessary for our theoretical results. We have revised our draft to show that our quantity of interest (training loss + model complexity) is upper-bounded by the prequential code length, and minimizing this upper-bound is a valid proxy to minimizing the objective itself. Such an approach of minimizing a bound is in line with well-established methods like variational inference, where the negative ELBO is minimized which upper bounds negative log likelihood. We kindly refer the reviewer to the updated Sections 2.2 and 2.3, which are now solely based on upper-bounds.

3) Connection between experimental results and theoretical claims

Our theory predicts that ICL trained to minimize next-token prediction error produces an in-context learning algorithm that jointly minimizes training error and model complexity. Crucially, by Occam’s razor, model complexity is most important to take into account when data is limited and generalization is difficult. This is in line with evidence from standard statistics stating that maximum likelihood estimator is a consistent estimator, implying that as the amount of data goes to infinity one uncovers the true model. It is only in finite and limited data settings that regularization strategies and choice of prior play a big role. This ascertains that the gap in performance below 10 data-points in-context is still significant.

The purpose of Section 3.1 was to test exactly this hypothesis that minimizing prequential code length results in better generalization for low-data regimes compared to minimizing training error alone and similar performance for high-data regimes, as our theory predicts given the relationship between prequential ICL and Occam’s razor. Our empirical observations (see “Findings” paragraph of Section 3.1) precisely corroborate our theoretical claims and are thus crucial in validating the hypothesis that we laid out. The fact that this gap in generalization performance sometimes exists only in very low data regimes for certain tasks just implies that for certain tasks like linear regression, less data is required for generalization.

Our other experiments are designed to test the abilities of current methods for ICL, i.e. to what degree do current methods (e.g., sequence model architecture) successfully minimize prequential code length in practice. These experiments were crucial in demonstrating that despite attractive theoretical properties, methods for ICL have significant room for improvement in terms of limitations we identified experimentally and interpreted using our theory.

For all experiments, we attempted to clearly identify in the “Findings” paragraphs what our theory predicts (in terms of Occam’s razor and model complexity) and then either confirmed or rejected these predictions based on experimental results, while highlighting the implications for current methods for ICL.

Minor suggestion

• We agree with the reviewer that the direct equivalence between next-token prediction loss and prequential code length could be clarified in Section 2.4. We added an equation in Section 2.4 to demonstrate this in line 229.
• Thank you for flagging the incorrect figure reference in section 3.4, we have updated the paper accordingly.
• The claim in line 1245 is indeed backward, thank you for pointing out this mistake. We have updated the paper accordingly.

We thank the reviewer for their time, feedback and positive evaluation of our work. We are encouraged that the reviewer found our work as innovative and providing a novel perspective, and address the questions and concerns raised by the reviewer below.

Generalization to non-iid data and complex real-world tasks

ICL is used as a general term for systems that learn conditioned on context without weight updates. As the reviewer astutely points out, this includes non-iid contexts such as natural language. In the majority of our experiments, we consider contexts as iid observations of supervised learning tasks, in line with a long history of studying such problems to better understand general ICL properties [1]. Our experiments on non-iid data in a task capturing temporal correlations akin to ones found in language (see Figure 3) suggest that the theory could be extended to non-iid data.

While we agree that extending the theory to non-iid data would be ideal, it is non trivial to do so as it would entail evaluating every possible ordering of the data to assess the best compression possible. That is, $K(D, p) < L_{preq}^{ordering}(D) + permutation\_matrix(ordering)$. Thus, a key challenge to extending our theory is an efficient way of computing these orderings.

Under-explored architectural dependencies

We agree with the reviewer that investigating the inductive biases and architectural choices that lead to better/worse minimization of prequential code length is a very interesting and important question. However, we believe that it is out-of-scope for our current work which is primarily intended as a theoretical result with confirmatory evidence that could subsequently drive future research on important questions like the role of architecture.

References

[1] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023, July). Transformers learn in-context by gradient descent. In International Conference on Machine Learning (pp. 35151-35174). PMLR.

After carefully considering the feedback from fellow reviewers and re-evaluating the paper, I have decided to lower my score from 8 to 5. While I initially found the paper's novel perspective on in-context learning and Occam's razor intriguing, the concerns raised about the theoretical rigor and applicability of the work are significant. The lack of clarity in the theoretical justifications and the limited extension to non-IID and real-world tasks diminish the paper's contribution to the field. Additionally, the underexplored architectural dependencies and issues with experimental validity suggest that the work is not yet ready for acceptance. I encourage the authors to address these criticisms in future revisions to strengthen the paper's impact.

Here are the concerns I share with my fellow reviewers

• Theoretical Rigor and Clarity. The approximations made, particularly regarding the tightness of bounds and the relation to Kolmogorov complexity, are not sufficiently justified. There is concern that key claims rely on loose bounds or approximate equalities without adequate explanation or empirical validation (Reviewers YT7Y, nR2b)

• Novelty of Contributions. The connection between next-token prediction loss and prequential coding is not novel and has been previously discussed in related work, such as Delétang et al. The paper does not sufficiently acknowledge or build upon these existing contributions (Reviewer YT7Y).

• Experimental Validity. I share the same concerns about experimental setups introducing ambiguities or unfair comparisons, such as differing input formats between models and lack of appropriate regularization techniques in baselines (Reviewers M5Ep, 1nBD)

• Applicability to Non-IID Data. Reviewer 1nBD shares the same concern as I initially did. The theory primarily addresses IID data and does not extend to non-IID or complex real-world tasks, such as language modeling with large language models (LLMs). This limitation reduces the practical relevance of the work.

We thank the reviewer for their time, feedback and for the positive assessment of our work. We address the key questions raised by the reviewer below.

1) Preventing $T_\phi$ from memorizing tasks

To put this question into context, memorization came into question because we need $K(T_\phi)$ – the complexity of the learner itself – to be small so that we minimize a tight bound of PCL. If the parameters $\phi$ memorize datasets, then $K(T_\phi)$ could be large. Thus, we avoid this pitfall by meta-learning the parameters $\phi$ on a sufficiently large distribution of datasets – with a fixed capacity, memorizing each dataset becomes impossible.

2) and 8) Validity of the approach for non-iid data

ICL is a general term for systems that learn conditioned on context without weight updates. In most of our experiments, we consider supervised iid tasks in line with a long history of studying such problems to better understand general ICL properties [1]. Our experiments on non-iid data in a task capturing temporal correlations akin to ones found in language (see Figure 3) suggest that the theory could be extended to non-iid data.

While we agree that extending the theory to non-iid data would be ideal, it is non trivial to do so as it would entail evaluating every possible ordering of the data to assess the best compression possible. Thus, a key challenge to extending our theory is an efficient way of computing these orderings, which is out-of-scope of the current work.

3) The author’s setting assumes the model is updated with every token seen, but isn’t the gradient only back propagated on the summed autoregressive loss?

We believe that there may be a small confusion. It is important to distinguish between (1) the updates of the sequence model that parametrizes the in-context learner and (2) the updates of the prediction model inferred from the context. During the pre-training phase, the sequence model $T_\phi$ is trained “with gradient steps over the cumulative loss” as the reviewer noted, which updates the in-context learning algorithm it implements. However, the “model” whose complexity our theory pertains to is the prediction function specified by the sequence model conditioned on context, i.e $T_\phi(D_{1:t-1})$. As explained in Section 2.4, this prediction function is parameterized by the latent activations after a forward pass of the sequence model given past tokens, and these latent activations do “update” as additional tokens are added to the context. The updates to this in-context model are not gradient updates, but are rather governed by whatever learning algorithm the sequence model implements after pre-training. Throughout the paper, we attempt to clearly indicate this distinction by calling the sequence model the “learner” and calling the implicit in-context model parameterized by latent activations the “model”.

4) Comparing a gradient-based learner and in-context learner

We would first like to clarify that we are comparing properties of learners which can be either (a) in-context learning based, or (b) standard optimization based. ICL learners have trainable parameters $\phi$ that are pre-trained once, after which they takes as input a dataset $D_{1:t-1}$ in-context to provide the model parameters $\theta = T_\phi(D_{1:t-1})$. Optimization based learners do not have trainable parameters and rely on optimization routines to provide model parameters (e.g. training a model with SGD using the training data $D_{1:t-1}$). Note that in this case, the dataset is not provided as context to a model.

To compare learners that minimize PCL to learners that do not, we train an MLP using SGD to minimize training risk. Our theory predicts that especially in low-data regimes where minimizing train risk leads to overfitting, ICL should yield lower test error. This prediction was borne out in our empirical studies.

5) Clarify why the expressivity of a sequence model scales as $N^2$, but the expressivity of the model it fits from context scales as $N$

These scaling trends were meant to be approximate based on linear layers with $N$ units and $N^2$ weights. As the reviewer points out, this is not always the case for all DNN layers, such as self-attention. Regardless, the “parameters” of the model implicitly fit from context do not include the attention weights, but only the latent layer activations which the model can use to modulate predictions based on context.

We thank the reviewer for their time, feedback and for the positive assessment of our work. We address the key questions raised by the reviewer below, and hope that it addresses all the concerns raised.

Why does the bottleneck model perform better than the model w/o bottleneck?

The reviewer hypothesized that the result seems at odds because the bottleneck model might lose information about the task. We’d like to clarify that a bottleneck model can actually capture all the task information since our problems involve task-latents of fixed size (e.g., regression weights, code in Mastermind), where the bottleneck is of sufficient dimension to encode it.

The reviewer also made an astute observation: unlike the bottleneck model, the vanilla Transformer receives separate tokens for $x$ and $y$ as input, to which positional token embeddings are added. This (and causal attention masking) is required to allow predictions at all $x$ tokens in parallel without letting them access their corresponding $y$, and is the common modeling choice for Transformers in prior work [1,2]. However, this choice makes learning more difficult. Based on the reviewer’s feedback, we decided to investigate if this detail explains the difference between models with and without bottleneck and show in Appendix E that it does: forcing the bottleneck model to use separate $x$ and $y$ tokens eliminates its performance advantage, leading to a fair comparison between the models.

For train-risk ICL, is the attention causal, and how is the query selected without inducing shortcuts?

For in-context learners that minimize train risk, we do not input all context subsequences to the bottleneck model. Using causal attention, we can process all context subsequences with a single forward pass.

The query for each subsequence is randomly chosen from the context points in that subsequence. This does not lead to shortcuts as the model does not know which point is chosen for the query, and hence must learn to model $\theta = T_\phi(D_{1:t-1})$ that would generalize well to any choice of context point.

Clarifying task difficulty and its impact on performance gap

The reviewer astutely observes that the performance gap is smaller for linear regression but larger for sinusoidal regression and Mastermind. This is because a complex task requires the algorithm to learn more complex functions to successfully minimize train risk. However, learning more complex functions with very limited data leads to overfitting, which is the basis for our hypothesis that as task complexity increases, simple predictors learned by minimizing prequential code length enjoy a bigger advantage over predictors learned by minimizing train risk.

The reviewer also makes a great suggestion to systematically analyze this by fixing the function class and varying the problem dimension. We conduct the suggested experiment for sinusoid regression and show that with increasing task complexity, the difference between prequential code solution and train-risk solution increases, as shown in Appendix C, thereby validating our initial hypothesis.

Comparing a gradient-based learner and in-context learner

As the reviewer notes, we sought to empirically compare Transformers – learners with parameters that are trained to minimize PCL – to SGD – arguably the most popular learner in ML. Despite using early stopping when training MLPs with SGD, we found that the predictors produced by SGD continue to overfit in low-data regimes.

However, we think that the reviewer makes an excellent suggestion to compare the prequential code lengths achieved by different out-of-the-box learning algorithms that regularize vanilla SGD (e.g., weight decay), which we now show in Appendix F.2. As expected from standard learning theory, these regularization techniques have substantial impact on prequential code length: regularization reduces overfitting in low-data regimes as it favors simple models, but at the expense of underfitting in high-data regimes.

References

[1] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023, July). Transformers learn in-context by gradient descent. In International Conference on Machine Learning (pp. 35151-35174). PMLR.

[2] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

Thank you for addressing my concern and adding multiple experiments in the appendix section.

Upon reading the added experiments, I notice that the Figure E.1 indicating that Transformer with a bottleneck architecture cannot learn the ICL at all (almost random results). This seems weird.

The comparison between prequential ICL and train-risk ICL (in Figure 2) should be done in a fair setting. With that, I mean in the main paper, I think it make more sense to put Figure E.1's result as the prequential ICL, instead of the Figure 2's truncated version of [x,y]. I understand this will downgrade the performance of prequential ICL, which is not what you want in the paper. Therefore I suggest the authors to look for ways to improve the prequential ICL performance in Figure E.1. Since the code is not released, I cannot give detailed suggestions. But maybe try with the curriculum learning applied in Garg et al.? This should attributes to the optimization issue of Transformer rather than the problem itself.

I am willing to raise the score if this could be addressed. Thanks for your time.

Dear reviewer,

As the discussion period is drawing to a close, we kindly ask if our latest response addresses your point which prevented you from strengthening your support. We believe that the paper is greatly improved thanks to your interventions, and we remain available to provide more details.