Does learning the right latent variables necessarily improve in-context learning?
Authors: Sarthak Mittal
Eric Elmoznino
Leo Gagnon
Sangnie Bhardwaj
Guillaume Lajoie
Dhanya Sridhar
Venue: ICML 2025 poster (2025)
Keywords: in-context learning
transformers
attention
latent variable
shortcuts
Primary area: general_machine_learning->transfer_multitask_and_metalearning
TL;DR: We test whether explicit latent-variable inference leads to better performance in in-context learning and identify the benefits and shortcomings of such methods.

Abstract: Large autoregressive models like Transformers can solve tasks through in-context learning (ICL) without learning new weights, suggesting avenues for efficiently solving new tasks. For many tasks, e.g., linear regression, the data factorizes: examples are independent given a task latent that generates the data, e.g., linear coefficients. While an optimal predictor leverages this factorization by inferring task latents, it is unclear if Transformers implicitly do so or instead exploit heuristics and statistical shortcuts through attention layers. In this paper, we systematically investigate the effect of explicitly inferring task latents by minimally modifying the Transformer architecture with a bottleneck to prevent shortcuts and incentivize structured solutions. We compare it against standard Transformers across various ICL tasks and find that contrary to intuition and recent works, there is little discernible difference between the two; biasing towards task-relevant latent variables does not lead to better out-of-distribution performance, in general. Curiously, we find that while the bottleneck effectively learns to extract latent task variables from context, downstream processing struggles to utilize them for robust prediction. Our study highlights the intrinsic limitations of Transformers in achieving structured ICL solutions that generalize, and shows that while inferring the right latents aids interpretability, it is not sufficient to alleviate this problem.

Does learning the right latent variables necessarily improve in-context learning?
Sarthak Mittal
Eric Elmoznino 122 Leo Gagnon 12 Sangnie Bhardwaj 123 Guillaume Lajoie 122
Dhanya Sridhar
Abstract
ence through in-context learning (ICL, Brown et al., 2020). Indeed, LLMs can learn from instructions and demonstra- tions provided as input, without requiring gradient- based learning. While ICL plays a key role in many L.M abilities (Lu et al., 2024), such as cdtructtionfollowiing (Wei et al., 2022) and chain-of-thought reasoning (Wei et al., 2023), the factors that influence its generalization particularly in out- of-distribution (OOD) settings remain poorly understood. Although ICL can leverage combination of instructions and demonstrations, our analysis focuses specifically on its ability to model predictions based on a task's training examples (demonstrations) provided in-context (Lampinen et al., 2024) beyond the modality of language.
Large autoregressive models like Transformers can solve tasks through in-context learning (ICL) without learning new weights, suggesting avenues for efficiently solving new tasks. For many tasks, e.g., linear regression, the data factorizes; exam- ples are independent given a task latent that gen- crates the data, e.g., linear coefficients. While an optimal predictor leverages this factorization by inferring task latents, is unclear if Transformers implicitly do so or instead exploit heuristics and statistical shorteuts through attention layers. In this paper, systematically investigate the effect of explicitly inferring task latents by minimally modifying the Transformer architecture with a bottleneck to prevent shortcuts and incentivize structured solutions. We compare it against stan- dard Transformers across various ICL tasks and find that contrary to intuition and recent works, there is little discernible difference between the two; biasing towards task-relevant latent variables does not lead to better -fittri-iitiiittionnn perfor- mance, in general. Curiously, we find that while the bottleneck effectively learns to extract latent task variables from context, downstream process- ing struggles utilize them for robust prediction. Our study highlights the intrinsic limitations of Transformers in achieving structured ICL solu- tions that generalize, and shows that while infer- ring the right latents aids interpretability, it is not sufficient to alleviate this problem.
A plausible hypothesis behind the success of ICL is that since many tasks are based on some low-dimensional latents (e.g.. complex games are described completely through their rules, linear regression through its underlying coefficients), Transformers might generalize to novel queries by inferring the task latents from the context (Hendel et al., 2023; Todd et al., 2024; Yang et al., 2025). This describes a parametric (Bishop, 2006) modeling approach that breaks the prediction mechanism into parts: 1) inferring the task latents (i.e. parameters) from the context, and then 2) using them to make predictions on novel queries. With the right prediction function, such an approach leads to systematic OOD generalization to new queries.
However, mounting evidence (Wang et al., 2023; Han et al., 2023; Song et al., 2024: Tang al., 2023b; Bhaskar et al., 2024; Crosbie & Shutova, 2024) suggests that Transformers instead often employ more shallow solutions which rely on direct comparison of the query to demonstrations (reminiscent of induction heads, Olsson et al., 2022). This is closer in spirit to paraamemttti approaches (such as nearest neighbors or kernel regression) which are known for their flexibility but poor generalization (Bishop. 2006). Since these solutions do not model the data generative process, they can be described as statistical shorteuts and risk poor performance on OOD context and queries e.g., learning the actual linear predictor for linear regression can generalize to any distribution over training and test points, but nearest-neighhour based interpolation might not. Interestingly, the functional form of attention operations is almost identical to that of kernel regression (Tsai et al.,
1. Introduction
Recent advancements in large language models (LLMs, Rad- ford et al., 2019) showease the Transformer architecture's s (Vaswani et al., 2017) ability to perform novel tasks at infer-
'Equal Contribution, Equal Supervision 'Mila Quebec AI Institute Université de Montréal Google DeepMind. Correspon- dence to: Sarthak Mittal <sarthmit@ gmail.com>, Guillaume La- joie <guillaume lajoie @ queb..ue
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Docs learning the right latent variables necessarily improve in-context learning?
Implicit model
Explicit model
po(y>|x* D)
z (D)
PO
Transformer
y+
Context Transformer
Prediction function
yo
Expressivity
x
Explicit latents
Dimensionality bottleneck
x1 Y1
x1 Y1
Ayp
Y'
Yn
X+
Performance
Interpretability
(xe, yo))
(x0w, yoo)
Figure 1. We compare the henefits of the implicit (left) and the explicit (right) model. Explicit models disentangle context aggregation and prediction into separate functions with inductive bias for inferring generative latent variables in order to solve the task. Implicit models are more expressive, but can learn non-parametric shortcut solutions that bypass latent variable inference.
2019; Han et al., 2023), making such solutions more natural for Transformers to express (Zhou al., 2023).
using a simple bottleneck a Transformer can improve interpretability and explicitly extract task-relevant latent variables, it also suggests that the limitations of Trans- formers in learning more structured and generalizable ICL solutions are not solely due to con-parametric shortcuts that skirt latent variable inference, but due to more fundamental architectural limitations. In sum, our contributions are:
In this paper, aim to test the hypothesis that biasing Transformers against non-paran aan-aarmmttri solutions improve their ICL performance by encouraging parametric modeling. We minimally modify the Transformer architecture to pre- vent such non-parametric shortcuts and compare the OOD performance of the resulting model to that of a traditional Transformer on large array of ICL tasks. We call this al-
Formalizing a framework to test whether parametric ICL solutions generalize better -uitt--diitributioo.
tered architecture an explicit model because of its inductive
Analyzing the benefits, lack thereof, of inferring the true latents explicitly.
bias of explicitly extracting structured latent variables to
solve the tasks, and call the traditional Transformer architec-
Highlighting flaws in the prediction function and down- stream latent utilization which hinders generalization.
ture an implicit model. Specifically, the explicit model pre- vents the query from directly attending to demonstrations in the context by introducing a bottleneck between the process- ing of the context and the query (see Figure 1), similar conditional neural process (Garnelo et al., 2018a). To study the impacts of this inductive bias favoring parametric solu- tions, we need to establish that explicit models successfully recover task latents. As such, we consider synthetic and real tasks for which the latent mechanisms are well understood, and systematically analyse the impact of task latents on generalization by comparing explicit and implicit models.
2. Notation
Throughout the paper, we denote datasets with the symbol D which consists of set of observations with inputs de- noted via E X' and their corresponding outputs as y E y. A task is defined by functional mapping g : X 2 y which maps observations to labels through some latents or parameters 2, eg. y = zT x for linear regression task, or y ~ Ni:: =T xx o²) for its stochastic counterpart. To ease readability, we will reserve x. E X for the query, i.e. the test time observation we want to generalize to, and y. E y its corresponding target. Finally, & denotes the parameters of the context aggregation component of explicit model, which inputs the dataset D and infers the corresponding parameters zy (D), and y the parameters of the prediction model which given a query T. and parameters E Z, pro- vides the prediction. For the implicit model, these operations are subsumed into a single model, with parameters yp
We find that the explicit model does not outperform the implicit OOD data, challenging the aforementioned hypothesis that avoiding non-pann con-aaraettric solutions enhances generalization. Our investigation into this lack of improve- ment reveals that the issue often lies in the explicit model's prediction function which has to leverage the inferred latent variables for downstream predictions on the query. Our controlled experiments and analysis on the interpretable nature of the bottleneck revealed strong evidence that while the explicit model often extracts relevant task latents, these are not properly utilized by the prediction function.
3. Implicit vs. Explicit Inference
We look at ICL the context of algorithmic problems where the task is predict the target y, from query x. when provided with some context examples D {(0,, 1:
While on one hand, our research demonstrates that
2

---

Does learning the right latent variables necessarily improve in-context learning?
a.
Synthetic regression tasks
Linear
Nonlinear (MLP)
Gaussian Process
Sinusoidal
Hodgkin-Huxley
03
0.2
10"
0.4
MSE
0.2
10"
0.1
0.1
02
10"
0.0
0.0
ID
GOD
OOD
000
ID
GOD
OOD
b.
Synthetic classification tasks
C.
Compositional tasks
Linear
Non Inear (MLP)
Raven's PM
Alchemy
Gene targeting
1.00
1.00
0.53
Accuracy (
Accuracy (f)
0.98
0.82
0.98
0.98
0.81
0.95
0.92
0.96
0.98
0.80
DD
0OD
OOD
000
Impliot
Expliott
x
Transfommer
Transformer
Transtomer
4
Figure 2. Comparison of implicit and explicit models in-distribution (ID) and out-of-distriuution (OOD) across various domains: (a) synthetic regression, (b) classification, and (c) compositional generalization tasks. Implicit models are in shown gray, explicit models with Transformer prediction in blue, and with MLP prediction in orange. Further details about tasks is provided in Appendix B.
sharing a common underlying structure defined by the task latent z and a functional form g. The goal of ICL to learn function that can utilize the context set D to provide predictions for new queries T... To achieve this, the model is trained on different draws of context sets (D1, ...) which share the same underlying functional mapping 9 x, z y but different realizations of the latent z, for example =( = z"'s x could be a linear regression system shared across different contexts D1, D2, but the underlying latents could be different, i.c. D1 is generated from 21 while D: from 22, similar to Von Oswald et al. (2023). We emphasize that this setup, we are not training models to do next-token prediction is done in language modeling; instead, given fixed context D that includes samples, we are attempting to make a prediction on a single novel query x., We therefore do not use a causal Transformer, and we allow all tokens to attend to each other.
Context Aggregation. This component deals with inferring the task-dependent latent variables from the in-context examples such that the downstream prediction becomes conditionally independent of the context, i.e. inferring from D such that p(y- x., 2, D) p(y. (2,2))
Predictive Modeling. This component refers to the process of estimating the predictive function that leverages context 'D to infer y. from a query x.. In the above example, it refers to learning the functional mapping g once z has been extracted from context aggregation.
As discussed, Transformers do not have clear incentive to make this explicit separation of context aggregation and predictive modeling. Instead, given context D, they implic- itly and jointly model both the function g along with D- dependent latent variable z inference to directly provide pre- dictions for the query x.. in contrast to separately estimating 9 and explicitly factorizing z. Thus, in order to enforce ex- plicit representation of 2, we propose a simple architectural modification where the query ac, cannot directly attend to the context, and the latent task representation is forced to summarize the context efficiently. Formally, we compare the following two models, which are illustrated in Figure 1.
Often, ICL solutions are learned via maximum likelihood:
argmax max ED....y (1o |ac., D)]
(1)
where Py represents the Transformer model and D is sam- pled from the parametric family defined through g. Thus, the transformer model Psp must not only learn the form of the prediction function g. but also how to efficiently aggregate information from the context D to infer z for downstream predictions on I.. Thus, this general framework can be naturally decomposed into two distinct parts.
Implicit Model. This refers to the traditional in-context learning computation performed by Transformer models. In this setup, given the set of observations D (context) and query x., the prediction y. is modeled directly
3

---

Does learning the right latent variables necessarily improve in-context learning?
Linear reg.
Nonlinear reg. (MLP)
Sinusoidal reg.
Linear cls.
Nonlinear cls. (MLP)
Raven's PM
0.20
Accuracy (?)
1.00
1.00
1,00
03
0.15
0.95
0.96
0.98
MSE
02
0.10
01
96
0.96
0.00
000
0OD
00D
OOD
OOD
0OD
Known
Implicit
Expliott
D Transformrr
Y4
D
Thanstomer
Tranaformer
Known
Xig
X
Figure 3. Performance on subset of tasks where the true latents and prediction function g are known. Implicit models are in shown gray, explicit models with Transformer prediction in blue, models trained with an auxiliary loss predict the true latents in purple and those using the true prediction function in green. Using the known prediction function leads to significantly better OOD performance.
as Pe (y- (x., D), where Pe is defined using a standard Transformer with parameters yp and is tasked with modeling both context aggregation and predictive modeling.
model is parametric and low-dimensional, but in case of a non- pooa-arrmetric very high dimensional parametric model, the implicit model should be better.
Explicit Model. This represents the architectural variation which minimally modifies the Transformer architecture by separating context aggregation and predictive modeling. It first constructs a task representation z (D) using the set of observations D and context model z with parameters b (context aggregation) and another network py to make prediction for a new point F, (predictive modeling) as z (D)). A key insight is that the task latents are invariant to the queries when modeling prediction. The context model is implemented with a Transformer zy with weights v and for the prediction function py, we study both Transformers and MLPs with weights 7. Importantly, the output of the context model z = (D) is a fixed-size vector with much lower dimensionality than the full context D. This information bottleneck prevents the query x, from attending directly to the context as in standard Transform- ers; instead, the context model must summarize D into underlying generative factors, thus ruling out potential shortcut solutions that bypass latent variable inference.
Finally, we note that our aim is not to construct the best possible explicit model architecture indeed, more sophis- ticated ones already exist based on amortized Bayesian inference (Garnelo et al., 2018h; Mittal et al., 2023). Instead, we are interested in investigating potential inductive biases for ICL by minimally modifying the standard Transformer architecture to remove certain shortcuts from the space of possible solutions. We leave the design of more performant architectures for future work and refer the readers to Appendix A for a detailed discussion of related work.
Py Py- [z., (y.|
4. Experiments
Our goal is to use both synthetic and real tasks that capture the key elements of ICL applications tease apart the ef- feets of implicit and explicit models both in-distribution (ID) and out-of-distribun
(OOD) generalization.
Task Setup. We conduct experiments across a variety of settings that admit task latents, from synthetic regression and classification to reasoning problems. For reasoning tasks that require compositional knowledge, we consider Raven's Progressive Matrices (Raven's 's PM) (John & Raven, 2003), Alchemy (Wang et al., 2021), Gene Targeting (Nor- man et al., 2019) and reusable mixture of experts. A brief description of our tasks is provided below, with a more thorough explanation in Appendix B.
Implicit vs. Explicit. Assuming Transformers do in fact favour shortcut-based solutions, we first hypothesize when each setup should perform better given different task characteristics. If the data is generated with a parametric model with few underlying parameters z (e.g. a linear model y = z" 2), the right predictor can be precisely described using the parameters 2, making the explicit model better suited. In contrast, if the data is generated with Gaussian Process (GP), the implicit model should be superior since by construction query prediction computes similarities with all points in the context. In this case, the task latents of Ga-based data with RBF kernel are infinite dimensional (i.e. a point in function space), which cannot be captured in the finiteddimensional bottleneck of the explicit model. In general, should expect the explicit model to be superior when the underlying true
Regression Tasks. We consider different data-generating pro- cesses, e.g., linear: prediction z"a and latents z, nonlinear (MIP)) prediction with neural network g(x,, z) and latents as its weights, sinusoidal: prediction as sum of sinusoids with different frequencies and the latents as its amplitudes.
Classification Tasks. Akin to the regression problems, we consider a linear and nonlinear (MLP) prediction for classifi- cation using an additional sigmoid activation on the output.
4

---

Docs learning the right latent variables necessarily improve in-context learning?
a.
Decoding the true latent Z
Linear reg.
Nonlinear reg. (MLP)
Sinusoidal reg.
Linear cls.
Nonlinear cls. (MLP)
0.00100
0.8
0.8
0.8
0.00075
0.004
0.4
0.6
W MSE
000000
0,002 02
0.4
0.2
0.00225
02
0 000000
0.0
0.000
0.0
0.0
Raven's PM
Gene targeting
Rule accuracy (T)
1:00
Target accuracy (T)
b.
Counterfactual interventions
0.15
0.75
Alchemy
0.50
0.10
1.0
Relative accuracy
0.25
0.05
0.5
0.00
0.00
Known
0.0
Implictt
Explict
Rempee
Graph
Potion map
Stone map
Tresaoommrr
3ig
Trerafommer
MLP
Latent component
xg
Figure Explicit models are interpretable as the bottleneck allows (a) linearly decoding the true latent, and (b) intervening to obtain correct counterfactual predictions. Implicit models shown in gray, explicit models with Transformer prediction in blue, and with MLP prediction in orange. Models green use the true prediction function g, while models in purple use an additional auxiliary loss based on true latents. To evaluate decoding performance in (a), we linearly decode the true latent directly from concatenated context examples, which yields significantly worse performance than decoding from the bottleneck. Baseline performances in units of the plots are Linear regression: 0.49, Nonlinear regression (MLP): 0.94, Sinusoid regression: 0.33, Linear classification: 0.86, Nonlinear classification (MLP): 0.97, Raven's 's PM: 0.5, and Gene targeting: 0.0. In (b), the "Relative accuracy" takes the baseline in account (details in C.4).
Raven's 's Progressive Matrices. A pattern-completin used in IQ tests where individual object attributes evolve according to different rules. The task is to complete sequence of frames satisfy the underlying rule, which is the latent variable and can be based on colors, shapes, etc.
task
the mean-squared error loss, while those based on classifi- cation use the cross entropy loss for training. Model archi- tecture details provided in subsection C. For ID evalu- ation, consider the underlying task latent x, context sam- ples D, and queries 2. to be sampled from the same distribu- tion as during training. The challenge in this case is simply to generalize from finite data. For OOD evaluation, we test two different cases depending on the kind of task. For our synthetic regression and classification tasks, the task latent and context samples D are sampled from the same distribu- tion as at training time, but the queries æ, sampled from Gaussian distribution with higher (3x) standard deviation, thus testing a form of out-of-domain generalization. For our reasoning- bbasd problems, we evaluate on task latents z that were not seen at training. The task latent in each of these reasoning- based problems is composed of parts (e.g., the Gene Targeting experiment, the latent is binary vector of targeted genes), which allows to test a form of composi- tional generalization (Wiedemer et al., 2023) in which all parts have full marginal support at training time, but novel combinations of those parts are evaluated at test-time.
Alchemy. Here, a latent causal graph describes how different stones and potions interact to generate new stones. The task infer the next stone given some transitions.
Gene Targeting. It represents real-world dataset of targeted gene knockouts and subsequent observations of gene expres- sions across many cells. The underlying latent variable is the set of genes that were knocked out in given experiment, and the task is to infer the gene expressions certain cells in an experiment given those of other observed ones.
Reusable Modular Mixture of Experts (MoE). In this task, apply sequence of operations gi on the input æ, where the choice of expert applied at layer is governed by cat- egorical latent z1 In particular, we apply five operations sequentially, leading to = 9xs Is. o... IR, (æ). Here 91. ... are shared across tasks and each task is uniquely identified by (2i...., 25) which are the task latents. This represents a reusable mixture of experts system with a hier- archical and compositional decomposition.
For all tasks, implicit and explicit models were trained from scratch over distribution of tasks latents, with control for the number of parameters to provide systematic comparison between the two. To better understand the shortcomings of different models, we also compare with
Training and Evaluation. Tasks based on regression utilize
5

---

Docs learning the right latent variables necessarily improve in-context learning?
normal distribution with wider standard deviation than was used at training (3x), effectively evaluating if the models could extrapolate to points further out along their domain despite only being trained within narrow distribution near the origin. For the compositional tasks in Figure 2 (c), we instead evaluated OOD performance by only training on certain configurations of the true latent variable z while evaluating on unseen ones. Importantly, at training time the models were shown data that included every possible value for each component of 2, but not every possible combination of these values were seen, thus evaluating a form of compositional generalization (Wiedemer et al., 2023). Similarly, for the reusable modular MoE task, we train the models on fraction of all possible combinations in the latent (z1....,=s) =s) and evaluate on all combinations.
Impict Explidt Transfomerr Expliit MLP
0.4
0.12
0.10
MSE (4)
0.08
0.06
0.04
Surprisingly, and counter to our predictions above, we found that the explicit model provided no significant benefit in OOD settings. In synthetic regression tasks shown in Figure 2 (a), all models failed to generalize and obtained substantially worse performance than during ID evaluation, with the implicit model actually being the one that had slightly lower degradation in performance. In classification and compositional tasks shown in Figure 2 (b, c), all models generalized fairly well OOD and with similar performance. Our results on reusable modular MoE task in Figure 5 further indicate that implicit models consistently outperform explicit ones across different proportion of latent combinations seen during training (X-axis). In summary, explicit models appear to provide no benefit across our tasks, both in terms of ID and OOD performance.
0.02
0.0
0.2
0.4
0.6
0.8 1.0
Ratio Permutations Seen
Figure 5. We conduct experiments on reusable modular MoE task where we train on a subset of combinations of experts, shown on the X-axis. Our results indicate that across different percentages of combinations seen during training, the implicit model consistently outperforms the explicit one at compositional generalization.
privileged oracles (known decoder using ground-truth g function, and known latent variable- using an auxiliary loss that includes the ground-truth 2).
Explicit models do not outperform the implicit models. The first evaluation setting that we considered was the ID performance. In this case, we should expect both implicit and explicit models to perform equally well, even if implicit models learn shorteuts rather than ground-truth task latents. This is because those shortcuts are tuned to minimize prediction error within the same data distribution that is being evaluated. Across all our tasks, the results indeed confirmed this prediction. Specifically, during ID evaluation, all models were capable of making highly accurate predictions (Figure 2). While the performance of the implicit model was generally slightly better than that of the explicit models, the benefits were marginal. Effectively, all models solved the tasks similarly well.
If the explicit model did learn the right latent variables in the bottleneck, it essentially implies that either the implicit model learns benign shortcuts (if at all) or that learning the right latent variables is not sufficient to improve generaliza- tion, both ID or OOD. In the following results, we see that the explicit model does indeed learn the right task latents.
Explicit models learn to infer the correct latent variable, but not how to use it. Why didn't the explicit model provide any benefit? Our initial hypothesis was that the im- plicit model could be susceptible to learning shortcuts that are sufficient to reduce the training loss and easy to express using self-attention between the query x. and context D. By summarizing the context in bottleneck zy., the explicit model should instead be forced to extract the true latent variable in order to minimize the training loss, thus learning solution that is closer to the actual data-generating process. There are then two possible explanations for the results in Figure 2: (1) the explicit models did not learn to extract the true latent variable despite inductive biases to do so induced by the bottleneck, or (2) they did extract the true latent variable but did not learn to use it in the correct way. We performed several experiments to test these two
While we expected that ID evaluation would be insufficient to demonstrate potential benefits of the explicit models, we expected to see differences in OOD settings. Both implicit and explicit models are sufficiently expressive to fit the training distribution. However, if an explicit model suc- cessfully learns the true task latents that generated the data while an implicit model learns statistical shortcuts that are specialized to the training distribution, we should expect the explicit model to generalize better OOD. As a reminder, for the synthetic regression and classification tasks in Figure 2 (a, b), OOD evaluation was done by sampling x. from a
6

---

Does learning the right latent variables necessarily improve in-context learning?
a.
Task performance (OOD)
2000
600
600
1500
300
400
400
MSE
300
200
200
100
100
6x
2x
6x100
2x10
Data Dimensionality
Context Length
Number of Parameters (Mi lions]
b.
Decoding the true latent z
0.6
0.5
0.5
Latent MSE
0.4
0.4
0.3
0.2
0.2
02
0.1
6x:10
2=
6x101
2 10
10'
Data Dimensionality
Context Length
Number of Parameters (Millions)
Inp
Expict
Turatoree
Trasffrme
Krenn
Figure 6. We analyze (a) Linear regression OOD performance and (b) latent variable linear decoding as a function of model and task parameters. Task performance scales similarly for implicit (gray) and explicit models (blue), while latent variable decoding scales similarly for the explicit model and models trained with the true prediction function g (green). possibilities, and found strong evidence for the second.
model was linearly projected to the same dimensionality as the true weights z, after which the prediction function took its dot product with queried input æ.. In this setting, if the explicit model extracts the correct latent variable, it should generalize perfectly both in and out of distribution. Our results in Figure 3 confirm that using the correct prediction function indeed provides substantially better OOD general- ization and effectively solves most tasks. This finding has significant implications: it suggests that while learning the true task latents may be necessary condition for gener- alization, this must also be supplemented with significant inductive biases in the prediction function for instance, through novel architectures - so that these task latents can be leveraged correctly. In the absence of such inductive biases, inferring the correct task latent appears to provide no benefits in practice. We do note that our nonlinear regres- sion tasks, where represents the weights of underlying MLP generating the data, were an exception to the results described here in that using an oracle prediction function performed poorly. In this case, we conjecture that the under- lying latent variable is too difficult to accurately infer from the context, while shorteut-based solutions would avoid la- tent variable inference altogether to provide robust solutions.
To test whether or not the explicit models failed because they did not extract the correct latent variable, we added an additional supervised loss term to Equation 1, || - W I=?, to force the explicit model to recover the true task latent z up to a linear transformation, where W is a learnable param- eter. The loss of this linear model was then backpropagated through the context model along with the task loss. Results in Figure 3 (purple) show that this auxiliary loss provided no improvement apart from minor increases in accuracy on Raven's ' PM, suggesting that incorrect latent variable inference does not explain the explicit model's 's suboptimal performance. Indeed, when we did not use the auxiliary loss as a training signal for the explicit model and just evaluated the quality of the learned task latents by training a separate linear predictor predict Zg., we found that could still accurately predict the true latent variable (see Figure 4 (a)). This means that in the absence of any explicit training signal to predict the true task latent, the explicit model accurately learns task latents up to linear reparameterizations..
Given that the explicit model correctly infers the true latent variable in its bottleneck, we study whether the prediction function is suboptimally learned. In other words, despite the explicit model having access to the true x, we hypothesize that P.y (v- F,, 2 does not learn the true lata-generating process y. = g( )), where is the true prediction func- tion- e.g., for linear regression g(x,, x) = xT x.. To validate this hypothesis, we trained explicit models with the prediction function g hard-coded as an oracle. For instance, in the linear regression task, the zn output by the context
Explicit models are easily interpretable. While explicit models do not provide performance benefits, the ability to extract the correct latent and summarize it in single bottleneck can still be useful for interpretability. On tasks with known underlying latent variables, we were able to linearly decode them from 2Φ with high accuracy, meaning that the information is not only present but also easily accessible (Figure 4 (a)). The exceptions were complex
7

---

Does learning the right latent variables necessarily improve in-context learning?
Queries sampled from Normal dist.
Queries sampled near context
0.100
MSE (1)
0.03
0.075
0.02
0.050
0.01
0.00
0.000
OOD (far)
O0D (fer)
Known 2
Implicit
Explicit
X4
Transtomer
Transtormer
Xig
Transtormer
MIP
Known
X4
Figure 7. We analyze sinusoid regression by sampling the query either from the normal distribution (leff) or close the context (right) during training hut always far from context during evaluation. Wc see that when queries are sampled close to the context, implicit models which can rely more on kemnel-regressonn based nearest- neighbtr solutions don't 't' generalize far from context, while explicit models do.
nonlinear regression tasks where the latents are difficult to infer and classification tasks where more context samples are needed to precisely identify the decision boundary. In contrast, finding such clear task-relevant latents is immensely challenging in an ordinary Transformer trained to do ICL, given that they can be distributed across mixture of many layers and token positions.
0.40
0.35
0.30
0.25
0.20
0.15
Furthermore, even when latent variables appear to be suc- cessfully identified using a linear decoder in some hidden layer of a Transformer, one often finds that the relationships merely amount to spurious correlations (Ravichander et al., 2021). For instance, one manually modifies the activations in this hidden layer such that different latent variable is predicted by the linear decoder, the model's predictions do not generally change in way that is "counterfactually" con- sistent with this new latent (i.e., the prediction is not what it should have been under the new latent variable). We there- fore used the Distributed Alignment Search (DAS) method from Geiger et al. 2023b (see subsection C.4) to search for units in the implicit and explicit models that can be manip- ulated obtain correct counterfactual predictions. For the explicit model, we limited this search to the bottleneck zy- We found that using the explicit model, we were indeed able to manipulate Zy and obtain correct counterfactual predic- tions, but we were not able to successfully do this using the implicit model, as shown in Figure 4 (b). Explicit models might therefore be useful for both mechanistic interpretabil- ity and scientific discovery (Geiger et al., 2023a), where we do not know the underlying task latents and want to be able to easily uncover them from the trained model, and subse- quently obtain a good predictor for an intervened system zero-shot given some knowledge about the intervention.
0.0
Mode
AG2 decoded
0.00
Figure 8. We compare using the auxiliary ground truuhd-truth task latent loss directly on the output of context aggregation, i.e. zy (D) (aux direct), linear decoding from it (aux decoding).
similarly, but that the implicit model reliably outperformed the explicit one unless used the known prediction function g. We also looked at the latent variable linear decoding accuracy in the explicit model and model properties (Figure 6 (b)). As expected, we found that the latent variable was easier to decode from the explicit model's 's bottleneck when there was less inherit uncertainty about its value (lower data dimensionality, longer context length) and when the explicit model was given more capacity. However, throughout the different settings, we see that while the explicit model does learn the true latent well, it is not sufficient to get performance boost over the implicit models. Further details on the setup of these scaling experiments is provided in subsection C.2.
function of these task
Impact of auxiliary loss on decoding from bottleneck. Additionally, we perform an experiment where instead of using an auxiliary loss obtained between the ground-truth task latents z and a decoding from the explicit model's bottleneck || - W I22 (called aux aux decoded), we instead force the bottleneck itself to be directly close to the ground-truth || - 1= (called aux direct). Since the prediction function relies on the bottleneck and not its decoding, removing this extra layer when providing additional supervision might allow the bottleneck to better
Scaling Trends across Different Properties. To better compare the implicit and explicit models, we investigated their OOD task performance on linear regression as we var- ied the different properties of the task (input dimensionality and context length) and the size of the model used (Figure 6 (a)). We found that performance in both models scaled
8

---

Docs learning the right latent variables necessarily improve in-context learning?
reflect the task latents and thereby aid prediction. Our results, however, indicate that doing does not lead any benefits OOD evaluation for linear regression, further strengthening the conclusion that effective task latent inference is not the biggest problem in such models.
not substantially aid generalization, they do come with in- terpretability benefits which are extremely useful when de- ploying Al systems at scale. Finally, we hope that our analysis sparks controlled experiments to understand the mechanisms behind in-context learning better as well as in- corporating and validating numerous inductive biases to see whether they do aid generalization and reduce the reliance on shortcuts.
Extreme Shorteut Injection. Finally, we test whether in- jecting extreme shortcuts during training pushes implicit models learn nearest-neighbor styled kemeelreggession so- lutions opposed to uncovering the underlying functional form. We consider two cases, where queries during training are sampled (a) randomly. or (b) near context points. The latter further incentivizes implicit models to learn nearest neighbour shortcuts. At evaluation, the queries are sampled far from the context. Our results on sinusoid regression in Figure 7 indicate that while implicit models perform well generally, they suffer considerably more in the presence of such injected shortcuts since explicit models distill the task latent from context independent of the query.
A Ackonwledgements
The authors would like to acknowledge the following re- searchers for valuable discussions and exchanges: Joao Sacramento, Johannes von Oswald. All authors acknowl- edge support from an unrestricted gift from Google inc. EE acknowledges support from Vanier Canada Graduate Schol- arship #492702. SM acknowledges the support of PhD Excellence Scholarship from UNIQUE. DS acknowledges support from NSERC Discovery Grant RGPIN-2023-04669, and Canada-CIFAR Al Chair. GL acknowledges sup- port from NSERC Discovery Grant RGPIN-2018-0444 2018-0422,, the Canada Research Chair in Neural Computations and Interfacing, and Canada-CIFAR AI Chair.
We further refer to Appendix D for a detailed analysis.
5. Conclusion
A commonly believed hypothesis is that Transformers do ICL through brittle statistical shorteuts rather than by in- ferring the underlying generative latent variables of the task, and that this explains their inability to generalize outside of the training distribution. Here, we empirically tested this hypothesis by minimally modifying the Transformer architecture through the use of a bottleneck that factorized the model into separate context aggregation and prediction functions, creating an inductive bias for explicit latent variable inference. While we confirmed that this model indeed learned to infer the correct latent variables across many ICL tasks, it surprisingly gave improvement in performance for either in-distribution out-of-distribution evaluation. Contrary to common belief, then, we showed that simply learning the correct latent variables for the tasks not sufficient for better generalization because end-to-end optimization does not learn the right prediction model to leverage these latent variables.
Impact Statement
This paper provides a comparative analysis to better understand the capabilities current in-context learners from a point of view of making them more robust and aligned with the true underlying models of the data. We believe that this is an important step towards understanding when and how machine learning models can rely shortcuts and spurious correlations, and understanding whether such correlations can be mitigated through the use of conditional independence assumptions and bottlenecks, as is investigated in this work.
While we show negative result that such bottlenecks do

---

Docs learning the right latent variables necessarily improve in-context learning?
References
Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F, Rezende, J., Eslami, S. M. A., and Teh, Y. W. Neural processes, 2018b
Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Architectures and algorithms, 2024.
Geffner, T., Papamakarios, G., and Mnih, A. Com- positional score modeling for simulation-baset infer- ence. In Krause, A., Brunskill, E., Cho, K., Engel- hardt, B., Sabato, S., and Scarlett, J. (eds.), Proceed- ings of the 40th International Conference on Machine Learning, volume 202 of Proceedings Machine Learn- ing Research, pp. 11098-11116, PMLR, 23-29 Jul 2023. URL https::: http:///pocccedeeeeeeee mlr press/ p v202 2002/geffe /202/gefffer223
Alain, G. and Bengio, Y. Understanding intermediate layers using linear classifier probes, 2018.
Beck, J.. Vuorio, R., Liu, E. Z., Xiong, Z., Zintgraf, L., Finn, C., and Whiteson, S. survey of meta-reinforcemmnt learning. 2023.
Bhaskar, A., Friedman, D., and Chen, D. The heuristic core: Understanding subnetwork generalization in pretrained language models, 2024. URL https://///// /tl:://arraa //txxiaxvvv..//// org/ abs/2403. 03942.
html.
Geiger, A., Potts; C., and leard, T. Causal abstraction for faithful model interpretation, 2023a.
Bishop, C. chine Learning. Springer, January 2006. URL
Pattern Recognition and Ma-
Geiger, A., Wu, Z.., Potts, C., Icard, T., and Goodman, N. D. Finding alignments between interpretable causal variables and distributed neural representations, 2023b.
https://www.icrrsset ar-us//reeerchh/publlaaaiio pettern--econttiooommmcchheeeeee
Geiger, A., Wu, Z., Potts, C., Ieard, T., and Goodman, N. D. Finding alignments between interpretable causal variables and distributed neural representations, 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Bren- del, W., Bethge. M., and Wichmann, F. A. Shortcut learning in deep neural networks. Nature Machine Intelli- gence, 2(11):665-67) November 2020. ISSN 2522-5839. doi: 10. 10.10088/042566-000002552--- URL htw.::tttpp// dx. doi org /10 100.0.113//222222 020 00257-2
1877-1901, 2020.
Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F. Data distri- butional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878-1888 1,
Gilbert, A., Horlbeck, M. A., Adamson, B., Villalta, E., Chen, Y., Whitehead, E. H., Guimaraes, C.,, Panning, B., Ploegh, H. L., Bassik, M. C., et al. Genome-scale crispr- mediated control of gene repression and activation. Cell, 159(3):647-661, 2014.
Crosbie, J. and Shutova, E. Induction heads as an essen- tial mechanism for pattern matching in in-context learn- ing, 2024. URL https: /////////ii/ org/ 2407. 07011.
Gordon, J., Bruinsma, W. P., Foong, A. Y., Requeima, J., Dubois, Y., and Turner, E. Convolutional conditional neural processes, arXiv preprint 1rXiv:1910.135:. 2019.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hermandez, D., Jones, A., Kermion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. https:/transsformer Sccuitspub/202//ftameeorr///
Grau-Moya, J., Genewein, T., Hutter, M. Orseau, L., Delétang, G., Catt, E., Ruoss, A., Wenliang, L. K., Mat- tern, C., Aitchison, M., and Veness, J. Learning universal predictors, 2024.
Guo, Y., Hao, Y., Zhang, R., Zhou, E., Du, Z., Zhang, X., Song, X., Wen, Y., Zhao, Y., Zhou, X., Guo, J., Yi, Q.. Peng. S., Huang, D., Chen, R., Guo, Q. and Chen, Y.
Garg, S., Tsipras, D., Liang, P., and Valiant, G. What can transformers learn in-context? case study of simple function classes, 2023.
Emergent communication for rules reasoning, 2023.
Hahn, M. and Goyal, N. A theory of emergent in- in-context learning as implicit structure induction, 2023.
Garell,, M., Rosenbaum, D., Maddison, C. J., Ramalho, T., Saxton, D... Shanahan, M., Teh, Y. W., Rezende, D. J., and Eslami, M. Conditional neural processes, 2018a.
Han, C., Wang, Z., Zhao, H., and Ji, H. Explaining emergent in-context learning as kernel regression, 2023.
10

---

Docs learning the right latent variables necessarily improve in-context learning?
Hastie, T. The elements of statistical learning: data mining, inference, and prediction, 2009.
Norman, T. M., Horlbeck, M. A., Replogle, J. M., Ge, A. Y., Xu, A., Jost, M., Gilbert, L. A., and Weissman, J. S. Exploring genetic interaction manifolds constructed from rich single-cell phenotypes. Science, 365(6455):786-79, 2019.
Hendel, R., Geva, M., and Globerson, A. In-context learning creates task vectors, 2023.
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey, A. Meta-learning in neural networks: A survey, 2020a.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen,
A., Conerly, T., Drain, D., Ganguli, D., Saatficd--Ddds
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey, A. Meta-learning neural networks: survey, 2020b. URL https: //ppsi///axvv.../ org/ abs 2004 05439.
Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark,
J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads, 2022. URL https:// arxiv. org/ abs 2209 11895.
John and Raven, J. Raven Progressive Matri- ces, pp. 223-237.
Springer US, Boston, MA,
2003. ISBN 978-1-4615-0153-4 178-1-46115-0153-4.11 URL htttt:::/ddi..rr 10. 10.1007/978 1 11188-----155-115----111
doi: 10.1007/
Pakman, A., Wang. Y., Mitelut, C., Lee, J., and Paninski, L. Neural clustering processes. In International Conference on Machine Learning, pp. 7455-7465. PMLR, 2020.
Kim, H., Mnih, A., Schwarz, J., Garnelo, M., Eslami, A., Rosenbaum, D., Vinyals, O., and Teh, Y. W. Attentive neural processes, 2019.
Pospischil, M., Toledo-Rodriguee M., Monier, C., Pi- wkowska, Z., Bal, T., Frégnac. Y., Markram, H., and Destexhe, A. Minimal hodgkin-huxley type models for different classes of cortical and thalamic neurons. Biolog- ical cybernetics, 99:427-444 2008.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv: 1412.698,, 2014.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Radev, S. T., Mertens, U. K., Voss, A., Ardizzone, L., and Köthe, U. Bayesflow: Learning complex stochastic mod- els with invertible neural networks. IEEE transactions on neural networks and learning systems, 33(4):1452-1446 2020.
Lampinen, A. K., Chan, S. C., Singh, A. K., and Shanahan, M. The broader spectrum of in-context learning. arXiv preprint arXiv:2412.03782, 2024.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Lu, S., Bigoulaeva, I., Sachdeva, R., Madabushi, H. T., and Gurevych, I. Are emergent abilities in large language models just in-context learning?, 2024. URL https: //arxiv.ores abs/ wta.//2332 2309. 01809.
Ravichander, A., Belinkov, Y., and Hovy, E. Probing the probing paradigm: Does probing accuracy entail task relevance?, 2021.
McCoy, T., Pavlick, E., and Linzen, T. Right for the wrong reasons: Diagnosing syntactic heuristies in nat- ural language inference. In Korhonen, A., Traum, D., and Mårquez, L. (eds.), Proceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics, pp. 3428-3448, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.18655//1/P119-1334 URL https::p////.m.. http//aaaannnnooog.oooooooo
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep gen- erative models. In International conference on machine learning, pp. 1278-122 1286. PMLR, 2014.
Ribeiro, M. T., Singh, S., and Guestrin, C. "why should i trust you?": Explaining the predictions of any classifier, 2016.
Mittal, S., Bracher, L., Lajoie, G., Jaini, P., and Brubaker, M. A. Exploring exchangeable dataset amortization for bayesian posterior inference. In ICML 2023 Work- shop on Structured Probabilistic Inference & Generative Modeling, 2023. URL. https: /ppnneveeeeeeeeeeeeeeeeee net./ forum? ? d= Zt 9A5LMNUG.
Singh, A. K., Moskovitz, T., Hill, F., Chan, S. C. Y., and Saxe, A. M. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation, 2024.
Nguyen, T. and Grover, A. Transformer neural processes: incertainty-awar meta learning via sequence modeling, 2023.
Song, R., Li, Y., Shi, L., Giunchiglia, F., and Xu, H. Shortcut learning in in-context learning: A survey, 2024. URL https: www.i.wwg/v/..ee.... 2211102018........

---

Docs learning the right latent variables necessarily improve in-context learning?
Tang, R., Kong, D., Huang, L., and Xue, H. Large lan- guage models be lazy learners: Analyze shorteuts in in-context learning. In Findings of the Association for Computational Linguistics: ACL 2023. Association for Computational Linguistics, 2023a. doi: 10. 18 2023. 1023..indings-aclee ac1.284. URL http: / doi. org. 10. 18653/66/// / vv/lv11/2202 11/2232... findings-asd acl..acc... 284.
Learning Research, pp. 35151-35174. PMLR, 23-29 Jul 2023. BRRLhttpps:/pprrceedii html.
r.rerss//////
Wang. J. X., King, M., Porcel, N., Kurth-Nelson, Z.,
Zhu, T., Deck, C., Choy, P., Cassin, M., Reynolds, M.,
Song, F., Buttimore, G., Reichert, D. P., Rabinowitz, N., Matthey, L., Hassabis, D., Lerchner, A., and Botvinick, M. Alchemy: A benchmark and analysis toolkit for meta- reinforcement learning agents, 2021.
Tang, R., Kong, D., Huang, L., and Xue, H. Large lan- guage models can be lazy learners: Analyze shortcuts in in-context learning. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 4645-4657, Toronto, Canada, July 2023b. Association for Computa- tional Linguistics. doi: 10.18663w11/20233mmmmmng 284. URL https: ac http://aaaaaoooooooorr//2222222 Aindings-acl.284.
Wang, L.., Li, L.,, Dai, D., Chen, D., Zhou, H., Meng, F, Zhou, J., and Sun, X. Label words anchors: An information flow perspective for understanding in-context learning, 2023.
Wei, J., Bosma, M., Zhao, V. Y., Guu, K., Yu, A. W. Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetumed language models are zero-shot learners, 2022. URL https:// arxiv. exxiv.org// wwgg/bsg/////oooooo./e//.eeeeabe aat/a/bss//00 01652.
Tejero-Cantero, A., Boelts, J., Deistler, M., Lueckmann, J.-M., Durkan, C., Gonçalves, J., Greenberg, D. S., and Macke, J. H. Sbi toolkit for simulation-based inference, 2020.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F., Chi, E.,, Le, Q., and Zhou, D. Chain-of- thought prompting elicits reasoning in large language models, 2023. URL https: /ttp///xxvvv.rg./e........... /bbs/ 2201.
Todd, E., Li, M. L., Sharma, A. S., Mueller, A., Wallace, B. C., and Bau, D. Function vectors large language models, 2023.
Wiedemer, T., Mayilvahanan, P., Bethge, M., and Brendel, W. Compositional generalization from first principles, 2023.
Todd, E., Li, M. L., Sharma, A. S., Mueller, A., Wallace, B. C., and Bau, D. Function vectors in large language models, 2024.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning implicit bayesian inference, 2022.
Tsai, Y.-H. H., Bai, S., Yamada. M., Morency, L.-P., and Salakhutdinov, R. Transformer dissection: unified understanding of transformer's attention via the lens of kernel, 2019.
Yang, L., Lin, Z.,, Lee, K., Papailiopoolos, D., and Nowak, R. Task vectors in in-context learning: Emergence, for- mation, and benefit. arXiv preprint arXiv:2501.092240 2025.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Zaheer, M., Kottur, S., Ravanbakhsh, S., Poezos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Ad- vances in neural information processing systems, 30, 2017.
Vettoruzzo, A., Bouguelia, M.-R., Vanschoren, J., Rögnvaldsson, T., and Santosh, K. Advances and chal- lenges in meta-learnng: A technical review, 2023.
Zhang, H., Li, L. H., Meng, T., Chang, K.WW, and den Broeck, G. V. On the paradox of learning to reason from data, 2022.
Vig, J., Gehrmann, S., Belinkov, Y., Qian, S., Nevo, D., Sakenis, S., Huang, J., Singer, Y., and Shieber, S. Causal mediation analysis for interpreting neural nlp: The case of gender bias, 2020.
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P. What algo- rithms can transformers learn? a study length general- ization, 2023.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J.. Mordvintsev, A., Zhmoginov, A., and Vladymy- rov, M. Transformers learn in-context by gradient descent. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine
12

---

Does learning the right latent variables necessarily improve in-context learning?
A. Related Work
In-Context Learning. In-Context learning (ICL) is an ability of certain trained models to take an entire task's dataset as input and parameterize solutions directly in their layer activations, which then condition subsequent computation on novel inputs from those same tasks. Generally, this ability is found in sequence models such as Transformers where the task dataset, "context", corresponds to an earlier part of the sequence. ICL was first observed in large-scale pre-trained LLMs (Brown et al., 2020), and is similar in many respects to meta-learning (Chan et al., 2022). These LLM findings were subsequently expanded to more controlled settings outside of the language modality, where Transformer models were directly trained on task distributions such as linear regression (Von Oswald et al., 2023; Garg et al., 2023), hidden Markov models (Xie et al., 2022), compositional grammars (Hahn & Goyal, 2023), regular languages (Akyürek et al., 2024) and Turing machines (Grau-Moya et al., 2024), with a set of task observations defining the "context". These works highlight that Transformers are indeed able to model many types of complex task distributions, approaching in many cases the performance of the Bayes-optimat predictor, the a-priori optimal solution (Xie et al., 2022). Our work lies along similar lines but using more complex tasks and systematic study into the differences between modeling the predictive space directly, or through a two-step process involving explicit inference of task latents.
Shortcuts in ICL. Shortcut learning is a phenomenon that has widely been observed in machine learning (Geirhos et al., 2020), and refers to where a model solves a task through statistical correlations that are accidental and thus not robust to even slight distribution shifts. A classical example of this in image classification is the usage of background cues to classify objects (Ribeiro et al., 2016). Similar mistakes are know be very common in NLP (McCoy et al., 2019), specifically in reasoning tasks (Zhang et al., 2022). Particularly relevant to our work, many authors have shown that has shown that Transformers are very prone to relying shortcuts when performing ICL (Tang et al., 2023a). For instance, (Olsson et al., 2022; Singh et al., 2024) have shown that induction heads play in important role in ICL by predicting that the continuation of a token will be the same as last time (i.e. [a][]... [a] - [b). As shown by (Von Oswald et al., 2023) this motif can be used to do linear regression, and can generally be seen as a form of kernel regression (i.e. p Tinm) x E, ((2,, ))k, (Han et al., 2023)). This observation draws link between those types of solutions and non panaamraettiic inference methods in statistics (Hastie, 2009), of which kernel regression is a member. In contrast (Hendel et al., 2023; Todd et al., 2023) have concurrently shown that in some cases Transformers encode "task vector" that they infer from the context and then use to do the prediction. There is therefore a need to better understand the nature of shortcuts in ICL and whether or not they can be easily avoided for better generalization. Our work explores this very question.
Neural Processes. The problem of solving new tasks in a zero-shot manner directly at inference is also closely tied to amortized Bayesian models (Kingma & Welling, 2013; Rezende et al., 2014; Radev et al., 2020; Geffner et al., 2023; Mittal et al., 2023). Conditional Neural Processes (CNPs) (Garnelo et al., 2018a) provide a framework akin to the explicit model, where the posterior predictive distribution is modeled through a bottleneck Zg. i.c. po(y_ z., D) = pe pe (x,, z (D)). However, CNPs do not look the relevance of zΦ to the true latent 2, and use the DeepSets (Zaheer et al., 2017) architecture to model ze, though recent research generalizes this setting to use Transformers (Nguyen & Grover, 2023) and other architectural backbones as well (Kim et al., 2019; Gordon et al., 2019). Our approach with the explicit model, however, is to precisely question whether task-specific latents are encoded via Z. which is now instead modeled using a Transformer architecture. Analogously, Neural Processes (NPs) (Garnelo et al., 2018b; Pakman et al., 2020) augment CNPs with probabilistic modeling, where z is now modeled explicitly as a latent-variable in the Bayesian sense, i.e. the likelihood is now modeled as pe (y) x.. D)- = f (ylyaa z) )e (z (D)dz, where = represents the latent variable and O the parameters of the likelihood model. The model is trained via the Evidence Lower-Bound (ELBO) with the amortized variational approximation (-ID)) Once trained, predictions for new datasets can be made by simply performing inference over the encoder 4e to obtain z, and then leveraging this latent variable to eventually give the predictions via po (y) (z., 2). Hence, while CNPs and the explicit model to share similarities in architecture, our goal is orthogonal in that we specifically use the explicit model to understand the impact of task-specific latent variable inference on ICL setups.
Meta-Learning. Meta-learning (Hospedales et al., 2020b;a) studies systems that can learn over two levels: rapidly through an inner-loop that is meta-learned using a slower loter- The goal in such methods is to learn good initialization common to the parameterized family of tasks, in manner that obtaining particular solution for new task is fast from this initial point. The inner loop provides an optimization trajectory for randomly drawn task from some initialization, which in itself is optimized in the outer loop to good solution applicable for the global set of tasks. Typically, evaluation is done on some meta-validation set of tasks not seen during training. Task distributions can for example be a set of different classfficatinn/eeressinn tasks (few shot learning, (Vettoruzzo et al., 2023)) or variations of a reinforcement learning (meta-RL, (Beck et al., 2023)). The goal is similar to ICL approaches in the sense that given a novel context D,
13

---

Docs learning the right latent variables necessarily improve in-context learning?
Counterfactual MSE (J)
Linear regression interventions
10.0
7.5
5.0
2.5
0.0
ID
OOD
Validation data
Implicit
Explicit
D
q
Transformer
Xq
Transformer
Transformer
x
Figure Same experiment as Figure 4b but with the linear regression task. Specifically, we use Distributed Alignment Scarch (DAS, (Geiger et al., 2023b), see Appendix C for details) to find the 10 dimensional subspace in cach model with best simulates counterfactual interventions on the task vector (in this case, the weight of linear regression). both explicit model, the subspace is taken at the bottleneck. In the implicit one, perform the DAS at all layer of the query token and report the best one. The reported metric is the MSE of the intervened model on the intervened regression problem (i.c. using the same query æ hut y's coming from the intervened linear regression weights, Appendix C for details).
one wants to make predictions for some query T.. However, a big difference is that ICL approaches bypass modeling a common initialization by working directly on the prediction side (implicit), or instead predict the optimal parameters directly zero-shot through inference on the context model (explicit).
Mechanistic Interpretability. Mechanistic interpretability is interested in understanding deep neural network's computations through interpretable abstraction, akin what computational neuroscience does with the brain. (Alain & Bengio, 2018) introduced the foundational technique of linear "probes"" which are linear models trained on the hidden state of network to predict an abstract feature of the input; the success of which suggests that such a feature is used by the model. Since then, this naïve approach has been criticized for being potentially misleading (Ravichander et al., 2021 I); in many cases a feature can be linearly decoded from a model without the model using it. More reliable methods grounded in causality (Vig et al., 2020; Geiger al., 2024) have now became the gold standard, and their use applied Transformers has been exploding in popularity (Elhage al., 2021).
B. Tasks
We consider the following tasks for our evaluations, specified by the data-generating prediction function g: x, z + which is used to generate the ICL dataset, where z represents the task-specific latent variable.
B.1. Regression Tasks
For regression tasks, since y E IR, we use the mean-squared-erroo loss to train the model.
Linear Regression. This refers to the task where y is obtained from an affine transformation on the input æ. In particular, y == g 2) 2T x. where z E Rixj No0 7). For our experiments, we set dim(x) and dim(y) 1.
Nonlinear Regression using MLPs. Here, the labels y are obtained from neural network which takes æ as an input. In particular, y g(xx fx (z), where fz is modeled as Multi-Layer Perceptron (MLP) network with 64 dimensional single hidden layer and ReLU nonlinearity. The distribution of the weights of the neural network is AO0 7). For our experiments, we set dim(x) ) = 2 and dim(y) = 1.
Sinusoid Regression. For this task, the label y is obtained as a summation of sine functions with different frequencies and
amplitudes, taking æ as an input. Mathematically, we structure the system as y = g(x, 2) Σ a; sin (2x),æ)) where A
denote the frequencies and as the amplitudes. The parameters for the system can be seen as = (a1:)} while di-K
14

---

Does learning the right latent variables necessarily improve in-context learning?
remains fixed throughout. Additionally, for our experiments we set K = 3, and consider the distributions - > ~ ()0,5) and 04 ((),1,) and set dam():) 2 and dim(y) =
Gaussian Process Regression. While the other tasks considered had parametric nature to it, this task on the other hand has more of non paraameettiic nature. Here, the task is that Y ~ (0, K (x. i.c. the set of labels is sampled from a joint Gaussian distribution, akin to drawing a random function through a Gaussian Process (GP) prior and then evaluating it at different points X; with K defining the kernel in the GP. In our case, we consider K ((x, x') exp
=-=="" 2n as the
RBF kernel and X (X:: X,), Y (Y,, Y₂) are the combined points for both the context and the queries, which split after this sampling. Here the latents has to store the kemel computations between the query and all the context points X., as well as the corresponding context labels Y, Storing this either involves storing the high-dimensional mapping of X, which is defined by the kernel K, or storing all the points X themselves. This is thus very high dimensional and weakly structured.
Hodgkin- ODE Prediction. This is an example of the task where the context 'D is not composed of iid entries, but instead observations from the Hodgkin-Huxley temporal dynamics model of neural activity unrolled through time:
dV di = 91 (EL-VE - V) + J (Ena - V) + JKn* (EK - V) + GMP (Ex - V) + Las + or (t)
Cm
Above, V represents the membrane potential which is the target of interest, represents the different points at which observations are provided, C. is the membrane capacitance, 9i is the leak conductance, Ej is the membrane reversal potential, de is the density of channels of type e (Na*. + Kt, M), E. is the reversal potential of c. (m, h, n. p) are the respective channel gating kinetic variables, and (t) is the intrinsic neural noise. The right hand side of the voltage dynamics is composed of a leak current, a voltage-dependent Nat current, a delayed-rectifier K* current, a slow voltage- dependent K* current responsible for spike-fre adaptation, and an injected current /ij. Channel gating variables g have dynamics fully characterized by the neuron membrane potential V. given the respective steady-state 9:0 (V) and time constant Ty (V) (details in (Pospischil et al., 2008)).
Importantly, in our experiments, we fix all parameters but (īNa; 9k) to values in Tejero- Carrero et al. 2020 and solve the differential equation for 6,400 pairs (āNo: āk) E [0, 40j² from t = 0 to t = 120 with 1000 time-steps. In other words, the Transformer has to regress to solutions of ordinary differential equations, where the task latents are x = (ōNa: [k]], the observations are æ and y = V, such that y = g(x, z). Here g represents the unrolling of the differential equation.
B.2. Classification Tasks
For classification tasks, since y is a categorical measure, we use a cross-entropy loss for training.
Linear Classification. Akin to linear regression, here we consider the case that y is obtained by an affine transformation of followed by a sigmoid function and a consequent sampling step. That is, y = g(x, z) Cateeorical(SSftmmxx where z E Rixi (N(0, I). For our experiments, set dim (x:) and y E {0, 1).
Nonlinear Classification Using MLPs. Here, the logits for the labels are instead obtained through a neural net- work taking æ as an input, and not an affine transformation. Mathematically, this can be seen as y =gx 2) ~ AaoggoricalSoftmaxs f (z)) where fs is modeled as a Multi-Layer Perceptron (MLP) network with a 64 dimensional single hidden layer and ReLU nonlinearity. The distribution of the weights of the neural network is z ~ (N, I). For our experiments, we set din(xx) = 2 and y E 10.
B.3. Compositional tasks
Reusable Modular Mixture of Experts (MoE). We consider a modular task which consists of sequential application of a choice of K experts 9K over the input æ. In particular, the computational graph consists of layers where at each layer I, expert zi operates on the output of the preceding layer to give the successive output, i.c. x' = 9x1 (a'-1) This task is compositional in nature because at each layer, any of the K experts can be called upon to perform unit of computation and the choice of the expert is defined by the underlying task latent 21, ....... 2r., each of which are categorical with K possibilities. In our specific implementation, we set T, to be 5, K to be 5, E 1 and y E TR... Fach expert 9i is parameterized as a linear layer followed by the tanh activation function. We enumerate all K- possible combinations and then only use a subset of them during training, while randomly sampling all for evaluation.
Alchemy, Alchemy is a meta-reinffrceement learning benchmark (Wang et al., 2021) where each environment is defined
15

---

Does learning the right latent variables necessarily improve in-context learning?
by set (GRAPH, POTION MAP, STONE MAP) of rules about how some set of potions transforms some stones. We extracted from it an ICL. classification dataset consisting of transformations transformations are compositional and symbolic; each potion affects only one of the three properties of stones (size, shape and color). An environment is specified by how observable stones and potions MAP to latent stones and potions, along with a GRAPH over these latent stones which specify the result of the Transformations. In total there is 109 GRAPH, 48 POTION MAP and 32 STONE MAPS, making for 167424 environments. We reserve 100,000 environments for evaluation and train of the remaining ones.
(STONE, POTIoN) y STONE. The
Raven's Progressive Matrices. Raven's Progressive Matrices (Raven's PM) reasoning task used for IQ tests (John & Raven, 2003). It consists of a 3x3 grid where each cell contains simple objects varying in small number of attributes (number, shape, size, color), but the bottom right cell is left empty. Subjects must notice a pattern in how the cells change from left to right in the first two rows of the grid, and then use that same pattern to complete missing cell in the bottom row. This is done by selecting one answer among N possible provided options for the missing cell. We use a symbolic version of the dataset that addresses bias in the original version (Guo et al., 2023). In this dataset, objects at a cell have 4 discrete attributes with 40 possible values each. our models, the context consists of the first two rows of the grid, the query consists of the last row with a masked out final cell, and the ground-truth latent variable is the underlying rule that generates a particular grid. Each rule is composed of set of sub-parts, and we evaluate unseen compositions.
Gene Targeting. We use Perturb-seq dataset collected by Norman et al. (2019) where researchers performed several genetic intervention experiments using CRISPR (Gilbert et al., 2014). In each experiment, either one or two genes were targeted and the resulting expressions across 5000 genes were observed across several cells. Here, we consider each CRISPR intervention experiment as a different context, the resulting cell genetic expressions as 5000-dimal observations, and a left-out cell with half of the genetic expressions randomly masked out as the query. The task is to predict the missing genetic expressions for the queried cell. We evaluate on held on held out CRISPR experiments with novel pairs of targeted genes.
C. Model Details
In the following section, we describe the standard architectural details used for all the tasks, as well as specific differences in the architecture used for the scaling experiments. Finally, we also provide details about the distributed alignment search mechanism.
C.1. General Details
For our implicit model, we use a standard Transformer with 8 layers. In the explicit model, for context uggregation we parameterize (D) using a standard Transformer with 4 layers, 256 dimensions latent, 512 dimensions MLP, and heads. For the predictor Py, we consider two options: a ReLU-actiavtion based MLP with three hidden layers of size 512 and a Transformer with the same configuration as 2 (D).
For the implicit model, we format the prompt for prediction as [_r,,1] ... [T,, Bn [ 0J. where every [:] represents a token. We use a distinct mask token ( to represent the target (which is the thing being predicted). For the explicit model, we first compute [211, y1] [ yn] to z (D) with the context Transformer, then we give [ (D) [a.,] to the predictor Transformer or [ (D), xy,] to the MLP.
For our experiments, the number of context points n is uniformly sampled from 16 to 128 for both training and evaluation. Training is done with new data being synthetically generated the fly, and evaluation either based the test set provided for real-world tasks simulated data of 1000 different contexts for synthetic tasks. All the models were trained with learning rate of 10-4 using the Adam optimizer (Kingma & Ba, 2014) for a 1000 epochs.
C.2. Scaling Experiments
For the scaling experiments, we only consider the linear regression case with base configuration of: (a) of dimensionality 100, (b) context size being sampled uniformly from (75, 125), and (c) 8 heads, 8 layers, 512 hidden dimensions and 256 bottleneck dimension for the transformer models.
From this base configuration, we changed only one of the configurations at each time to test for scaling trends for each property independently. In particular, we ablated over (50, 100, 250) for the dimensionality of æ. (50, 100, 250) for the context length which was sampled from +25 range and the model size. The smallest model size considered had heads, 4
16

---

Docs learning the right latent variables necessarily improve in-context learning?
Loss
E(21, 22, E3, 24)
Replace subspace activation
Model
Model
{I,,, (1))
E(21, 22, 23, 24)
200
{G α)) Ji=1 E(21, 32, 23, EA)
Figure Illustration of DAS training procedure
layers, 256 hidden dimensions and 128 feature dimensions. The medium multiplied each of these properties by 2x, and the biggest model subsequently multiplied it by 2 again. For the explicit models, considered the same scaling paradigms with the number of layers being split by half to accommodate a separate context model and prediction model.
All the models were trained with a learning rate of 10-6 using the Adam optimizer for 5000 epochs.
C.3. Compute Details
We train most of our models on single RTX8000 NVIDIA GPUs, where it takes roughly 3-6 hours for each experiment to run. Our scaling experiments on the other hand often required 1-2 days on single GPUs for training each model.
C.4. Distributed Alignment Search details
To find subspace causally associated with task latent in Alchemy, we use method based on Distributed Alignment Search (DAS) by (Geiger et al., 2023b). This procedure is performed for a location IRd (e.g. the bottleneck) and latent E {1, 2, 3} (GRAPH, STONE MAP, POTION MAP).
First, we run with the model on D and D& for every possible query Z., We call the base and the source and only differ by the ith latent. For every run, record the activity of the source model at the location T.. Then, the base model again but this time fixing the subspace of defined by the orthogonal projection 11 E dd to it's value in 12, A single projection II is learned over all possible combination 2, 2 and x. with cross-entropy loss between the prediction of the base (intervened) model and the true counter-factual result of changing the latent to ž, See Figure 10 for an illustration of the process. A subspace evaluated by looking at the accuracy of the counterfactual interventions over a dataset held-out 2, pairs; a quantity called the Interchange Intervention Accuracy (IIA). In Figure Figure 4 (b) we report the validation IIA relative to a baseline corresponding to the counterfactual accuracy if we don't 't perform any BASELNN intervention (because changing the latent sometimes doesn't change the prediction) 1 BASELINE
D. Analysis of Experiments
Based on the empirical evidence presented in section 4, we finally provide details and analysis into the results to further the understanding of the conclusions. In particular, our key analysis includes
Explicit Models sufficiently uncover task latents. We see that in problems where the context provides enough evidence
17

---

Does learning the right latent variables necessarily improve in-context learning?
Sruezitsl
Uheer
Situeidss
17.
15.
12.
NA0E L-)
18
14-
Implictt Proxy
Explicit
Transfomer
A
MLP
Yq
D Transformer
xg
MLP
-gg
Figure 11. To further understand the difference between the explicit and implicit model, we utilize an implicit proxy model which shares the same architecture as the explicit model with MLP prediction with just one key difference: the task latent depends on the query 2y as well. This task latent can be understood as the final attention layer output of the implicit model, after which an additional MLP is utilized to provide prediction. Our findings on linear and sinusoidal regression demonstrate that as we move further and further aut-of-distribution, the implicit proxy model performs better than the explicit model (left figures), but recovers the underlying true task latents worse (two right tigures). This provides additional validation of our hypothesis.
to uncover the true task latents, explicit models are able to do so. In particular, this hints at the fact that explicit models do perform downstream prediction based on true task latents whenever these latents can be sufficiently identified from the context examples.
Explicit Models do not generalize better than implicit ones. Our analysis also reveals that while explicit models often do uncover the right task latents, they are still not able to surpass implicit models even on OOD generalization. This could be due to implicit models also uncovering the true underlying prediction function but in distributed fashion, or explicit models not being able to leverage the learned latents in downstream prediction.
Learned downstream prediction is often sub-optimal. Our results indicate that it is indeed the case that while the explicit models do uncover the right latents, they fail to generalize well OOD because the downstream prediction function fails to generalize.
This is further strengthened by Figure 11 where additionally leverage the query in context aggregation, thus interpolating between explicit and implicit model while maintaining bottleneck. Our results indicate that despite worse latent variable inference far from the query, it still has better predictive generalization when compared explicit models. This further strengthens the claim that better latent variable inference is not the sole problem, and learning the right downstream prediction is as important.
Classification tasks vs. regression tasks. OOD performance is generally strong (across all models) for classification because decision boundaries are within the training domain and do not change beyond it. In contrast, for regression tasks, the function continues to change beyond the observed training domain, making OOD prediction more difficult. This also why known prediction functions give little benefit in classification tasks: they already solved well OOD with ordinary implicit and explicit models.
The explicit model with known prediction function does not give benefits in nonlinear (MLP) regression. This is because the problem of inferring an MLP's weights given some context examples is too difficult, so the explicit model opts for a different, non- parameamettiic solution. This is supported by the latent variable decoding results in Fig. 5 (previously Fig. 4), which show that even with a known prediction function the explicit model does not learn to infer the correct latent variable for the nonlinear (MLP) regression task.
Impact of Output Dimensionality. In addition to standard experiments, we consider linear regression with varying output dimensionality as another measure of difficulty of the task. Our results in Figure 14 showcase that the implicit model fares much better than the explicit model, however having known prediction function leads to much better performance.
Size of Context Aggregator. While we primarily controlled for the total number of parameters, it is possible that context aggregation requires more parameters and thus equally splitting parameters between this aggregation and prediction may be the cause for suboptimal performance. To study this properly, we conduct experiments where the context aggregator the
18

---

Docs learning the right latent variables necessarily improve in-context learning?
a.
Synthetic regression tasks
Linear
Nonlinear (MLP)
Gaussian Process
Sinusoidal
Hodgkin-Huxkee
0.4
0.2
0.3
MSE
10"
0.2
10
0.
0.0
ID
OOD
0OD
OOD
ID
0OD
OOD
b.
Synthetic classification tasks
C.
Compositional tasks
Linear
Non inear (MLP)
Raven's PM
Gene targeting
1,00
0:83
Accuracy (T)
Accuracy (t)
0.98
0.82
0.84
0:81
0.98
0.82
0.96
0.90
0.80
GOD
GOD
Impliiit
Explicit
Explicit, aggregator Implicit model
Transformer
Transformrr
Tranalormer
MUP
-gg
kg
Figure 12. Similar to Figure 2, but with additional explicit model where the context aggregator is identical in size and hyperparameters (M) as implicit model. Implicit models are in shown gray, explicit with Transformer prediction blue (light blue total model size identical to implicit as in main paper, dark blue context aggregator identical to implicit), and with MLP prediction in orange.
same size as the implicit model and we see similar results in Figures 12 and 13. addition, ablate with differently sized context aggregators and prediction models for different tasks in Tables to
Visualization of Explicit Predictions. We visualize the failure of explicit model in learning the right prediction function, which can be seen in out of distribution for sinusoid regression in Figure 15.
E. Mathematical Formalism
In this section, we provide a formal distinction between the implicit and explicit model. In both the approaches, the goal is model the true posterior predictive payla D); however the methods model through different conditional independence setup.
Implicit Model. In this setup, we model the predictive distribution as Po xx, D), where the training is done as
arg max Em,,,, ( D))
and then given a query x and datasct D, the inference is done simply by sampling or estimating the mean of p(y],, D).
Explicit Model. Contrary the implicit model, the explicit model parameterizes the predictive distribution as pyy (y 2i (D)), with similar training procedure as above. Note that the predictive distribution only interacts with the dataset D through the latent z (D) while the implicit model allows unconstrained access to D.
Implicit Proxy Model. To better understand the differences that play a role from architectural differences and parameteriza- tions, we use exactly the same architecture as the explicit model to obtain a version of the implicit model. Such a model parameterizes the predictive distribution as Py (y x z "D, x)), with a similar training procedure as above. Note that the only difference with the explicit model here is that the conditional dependence of the query and the task latents is broken.
We refer the reader to Figure 16 for plate diagram of the corresponding architectures.
19

---

Docs learning the right latent variables necessarily improve in-context learning?
Linear reg.
Nonlinear reg. (MLP)
Sinusoidal reg.
Linear Nonlinear (MLP)
Raven's PM
1.00
Accuracy (T)
0.
0
0.98
MSE
02
0.1
0 96
0.96
Known
Explictt
1 Bidipes
Explicit context aggregator Transformer same size as for Implicit model
89
Nq"
Figure 13. Similar to Figure 3 but with an additional explicit model where the context aggregator is identical size and hyperparameters (M) the implicit model. Performance tasks where the latents z and prediction function known. Implicit models in gray, explicit models with Transformer prediction in blue, models trained with an auxiliary loss to predict true latents in purple and those using the true prediction function in green. Using known prediction function leads to significantly better OOD performance. Lighter color indicates total model size identical to implicit as in main paper, darker context aggregator identical to implicit.
MSE (1)
2-
10
3 10
10
10
Output Dimensiona ity
Implicit
Explicit
D
Transformer
Nig
D
Transformer
Transforer
Known
xg
Xqq
Figure 14. We study the impact of output dimensions on the performance explicit and implicit models for linear regression with 8-dimensional inputs. The output dimensionality considered he and dimensional, and results indicate that implicit methods (gray) outperform explicit ones (blue), but the same explicit models with known prediction function (green) scales much better.
20

---

Docs learning the right latent variables necessarily improve in-context learning?
Thee Contest
Thue Contoct: Quatry
Thue Contictt
The Content Queryy
Querry
Querry
WWa M Amm
MAA MMAmMm
MmMMM
Figure 15. Illustration of explicit model with MLP prediction the sinusoid task OOD. True function is shown in black, model output in green, context points blue and query points in red. Our results indicate the failure learned prediction function away from context.
y4
yo
Z
yx
N
N
N
Implicit Model
Implicit Proxy Model
Explicit Model
Figure Plate diagram the implicit model (left), implicit proxy model (middle) and the explicit model (right), where gray blocks refer to observed variables and white refers to unobserved variables. Trainable parameters are indicated without circles. In the explicit model case, z currently modeled as a dirac measure defined the trainable parameters th. One can see the implicit proxy model as very similar to the implicit model where output of the last attention layer corresponding to the query token is further processed to give prediction. similarity the explicit model is also clear shares exactly same parameterization.
21

---

Docs learning the right latent variables necessarily improve in-context learning?
Ly Loss (4)
Model
# Context Layers
Prediction Layers
Parameters (M)
IID
OOD Query
OOD Latent
0.011
0.016
0.782
Explicit-Known
6
3.2
0.012
0.019
0.923
8
4.2
0.013
0.023
1.391
3.0
0.015
0.169
0,883
4
6
3,6
0.017
0.174
0.728
4
8
4.1
0.015
0.082
0.929
4.1
0.011
0.201
1.047
Explicit-MLP
6
4.6
0.013
0.131
0,509
6
8
5.1
0.012
0.077
0.439
8
4
5.1
0.012
0.165
1.043
6
5.7
0.017
0.171
1.020
6.2
0.016
0.082
1.091
4
4
4.2
0.012
0.267
0.856
5.3
0.014
0.335
1.014
8
6.3
0.011
0.314
0.970
6
4
5.3
0.025
0.389
1.128
Explicit-Tsf
6
6
6.3
0.011
0.379
1.091
8
7.4
0.012
0.289
0.981
6.3
0.012
0.354
1.171
6
7.4
0.013
0.311
1.134
8
8
8.4
0.012
0.369
1.373
2.1
0.013
0.258
0.817
Implicit
6
3.2
0.011
0.261
0.638
4.2
0.011
0.287
0.661
Table Analysis different design choices for linear regression where implicit and explicit models tested with different number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization as well as out distribution generalization where either the query is OOD or underlying latent is.
22

---

Docs learning the right latent variables necessarily improve in-context learning?
Accuracy (t)
Model
Context Layers
Prediction Layers
Parameters (M)
IID
OOD Query
OOD Latent
2.1
96.381
97.779
96.288
Explicit-Known
3.2
96.331
97.746
96.324
4.2
96.426
97.788
96.299
4
3.0
96.297
97.646
96.150
6
3.6
96.176
97.410
95.989
8
4.1
96.338
97.772
95.964
4.1
96.289
97.679
96.125
Explicit-MLP
6
4.6
95.885
97.154
95.803
8
5.1
96.167
97.503
96.039
4
5.1
96.194
97.496
95.761
5.7
96.158
97.508
96.224
6.2
95.915
97.006
95.758
4
4.2
96.386
97.760
96.239
5.3
96.288
97.676
96.163
8
6.3
96.453
97.874
96.349
4
5.3
96.343
97.722
96.228
Explicit-Tssf
6
6.3
96.415
97.853
96.365
8
7.4
96.019
97.335
96.074
4
6.3
96.288
97.733
96.306
6
7.4
95.811
96.976
95.976
8
8.4
96.304
97.729
96.068
2.1
96.525
97.996
96.149
Implicit
3.2
96.557
98,036
96.440
4.2
96.554
98.147
96.357
Table Analysis different design choices for linear classification where implicit and explicit models are tested with differrnt number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization as well as out distribution generalization where either the query is OOD or underlying latent is.
23

---

Docs learning the right latent variables necessarily improve in-context learning?
La Loss (4)
Model
Context Layers
Prediction Layers
Parameters (M)
IID
OOD Query
2.1
0.001
0.002
Explicit-Known
3.2
0.000
0.001
4.2
0.001
0.002
3.0
0.001
0.295
6
3.6
0.001
0.254
8
4.1
0.001
0.208
4
4.1
0.001
0.298
Explicit-MLP
4.6
0.001
0.307
8
5.1
0.001
0,229
4
5.1
0.001
0,399
6
5.7
0.002
0,266
8
6.2
0.001
0.222
4
4.2
0.001
0.179
5.3
0.001
0.196
6.3
0.001
0.186
4
5.3
0.001
0,212
Explicit-
6
6.3
0.001
0,230
8
7.4
0.001
0,211
6.3
0.002
0.201
7.4
0.001
0.209
8
8.4
0.001
0,203
2.1
0.001
0.186
Implicit
3.2
0.000
0.143
4.2
0.000
0.153
Table Analysis different design choices for sinusoid regression where implicit and explicit models tested with different number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization as well distribution generalization.
Lg Loss (1)
Model
# Context Layers
# Prediction Layers
# Parameters (M)
IID
OOD Query
4
3,0
0.039
0.205
4
6
3.6
0.033
0.185
4.1
0.037
0.181
4.1
0.034
0.186
Explicit-MLP
4.6
0.030
0.183
8
5.1
0.029
0.174
4
5.1
0.030
0.186
0.032
0.176
8
6.2
0.031
0.175
0.028
0.181
5.3
0.027
0.172
4
8
6,3
0.027
0.172
6
4
5.3
0.025
0.170
Explicit-Tsff
6.3
0.024
0.169
0.027
0.171
6.3
0.025
0.171
8
6
7.4
0.028
0.173
8
8
8.4
0.030
0.177
4
2.1
0.011
0.148
Implicit
3.2
0.010
0.148
0.010
0.147
Table Analysis of different design choices for GP regression where implicit and explicit models are tested with different number of parameters different parameter split between context aggregation and prediction. We conduct experiments in in-distribution generalization as well as out of distribution generalization.
24

---

Docs learning the right latent variables necessarily improve in-context learning?
Ly Loss (4)
Model
# Context Layers
# Prediction Layers
# Parameters (M)
IID
OOD Query
0.084
1.323
Explicit-Known
6
3.2
0.074
1.168
8
4.2
0.073
1.135
4
3,0
0.051
2.658
4
3,6
0.051
2.708
4
4.1
0.057
2,054
6
4.1
0.051
2.814
Explicit-MI.. PL.P
6
4.6
0.052
2.256
6
5.1
0.058
2.395
8
5.1
0.047
2.465
8
5.7
0.053
2.061
6.2
0.054
1.906
4
4.2
0.039
2.052
5.3
0.044
2.185
4
6.3
0.046
2.331
6
5.3
0.042
2.154
Explicit-Tsf
6
6,3
0.050
2.118
6
7.4
0.052
2.351
8
0.043
1.945
8
7.4
0.044
2.148
8
8.4
0.048
2.178
4
2.1
0.024
1.443
Implicit
6
3.2
0.024
1.487
8
4.2
0.024
1.509
Table Analysis different design choices MLP regression where implicit and explicit models are tested with different number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization as well as out distribution generalization.
25

---

Docs learning the right latent variables necessarily improve in-context learning?
Accuracy (t)
Model
# Context Layers
# Prediction Layers
# Parameters (M)
I1D
OOD
94.426
90.753
Explicit-Known
6
3.2
94.633
91.086
8
4.2
94.492
91.443
3,0
94.888
92.775
4
3,6
94.971
92.861
4
4.1
94.465
92.342
6
4.1
94.878
92.807
Explicit-MLPP
6
94,933
93.024
6
5.1
94,851
93.097
8
4
5.1
94.818
93.036
8
5.7
94.822
93.024
94.878
92.985
4
4.2
95.074
92.990
95.060
92.988
4
6.3
95.157
93.190
6
4
5,3
94.989
93.010
Explicit-Tsf
6
6.3
95.060
93.275
6
7.4
95.078
93.119
95.010
93.133
8
7,4
95.058
93.121
8
8
8,4
94.969
92.954
4
2.1
95,285
93,394
Implicit
6
3.2
95.325
93.336
8
4.2
95.324
93.351
Table Analysis different design choices MLP classification where implicit and explicit models are tested with different number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization as well as out distribution generalization.
26

---

Docs learning the right latent variables necessarily improve in-context learning?
Accuracy (t)
Model
# Context Layers
# Prediction Layers
Parameters (M)
III
OOD
4
2.1
99.864
99.844
Explicit-Known
99.876
99.875
8
4.2
99.545
99.539
97.929
97.894
Explicit-MLE
4
4.0
97.551
97.532
4
8
4.5
97.804
97.802
6
4
4,5
97.394
97.383
5.0
97.524
97.517
6
98.527
98.514
8
5.5
29.883
29.587
8
6.1
98.240
98.216
8
6.6
96.726
96.677
4
4
4.2
98.522
98.513
4
5.3
98.467
98.455
98.006
98.005
6
5.3
98.494
98.459
Explicit-Tsf
6
6,3
97.930
97.874
6
7,4
98.566
98.591
8
6.3
98.280
98.260
97.750
97.744
8
8.5
98.468
98.456
98.359
98.334
Implicit
6
3.2
98.332
98.309
8
4.2
97.164
97.176
Table Analysis of different design choices for RAVEN's progressive matrices where implicit and explicit models tested with different number of parameters or different parameter split between context aggregation and prediction. We conduct experiments on in-distribution generalization well as compositional out of distrihution generalization.
R² (t)
Model
# Context Layers
# Prediction Layers
# Parameters (M)
Performance
4
9.5
0.815
6
10.0
0,827
10.5
0.814
4
0.814
Explicit-MLP
6
11.0
0,818
8
11.6
0.819
4
11.6
0.828
12.1
0.820
8
12.6
0.817
4
8.1
0.829
6
9.1
0.824
8
10.2
0.823
4
9.1
0.808
Explicit-Tsf
10.2
0.812
8
11.3
0.812
4
10.2
0,810
6
11.3
0.818
8
12.3
0.821
4.7
0.819
Implicit
5.7
0.823
6.8
0.817
Table Analyyis of different design choices for gene targeting experiments where implicit and explicit models are tested with different number of parameters different parameter split hetween context aggregation and prediction. We conduct experiments compositional out of distribution generalization.
27


Reviews:
[{"review_id": "A5rnMuTexk", "reviewer": "Reviewer_1Akw", "summary": "This paper investigates whether explicitly inferring latent variables of an underlying task improves in-context learning performance in transformers. They find that explicit modeling of latent variables does not necessarily improve performance compared to standard implicit models. They also find that while the explicit model does learn latent variables, the main problem with generalization is its prediction function is not properly trained.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the value of our work and providing constructive cristicism.\n\n> Additional References\n\nWe appreciate the reviewer bringing two relevant papers to our attention and will include them in the final version.\n\n- 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.\n- 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. \n\nWhile 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.\n\n> why the prediction function is not trained well enough even though it appears like the latents are being modeled properly\n\nOur 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.\n\n> algorithmically compare the learned prediction function to the optimal one\n\nWe 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.\n\n> 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?\n\nWe 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].\n\n> Additional insights into inductive biases to improve prediction function\n\nWe 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.\n\nWe hope that our response has resolved the reviewer's concerns and would be happy to provide further clarifications.\n\n[1] Hollmann, Noah, et al. \"Tabpfn: A transformer that solves small tabular classification problems in a second.\" arXiv preprint arXiv:2207.01848 (2022).\n\n[2] Müller, Samuel, et al. \"Transformers can do bayesian inference.\" arXiv preprint arXiv:2112.10510 (2021).", "claims_and_evidence": "Yes, the authors are careful to make claims that are backed by sufficient evidence.", "methods_and_evaluation_criteria": "Yes, there are several appropriate datasets the authors use to show the generality of their findings.", "theoretical_claims": "I did not check any proofs, but also did not see any theoretical claims in the paper.", "experimental_designs_or_analyses": "The experimental design is solid. The experiments that are clearly explained, and questions I had about details while reading I was able to find in the appendix.", "relation_to_broader_scientific_literature": "This work finds that causing the transformer architecture to explicitly infer latent variables of few-shot prompts is not sufficient for good ICL generalization. There has been a few works that suggest LLMs infer latent variables while doing ICL [1,2]. But this work showed little difference in performance between standard transformer and the explicit model. This is an interesting data point among previous and related findings. For instance, while [3] & [4] provide evidence suggesting that LLMs trained on natural text seem to do implicitly model latent variables, one contribution of this paper might be that it provides a way to perhaps understand the failure modes of these models - that it can similarly be attributed to a poor \"prediction function\" which seemed to be the problem with the failure cases of the explicit model in this work. There are some nice contributions of this work that might help us understand the architectural and algorithmic shortcomings of the models we use.\n\n___\n[1] Xie, et al. An explanation of in-context learning as implicit bayesian inference. ICLR 2022. (https://openreview.net/forum?id=RdJVFCHjUMI)\n\n[2] Wang, et al. Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning. NeurIPS 2023. (https://proceedings.neurips.cc/paper_files/paper/2023/hash/3255a7554605a88800f4e120b3a929e1-Abstract-Conference.html)\n\n[3] Hendel, et al. In-Context Learning Creates Task Vectors. EMNLP 2023. (https://aclanthology.org/2023.findings-emnlp.624/)\n\n[4] Todd, et al. Function Vectors in Large Language Models. ICLR 2024. (https://openreview.net/forum?id=AwyxtyMwaG)", "essential_references_not_discussed": "This work is well-placed among previous and concurrent work. I wanted to point out two works that I think are related, but were not cited, which study the idea of in-context learning centering around latent variables [2,5].  [2] support the bayesian inference of an internal latent view of in-context learning.\nThe work in [5] shows an example of how simple implicit transformer models seem to learn latent variables like others suggest is happening in larger LLMs. Maybe a discussion of how [5] relates to this paper's study of implicit vs. explicitly modeling latent variables would be helpful to provide more context to how we might interpret both results.\n\n___\n[2] Wang, et al. Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning. NeurIPS 2023. (https://proceedings.neurips.cc/paper_files/paper/2023/hash/3255a7554605a88800f4e120b3a929e1-Abstract-Conference.html)\n\n[5] Han, et al. Emergence of Abstractions: Concept Encoding and Decoding Mechanism for In-Context Learning in Transformers. (https://arxiv.org/abs/2412.12276)", "other_strengths_and_weaknesses": "This is a solid paper. The proposed setting of explicitly modeling latent variables is nice, and the experiments are carefully designed to test specific hypotheses. The results are thoughtfully presented without overclaiming and are supported by evidence across a variety of tasks.\n\nThe main weakness I'd say of the paper is that since prediction function failure was the main reason the explicit model didn't work, there could be more discussion about how to improve this during training beyond the current text: (i.e., \"supplemented with significant\ninductive biases in the prediction function\"). Could you provide some examples of what this might look like for different tasks?  Does next-token prediction (e.g. language modeling) have sufficient inductive biases that might mediate some of these problems, or do you think it's purely architectural?", "other_comments_or_suggestions": "Here's a list of minor typos I found while reading through:\n- Line 131: \"shown gray\" -> shown in gray\n- Line 782: \"Hodgkin-Hoxley\" ->  Hodgkin-Huxley?\n\n\n___\nNote after Rebuttal period:\nAs before the rebuttal, I am leaning towards accepting this paper. This paper does have its limitations, but the things I learned from this paper I think outweigh any reservations I may have had. The rebuttal by the authors answered my questions, and I feel they have also addressed the concerns of the other reviewers as well.", "questions_for_authors": "My main question is related to the learning of a good prediction function. \n1. Do you have any intuition as to why the prediction function is not trained well enough even though it appears like the latents are being modeled properly? Is it because the ID training data is sufficiently different from the OOD data? Perhaps it depends on the task.\n\n2. It's a bit surprising that explicit training does not learn a \"good\" prediction function. Is there some way to algorithmically compare the learned prediction function to the optimal one? This might be a nice way of characterizing the failure modes. Would freezing the first half of the explicit model after a while and training only the predictor function approach the optimal one? (It seems like they should be mutually reinforcing in a sense)\n\n3. In Line 136-142, left 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fYuGqNOUz5", "reviewer": "Reviewer_1soF", "summary": "The paper addresses a key question in in-context learning: whether explicit latent variable learning leads to better generalization, especially out-of-distribution (OOD). The conclusion is that the explicit bottleneck architecture does not help in terms of generalization.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n> Is there any difference in the result when using an encoder (as in the paper) vs. using a decoder?\n\nThe 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].\n\n> Could you check the other OOD task for the regression problem, especially the ones that have OOD on the latent variable?\n\nWe 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).\n\nInspired 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}$.\n\n> Could you explain how many layers of Transformer are used in the implicit and explicit model\n\nWe 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.\n\n> I wonder if the effect of \"dimension\" in the output also matters\n\nWe 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.\n\n> why is the classification task not beneficial from predicting the latent variable directly?\n\nFor 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.\n\nWe hope that our response has resolved the reviewers' concerns and would be happy to provide further clarifications.\n\n[1] Hollmann, Noah, et al. \"Tabpfn: A transformer that solves small tabular classification problems in a second.\" arXiv preprint arXiv:2207.01848 (2022).\n\n[2] Müller, Samuel, et al. \"Transformers can do bayesian inference.\" arXiv preprint arXiv:2112.10510 (2021).", "claims_and_evidence": "Yes, the paper design experiments test the intended hypothesis well.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "1. Is there any difference in the result when using an encoder (as in the paper) vs. using a decoder? Why did the author choose to use an encoder in particular?\n2. For the regression problem, why did you choose the OOD task to be scaled x, instead of scaled w, since by bottlenecking the latent variable, the explicit model should probably be better at any OOD of the latent variable? Could you check the other OOD task for the regression problem, especially the ones that have OOD on the latent variable?\n3. Could you explain how many layers of Transformer are used in the implicit and explicit model (including the context model and the predictor head)? Would there be any impact regarding the depth of the Transformer itself on the OOD generalizability? For instance, the deeper model may be better at OOD than the shallower model.\n- I see from the appendix that you consider two options in the predictor, one is a 4-layer TF, another is the MLP. Could you explain clearly which model you used to report the result? And why do you think using a shallower Transformer for the context extraction is sufficient for the explicit model to generalize?\n4. By the ablation study on \"explicit models learn to infer the correct latent variable, but not how to use it\", the author suggests that by forcing the model to explicitly predict the correct latent variable, the explicit model generalizes better. I wonder if the effect of \"dimension\" in the output also matters. If I understand correctly, the y_q is one-dimensional, while the latent variable is 1 dimension for the linear regression, and 2 dimensions for the other tasks. Would it be possible to put the nonlinear & sinusoidal to also 1 dimension in x when conducting the Figure 3 experiment? This could potentially rule out the error in terms of the dimensionality.\n5. Also in Figure 3, why is the classification task not beneficial from predicting the latent variable directly?", "relation_to_broader_scientific_literature": "This paper discuss the key hypothesis that learning the latent variable is not beneficial to the overall OOD performance of the transformer. This is an interesting observation to understand how transformer learn to solve the synthetic problems.", "essential_references_not_discussed": "Not that I aware of", "other_strengths_and_weaknesses": "Strengths: The paper uses extensive experiments designed to investigate the hypothesis of whether learning an explicit latent variable helps in the transformer's generalization ability.\n\nWeakness: It is in general hard to draw a fair comparison on the OOD experiments. For instance, as I mentioned in the experiment section, the author chooses 8 layers for implicit TF, and 4 layers for explicit TF when extracting the latent z, and MLP / 4 layer TF for the predictor. However, as far as I know, the depth of the model strongly affects the model's OOD performance in the regression task. It is not justified why the explicit model is designed like that (with 4 layers allocated for context extraction), and could be beneficial if there are more explicit model designs, for instance, the context extractor = implicit model depth, is tested to accompany the results.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Update after the rebuttal, and after the direct message from the author to AC:\n\nI appreciate the authors for conducting extensive experiments varying the number of parameters. *I apologize for initially overlooking the additional results provided in the attached files*. I have now carefully reviewed the full set of results and summarized the key information below:\n\n* For linear regression: result is mixed\n| Implicit Model         | Matching Explicit Model in terms of param          | OOD Query (which is better) | OOD Latent (which is better) |\n|------------------------|-----------------------------------|-----------------------------|-------------------------------|\n| 4 layers               | N/A                               | N/A                         | N/A                           |\n| 6 layers               | (explicit-MLP) L_context=4, L_prediction=4 | explicit                    | implicit                      |\n| 8 layers               | (explicit-MLP) L_context=4, L_prediction=8 | explicit                    | implicit                      |\n| 8 layers               | (explicit-MLP) L_context=6, L_prediction=4 | explicit                    | implicit                      |\n| 8 layers               | (explicit-Tsf) L_context=4, L_prediction=4 | explicit                    | implicit                      |\n\n* For linear classification, sinusoid, MLP regression, MLP classification: implicit is better\n| Implicit Model         | Matching Explicit Model in terms of param          | OOD Query (which is better) | OOD Latent (which is better) |\n|------------------------|-----------------------------------|-----------------------------|-------------------------------|\n| 4 layers               | N/A                               | N/A                         | N/A                           |\n| 6 layers               | (explicit-MLP) L_context=4, L_prediction=4 | implicit                    | implicit                      |\n| 8 layers               | (explicit-MLP) L_context=4, L_prediction=8 | implicit                    | implicit                      |\n| 8 layers               | (explicit-MLP) L_context=6, L_prediction=4 | implicit                    | implicit                      |\n| 8 layers               | (explicit-Tsf) L_context=4, L_prediction=4 | implicit                    | implicit                      |\n\n* For RAVEN: result is mixed\n| Implicit Model         | Matching Explicit Model in terms of param          | OOD (which is better) | \n|------------------------|-----------------------------------|-----------------------------|\n| 4 layers               | N/A                               | N/A                         | \n| 6 layers               | (explicit-MLP) L_context=4, L_prediction=4 | implicit                    | \n| 8 layers               | (explicit-MLP) L_context=4, L_prediction=6 | explicit                    | \n| 8 layers               | (explicit-Tsf) L_context=4, L_prediction=4 | explicit                    | \n\n* For gene, result is not interpretable since no explicit runs have similar number of params than implicit runs.\n\nThat said, under a similar total number of parameters, the implicit model generally performs better, though there are a few exceptions.\n\nThe following concern still remains, and I hope to see results averaged over more trials to clarify this point:\n> Additionally, I am somewhat confused. In the paper, the OOD task for linear regression is defined as OOD query. However, the new results suggest that for OOD query, the explicit model performs better when the latent parameters are learned explicitly under a similar parameter count. Specifically, in Figure 2, the blue block (representing the explicit model as Transformer) should correspond to the Explicit-Tsf model with L_context=4 and L_prediction=4. However, this seems to lead to a conflicting conclusion with what is now presented. I hope I’m not misunderstanding or misinterpreting anything here—perhaps this result is based on a single run, and the variance might explain the inconsistency. If that’s the case, it would be helpful to clarify whether these results are averaged over multiple runs or represent individual trials.\n\nIn addition, I would like to see more fine-grained explicit model configurations that match the number of parameters of the implicit models. Among the results provided during the rebuttal, only a few explicit setups are comparable in parameter count—one setup roughly matches the 6-layer implicit model, and one or two setups roughly match the 8-layer implicit model.\n\nOverall, I see promising signals that the implicit model outperforms in general, especially on ID tasks. However, I would also encourage the authors to investigate fairer comparisons by matching ID performance—e.g., tuning the parameter count so that both models achieve similar ID results, and then evaluating on OOD. Currently, a lot of settings (GP, MLP regression) have implicit model with better ID than the explicit model, and the the source of the OOD gap needs to be investigated.\n\nLastly, I would suggest that the authors moderate the strength of their conclusions. Since it is difficult to cover all possible setups, it would be helpful to explicitly limit the scope of the findings and include a discussion paragraph acknowledging that the results are task- and architecture-specific.\n\nAt this point, I fully understand the authors’ efforts and frustrations during the intense rebuttal phase. I have updated my score to 4, and I hope the authors will seriously consider the comments above. Once again, I sincerely appreciate the authors’ engagement and detailed responses throughout the discussion.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wZgHYxd3P0", "reviewer": "Reviewer_gppE", "summary": "This paper notes that when we do in-context learning, it is likely that the network is, in some sense, learning about the structure of the task. This paper considers task spaces that are explicitly low-dimensional, such as linear regression, where you can use in-context learning to give information about the linear regression.  To encourage the network to use this low-dimensional structure, they add a bottleneck to the transformer.  However, they find that this does not improve performance.  Thus, the paper is in many ways presenting a \"negative result\".", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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].\n\n> improve in-context learning by introducing a bottleneck, … very rare view in the field.\n\nWe 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.\n\nEven 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\n- MoE architectures [2] use only a subset of parameters\n- Perceiver models [3] introduce bottlenecks through learned latent variables\n- Retrieval and memory augmented methods [4] introduce bottlenecks by selecting a subset of data for context\n- Information bottleneck [5] are well studied for improving generalization\n- 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).\n\nIn 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.\n\n> transformer based LLMs work really well.\n\nWe agree with the reviewer that they do, but this should not disincentivize research into improving or analyzing them.\n\n> limited to settings with known latent variables\n\nOur 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.\n\n> interpretability is most interesting when we don't know the latent variables.\n\nWe 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.\n\nWe 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.\n\n[1] Karl, F., Kemeter, L. M., Dax, G., & Sierak, P. (2024). Position: embracing negative results in machine learning. arXiv preprint arXiv:2406.03980.\n\n[2] Liu, Aixin, et al. \"Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.\" arXiv preprint arXiv:2405.04434 (2024).\n\n[3] Jaegle, Andrew, et al. \"Perceiver io: A general architecture for structured inputs & outputs.\" arXiv preprint arXiv:2107.14795 (2021).\n\n[4] Lewis, Patrick, et al. \"Retrieval-augmented generation for knowledge-intensive nlp tasks.\" Advances in neural information processing systems 33 (2020): 9459-9474.\n\n[5] Tishby, Naftali, and Noga Zaslavsky. \"Deep learning and the information bottleneck principle.\" 2015 ieee information theory workshop (itw). Ieee, 2015.\n\n[6] Doshi-Velez, Finale, and Been Kim. \"Towards a rigorous science of interpretable machine learning.\" arXiv preprint arXiv:1702.08608 (2017).", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "N/A (no such theoretical claims).", "experimental_designs_or_analyses": "I am satisfied that the evidence they present supports their claims.", "relation_to_broader_scientific_literature": "Good connections drawn in the Introduction. Related work is in the Appendix, which I don't really mind, but some might object to.", "essential_references_not_discussed": "None to my knowledge.", "other_strengths_and_weaknesses": "My central issue with this paper is that it is --- ultimately --- a negative result.\n\nThat can be fine if the negative result is sufficiently interesting, surprising and convincingly argued.\n\nHowever, I don't think that's the case here.  In particular, while the authors clearly expected to improve in-context learning by introducing a bottleneck, I believe that this would be a very rare view in the field.  That's basically because transformer based LLMs work really, _really_ well.  In particular, they work really well at in-context learning, and they work really well at a broader array of tasks that share character with in-context learning (even completing the next token using a pre-trained model requires bringing in lots of information from the context).  Now, perhaps _the_ essential component of transformers is self-attention.  And of course, self-attention is the _opposite_ of a bottleneck, as it allows each token to attend to any previous tokens.  So I'm really not sure who would expect introducing a bottleneck to improve performance.  Indeed, whenever we introduce bottlenecks of any form into attention (sliding window, quantised KVs etc.) you get worse performance.\n\nAdditionally, the interpretability results aren't that interesting as:\n* They're restricted to a network with bottlenecks that no one (presumably, not even the authors) would use in practice.\n* Their experiments are limited to settings with known latent variables (interpretability is most interesting when we don't know the latent variables).\n\nNow, if you buy that the result is expected, then I'm really not sure that the result is suitable for ICML.  I would instead recommend that the authors consider a venue such as TMLR, which has two key criteria for acceptance:\n\"Are the claims made in the submission supported by accurate, convincing and clear evidence?\"\n\"Would some individuals in TMLR's audience be interested in the findings of this paper?\"\nThe paper does clearly meet these thresholds.  But ICML requires something more (which was the entire point of setting up TMLR).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3grUs9k25q", "reviewer": "Reviewer_TjQ5", "summary": "This paper delves into the mechanism for Transformers to do In-Context Learning (ICL). A common belief is that TF do CL through some statistical shortcuts and hence can not generalize well in OOD tasks. The authors test this hypothesis by minimally modifying the architecture to encourage the model to explicitly aggregate the information in the context to learn the task representation before performing the final prediction to the label fo the query. By modifying the architecture, the authors are able to compare a usual TF (the so-called implicit model) and the modified one (the explicit model). They conduct a series of experiments on various ICL tasks, including linear regression, non-linear regression and classification, and some reasoning tasks. They showed that 1. The explicit model does not outperform the implicit model in various tasks and in both ID and OOD cases. 2. They prove that the explicit model can learn the task representation very well, so the reason why they are poor in OOD cases is that they did not learn the final prediction part well. Replacing the final prediction function with an oracle leads to a much better performance. \n\nIn general, the goal, experiments, and results are pretty well presented in this paper. The experiments are sufficient and clear. I like the style of this paper. Although they focused on a small question, they studied it deeply.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\nTo 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).\n\nFinally, 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}$.\n\nWe hope that our response has resolved the reviewers' concerns and would be happy to provide further clarifications.", "claims_and_evidence": "I have one suggestion.\n\n1. The OOD task for linear/nonlinear regression/classification only relies on sampling the query input from a distribution with a larger variance. This is not very sufficient since this type of 'query shift' can be well tolerated even by a single layer of the linear self-attention [1] model. So I will suggest trying more complex distribution shifts and see what happens for both models.\n\n[1]. Trained Transformers Learn Linear Models In-Context.", "methods_and_evaluation_criteria": "/", "theoretical_claims": "/", "experimental_designs_or_analyses": "/", "relation_to_broader_scientific_literature": "/", "essential_references_not_discussed": "/", "other_strengths_and_weaknesses": "/", "other_comments_or_suggestions": "/", "questions_for_authors": "/", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "/"}}, {"review_id": "vObkLdYUp9", "reviewer": "Reviewer_vn7o", "summary": "This paper investigates whether explicitly inferring task-relevant latent variables improves in-context learning (ICL) in Transformer models. The authors introduce an explicit model that enforces structured inference of latent variables and compare it with a standard implicit model that learns ICL end-to-end without explicit latent variable inference.\n\nThe authors find that explicitly inferring task latents does not improve generalization. The explicit model effectively extracts latent variables but struggles to utilize them for robust predictions. This suggests that the challenge in ICL is not just learning task latents but also correctly leveraging them in downstream prediction.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing valuable and constructive feedback.\n\n> Single latent variable can fully summarize context information independently of $x_{query}$ seems unnatural\n\nIt 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\n\n- 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}$. \n- The same assumption is also used in representation and unsupervised learning.\n- Multiple works on ICL also claim this to be the underlying mechanism (Eq. 1 in [5,6])\n\nSince 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].\n\n> restricting direct attention from x_{query} to other context tokens may artificially constrain the model’s behavior rather than isolating a meaningful causal mechanism\n\nIn 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$.\n\n> Implicit models may also recover task-relevant latents\n\nIndeed 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.\n\n> failure of the explicit model to leverage inferred latents for prediction\n\nWe 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.\n\n> bottleneck structure itself may limit information flow, rather than revealing a true failure of latent inference\n\nMost 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?\n\n> why does training another classifier on the same inferred latents yield better results?\n\nIn 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.\n\nWe hope that our response has resolved the reviewer's concerns and would be happy to provide further clarifications.\n\n[1] Shalev-Shwartz, Shai, and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.\n\n[2] Andreas, Jacob, et al. \"Neural module networks.\" Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.\n\n[3] Tishby, Naftali, and Noga Zaslavsky. \"Deep learning and the information bottleneck principle.\" 2015 ieee information theory workshop (itw). Ieee, 2015.\n\n[4] Sun, Yu, et al. \"Test-time training with self-supervision for generalization under distribution shifts.\" International conference on machine learning. PMLR, 2020.\n\n[5] Müller, Samuel, et al. \"Transformers can do bayesian inference.\" arXiv preprint arXiv:2112.10510 (2021).\n\n[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).", "claims_and_evidence": "I am unsure whether the chosen setting and the assumption that a single latent variable can effectively summarize all context information is the best way to explain how Transformers behave. In particular, enforcing this assumption by preventing x_query from attending to other context tokens may impose an artificial constraint that does not align with how Transformers naturally process information.", "methods_and_evaluation_criteria": "See previous section.", "theoretical_claims": "No theory in this paper.", "experimental_designs_or_analyses": "To some extend see \"Claims And Evidence\" section.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The paper explores a potentially interesting question regarding the role of explicit latent variable inference in in-context learning. However, in its current form, the study feels **incomplete** in several ways.  \n\nFirst, the assumption that a **single latent variable can fully summarize context information independently of \\( x_{query} \\)** seems **unnatural**. Transformers process information dynamically, and restricting direct attention from \\( x_{query} \\) to other context tokens may **artificially constrain** the model’s behavior rather than isolating a meaningful causal mechanism. This setup may not fully capture how in-context learning operates in standard architectures.  \n\nAdditionally, while the negative result is valuable, the study **does not sufficiently rule out alternative explanations**. For example:  \n- **Implicit models may also recover task-relevant latents** in their final-layer representations, but this is not systematically tested.  \n- The failure of the explicit model to leverage inferred latents for prediction is **not fully explained**—is this due to architectural constraints, optimization dynamics, or a fundamental limitation of explicit inference?  \n- The **bottleneck structure itself may limit information flow**, rather than revealing a true failure of latent inference to improve generalization.  \n\nThese concerns are further detailed in the **\"Questions for Authors\"** section. Addressing them could significantly **strengthen the clarity and impact of the paper.**", "other_comments_or_suggestions": "I find it difficult to pinpoint a clear takeaway from this paper. While it presents several interesting observations, the main message remains unclear.", "questions_for_authors": "Thank you for your insightful work on the role of latent variable inference in in-context learning. I have a few questions regarding your methodology and findings:\n\nRecovering True Latents in Implicit Models\n\n1. In your experiments, you show that the explicit model successfully extracts the correct task latents. However, have you considered whether a standard Transformer (implicit model) might also recover similar latent representations, perhaps in the final-layer embeddings or through a linear probe analysis? This would help clarify whether the explicit model is truly unique in this regard.\nDependence of Latents on Query Input\n\n2. The explicit model prevents the query x_query from directly attending to the context. However, does this design fully isolate the necessity of learning the correct latent variables? If the true latent is inherently dependent on x_query (e.g., in tasks where the latent determines the relationship between query and context), wouldn’t this architectural constraint potentially hinder optimal inference?\nFailure to Leverage Learned Latents for Prediction\n\n3. One of your key findings is that even when the explicit model correctly infers task latents, it does not generalize better. Could this be an issue with optimization dynamics (e.g., the choice of Adam or gradient descent) rather than an inherent failure to use the latents? Given that the last layer acts as a simple classifier, why does training another classifier on the same inferred latents yield better results? This suggests a mismatch between latent inference and prediction that is not fully explained.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper initially received mixed reviews. 
- Reviewers 1Akw and TjQ5 were strongly supportive. However, TjQ5's review lacked detail, and the reviewer did not participate further, such as by providing clarifications or joining the discussion. Therefore, I downweighted Reviewer TjQ5's score accordingly.
- Reviewer 1soF is leaning positive but critical of certain comparisons. 
- Reviewer gppE found the results unsurprising and despite the rebuttal thinks that the paper does not meet the ICML standards.
- Reviewer vn7o found the setup and main message unclear, even after the authors' rebuttal.

In the reviewer discussion phase, reviewers 1soF and vn7o again expressed their concerns. Following a note by the authors indicating that 1soF might have overlooked recent results, the reviewer updated their score and now recommends acceptance.

While reviewers gppE and vn7o have expressed concerns, the majority recommend acceptance. I suggest that the authors address the concerns raised by these reviewers in the final version of the paper, should it be accepted.
Decision: Accept (poster)