Does learning the right latent variables necessarily improve in-context learning?

We thank the reviewer for providing valuable and constructive feedback.

Single latent variable can fully summarize context information independently of $x_{query}$ seems unnatural

It is important to note that the dimensionality of the latent variable in explicit models is kept sufficiently large to encode the true latents for all the tasks considered (except GP). The independence assumption of latent variable and $x_{query}$ is standard across the field

• Deep learning models learn parameters $\theta$ of a neural network from training data $\mathcal{D}$ to generalize to test observations $x_{query}$. This is the setup for vision, natural language and tabular tasks where $\theta$ is kept independent of $x_{query}$.
• The same assumption is also used in representation and unsupervised learning.
• Multiple works on ICL also claim this to be the underlying mechanism (Eq. 1 in [5,6])

Since most machine learning approaches train a model independent of the test set, they follow this independence assumption. In contrast, test-time training methods can be seen as not following this assumption [4].

restricting direct attention from x_{query} to other context tokens may artificially constrain the model’s behavior rather than isolating a meaningful causal mechanism

In all the tasks (except GP Regression), the correct causal mechanism is to infer the true latents solely from context and then leverage it for prediction. This indeed constrains the hypothesis class and theoretical results suggest that reducing the hypothesis class leads to tighter bounds as long as the true solution is realizable [1], which it is as the independence assumption is satisfied in the true data generating distribuion, i.e. $y_q \perp \mathcal{D} \,|\, x_q, z$.

Implicit models may also recover task-relevant latents

Indeed they can, but the search space is larger. Given $T$ tokens and $L$ layers, it is not clear which combination of $𝑇\times𝐿$ representations should you decode latents from. Different places to probe can provide vastly different inferences; whereas in explicit model there is only one natural place to probe as it is precisely trained to infer the latents. We refer the reviewer to Figures 4(b) and 9 where we investigate this with probes in both implicit and explicit models and show that the counterfactual performance is worse in implicit models, showing that either the latents are not sufficiently encoded or just difficult to find.

failure of the explicit model to leverage inferred latents for prediction

We conjecture that this is due to optimization dynamics or lack of inductive biases since our prediction model is expressive enough to represent a linear mapping, but is not able to do so when combined with the problem of inferring latents, in OOD settings. Note that this is similar to standard studies in OOD generalization that highlight failure cases of deep learning methods.

bottleneck structure itself may limit information flow, rather than revealing a true failure of latent inference

Most works on inductive biases limit information flow and can aid generalization if this is reflected in the data too, eg. modular systems [2] limit information flow by blocking non-activated experts, information bottlenecks through KL bounds [3], etc. In our work, we consider tasks where inferring the correct latents should block additional information flow from the context to the query. Could the reviewer clarify if they meant something else by true failure of latent inference?

why does training another classifier on the same inferred latents yield better results?

In one case we keep the last layer fixed and learn just the latent inference (known prediction) while in the other, we learn both latent inference and prediction jointly. Our experiments demonstrate that the latter leads to sub-optimalities, i.e. joint training of prediction parameters and latent variable inference causes problems.

We hope that our response has resolved the reviewer's concerns and would be happy to provide further clarifications.

[1] Shalev-Shwartz, Shai, and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.

[2] Andreas, Jacob, et al. "Neural module networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

[3] Tishby, Naftali, and Noga Zaslavsky. "Deep learning and the information bottleneck principle." 2015 ieee information theory workshop (itw). Ieee, 2015.

[4] Sun, Yu, et al. "Test-time training with self-supervision for generalization under distribution shifts." International conference on machine learning. PMLR, 2020.

[5] Müller, Samuel, et al. "Transformers can do bayesian inference." arXiv preprint arXiv:2112.10510 (2021).

[6] Han, Seungwook, et al. "Emergence of Abstractions: Concept Encoding and Decoding Mechanism for In-Context Learning in Transformers." arXiv preprint arXiv:2412.12276 (2024).

We thank the reviewer for their feedback and appreciate that they found the paper to be a clear, well written, in-depth analysis of the subject.

To alleviate their concerns regarding OOD generalization, we refer them to Figures 2(c) and 5 where we test for compositional generalization instead of just shifting the query. In these tasks, novel combinations of underlying latents are provided during inference as opposed to solely changing the query (see Training and Evaluation on Page 5).

Finally, we provide additional analysis where we test for OOD generalization in linear regression and classification by changing the distribution of the underlying weight vectors $w$ to be sampled from a normal distribution with larger variance. Our results, provided in Table 1 here (https://anonymous.4open.science/r/explicit-implicit-rebuttal-B263/explicit-implicit-rebuttal.pdf), indicate that it is not the case that explicit models are able to generalize better than implicit ones in such OOD scenarios. Even further, we see that while known prediction is a bit better, it still lags behind the implicit models because this setting is OOD for the context aggregator while maintaining in-distribution $x_{query}$.

We hope that our response has resolved the reviewers' concerns and would be happy to provide further clarifications.

We thank the reviewer for their feedback, but strongly disagree that our paper is a negative result that is “not sufficiently interesting, surprising or convincingly argued”. Besides the fact that several reviewers (1Akw, 1soF, TjQ5) found our motivation (detailed below) and study interesting and convincingly argued, ICML has a track record of accepting negative results [1].

improve in-context learning by introducing a bottleneck, … very rare view in the field.

We believe the reviewer has misunderstood the motivation of our work. Our goal is not to provide a new architecture but to investigate the hypothesis that Transformers suffer from sub-optimal performance primarily due to insufficient latent variable inference. To systematically test this hypothesis - which is an open debate (Lines 19-52 RHS provide ample citations on both sides) - we use an architecture biased towards latent variable inference using a bottleneck (Lines 52-89 LHS). If Transformers' performance is linked to explicitly inferring task latents, then biasing the model towards this solution ought to improve OOD generalization. If it doesn't, we can conclude that explicit latent variable inference isn't sufficient to improve ICL. We do not suggest that our architectures, or bottlenecks, are the answer; they are simply minimal interventions to test the importance of correct latent variable inference in ICL.

Even though our goal is not to introduce new architectures, it is not a "rare view" in ML to improve generalization via bottlenecks for which ample evidence exists

• MoE architectures [2] use only a subset of parameters
• Perceiver models [3] introduce bottlenecks through learned latent variables
• Retrieval and memory augmented methods [4] introduce bottlenecks by selecting a subset of data for context
• Information bottleneck [5] are well studied for improving generalization
• Parametric assumptions (eg. training a neural network on data and then throwing away the data) introduces a bottleneck as opposed to methods like kNNs which retain entire datasets (note the bottleneck here is the trained model, which has strictly less information than the data it was trained on).

In our own study, explicit models with known prediction function — which is more bottlenecked than both the explicit and implicit model — outperforms both on various tasks. In light of these works, we believe that the story is more nuanced than introducing bottlenecks (or inductive biases) reduces performance.

transformer based LLMs work really well.

We agree with the reviewer that they do, but this should not disincentivize research into improving or analyzing them.

limited to settings with known latent variables

Our goal was never to provide a method for interpretability. To investigate the hypothesis that latent variable inference is not the problem, we need tasks where we can evaluate whether we are inferring the latents well. We rely on counterfactual predictions, a commonly used interpretability tool, as the metric to evaluate the extent to which the models infer task latents (hence the requirement for ground-truth latents). Our analysis highlights the difficulty of inferring task latents from implicit models, even though they perform well. This study thus contributes to a body of empirical evidence that allows us to conclude that improving task latent inference by itself is not the key to improved ICL generalization.

interpretability is most interesting when we don't know the latent variables.

We disagree because in such cases, a metric for what is more interpretable is either unavailable or there is a noisy, potentially incorrect, proxy for it which leads to mis-interpretations [6]. Thus, to rigorously test our hypotheses, we relied on tasks with known latent variables.

We hope that our detailed response has addressed the reviewer’s concerns and we would be happy to engage in further discussion to understand and resolve further questions. We hope that our response sufficiently highlights why the hypothesis we study is well motivated and interesting.

[1] Karl, F., Kemeter, L. M., Dax, G., & Sierak, P. (2024). Position: embracing negative results in machine learning. arXiv preprint arXiv:2406.03980.

[2] Liu, Aixin, et al. "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model." arXiv preprint arXiv:2405.04434 (2024).

[3] Jaegle, Andrew, et al. "Perceiver io: A general architecture for structured inputs & outputs." arXiv preprint arXiv:2107.14795 (2021).

[4] Lewis, Patrick, et al. "Retrieval-augmented generation for knowledge-intensive nlp tasks." Advances in neural information processing systems 33 (2020): 9459-9474.

[5] Tishby, Naftali, and Noga Zaslavsky. "Deep learning and the information bottleneck principle." 2015 ieee information theory workshop (itw). Ieee, 2015.

[6] Doshi-Velez, Finale, and Been Kim. "Towards a rigorous science of interpretable machine learning." arXiv preprint arXiv:1702.08608 (2017).

We thank the reviewer for acknowledging the value of our work and providing constructive cristicism. Throughout this comment, we will refer to additional experiments that are provided here: https://anonymous.4open.science/r/explicit-implicit-rebuttal-B263/explicit-implicit-rebuttal.pdf

Is there any difference in the result when using an encoder (as in the paper) vs. using a decoder?

The only difference is that a decoder would use a causal mask as opposed to our setting. Since we feed tokens as [$(x_1, y_1), (x_2, y_2), \ldots (x_n, y_n)$] where $(\cdot)$ defines a token, we cannot leverage a causal mask (refer to our response to Reviewer 1Akw). Our training loss, however, is an unbiased estimator of a similar next-token prediction decoder loss (modeling $y_{i+1}$ given [$(x_1, y_1), (x_2, y_2), \ldots (x_i, y_i), (x_{i+1}, \phi)$] for all $i$ in parallel) but more expressive since context points can attend in anti-causal direction as well. We chose this setup as it provides more supervision than feeding a token as either $x$ or $y$, and is a common choice for many related works [1,2].

Could you check the other OOD task for the regression problem, especially the ones that have OOD on the latent variable?

We thank the reviewer for this great suggestion and already conduct similar experiments with OOD latents in Figures 2(c) and 5, where we test for compositional generalization instead of OOD queries. In these tasks, novel combinations of underlying latents are provided during inference as opposed to solely changing the query (see Training and Evaluation on Page 5).

Inspired by the reviewer's suggestion, we conduct additional analysis where we test for OOD generalization in linear regression and classification by changing the distribution of the underlying weight vectors $w$ to be sampled from a normal distribution with larger variance. Our results, provided in Table 1 of additional experiments, indicate that it is not the case that explicit models are able to generalize better than implicit ones in such OOD scenarios. Even further, we see that while known prediction is a bit better, it still lags behind the implicit models because this setting is OOD for the context aggregator while maintaining in-distribution $x_{query}$.

Could you explain how many layers of Transformer are used in the implicit and explicit model

We use 8 layers in our experiments, where the explicit model splits prediction and context aggregation evenly (4 layers each). We conduct additional experiments where we compare an 8-layered implicit model to an explicit model with 8-layered context aggregation, and refer the reviewer to Figures 1 and 2 in additional experiments. Our results indicate that even with a larger context aggregation model, the same results hold.

I wonder if the effect of "dimension" in the output also matters

We point the reviewer to Figure 6 in the main paper which studies the role of dimensions. In general, the trend is consistent with increasing the complexity of the task, whether it be through the dimensionality of the input, the latents or the output. We also refer to the Figure 4 of additional experiments where we study the role of size of output dimensions in linear regression.

why is the classification task not beneficial from predicting the latent variable directly?

For the linear classification task, we believe that all the models have saturated their performance (note > 97%). For the nonlinear classification, when we fix the prediction function, there is an infinite set of latents that could lead to the same functional form but the space of the solution space is quite entangled and convoluted (note that this refers to permutation and scaling symmetries that leave the functional form unchanged). In contrast, it might be easier to explore alternate solutions by changing the prediction function to have a smoother landscape of possible latent variables.

We hope that our response has resolved the reviewers' concerns and would be happy to provide further clarifications.

[1] Hollmann, Noah, et al. "Tabpfn: A transformer that solves small tabular classification problems in a second." arXiv preprint arXiv:2207.01848 (2022).

[2] Müller, Samuel, et al. "Transformers can do bayesian inference." arXiv preprint arXiv:2112.10510 (2021).

We thank the reviewer for acknowledging the value of our work and providing constructive cristicism.

Additional References

We appreciate the reviewer bringing two relevant papers to our attention and will include them in the final version.

• Wang et al. (2023) shows that inferring latents ($\theta^d$ in their setup) is helpful but relies on a finite set of tasks and latents that are shared across different batches of observations. In contrast, our approach considers an uncountable set of tasks. They also assume data beyond the current context that belongs to the same task, while we only consider the current context as defining the latent.
• Han et al. (2024) show that performance on in-distribution (ID) tasks highly depends on how well the task latent $\theta$ is encoded. However, they also consider only finite set of tasks and do not focus on OOD generalization. In their language, we show that if the concept decoding part of the model doesn't have the proper inductive biases, it will only learn how to use the encoded concepts ID.

While we show that explicit models do not outperform implicit ones even though they correctly infer the latents, it could mean either the implicit model uses a different mechanism or infers the latents in a potentially distributed, uninterpretable way depending on the task. Answering this is an important future work, and the reviewer points out relevant works that aim to do so in certain settings.

why the prediction function is not trained well enough even though it appears like the latents are being modeled properly

Our hypothesis is that a learnable prediction function doesn't have appropriate inductive biases, either coming from its architecture or the training data. For example, a MLP prediction function can learn to be linear in the training regime, while being arbitrarily nonlinear outside, leading to suboptimal OOD performance. A similar argument can be made for other tasks. We will add a discussion about this in the draft.

algorithmically compare the learned prediction function to the optimal one

We refer to Figure 11 for an example of a learned prediction function evaluated away from the training distribution. In addition, we refer to Figure 4 (https://anonymous.4open.science/r/explicit-implicit-rebuttal-B263/explicit-implicit-rebuttal.pdf) which illustrates the performance of explicit models with MLP prediction OOD on the sinusoid task.

you mention you do not train with next-token prediction. Is there a reason this wouldn't work? Do you think this limits your results/claims in any way?

We could have trained our models using next-token prediction by feeding in tokens as [$(x_1), (y_1), (x_2), (y_2),\ldots (x_n), (y_n)$], where arguments inside $(\cdot)$ constitute a token. In this case, we could have done an augmented version of next-token prediction where we only consider losses corresponding to $x_i$ tokens (predicting the next query point is non-informative since iid samples). However, this requires the model to additionally learn which $y$ corresponds to which $x$. We instead feed [$(x_1, y_1), (x_2, y_2), \ldots (x_n, y_n)$] so the model does not need to learn that information, it is provided as input. However, this prohibits learning via next token prediction solely computationally and not algorithmically. In practice, our training loss is an unbiased estimator of predicting $y_{i+1}$ given [$(x_1, y_1), (x_2, y_2), \ldots (x_i, y_i), (x_{i+1}, \phi)$] for all $i$ in parallel, while being more expressive by allowing anti-causal communication within the context. This choice of modeling is common across a number of related works [1,2].

Additional insights into inductive biases to improve prediction function

We believe that there are a number of directions for inductive biases that could be worth pursuing. One direction is the architectural design of the prediction function, with or without task-specific knowledge baked in (e.g. using $\sin(\theta x)$ as the predictor for sinusoidal regression leads to perfect OOD generalization). Alternately, one could also look at optimization strategies (eg. alternate optimization instead of jointly optimizing, freezing one part of the network, etc.) that could lead to better convergence of the prediction function. As the reviewer rightly points out, next token prediction also provides an inductive bias towards this goal. We defer an in-depth analysis into the inductive biases as well as improving the prediction function as future work.

We hope that our response has resolved the reviewer's concerns and would be happy to provide further clarifications.

[1] Hollmann, Noah, et al. "Tabpfn: A transformer that solves small tabular classification problems in a second." arXiv preprint arXiv:2207.01848 (2022).

[2] Müller, Samuel, et al. "Transformers can do bayesian inference." arXiv preprint arXiv:2112.10510 (2021).