Reconciling In-Context and In-Weight Learning via Dual Representation Space Encoding
Authors: Guanyu Chen
Ruichen Wang
Tianren Zhang
Feng Chen
Venue: Accepted by TMLR ()

Abstract: In-context learning (ICL) is a valuable capability exhibited by Transformers pretrained on diverse sequence tasks. However, previous studies have observed that ICL often conflicts with the model’s inherent in-weight learning (IWL) ability. By examining the representation space learned by a toy model in synthetic experiments, we identify the shared encoding space for context and samples in Transformers as a potential source of this conflict. To address this, we modify the model architecture to separately encode the context and samples into two distinct spaces: a \textit{task representation space} and a \textit{sample representation space}. We model these two spaces under a simple yet principled framework, assuming a linear representational structure and treating them as a pair of dual spaces. Both theoretical analysis and empirical results demonstrate the effectiveness of our proposed architecture, CoQE, in the single-value answer setting. It not only enhances ICL performance through improved representation learning, but also successfully reconciles ICL and IWL capabilities across synthetic few-shot classification and a newly designed pseudo-arithmetic task. The code is available at: \url{https://github.com/McGuinnessChen/dual-representation-space-encoding}.

Published in Transactions on Machine Learning Research (02/2026)
Reconciling In-Context and In-Weight Learning via Dual Representation Space Encoding
Guanyu Chen
hhen-
Deportment of Automation, Tsinghua University
Ruichen Wang
mong-re22 edu.ccm
Department of Automation, Tsinghua University
Tianren Zhang
trahanng@mailltsiinghu... edu.ccn..e.ccc...
Department of Automation, Tsinghua University
Feng Chen
Shenfeng@@maill.sinnguu en..cn
Department of Automation, Tsinghua University
Reviewed on OpenReview: https: // openeevie.......... net/ forun? id- DJK7VIOKAU
Abstract
In-context learning (TCL.) is valuable capability exhibited by Transformers pretrained on diverse sequence tasks. However, previous studies have observed that ICL often conflicts with the model's inherent in-weight learning (IWL) ability. By examining the representation space learned by a toy model in synthetic experiments, we identify the shared encoding space for context and samples in Transformers as a potential source of this conflict. To address this, we modify the model architecture to separately encode the context and samples into two distinct spaces: task representation space and sample representation space. We model these two spaces under simple yet. principled framework, assuming linear representational structure and trcating them as pair of dual spaces. Both thcoretical analysis and empirical results demonstrate the effectiveness of our proposed architecture, CoQE, in the single-val answer setting. not only enhances ICL performance through improved representation learning, but also successfully reconciles ICL and IWL capabilitics across synthetic few- shot classification and newly designed pseudo-arithmelic task. The code is available at:
1 Introduction
In recent years, large-scale models based on the Transformer architecture have demonstrated remarkable capabilities across language (Brown et al. 2020 Guo et al. 2025), vision (Achiam et al. 2023 Maaz et al. 2024), and robotics (Driess et al. 2023; Zitkovich et al. 2023 Among these capabilitics, the in-context learning (ICL) ability has drawn increasing attention, as it ollers a general paradigm for task generalization. ICL refers to the capability of pretrained Transformer model to solve previously unseen tasks by using demonstration examples in the prompt-withe
updating its parameters. In contrast, in-weight learning (IWL) characterizes the conventional ability of model to recall the memory stored in weights. An ideal model would seamlessly integrate both capabilities: relying on memory to handle training tasks, while adapting to new tasks through contextual cues.
However, recent studies suggest that there exists an inherent conflict between ICL and IWL Park et al. 2025 Nguyen Reddy 2025). This leads to notable performance degradation when the demonstration examples deviate from the training distribution (Chan et al. 2025), thereby limiting the generalization ability
1 "Co-corresponding authors.

---

Published in Transactions on Machine Learning Research (02/2020)
of ICL. This conflict is closely related to training settings such as data distribution CCaa et al.) 2022 model size (Shi et al. 2024), and training time (Singh et al. 2025 Singh et al. (2023) hypothesize that it arises from the competition between the two capabilities for the shared model circuits. Nguyen & Reddy (2025 on
the other hand, attributes this to the different relative learning rates of ICL and IWL. Furthermore, some
studies (Chan et al., 2022 Singh et al. 2023 Anand el al.] 2025) have attempted to alleviate the lension between ICL and IWI. but can only achieve a tradooff under particular Zipfian data distributions. Hence, it remains a valuable question to understand and eliminate the conflict between ICL and IWL.
In this work, we aim to understand why models tend to oscillate between ICL and IWL capabilities under different settings from the perspective of representation learning. We observe that in models with strong IWL performance, the learned representation space provides a sufficiently rich representation of the query samples themselves. In contrast, models with strong ICL performance learn representation spaces that effectively encode contextual information. However, these two desirable structures of representation are difficult to acquire simultancously, making it challenging to achieve strong ICL and IWL performance at the same time (see Figure | bottom). Motivated by this observation, we hypothesize that the ICL-IWL, conflict can be alleviated by encoding the context and query samples into separate spaces, thereby preventing interference between the two types of information. We refer to these spaces as the task representation space and the sample representation space, respectively. To provide a principled modeling of these two spaces, we build upon the widely accepted linear representation hypothesis (Mikolov et al., 2013 Nanda et al. 2023 Park et al. 2024) and propose to model the task representation space as the dual space of the sample representation
space. We prove the completeness of a sample representation space under sufficient training tasks, which could facilitate task generalization by ICL. We also formalize the entangled nature of Transformers' encoding process standard softmax attention does not support such dual-space structure, highlighting contrast with linear attention mechanisms commonly adopted in recent theoretical analysis.
Furthermore, we propose a straightforward architectural design, CoQE, as an algorithmic exploration. Unlike standard Transformers, CoQE employs separate pathways to encode context and query samples, aiming to learn the lask and sample representation space, respectively. The final model output obtained by computing the inner product between elements from the two spaces according to the Riesz representation theorem. Our synthetic experimental results show that CoQE not only achieves lower ICL error in both in-distribution and scenarios, but also reconciles ICL and IWL across synthetic few-shot classification task and our newly designed pseudo-arithmetic task.
of Chan et al. b Singh et al. 2023
2025
2 Related Work
Investigations on ICL Mechanisms. From the theoretical perspective, existing works have examined
how Transformers perform ICL across various scenarios (Zhang et al. 2024a Li ct al. 2023 Tian ct al. [023)]
Nichani et al. 2024 Chen et al. 2024 Wu et al. 2024 Huang & Ge 2024 Oko et al. 2024 Liang el al. 2025
These studies typically analyze simplified architectures such as linear self-attention or query-key-coml kembbomeddm
formulations, and many of them conduct analyses from the perspective of Bayesian model averaging (Zhang
ct al. 2025b Ye et al.
2024 From the empirical perspective, Garg et al. 2002 firstly demonstrated that
seansformer-baseee based TCL. can generalize elfectively to out-of- ributtion (OOD) tasks, leading to a surge of
interest in exploring its generalization behavior (Ahuja & Lopez- 2023
Kossen et al. 2024 Pan et al.,
2023
Fan ct al.) 2024 Xiong ct al. 2025 Yadlowsky et al.] 2023 Bu ct al. 2025). Some work has shown
that LLM can implicitly encode task vectors during ICL (Hendel et al.) 2023
Todd et al. 2024 Guo et al.
2024
Yang et al., 2025 Han et al. 2025). Unlike these works that focus only on task information encoded in the model's representation space, we examine the relationship between task information and the intrinsic information of the samples. From this perspective, we seek to explain the conflict between ICL and IWL.
Relationship between ICL and IWL. Beyond investigations on the ICL mechanisms, some studies have found that ICL is not a guaranteed and stable capability of Transformers; rather, it competes with the model's inherent in- weight learning (IWL) ability, which relies on information stored in the weights (Chan et al. 2022 Singh et al.] 2023 Reddy 20044) Panwar et al., 2024).. Chan et al. (0022) examined the impact of different training data distributions on both abilities, finding that only when the training data follows certain Zipfian distribution can both abilities coexist. [Singh et al. (2023) further confirmed the transient
2

---

Published in Transactions on Machine Learning Research (02/2026)
nature of ICL, observing that it always fades after emerging and gives way to IWL. Nguyen & Reddy (2025) attributes this to the different relative learning rates of TCL and IWL, and conducted an analysis on a simplified one-layer transformer model. [Chan et al. (0025) proposed simple theoretical model, which is a linear combination of an in-weight learner and an in-context learner. [ingh et al. (2025) empirically discovered a more complex coopetition relationship between ICL and IWL. However, to date, no work has truly achieved coexistence between ICL, and IWL across varied training conditions.
Linearization in Latent Space. Beyond task-specific vectors, line of work has examined how models
internally encode variety of abstract concepts as linear vectors in latent space, giving rise to the commonly
accepted linear representation hypothesis (Mikolov et al. 2013 Nanda et al. 2023 Park et al.
2024). Several
studies have shown that concepts such as truthfulness Marks & Tegmark, 2024). time and
space (Gurnee
& Tegmark 2024), and other semantic properties (Dalvi et al. 2022 Merullo et al., 2024
Ye et al. 2025)
can emerge in the model's latent space, using lincar probes as the primary tool. Additionally, larger models
tend to yield more disentangled and interpretable internal representations (Bricken al. 2023 Cunningham
et al. 2023). and this be regarded evidence of world model within large models (Zhang al. 2025a In this work, to provide principled treatment of the encoding spaces of context and query samples, introduce the notion of a linear task representation space, grounded in the linear representation hypothesis.
3 Preliminaries
In-context learning setup. The basic sctup for analyzing ICL was first introduced by Garg et al. (2022) and has since been widely adopted (Yadlowsky et al. 2023 Pan et al. 2023). Consider distribution Dx over an input space X' S IRd,, and let F denote a class of functions over a distribution Dy. For cach prompt, first sample task f Dr, then draw set of input-output pairs ((x,,;;) where x; y; = fxx).. These sample pairs serve context demonstrations. Then, we independently generate a query input Xg Dx. The final prompt is gathered as a sequence:
i.d.. Dx and
P (x1,N11... Xn. Nno X,).
The ICL capability of pretrained model Mi, refers its accuracy produce predictions 8, (ll(P) for Y, f(x,), without having explicit knowledge of the current task and without updating its parameters.
Chan et al. 2022) extend this setting by introducing few-shot image classification tasks. In this setup. x represents an encoded image, and F, as set of classifiers, maps X' to finite label set y. The ICL capability refers to the model's ability to correctly classify a query image Xg based on the demonstrations.
Transformer model. A standard single-head self-attention (SA) layer Vaaswan et al. 2017 operates on an input matrix Z E Rd.XL, where L is the sequence length and de the embedding dimension. Let Q WoZ, K Wx Z, V Wy Z with Wo, WK e Rd xd, and Wv Rd.xde The attention output is
KTQ
SA(Z) = Z + Wo V softmax
dd
where Wo xd and the softmax is applied column-wise.
For the theoretical analysis of ICL, the prompt P is typically re-organized into an embedding matrix:
Xg X1 Xn yn 0 %))) E
Z =
V1
where de is the input feature dimension. Moreover, they often use linear self-attention (LSA) variant obtained by removing the softmax and merging parameters:
LSA(Z) = Zz WovZZ"WkoZ,
where Wov WoWy, WKQ WWWg E [d are trainable, and 1/n scaling constant. The model prediction 8, for the query is taken the bottom-right entry of LSA
3

---

Published in Transactions on Machine Learning Research (02/2026)
1.0
(64, (8,012))
Transformer or CoQE
0.8
Train sequence:
(64, 1)
IWL accuracy
b
b 621
6
216
0.6
621
216
45
216
579
621
ICL eval:
0.4
b
a
b
b
O
b
Q
IWL eval:
0.2
b
e
Q
(64.12.0.9.0)
2
Transformer
967
0.0
0.5
0.6
0.7
0.8
0.9
1.0
context
query answer
ICL accuracy
labeled as context
labeled context
no supportinn context
Figure (Top left) Overview of our synthetic task setup. (Top right) ICL and IWL performances under different training settings (E, L, Poursty, a). The Transformers fluctuate between ICL and IWL capabilities, whereas CoQE models robustly reconcile the two capabilities. (Bottom) Representation visualization of ten classes on distinct context conditions. We observe that good clusters for samples and good clusters for contexts are hard to achieve simultancously. Detailed discussions are presented in Scction
Dual space. Before formally introducing our dual-space modeling framework, we first present the general mathematical definition of the dual space.
Definition 3.1 (Dual space). Let V be a finite- dimensional inner product space over field F (typically R C) with inner product (.-)). The dual space of V, denoted V*, is the set of all linear functionals from V to F:
A (f: V +F is linear).
(1)
For every f E V there exists a unique vector w E V, called the Riesz representation of f, such that f(v) = (w, v), V. Let {ei,..., en) the basis V. The dual basis (e',...,,e") e". C V is defined by e'(e,) = 6 1 < i.j < n, where 5,y the Kronecker delta.
In the following, we will show that this dual-space formulation can be used to model the relationship between task representation space and the model's sample representation space. Moreover, by the Riesz representation theorem, elements from the two spaces can be composed via inner product.
4 Representation Space Analysis
In this section, we train and evaluate Transformers on synthetic datasets (Figure top left.) under a range of training conditions. By examining the learned representation spaces, we investigate how models with good ICT performance differ from those that excel in IWL.
4

---

Published in Transactions on Machine Learning Research (02/2026)
training serning
0.8
ICL accuracy
(64,1200..22
(84,12009.1)
0.7
(64.12.0..0)
70
70
0.6
1OL
80
60
0.5
50
50
0.0
0.1
0.2
0.3
Context silhouette coefficient
100
100
1.00
IWL acouracy
IWL accuracy
80
80
0.75
80
60
0.50
40
40
0.25
20
20
0.00
0.0
0,2
0,4
0.6
64
256
512
4
a
12
Sample silhouette coofficient
ResNet embedding dimension (E)
Transformer layer depth (L)
Figure 2: (Left) For three training settings and four training checkpoints (3k, 10k, 50k, 100k steps), there are clear positive correlations between ICL performance and CSC, and between IWL performance and SSC. (Right) Through various experiments, we observe that larger improves IWL but causes ICL to disappear, and larger L consistently enhances ICL while slightly harming IWL. Different colors and line styles represent different training settings. Detailed results are provided in Table
4.1 Synthetic Task Setup
We use synthetic few-shot classification task based on the Omniglot dataset (Lake et al. 2015 The dataset contains 1, 623 character classes, each with 20 samples. We construct prompt sequences that each consists of eight image-label pairs followed by a query image (Chan et al.) 2022 Singh et al., 2023 2025). We use a ResNet to encode images input, and the training objective minimizes the cross-entropy loss between the modcl's prediction and the correct label for the query image.
Training sequences have two key properties that affect the tradeoff between ICL and IWL: burstiness and Zipfian exponent. In bursty sequences, three out of the eight image-label pairs in the context share the same class as the query sample. This setup allows the model to infer the correct label based on context alone, which has been found to incentivize ICL while suppressing IWI. (Chan et al. 2022). To avoid repetition biases, bursty sequences additionally include three image-label pairs from a distinct distractor class. Pbussy denotes the proportion of bursty soquences in the training set, while the rest are generated via random sampling. The second factor the Zipfian exponent, which controls the frequency distribution of different classes. Under the Zipfian distribution, the class probability is defined as p(R = r) 1/ra, where R is the rank of the class, and is the Zipfian exponent. When O, the distribution becomes uniform. Chan et al.) (2022) observe that when = 1, a sweet spot emerges, where the model reaches tradeoff for both ICL and IWL. During training, we keep image-label mappings fixed.
Test sequences are divided into two kinds, corresponding respectively to the evaluation ICL and IWL capability. For ICL, we use sequences with four images from each of two classes, and we set the class labels either or randomly for each sequence. Accuracy on this evaluator measured across and 1 possible outputs, and chancc-level accuracy is 50%. As these labels not associated with these images during training, the only way to achieve above-chance accuracy is to refer back to the context. For IWL, we use sequences where none of the context images come from the same class as the query, but all of the image- label mappings are the same as during training. In this case, ICL is not useful, as there are no matching images in context, so the model must rely on mappings stored in weights.
4.2 Observations and Analysis
During training, we consider four factors: the ResNet embedding dimension E, the number of Transformer layers I,, the bursty proportion Poursty; and the Zipfian exponent a. The training outcomes under different
5

---

Published in Transactions on Machine Learning Research (02/2026)
combinations of these factors are shown in the red region of Figure top top right. We also collect and visualize the query sample representations of three models on the boundary of the red region, corresponding to the configurations (64, 12, 0.9, 2), (64, 12, 0.9, 0), and (64, 0.9, 1). Specifically, we examine the token embed- dings of x2 under three different context conditions: (1) when four images of the same class as X2 appear in the context and are labeled as 0, (2) when they are labeled as 1, and (3) when no image of the same class appears in the context. For the choice of token embedding layer, we empirically find that different layers yield consistent qualitative results; therefore, we directly use the final layer embeddings. We focus on the ten most common query classes, and the UMAP visualizations are shown in the Figure bottom.
Observation 1: ICL corresponds to good representations for context, while IWL corresponds to good representations for samples. The two are difficult to achieve simultaneously. As shown in Figure [ bottom, the first model forms well-clustered representations of query samples from different categories (points with different colors), but fails to capture the label distinctions provided in the context demonstrations (the squares and circles of the same color overlap). This indicates that the learned repre- sentation space sufficiently encodes intrinsic information of samples but lacks a clear encoding of the task information contained in the demonstrations. Consequently, the model achieves high IWL accuracy but ex- hibits almost no ICL capability. In contrast, the second model shows the opposite behavior: it fails to form distinct clusters for samples from different categories but demonstrates sensitivity to contextual information (the squares and circles of the same color are shifted apart). So the model has little IWL ability but possesses some ICL capability. The third model achieves a tradeoff between the sample information and contextual information in its reprosentation space, thereby realizing a tradcoff between ICL and IWL performance as well. However, the training conditions are rather restrictive, and both abilities remain suboptimal.
To further support our observation, we employ two quantitative metries derived from the standard silhouette coefficient, which evaluates cluster cohesion and separation. The silhouette score s(x) mcasures how similar sample x is to its own cluster compared lo other clusters. Using Euclidean distance d(.;)) +), we define a(x)
as the average distance between x and all other points in its own cluster Cz, and b(z) as the average distance to the nearest neighboring cluster Ck (where k + x).
a(x) = C2l
1
Σ
d(rxx',, b(x) = min (
Σ d(x,,a)) x'ECs
(2)
C.fC.
|Cal
of E'EC xx
The silhouctte score is then given by:
b(x) - a(x)
s(x) max(a(x), b(x)
(3)
The first metric we employ is context silhouette cocfficient (CSC). It measure whether the model learns 'good representations for context'. For a fixed object class c [I....,N} (with N = 10), we treat the samples as two distinct clusters based on their contextual labels (0 or 1). Here, a(z) measures the distance to samples with the same context label, while b(=) measures the distance to the opposite context label. We average the scores within each class and then across all N classes:
CSC = N
Σ S(2).
(4)
ZEZo
The second metric is sample silhouette coefficient (SSC). SSC adapts the metric to measure "good represen- tations for samples' without context support. In this setting, the clusters are defined by the N ground-truth object classes. For a sample i with class Vi, a(i) is the intra-class distance, and b(i) is the distance to the closest distinct class. The final metric is the average across all M samples of all classes:
M
SSC = M
Σ (1).
In Figure leff,, we report how ICL performance correlates with CSC, and how IWL performance correlates with SSC. Both show a clear positive correlation, providing quantitative evidence for our observation that the structure of the learned representation space is tightly linked to the model's ICL and IWL behavior.
6

---

Published in Transactions on Machine Learning Research (02/2026)
Observation 2: Model size also affects the ICL-IWL tradeoff. We make another new finding that model size also strongly affects the ICL-IWI. tradeoff in standard Transformers (Figure 2 right). We observe that, under the same conditions, a 12-layer Transformer exhibits stronger ICL but weaker IWL compared to a 4-layer Transformer. We suppose that this is due to the Transformer's 's inductive bias toward attending to context, compared to just memorizing context-irrelevant sample information. We also observe that increasing the ResNet embedding dimension from 64 to 512 nearly eliminates the model's ICL ability while substantially enhancing IWL. Notably, we connect fully connected layer after the ResNet to reduce the dimension back to 64 before imputting to the Transformer, ensuring that the latter's role remains unchanged. We speculate that the larger ResNet increases the expressivity of individual sample tokens, and when single token is sufficiently expressive to solve the task, the model tends to ignore the context. This is consistent with Observation that sample information and contextual information are difficult to coexist during training.
5 Dual-Space Modeling of Task and Sample Repressentations
We have shown that the conflict between ICL and IWL originates from the difficulty simultaneously en- coding sample information and contextual information within the same representation space of Transformers. We aim to reconeile ICL and IWL by encoding the context and the query samples into separate represen- (ation spaces, thereby addressing their incompatibility within a single space. This raises new question: how can we provide a principled modeling of these two spaces such that they can interact in a coherent and meaningful way? To this end, we start from the simple yet powerful assumption of a linear representational structure and propose our dual-space modcling framework. All proofs are provided in the Appendix A
5.1
Task Representation Space
We begin with the widely acknowledged linear representation hypothesis (Mikolov et al. 2013 Nanda et al. 2023 Park et al. 2024), which is the idea that models internally encode abstract semantics samples as linear vectors in latent space. And because tasks are inherently semantic, this hypothesis can be formalized from a measurement perspective according Park et al. (2024).
Definition 5.1 (Lincar sample representation space). Let X C C denote the input space and y the label set. A linear sample representation space M C Rd is a finite- penite-dimensionall inner product space equipped with a mapping o: X > M, such that
1. (Learnability) o is parameterized by a model M. and can be learned from data;
2. (Linear Measurement) In the case of regression with y S R, there erists a linear transformation w such that, given any (x,y) pair, the label can be expressed
y (ws, o(x))-
(6)
In the case of classification with y {0, 1}, the label probability is given by
logit.(( = x) = (w, ((x)).
(7)
Definition 5.1 formalizes the notion of a sample representation space under the linear representation hypoth- esis in the single-task setting. can be straightforwarde extend to the multi-task case, assuming that there exists shared linear sample representation space across different ICL tasks. Note that this assumption has been implicitly embedded in wide range of thcoretical and algorithmic work (Garg al. 2022 Hu et al.,
2023
Kim & Suzuki 2024 Zhang et al. 2024b). Based on this assumption, we define the corresponding linear task transformation space. Without loss of generality, we consider only the regression case.
- R) denote a task function space defined over the input space X. We assume that there exists a sample representation space Mr S Rd, together with mapping o5, such that My is linear with respect to X and each label set 3/ (f(xx) x E X), Vf F. A linear lask transformation space is then defined as linear functional space T (1: M, + R), equipped with mapping 1b : F + T such that for any f E F, +(F)) =t satisfying
Definition 5.2 (Linear task transformation space). Let F ==
f(x) = [(r(()), x
(8)
7

---

Published in Transactions on Machine Learning Research (02/2026)
Building upon this foundation, our key theoretical contribution is that find the task transformation space can be modeled as the dual space of the sample representation space. See Appendix A.1 for the proof.
Proposition 5.3 (Task-sample duality). Lel X be the inpul space and Yf the multiple label sels corre- sponding to each task f F. Under Definition 5.2 there ewists a linear sample representation space My and a linear task transformation space T, where T is the dual space of MMx i.c. M
Definition 5.4 (Task representation space). Under Proposition 5.3 for each task E +(f) E T admits unique Riesz representation wy. The task representation space WF is defined the set of all such Riesz represenlalions. Then for any J E F, we have
f(x) (F)) VXEX
(9)
In summary, we map various nonlinear ICL tasks vectors in the task representation space, leveraging the linear representation hypothesis, the dual-space formulation, and the Riess representation theorem. From the above formulation, can further define basis representations and basis transformations.
Definition 5.5 (Basis task representations). Under Proposition 5.3, let (m1...., ,,a} be a basis of the sample representation space My, and let {t),...,ta) be the corresponding dual basis of the task transformation space T. The basis task representations are defined as the Riesz representations of (ti..... :ta) denoted by (wi....,wa)) which satisfy
(wi, dejs 1 isijsd.
(10)
Thus, every sample representation (x) can decompose uniquely as dx(x) E ax(x)m, and every task representation wy can decompose uniquely as wy = Σÿ= Brwj. The output can given by
d
(ws, 0x(x))
Σ (x(x),
Our next Theorem 5.6 shows that, under the dual-space modeling framework, a sufficient set of ICL tasks guarantces a basis-covering sample representation space. See Appendix A.2 for the proof.
Theorem 5.6 (Completeness of basis representations under task traversal). Under Proposition 5.3
assume that learner with sample representation mapping de is presented with task traversal curriculum C such that: span(t teC = T. Then, if the learner achieves zero empirical error, the learned representation x} My.
mapping de satisfies: spanf
It can be expected that, by promoting basis- coder-co sample representation space, our dual-space modeling framework could enhance ICL capability. We empirically validate this in Section 7.1
5.2 ICL Formulation under Dual-Space Modeling
In this section, we formulize ICL under our dual-space modeling framework.
Definition 5.7 (Context-induced task representation in ICL). the ICL setting, the task representation can be specified jointly by two components: (1) contert demonstration of labeled ezamples Zi,n .2.....:2....................... Zn) with z, (x,,,,) X V, and (2) representation mapping @ X &+RRR Rd. That is
wy
A wy(Zisms Φ)
(11)
Definition 5.7 formalizes that in the ICL setting, the task specified by a prompt is determined by its context portion. Thus, the encoding space of context naturally serves as the task representation space. We further show that existing theoretical analyses of ICL based on LSA (Zhang et al. 2024a Kim & Suzuki 2024 Wu et al. 2024) are fully compatible with our framework, from which we can derive closed form of wy.
Proposition 5.8 (Closed form of Wy under simplified LSA). Consider LSA layer applied after feature encoder Φ X + Rd implemented by an MLP. Suppose the LSA projection matrices WKQ and Woa are initialized such that
Wor = 1, O0
1),
1
0g
WkQ = 00
*
8

---

Published in Transactions on Machine Learning Research (02/2026)
output
output
task representation space
Task Encoder (contextual)
inner product
feed-forwardd (token -wise)
dual
Nx
Multi-head Attention (contextual)
xxx)
sample representation space
Sample Encoder (token-wise)
Embedding
context
query
context
query
(a) Transformer
(b) CoQE
Figure Comparison of Transformer and CoQE architectures. Unlike the Transformer, which encodes both context-level and sample-lovel information into the same representation space, CoQE implements dual representation spaces encoding that explicitly distinguishes between context and samples.
Then the final prediction takes the form ÿ = (Wy((1n 6), o(xy)), where
(y Φ) n y (xx:).
(12)
Proposition 5.8 5.8ccan can explain the effectiveness of LSA simplification in analyzing ICL: it implicitly performs our dual-space modeling between the task representation space and the sample representation space. However, we argue that it fails capture the entanglement of standard Transformers encoding progress, which use the original, unsimplified SA. As we will show in the next theorem, SA cannot realize such dual-space modeling,
Theorem 5.9 (Entangled structure under general SA). For standard SA model with softmaa-based atten- tion weights, there does NOT erist pair of do and wo(Zinn do), such that the model prediction admits the following decomposition:
8
do), (x())))
(13)
From our dual-space modeling perspective, Theorem 5.9 formalizes the entangled nature of how Trans- formers encode context and sample-level information. See Appendix A.4 for the proof. We posit that this entanglement is the underlying reason for the observed conflict between ICL and IWL.
6 CoQE: A Transformer with Separate Context-Query Encoding
We have established principled modeling framework for the task representation space and the sample representation space. Next, we propose straightforward yet effective architectural modification: CoQE, a Transformer with separate Context-Query Encoding.
The CoQE model consists of two modules: shared sample encoder de and a dedicated task encoder wi. as shown in Figure 3 (b). The sample encoder generates general-purpont representations for all samples, including the query. We implement it with token-wise module, for it should process samples independently without considering context. The task encoder, the other hand, operates on the general representations of the context and focuses on producing the representation of the current task. Thus this module should be contextual and has the capability to condense sequential information. Finally, the prediction output is ob- tained by computing the inner product between the task representation and the query sample representation. Taking the regression task as an example, the formalization of CoQE output is follows:
8q { wi ( (31,2)) ), 2o (xx) }-
(14)

---

Published in Transactions on Machine Learning Research (02/2026)
Figure X compares the architectures of the Transformer and CoQE. The Transformer also contains token- wise components like feed-forward networks, and contextual components like multi-head attention modules. When stacked, these modules collectively exhibit contextual behavior, and the final token output intertwines with the context information in a complex manner during the forward pass. In contrast, CoQE explicitly separates the contextual and token-wise parts, which are responsible for learning the task representation space and the sample representation space, respectively. The two spaces interact through a well-defined inner product according to the Riesz representation theorem.
We aim to evaluate model on regression tasks for ICL performance, and few-shot classification tasks for ICL-IWL performance. In the following, we will give the specific implementation of COQE. Notably, due to their different properties, the task encoder constructs the task representation space in distinct ways.
6.1 Implementation for Regression
We employ a two-layer ReLU network as the sample encoder of CoQE, and a GPT-2 style Transformer the task encoder. We take the final output token of the task encoder directly the lask representa- tion induced by the context. The regression output is then computed as the inner product between it and the query sample representation. For fair com- parison, the baseline Transformer also equipped with the same two-layer ReLU embedding module. See Appendix B.2 for more implementation details.
class vectors
label 0
Lonig
+
class vectors (modified)
noise
mod
X1
$2
X
22
X
Task Encoder
6.2 Implementation for Few-shot Classification
static
Transformer
We use ResNet CoQE's sample encoder, and Transfromer as the task encoder. The task repre- sentation space is constructed in a distinct way from regression. A multi-class classification task can be regarded as a collection of sub-tasks that identify each class. Thus, we let it correspond to a set of task representations, each of which is associated with one class. ICL requires producing the task representations corresponding to the classes in the context, whereas IWL requires static memorization of all classes. To construct task representation set compatible with both, we assign a parameterized vector to each class, representing a static version of its task representation.
dynamic
context embed
Figure 4: Construction and training of the task Rep- resentation space for few-shot classification.
In each forward pass, the classes appearing in the context are encoded by the task encoder to obtain their corresponding task representations, as illustrated in Figure ! These dynamic vectors replace the corresponding static class vectors, and modificd class voctors are used to compute logits for prodiction. The resulting training loss is denoted as Lmod- Additionally, to accelerate the training of the static class vectors, compute additional set of logits from the unmodified class vectors during training, with the resulting classification loss denoted as Corig. These logits are not used during testing. Therefore, the total training loss is Emod Corig. During experiments, we observed that Lmod tends to converge to Corig. which means the task encoder fails to dynamically encode the context over training, and the learning of the task representation space restricted to the set of static class vectors. To prevent the task representation space from collapsing during training, we add Gaussian noise to the modified logits during training. with the variance increasing over training steps. The initial noise follows (Mo, 1). This trick can be interpreted as indirectly performing random sampling in the task representation space, that is, (wy, 0(x)) + (wy + e', ((x)). Through it we maintain dynamic task representation space. See Section 7.2 for experimental results and ablation studies.
7 Experiments
In this section, we first evaluate the ICL performance of CoQE on regression tasks. Then we evaluate ICL-IWT. performance on the synthetic classification task and newly designed pseudo-arithmetic task.
10

---

Published in Transactions on Machine Learning Research (02/2026)
Linear
parse linear
2-layer NN
Combination
6
20
10
CoQE
DoCE
CoQE
CoCE
5
Transformer
Tranaforaar
Transformer 2-layer NN. GD
Transformer
Least Squares
Least Squares
4
Lasseo (alp) 01)
MSE
3
0.5
0.0
5
10
15
20
10
20
30
40
50
10
20
30
40
50
# In-context examples
# In-context examples
# In-context examples
# In-context examples
(a) In-distribution testing
Scaled X
Scaled X
Different orthants
Scaled
3
CoQE
50
Transfomer
Least Squares
8
40
MSC
30
20
2
5
10
15
20
20
15
20
In-context examples
# In-context examples
# In-contexi examples
# In-context exampies
(b) Out-of-distribuiie testing
Figure 5: Results of ICL regression. We provide optimal baselines for test settings except for combination functions. CoQE, consistently achieves lower TCL error than the Transformer in both ID) and OOD scenarios.
7.1 Regression
Setup. We adopl general framework for training models to perform ICL over a function class F. To construct training prompts, we first sample a task function f ~~ Digain, then draw k i.id. inputs ,I. Xx Diain The prompt is formed as P = (x1, fxxi)..... ..., Xk, Xk; fxx)). Let pi denotes the prefix containing the first input-output examples and the (i+1 1)th imput: p = (x1, (xx1)...,, x.. f(x,), The training objective of a model Mg minimizes the expected loss over all possible prefixes:
k-1
((Mo(P)), s(xCHi)) K Σ
min Ep
where .(..)) is a mean squared error (MSE) loss function. At test time, we first sample a test function f ~ Disat then draw j S k - inputs XI,...,X, Xj ~ Diiett and Xg from Dquery to construct the test prompt: pi = [x1,,(xx),--.... X,, fxx]: x²). Then we measure the MSE between M. and f(x,)-
To compare ICL capability of CoQE with the standard Transformer, we consider two major evaluation scenarios: in-distribution (ID) testing and dustr--disttiiuutin (OOD) testing. For ID testing, we set D Distt Pquery: and Dirain Disat Specifically, we use the following four classes of functions F: linear functions, sparse linear functions, two-layer ReLU networks and combination functions. The latter two classes of noulinear functions allow the model to reduce ICL difficulty by learning Lask-invariant representations. Through them, we can empirically validate Theorem 5.6 which shows the benefits of dual-space modeling for representation learning. For OOD testing, we consider four different cases of distribution shifts under linear functions. See Appendix B.2 for more setup details.
Results. In the ID scenario, CoQE consistently achieves lower ICL error than the Transformer (Figure 5 (a)). For regression on more challenging combination functions, the Transformer exhibits substantial fluc- tuations, whereas CoQE attains much smaller error variance. We attribute this to CoQF's more effective
11

---

Published in Transactions on Machine Learning Research (02/2026)
ICL
IWL
ICL
IWL
1.0
1.0
1.0
1.0
0.9
0.8
0.9
0.8
Accuracy
0.6
0.7
0.7
0.6
0.2
0.6
0.2
0.5
0.0
0.5
0.0
20
20
40
60
80
100
20
60
20
40
60
80
100
Steps (k)
# Steps (k)
Steps (k)
# Steps (k)
CoQE Transformer
burdy 0.9
Poursty 0.5
Bunsty 0.1
CoQE Transformer
0=
P Puursty 0.9
Pbursty 0.5
Paursiy 0.1
-0=
a=2
(a) Results under different burstiness
(b) Results under different Zipfian distribution
Figure 6: Learning curves of synthetic few-shot classification under different data distribution factors.
learning of the sample representation space, and present further results in Appendix B.2 In the OOD scenario, CoQE also achieves substantially lower error than the Transformer across all four tested (Figure 5 (b)). Notably, the second case is adapted from Mittal et al. (2025), who similarly aims to enforce the model to explicitly learn task variables. However, they found no improvement in OOD performance, contrary to our results. This indicates that simply introducing task variables is insufficient and highlights the value of our proposed dual-space modeling and corresponding architecture design.
7.2 Synthetic Few-shot Classification
Baselines. Besides standard Transformers, we introduce two more baselines that were proposed alle- viate the tension between ICL and IWL: L2 regularization (Singh et al. 2023 and probabilistic temporary forgetting (Anand et al. 2025
Results. Figure top right presents the ICL and IWL accuracies of Transformers and CoQE under various factors after 100k training steps. The Transformers fluctuate between ICL and IWL, capabilities across different conditions, whereas our models robustly occupy the upper-right blue region, indicating Parcto improvement in both abilities. The other two baselines perform even worse than the standard Transformer; for clarity, we omit them from the figure, and report the detailed results in Appendix B.3] Figure G shows the learning curves under different values of Poursty and Zipfian exponent. We could observe that CoQE's ICL accuracy rises rapidly at the beginning, but declines slightly between 10k and 30k steps. This behavior aligns with prior findings ICT strategy: it emerges quickly and then gradually fades (Singh al. 2023 However, under our algorithm, the model quickly restrains this fading trend and continues to recover steadily, We also discuss the effect of parameter scale, showing that CoQE outperforms the standard Transformer under the same parameter budget, and its performance continues to improve as the model size increases. All detailed results are presented in Appendix B.3
Ablation study. We study the effect of Gaussian noise on the model performance, as shown in Table Without any noise, the model's ICL ability ultimately yiclds entirely to IWL. When Ao = 5, the model achieves maximal ICL performance while retaining high IWL capability. This is the default noise magnitude used in our experiments. See Appendix B.3 for more analysis and details.
7.3 Extending to Language Model Token Embeddings
To further verify the generality of our approach and explore its potential extension to natural language pro- cessing (NLP) (asks, we are inspired by Singh el al. (2023) and exlend our few-shot classification experiments from images to token embeddings of large language models (LLMs). we construct a fixed token embedding dataset through selection and clustering for token embedding matrix of Llama series (Dubey et al.) 2024)). see Appendix(B.4 details. Compared with Omniglot images, this dataset has larger intra-class variability and preserves meaningful semantic information relevant to NL.P tasks. We fix the training data distribution
12

---

Published in Transactions on Machine Learning Research (02/2026)
Table 1: Results under different noises. We
Table 2: Results of using language token embeddings.
use Ao 5 default noise magnitude.
Token Embedding
Model
ICL
IWL
ICL
IWL
Transformer
50.12
96.65
Llama-3.1-8B
frese-free
55.12
99.62
CoQE
65.35
93.96
Mio 3
81.91
95.31
Transformer
49.33
98.81
Aio 91.15
89.30
Hamma-322-B
CoQE
75.01
96.57
μt 7
88.22
77.70
9
86.01
72.62
Llama-3.2-1B
Transformer
49.95
99.12
CoQE
80.08
97.34
Conditional Pseudo-Arithmett Task
ICL
IWL
Prompt of Fine Tuni-nnnng and IWL Eva!
1.0
Condition: 15 23 38. 43 43 = 86, 17 5 22.
Question: Chef cooked 17 steaks and burgers. How many dishes did he cook? Answer:
Accuracy
0.6
Generate weihtss
22
0.4
0.4
Prompt of ICL Eval
CoQE
Condition: 15 = Y
0.2
0.2
Venilla Fine-Tuning
Question: Chef cooked Y steaks and burgers.
Regulsrization Temporary Forgetting
many dishes did he cook? Answer:
Generate prior info weights
300
300
C
Fine-tuning steps
Fine-tuning steps
Figure 7: (Left) Overview of conditional pseudo-arithmetic lask. (Right) Training curves of CoQE and three baselines. The IWL accuracy of CoQE improves more slowly, because its learning paradigm differs from the model's original autoregressive pro-training. However, after fine-tuning, all methods achieve nearly 100% IWL performance, and CoQE achieves substantially higher ICL accuracy.
parameters at Auurty- = 0.9 and Zipfian exponent a The results in Table 2 show that the Transformer almost fails to acquire ICL ability, consistent with [Singh] et al. (2023). In contrast, CoQE achieves substantially better ICL capability while maintaining strong IWL performance. We believe these results demonstrate that our dual spaces encoding can enhance ICL capability and mitigate the ICL IWL conflict in NLP scenarios where semantic information is richer.
approximate natural language distributions.
7.4 Conditional
Task
To further verify the generalization ability of our algorithm on generative tasks and pretrained models, we designed conditional pseudo-arithmetic task, inspired by Anand et al. (2025) and Chen et al. (2025). We construct sequences consisting of arithmetic conditions followed by a word problem, as shown in Figure 7 Here, the question can be answered solely based on the given conditions. We fine-tune GPT-2 on 5,000 such sequences and evaluate accuracy on the same distribution, which we interpret as its TWL, performance. After fine-tuning, we then construct soquences that mix standard arithmetic with "pseudo-arithmetie"" over letters. In this case, the model must rely on the contextual relation given in the sequence (rather than memorized weights) to answer correctly, which we to measure ICL performance. In this setup, x, and X₂ are generalized from single tokens to sequence-level inputs.
We exlend CoQE lo the pretrained GPT-2 as follows. During fine-tuning, we modily the attention mask so that tokens in the Question field cannot attend to the Condition field, thereby implementing separate Context-Query Encoding. On the output side, we largely follow the implementation for few-shot classifica- tion. The difference that we do not introduce forig or noise injection, because we observe that fine-tuning pretrained GPT-2 this dataset is fast and stable. Besides standard fine-tuning, we also introduce T
13

---

Published in Transactions on Machine Learning Research (02/2026)
regularization and probabilistic temporary forgetting baselines. All models fine-tuned for 300 steps. The training curves are shown in Figure 7 right.
After fine-tuning, all methods achieve nearly 100% IWL performance, while CoQE achieves substantially higher ICL accuracy than the other three baselines. typical failure mode we observe is using the wrong condition, for example: 'Condition: 42 + 46. 31 + 20 51. I. + Y = F. Question: T read T. pages yesterday and Y pages today. How many pages did 1 read? Answer: 51." This indicates that the model is still limited in leveraging the correct in-context relation in this scenario. Moreover, L2 regularization and probabilistie temporary forgetting perform even worse than vanilla fine-tuning, This demonstrate limitations of their design: regularization can take effect at very large training steps (e.g., 10"), and probabilistic temporary forgetting largets ICL settings that involve random-vector inputs, which differs from our pseudo- arithmetic setting. Overall, this simple autoregressive experiment further validates the effectiveness and transferability of CoQE on generative tasks and pretrained models.
8 Discussion
Our analyses and experiments primarily focus on small models. A natural question is whether large models exhibit similar imbalance between ICL and IWL. Piantadosi (2014) showed that a Zipfian distribution of a 1 closely approximates the empirical distribution of natural language, which serves as sweet spot for the tradeoff between ICL and IWL (Chan et al., 2022). Nevertheless, Chan et al. (2025) pointed out that LLMs can still experience conflicts between ICL and IWL when demonstration examples deviate from the training distribution. Beyond this, the growing demand for multimodal large models (e.g., VLMs, VLAs) trained on increasingly diverse data distributions, as well as the emergence of new architectures, makes it less tenable to rely on the fortunate coincidence of natural language statistics to balance ICL and IWL. These considerations underscore the importance of our exploration in this direction.
Limitations. Theoretically, a limitation of our work that our modeling is built upon the linear represen- tation hypothesis. Although this assumption also underlies many other theoretical and empirical studies, its applicability to complex NLP tasks like multi-step reasoning still requires further evidence. Another limitation is that our dual-space formulation currently assumes that the answer Yg is a single value rather than a sequence. Extending the theoretical modeling to handle open-ended generation tasks is valuable.
Emperically, the main limitation is that the evaluation of our algorithm has so far been confined to structural synthetic experiments. For tasks where such a context-query query decomposition is not available or not well defined, our current algorithm docs not yet provide an effective way to apply dual-space modeling. Extending our algorithm broader general NLP multimodal settings is also an important direction for future work.
Broader Impact Statement
Reconciling ICL and IWL under diverse conditions is crucial for the robustness of large models as data distributions grow more varied and architectures evolve. This work aims to enhance goneralization and reliability of large models, contributing safer and more equitable deployment in real-world applications.
Acknowledgments
This work was supported in part by the National Key Research and Development Program of China (No. 202400...0006), in part by the National Key Research and Development Program of China under STT 2030 Major Projects (No. 20212D0200300), and in part by the TSinghua-Fuuzou Data Technology Joint Research Institute (Project No. JIDT2024013).
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Igge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint Xrx::2303..00 08774, 2023.
14

---

Published in Transactions on Machine Learning Research (02/2026)
Kartik Ahuja and David Lopez-Paz. A closer look at in-context learning under distribution shifts. arXiv preprint or Xiv:23305 16704, 2023.
Suraj Anand, Michael A Lepori, Jack Merullo, and Ellie Pavlick. Dual process learning: Controlling use of in-context vs. in-weights strategies with weight forgetting. In The Thirteenth International Conference on Learning Represenlations, 2025.
Trenton Bricken, Rylan Schaeller, Bruno Olshausen, and Gabriel Kreiman. Emergence of sparse representa- tions from noise. In International Conference on Machine Learning, pp. 3148 3148-3191. PMLR, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda, Qingfu Zhang, Hau-San Wong, and Taiji Suzuki. Provable in-context vector arithmetic via retrieving task concepts. In Forty-second International Conference on Machine Learning, 2025.
Bryan Chan, Xinyi Chen, András György, and Dale Schuurmans. Toward understanding in-context vs. in-weight learning. In The Thirleenth International Conference on Learning Represenlations, 2025.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in trans- formers. Advances in neural information processing systems, 35:18878 18891, 2022.
Guanyu Chen, Peiyang Wang, Yizhou Jiang, Yuqian Liu, Chujic Zhao, Ying Fang, Tianren Zhang, and Feng Chen. Exploring the hidden reasoning process of large language models by misleading them. In Findings of the Association for Computational Linguistics: EMNLP, 2025.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. In Advances in Neural Information Processing Systems, volume 37, pp.
Curran Associates, Inc., 2024.
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoencoders find highly interpretable features in language models. arXiv preprint ar arXiv:2309.00 086000 2023.
Fahim Dalvi, Abdul Rafae Khan, Firoj Alam, Nadir Durrani, Jia Xu, and Hassan Sajjad. Discovering latent concepts learned in bert. In International Conference on Learning Representations, 2022.
Danny Dricss, Fei Xia, Mchdi SM Sajjadi, Corcy Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: an embodied multimodal language model. In Proceedings of the 40th International Conference Machine Learning, pp. 3469-8488, 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishck Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv e-prinds, pp. Xiv 2407, 2024.
Ying Fan, Steve Yadlowsky, Dimitris Papailiopoulos, and Kangwook Lee. Transformers can learn meta-skills for task generalization in in-context learning. In NewrIPS 2024 Workshop on Compositional Learning: Perspectives, Methods, and Paths Forward, 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. In Advances in Neural Information Processing Systems, 2022.
URL httpss// net//forummt---11222aaaa.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in 1mms via reinforcement learning. Xix preprint arXi::201 12948, 2025.
15

---

Published in Transactions on Machine Learning Research (02/2026)
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do trans- formers learn in-context beyond simple functions? a case study on learning with representations. In The Twelfth International Conference on Learning Representations, 2024.
Wes Gurnee and Max Tegmark. Language models represent space and time. In The Twelfth International Conference Learning Representations, 2024.
Seungwook Han, Jinyeop Song, Jeff Gore, and Pulkit Agrawal. Emergence and effectiveness of task vectors in-context learning: An encoder decoder perspective. In Forty-second International Conference on Machine Learning, 2025.
Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 9318 9333, 2023.
Yihan Hu, Jiazhi Yang, Li Chen, Keyu Li, Chonghao Sima, Xizhou Zhu, Siqi Chai, Senyao Du, Tianwei Lin, Wenhai Wang, et al. Planning-oriented autonomous driving. In Proceedings of the IEEE/CVF conference on computer vision and pottern recognition, pp. 17853-17881 2023.
Ruomin Huang and Rong Ge. Task descriptors help transformers learn linear models in-context. In ICML 2024 Workshop on In- Content Learning, 2024. URL http::://pporrevveeeeee///orooop
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dy- namics on the attention landscape. In International Conference on Machine Learning, pp. 24527-24561. PMLR, 2024.
Jannik Kossen, Yarin Gal, and Tom Rainforth. In-context learning learns label relationships but is not conventional learning. In The Twelfth International Conference Learning Representations, 2024. URL https:///oenrreview wtt.e/wwwweeooooo//aa/a//tttoroa
Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6666) 2015.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mech- anistic understanding. In Proceedings of the 40th International Conference on Machine Learning, pp. 19689-19729, 2023.
Haodong Liang, Krishna Balasubramanian, and Lifeng Lai. Transformers handle endogeneity in in-context linear regression. In The Thirteenth International Conference Learning Representations, 2025.
Muhammad Maaz, Hanoona Rasheed, Salman Khan, and Fahad Khan. Video-chatgpt: Towards detailed vidco understanding via large vision and language models. In Proceedings of the 62nd Arnual Meeting of the Association for Computational Linguistics (Volune 1: Long Papers), pp. 12585 12602, 2024.
Samuel Marks and Max Tegmark. The geometry of truth: Emergent lincar structure in large language model representations of true/falso datasets. In First Conference on Language Modeling, 2024. URL
https://openrevieet wttpeeewwweee/oooeeeeeeeeeeeeeee met//forrr//dddraaayyyyij
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vec-style vector arithmetic. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Longuage Technologies (Volume 1: Long Papers), pp. 5030-5047,
2024.
Tomá Mikolov, Wen-tau Yih, and Geoffrey Zweig. Linguistic regularities in continuous space word rep- resentations. In Proceedings of the 2013 conference of the north american chopter of the association for computational linguistics: Human language technologies, pp. 746-751, 2013.
Sarthak Mittal, Eric Elmoznino, Leo Gagnon, Sangnie Bhardwaj, Guillaume Lajoic, and Dhanya Sridhar. Does learning the right latent variables necessarily improve in-context learning? In Forty-second Interna- tional Conference on Machine Learning, 2025.
16

---

Published in Transactions on Machine Learning Research (02/2026)
Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models. EMNLP 2023, pp. 16, 2023.
Alex Nguyen and Gautam Reddy. Differential learning kinetics govern the transition from memorization to generalization during in-context learning. In The Thirteenth International Conference on Learning Represenlations, 2025.
Eshaan Nichani, Alex Damian, and Jason Tee. How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pp. 38018-38070, 2024.
Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained transformer efficiently learns low- dimensional target functions in-context. In The Thirtg-eightt Annual Conference on Neural Information Processing Systems, 2024. URL httpps// wwt//eeeuedddduiieneoooooa
Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen. What in-context learning "learns" in-context: Disentangling task recognition and task learning. In The 61st Annual Meeling Of The Associalion For Computational Linguistics, 2023.
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2024.
Core Francisco Park, Ekdeep Singh Lubana, and Hidenori Tanaka. Competition dynamics shape algorithmic phases of in-context learning. In The Thirleenth International Conference on Learning Represenlations, 2025.
Kiho Park, Yo Joong Choc, and Victor Veitch. The linear representation hypothesis and the geometry of large language models. In Proceedingss of the 41st International Conference on, Machine Learning, volume 235, pp. 39643 39666, 2024.
Steven T Piantadosi. Zipf's word frequency law in natural language: A critical review and future directions. Psychonomic bulletin & review, 21(5):1112 1130, 2014.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training- OpenAT blog, 2018.
Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference on Learning Representations, 2024.
Zhenmei Shi, Junyi Wei, Zhuoyan Xu, and Yingyu Liang. Why larger language models do in-context learning differently? In Proceedings of the 41st International Conference on Machine Learning, pp. 44991 45013, 2024.
Aaditya Singh, Stephanie Chan, Ted Moskovitz, Erin Grant, Andrew Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. Advances in neural information processing systems, 36:27801-27819, 2023.
Aaditya K Singh, Ted Moskovitz, Sara Dragutinovié, Felix Hill, Stephanic CY Chan, and Andrew M Saxe. Strategy coopetition explains the emergence and transience of in-context learning. In Porty-second Inter- national Conference on Machine Learning, 2025.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems, volume 36, pp. 71911-71947, 2023.
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron Wallace, and David Bau. Function vectors in large language models. In The Twelfth International Conference on Learning Representations, 2024. URL http:///pppppeeeeeeeeeeerrt/aaaa
17

---

Published in Transactions on Machine Learning Research (02/2026)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: Statc-of-the-art natural language processing. EMNLP 2020, pp. 38, 2020.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining lasks needed for in-context learning of linear regression? In The Twelfth International Conference Learning Representations, 2024.
Zheyang Xiong, Ziyang Cai, John Cooper, Albert Ge, Vasilis Papageorgion, Zack Sifakis, Angeliki Giannou, Ziqian Lin, Liu Yang, Saurabh Agarwal, et al. Everything everywhere all at once: Llms can in-context learn multiple lasks in superposition. In Forty-second International Conference Machine Learning. 2025.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint (Xiu:231100087) 2023.
Liu Yang, Ziqian Lin, Kangwook Lee, Dimitris Papailiopoulos, and Robert Nowak. Task vectors in in in-context learning: Emergence, formation, and benefit. arXic preprint orXiv:2501.09244 2025.
Naimeng Ye, Hanming Yang, Andrew Siah, and Hongscok Namkoong. Pre-training and in-context learning is bayesian inference la de finetti. arXiv preprint arviv: 033800 2024.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu Zhu, Physics of language models: Part 2.1, grade- school math and the hidden reasoning process. The Thirteenth International Conference on Learning Representations, 2025.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1 2024a.
Tianren Zhang, Chujie Zhao, Guanyu Chen, Yizhou Jiang, and Feng Chen. Feature contamination: Neural networks learn uncorrelated features and fail to generalize. In International Conference on Machine Learning, pp. 60446-6048 60495. PMLR, 2024b.
Tianren Zhang, Guanyu Chen, and Feng Chen. When do neural networks learn world models? In Forty- second International Conference Machine Learning, 2025a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learn- ing learn? bayesian model averaging, parameterization, and generalization. In The 28th International Conference on Artificial Intelligence and Statistics, 2025b. URL http::///openrevveee.ee/ttrooree |500FOFc60
Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Tod Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, al. Rt-2: Vision-languaag--ctionn models transfer web knowledge to robotic control. In CoRL, 2023.
18

---

Published in Transactions on Machine Learning Research (02/2026)
A Proofs of Theoretical Results
A.1 Proof of Proposition 5.3
For case of presentation, first restate the proposition and then introduce its proof.
Proposition A.1 (Task-sample duality). Let X be the input space and Vj the multiple label sets corre- sponding to each task E F. Under Definition 5.2. there evists linear somple represenlation space My and a linear task trunsformation space T, where T is the dual space of My, i.c. TMM
Proof. To prove the proposition, we must show that the linear task transformation space T is equivalent to the dual space of the linear sample representation space My. denoted as Mÿ- The proof proceeds by demonstrating mutual inclusion: (1) T S M and (2) M CT.
Step 1: Proof of T S Mz. Let be an arbitrary element in the task transformation space T. According to Definition 5.2 is a linear function such that
t(m) = ().,m). Vm E My.
Since & is a linear functional on My, it is by definition element of Mý- As was arbitrary element of T, it follows that every element in T corresponds to unique linear functional in M Thus, we have establisher that T M'
Step 2: Proof of M2 CT. Conversely, let l' be an arbitrary linear functional in the dual space Mý. By Definition 5.1 M, is finte-diiennsioal inner product space. By the Riesz representation theorem, for any linear functional t' E M2, there exists a unique vector, let's call it we E _Mx, such that for all m E M,F
t'(m) (we m).
Now, let us define a function fe X -> using this functional t':
(x = (x(xx(x)) (wer, (x(x)).
This function Ser has the exact mathematical form of task function as specified in Definition 5.2 Therefore, fo can be considered valid task belonging to the task function space F. Definition 5.2 states that for any such task fr E F, there exists a unique linear task representation vector, which we denote wy, that represents it. This means:
(x) = (wy, (x()))
By equating the two expressions for (x(x), we obtain:
(we, ox()) =
Vx E X
This implies that (wy -yy, m) = 0 for all m in the image dy. Since the sample representation space My is spanned by the image of Φx, this condition holds for all m E My. The only vector orthogonal to every vector in an inner product space is the zero vector. Therefore:
(wy - -wy =0 - We wg.
Since wy corresponds to an element of T, it follows that we is also corresponds to an element of T, and thus t' E T. As our choice of t' was arbitrary, we have shown that every linear functional in M corresponds to an element in T. Thus, we have established that M CT.
A.2 Proof of Theorem 5.6
For of presentation, we first. restate the theorem and then introduce its proof.
19

---

Published in Transactions on Machine Learning Research (02/2026)
Theorem A.2 (Completeness of basis representations under task traversal). Under Proposition 5.3
assume that learner with sample represenlation mapping dh is presenled with lask traversal curriculum C such that: spanft t t = T. Then, if the learner achieves zero empirical error, the learned representation mapping the salisfis:: span [x((x) x x MM.
Proof. By Proposition 5.3 fix a basis {m; of My and its dual basis ft: T, satisfying tz(m;)
(w), mj) dij. For any m MM write the unique decomposition Ei= (m(m)) my and for any T write 1 = Ei= B:(1)4: 4. The bilinear pairing then reduces to
{(m) = Σ (m) (c())
(15)
Let OF : X My denote the sample representation guaranteed by Definition 5.2 and define the coordinate vectors
a (x) - (an
(x(()(x))) R", a²(x) _
(x() ER.
Zero empirical error every curriculum task t E C means
f()o(x)) t(()(xx) for all E C and all training x.
By linearity of Equation 15 his equality holds for any linear combination of curriculum tasks; hence holds
for all t span(C) T:
t()o(x)) = t(os(x)),
AtET.
(16)
Take in Equation 16 the particular choice t = t, (the i-th dual basis functional). Using ((m) = (m) from Equation 15 we obtain for every x and every E [d].
ax((b(x))) (x)) (x(xx(x)) (x)))(x))
Thus a (x) = a" (x) pointwise for all (training) x. Consequently (x(x) and @x(x) have identical coordinates in the basis {m for all x, so
(pon( x X} -span x E X}.
Without loss of generality, take My span(o=(x) E X}. Therefore span( x E X} = MM. proving the claim.
A.3 Proof of Proposition 5.8
For ease of presentation, we first restate the proposition and then introduce its proof. Proposition A.3 (Closed form of wy under simplified LSA). Consider an LSA layer applied after a feature encoder Φ + Rd implemented by an MLP. Suppose the LSA projection matrices WKQ and WoA are initialized such that
Wov = Oo 1). WKQ =
-
+ 0a
0
*
Then the final prediction takes the form ÿ = (Wy(Zen Φ), ((xy)), where
wy (2 Φ) Σ y ((x,).
(17)
Proof. According to Kim & Suzuki (024). under the conditions of Proposition 5.8 given as follows:
the expression of ý is
8
Σ vx)(x)) T
(18)
1=1
Hence, Proposition 5.8 is readily provod.
20

---

Published in Transactions on Machine Learning Research (02/2026)
A.4 Proof of Theorem 5.9
For ease of presentation, we first restate the theorem and then introduce its proof. Theorem A.4 (Entangled structure under general SA). For a standard SA model with softmar-base based at- tention weights, there does NOT exist a pair of do and wo(Zine Φo), such that the model prediction admits the following decomposition:
8, (wo(Z Qo), oo(x)))
(19)
Proof. We argue by contradiction. Assume there exists a finite-dimensional feature map do and context- only cocfficient vector wo( (u(Zim Φ0) such that the identity holds for all contexts and queries.
From SA equations to ratio of exponentials in Xy. Let the sequence length be I, = n-11 Stack token embeddings as Z Bi....,, Zn Zn] R2xL, A single-head self-attention (SA) layer computes
Step 1:
Q = W K = WKZ, == WvZ.
with Q. K E Rd Ve Rd. Denote the i-th key/ key/value /vauee columns by k == K, WKZ, U2 u; V₂2 = Wyz, and the query column at position q by 9: Q. Wozz The attention weights for the query position form probability vector alq) E A with coordinates
exp ((ki, (li.g)//dd
ax(1) Σ/= exp ((kj, 93 /dz)
i= L L.
(20)
In the theoretical analysis of ICL, is common to set Zy = [xg. 0]. Without loss of generality, we assume that the query token embedding depends affinely on the input feature Xg E Rdz:
Bg ExX, + Tg:
where Eg E Rdxd, a fixed embedding matrix and rq Rd could collect position encodings and other context-independent parts at position q. Then the query vector also affine in Xgi
Woug = WoF X, + WQrs = Ux, + no:
with U := Eg Rddxdd and uo = Worg E Rddx. Plugging q = Ux, + uo into the logits in Equation 20 yields, for each key i,
(ks.9)
(ks,, Ux,) + (k,, uo)
= TTxx, + (z))
ddl Vdh.
vdl
where we define the (query input) slope and the (context) offset by
U
(Kis uo) vdk
0i Vdl
E Rd,,
bı(z)
R.
Hence, for fixed context 2im (which fixes all hy and uo), the attention weights are softmar of offine functions of xgi
z) = exp (a Xg b,(())
i 1 1..,.., L.
(21)
Σj=1 exp (a] X, + bz))))
The SA output at the query position is hy = 2y + Wo Σte az (Xg: 2) v. For a fixed linear predictor w E Rd (or equivalently choosing a fixed output coordinate), the scalar prediction is
(x,))) w Th, h4 w Zu
w Wooz) (x(x)
(22)
affine in Xy
(2)))
If we choose w orthogonal to Im (Ez) (always possible unless E, = 0), then T Zg = w TT, rg) w rg is a context-only constant; denote c(z) = w rg. With 7(z) = (71 (z) ... 7 (z)) Equation 22 simplifies
8 (x,x) c(z) + (xx;;2))
(23)
where al=zz) given by the ratio- exponenntss form Equation 21 This exhibits the claimed dependence of Sg on X2 through softmax affine functions of Xg-
21

---

Published in Transactions on Machine Learning Research (02/2026)
Step 2: A two-key reduction yields linearly independent logistic family. Specialize to dN = and one context keys (n= = 1) with ai # 0z. Choose Wo, V so that c(z) = and m(x) = 1,72(2) 0. Then Equation 23 reduces to
exp(aix, + b1
(x exp(a1x, + (z)) exp(a2x, + b2(x))
1 ((x) e ax,"
where a a= a2 # 0 and t(z) == exp (bz(2)) b1(z)) > 0. As the context varies, t(z) can take arbitrarily many distinct positive values, so SA realizes the ono-parameter family of functions
F = (1(x) = 1 ax x>0].
Fix distinct .1,...,bo im O. Suppose there exist scalars A Am with E.H Acfe, = 0 for all x R. Multiplying both sides by I txxx and letting axx gives the polynomial identity
Σ X I
for all 0.
A polynomial that vanishes on an infinite set is identically zero; hence the identity holds for all s R. Evaluating s -1/tp yields
R II
II j
1
1 1s
0.
Since the t; are distinct, each product is nonzero, forcing A 0 for all k. Thus fe....., ftm are linearly independent. Consequently, the linear span of J is infmite-dimensional.
Step 3: Contradiction with any unite-dimensional bilinear decomposition. If the bilinear decom- position 8, (x,) = (z), (x(xx)) held with a fixed feature map 60 : RR - Rd (independent the context), then for all contexts the functions X, + 8x,g)) would lie in the d-dimensional linear span of the coordinate functions of Φn- However, Step shows that by varying the context, SA realizes an infinite set. F of pairwise linearly independent functions in x, which cannot be contained in any finite-dimensional linear subspace. This contradiction rules out the existence of such (dp,wa).
B Experimental Details and Additional Results
In this part of the appendix, we provide detailed descriptions of the experiments in the main text and include additional experimental results.
B.1 Representation Visualization
To further illustrate the training dynamics in Figure 2 left, we visualize the evolution of representations for three common classes during training under training settings (64, 12, 0.9, 2) and (64, 0.9, 0), as shown in Figure 8 We find that under different data distributions, the model's representation space diverges significantly from the very beginning of training and eventually converges to distinct structure. This clearly reflects that the balance of sample-level and context-level information is highly sensitive to the training distribution.
B.2 Regression
Implementation details. We use Transformer architectures from the GPT-2 family RRdford el. al. 2018) as implemented by HuggingFace (Wolf et al. 2020). Specifically, the Transformer baseline we use is configured with an embodding dimension of 64, layers, and attention heads, resulting a total of 0.2M parameters. The task encoder CoQE uses the exact same Transformer configuration. The representation encoder of CoQE consists of two-layer ReLU network, implemented linear projection, followed by
22

---

Published in Transactions on Machine Learning Research (02/2026)
steps = =3k
steps = 10k
steps =50k
Piursty 0.9 a=2
hurnny = 0.9 a=
Phursey = 0.9 a=1
labeled 0 in context
labeled as in context
support in context
Figure 8: The uncropped visualization results of training dynamics.
ReLU activation, a LayerNorm, and a second linear layer. For fair comparison, the baseline Transformer's ss embedding module uses the exact same two-layer ReLU network. During training across the four classes of functions, we use a batch size of 64 and a learning rate of 5e - 5. For the three tasks except combination functions, models are trained for 1 x 10 steps, while the combination task is trained for x 10s steps due its increased difficulty. All experiments are conducted on an NVIDIA RTX 4090 GPU,
Setup details. We consider two major evaluation scenarios for regression: in-distribution (ID) testing and anttof-distribution (OOD) testing. In the ID scenario, we set D Dtest D Dquery, and D D Specifically, the following four classes of functions F:
23

---

Published in Transactions on Machine Learning Research (02/2026)
Linear functions: == (f f(x) w w R*), where = 10. We sample xI....,, X,.Y, and w independently from the isotropic Gaussian distribution (0, Ta), then compute f(z;) = w x; to construct the prompt. In this setting we use the least squares estimator as the optimal baseline.
Sparse linear functions: == {f
= = w R Iwwloo S Ss), where d = 10 and s = 3. We also sample xi...., xj, X2 and w independently from (0, Ia), and then zero out all but coordinates of w uniformly at random. We use the least squares estimator and Lasso, which leverages sparsity with an 4-orr regularizer as baselines.
Two-layer ReLU neural networks: F= = - a [w][( a E E,R w, RR), where (=)) [0, =} is the ReLU activation function, and d= 5, h 10. We sample x,s and X, from Ta), along with network paramcters a,s from 2/0.)/h). We sample w,s from W00 Ia), and share them across all lasks in F. The baseline is two-layer neural nelwork of the same archilecture trained on in-context examples using Adam.
Combination functions: F= V [(x) w (x), E RR), where + is element-wise combina- tion function. For X = X2, X3. X4, J. (x) = [|x]], x3, x cos(TX4), c0.2xa JT. We sample x,S, Xg and w from (0, Is) independently. In this setting, there is no naturally optimal bascline, so we compare only with the Transformer.
The latter two classes of nonlinear functions allow the model to reduce ICL difficulty through representation learning, by learning lask-invariant W,s or Φ.
In the OOD scenario, we consider four different cases of distributional shifts under linear functions.
Dipain A Di = Queery- We consider the setting where the prompt inputs x,s scale between training and testing is different. We scale them by a factor of 0.8 1.2.
Diain D + queery We sample the context examples from the same distribution training time, but sample X, from a Gaussian distribution with 3x higher standard deviation.
Disat # D Dquery. We fix the sign of each coordinate to be randomly positive or negative for all prompt inputs x,s, and draw Xg from A/(0,)) before.
Dirain * Djest. We consider scaling the weight vector by factor of 0.8 1.2, to capture shifts of task functions.
Through the above diverse evaluation settings, we comprehensively demonstrate that CoQE consistently exhibits stronger ICL capability than standard Transformer of comparable size regression tasks.
Additional results on representation learning. Our Theorem 5.6 shows that under the dual-space modeling framework, a sufficient set of tasks guarantees a basis-covering sample representation space that the model learns. For empirical validation, we design the task type two-layer ReLU networks and combination functions, whose different task functions share common sample representation space their construction. Figure (a) shows that CoQE indeed achieves smaller ICL error on these tasks. Furthermore, under the combination functions task, set the sample representation space dimension of CoQE to 5, matching that of 4 and directly visualize the 5-dimensional sample representation space learned by CoQE after training (Figure D.. From the figure, we can observe that the five dimensions appear to differentiate in a manner close to the respective transformations of the five dimensions of Φ. Although, due to the equivalence of sample representation spaces under linear transformations, ie., f = w ((x) = w -1 H #(xx) where H denotes an arbitrary invertible matrix, it essentially impossible for the model to learn t with perfectly identical scale and shape. The current differentiation can be regarded as another empirical proof of Theorem 5.6 that our dual-modeling could facilitate learning of the basis-covering sample representation space.
24

---

Published in Transactions on Machine Learning Research (02/2020)
dim=0
0.8
dim=1
dim=2
0.6
dim=3
x
dim=4
0.4
0.2
c.2x
Output
0.0
COS
-0.2
-0.4
-0.6
x²
-2.0
-1.5
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
Input
Figure 9: The 5-dimensional sample representation space learned by CoQE for the combination functions.
B.3 Synthetic Few-shot Classification
Implementation details. In our experiments, employ ResNets of three sizes (with embedding di- mensions E 64, E 256 and F 512) encode images. All architectures consist of four groups, each containing two residual blocks. The difference lies in the embedding dimensions of each group: for the E = 64 ResNet, the four groups produce embeddings of sizes 16, 32, 32, and 64, respectively; for the E = 256 ResNet, the sizes are 64, 128, 128, and 256; for the E = 512 ResNet, the sizes are 64, 128, 256, and 512. Then a fully connected layer appended to the E = 256 and E 512 ResNet to project the final embedding dimension back to before feeding into the Transformer. Larger embedding dimension of Resnets clearly possesses a stronger capacity for extracting visual sample representations, As result, the resulting embedding tokens for the sample encoder. We also employ two Transformer configurations with different layers: L 4 and L= 12. Both variants use embedding dimension of 64 and 8 aotention heads. We have shown that CoQE with only an L = 4 Transformer in the task encoder, can match the ICL and IWL performance of an L = 12 Transformer. In our experiments, a baseline Transformer with = 64 and L 12 contains approximately 0.9M parameters, CoQE with E 256 and L has 1.0M parameters, while CoQE with 512 and - has 2.0M parameters.
expressive. For CoQE, we find that the E 64 ResNet is insufficient
When training CoQE, we add Gaussian noise to the modified logits to prevent the task encoder's output from collapsing lo static vector. Specifically, the initial noise is sampled from ((Ho 1), and both the mean and standard deviation are incremented by 1 every 10a training steps. During training of baseline Transformers and COQE, we use a batch size of 24, a learning rate of le - 4, and train for 1 x 10 steps. All experiments are conducted on 8x NVIDIA V100 GPUs
Detailed results of all settings and baselines. In Table we report the ICL and IWL performance across all model scales, training settings, and algorithms. In Table report the effect of appropriately scaling both the Transformer and CoQE. These results demonstrates that CoQE can robustly reconcile ICL and IWL
Ablation study. We present the training curves CoQE under different levels of noise (Figure 10p. is evident that, the absence of noiso, the model's ICL capability rapidly decays after an initial emergence, accompanied by a similarly rapid increase in IWL performance. Although this observation is not made under standard Transformer model, hypothesize that the underlying phenomenon extends beyond
25

---

Published in Transactions on Machine Learning Research (02/2026)
Table 3: Performance comparison of different Transformer variants and baselines. All values are presented as
percentages. The maximum and second-largess values are indicated by boldface and underline, respectively.
Settings (Buursty; 0.)
Model
Config
Metric
Avg.
(0.9, 1)
(0.9,
(0.9,2) 2)
(0.5,
(0.1, 0)
= 64, L 12
ICL
84.2
73.1
49.7
64.5
50.5
64.4
# #Param: 0.9M
IWL
65.1
4.1
98.0
17.8
51.5
47.3
== 6 #Param: 0.6M
ICL
74.4
67.2
50.2
57.5
50.3
59.9
IWL
65.2
8.1
97.9
21.4
57.9
50.1
= E 4
ICL
50.0
50.5
50.1
50.2
50.4
50.2
#Param: 0.4M
IWL
66.2
15.1
98.1
27.3
63.1
54.0
E = 256, L = 12
ICL
49.8
51.4
50.0
50.2
50.5
50.4
#Param: 1.4M
IWL
72.3
20.1
98.4
58.1
75.4
64.9
Transformer
= 256, L = 8
ICL
50.8
50.7
50.2
49.9
50.7
50.5
#Param: 1.2M
IWL
73.1
28.7
98.4
60.2
78.5
67.8
= 252 L
ICL
50.0
50.3
50.6
49.9
50.7
50.3
#Param: 1.0M
IWL
74.4
38.8
98.6
63.4
81.1
71.3
E 512, L 12
ICL
50.5
50.5
50.3
50.1
50.5
50.4
#Param: 2.4M
IWL
79.6
36.0
98.9
98.4
99.0
82.4
E = 512, L = 8
ICL
50.4
49.8
50.3
50.2
51.0
50.3
#Param: 2.2M
IWL
81.0
49.3
99.0
98.9
99.1
85.5
= 512, I. 4
ICL
49.9
50.0
51.0
49.5
51.0
50.3
#Param: 2.0M
IWL
82.5
62.5
99.1
99.4
99.2
88.5
Transformer
E = 64, L
ICL
80.0
70.3
49.9
59.8
50.1
62.0
+ Noise
#Param: 0.9M
IWL
60.4
1.2
91.3
8.9
49.9
42.3
Transformer
E 6,, 12
ICL
80.5
64.2
50.2
67.8
53.1
63.2
+ Regularization
# #Param: 0.9M
IWL
64.9
3.1
63.2
36.4
98.5
53.2
Transformer
E 64, L 12
ICL
50.4
51.0
49.8
50.5
51.0
50.5
+ Temporary Forgetting
# Param:: 0.9M
IWL
54.8
0.8
97.1
7.2
38.5
39.7
E = 64, L = 12
ICL
51.3
51.0
50.0
50.1
50.3
50.5
Linear Transformer
#Param: 0.8M
IWL
68.3
8.0
98.4
21.0
53.5
49.8
E = 256, L = 4
ICL
90.3
89.2
85.8
88.0
90.3
88.7
# #Parmm: 1.0M
IWL
82.0
81.0
63.1
82.0
87.9
79.2
CoQE
512, L 4
ICL
92.1
91.2
88.0
90.2
92.0
90.7
# #Param: 2.0M
IWL
90.0
89.4
67.9
91.0
95.4
86.7
26

---

Published in Transactions on Machine Learning Research (02/2026)
Table 4: The effect of appropriately scaling both the Transformer and CoQE.
Model
Config
ICL
IWL
Model
Config
ICL
IWL
L = 12 81.2
63.1
I = 64 84.2
65.1
Trans. (E 64, 64)
L 86.5
62.9
Trans. (E 64, 12)
H
51.0
67.4
L = 16
73.0
62.9
H = 80
51.0
68.1
L 90.3
82.0
H 64
90.3
82.0
CoQE (E 256, H = 64)
L =
89.3
84.7
CoQE (E = 256, = 4)
H = 72 92.6
83.0
L 89.7
85.2
H 80 90.1 82.2
model architecture, reflecting the intrinsic properties of the two strategies. 1CL is lightweight, dynamic strategy, whereas IWL is more training-intensive but ultimately more stable. In standard Transformers, where the two strategies are difficult to co-exist, training often leads to a transition from ICL to IWL. In contrast, CoQE enables robust cocxistence both stratogics through explicit modeling and learning of the task representation space, as well as the use Gaussian noise to isolate the bask-Iransformations associated with each strategy.
1.0
0.9
0.8
Accuracy
0.8
Acourapy
0
0.7
3
5
0.8
0.5
100
Steps (k)
Steps (k)
(a) ICL accuracy
(b) IWL. accuracy
Figure 10: Learning curves under different noise levels
B.4 Llama Based Synthetic Few-shot Classification
Datasets construction. We saved the token embedding matrices of the Llama-3.2-1BB Llama-3.2-3B, and Llama-3.1-8B models. Each embedding matrix has shape of 128,256 dmodel, where dmadel = 2048, 3072, and 4096 for the IB, 3B, and 8B models, respectively. To construct "classes" for our few-shot classification experiments from these raw token embeddings, we apply the following three processing steps:
1. Subselection: Among the 128,256 tokens, many correspond to special symbols or non-English text. We selected all tokens consisting solely of English letters A-Z, a-z and underscores, resulting in 33,588 tokens.
Clustering: We applied spherical clustering using FAISS, performing 2,400-way clustering and re- taining all clusters containing more than 10 tokens. From these, we randomly selected 1,200 clusters.
Cluster Sampling: Many clusters contained more than 10 tokens. To maximize intra-class variability, we selected the 10 tokens farthest from each cluster center. Consequently, we obtained 1,200 classes, cach containing 10 samples.
Compared with Omniglot images, this construction produces classes with greater intra-class variability while reflecting semantic relations more relevant NLP tasks. Some sample clusters from 1lama-3.2-33 are shown below:
27

---

Published in Transactions on Machine Learning Research (02/2026)
English
edish
apancse
French
Chinese
orean
California
Russian
frican
ustralian
Class
_class
addClass
_CLASS
className
removeClass
classList
Classifier
hasClass
getClass
ellow
green
Green
iolct
orange
urple
agenta
purple
Pink
greens
YELLOW
Implementation details. We trained both the Transformer and CoQE models on the constructed dataset following the same procedure described the main text, fixing the data distribution parameters to Pbursty 0.9 and Zipfian exponent a architecture, for both Transformer and CoQE, used threo-layer MLP instead of ResNet the sample encoder, with hidden dimensions of 512, 256, and 64, respectively. The Transformer consists of 12 layers, while the task encoder in CoQE contains 4 layers, resulting in fewer parameters for CoQE compared with the Transformer. All other training hyperparameters such as learning rate and batch size, are set identical to those used in the image few-shot classification experiments.
approximate the natural language distribution. In terms of model
28


Reviews:
[{"review_id": "R6mq1SngCE", "reviewer": "Reviewer_3mWf", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper examines the ICL–IWL trade-off in Transformer-like models. It introduces a dual-space view where a task is represented as a linear functional over sample representations (via the Riesz theorem), argues that softmax self-attention entangles context and query encoding, and proposes CoQE to separate them by encoding context into a task vector and queries into sample vectors combined via an inner product. Experiments on regression and few-shot classification (plus a token-embedding classification proxy) show improved ICL/IWL trade-offs.\n\nStrengths:\n\nS1. The dual-space viewpoint provides a clean conceptual lens connecting task inference from context to linear readout over sample representations.\n\nS2. The proposed architecture is simple and interpretable, with a clear mapping from the motivating framework to an implementable design.\n\nS3. Results suggest an improved ICL/IWL trade-off on the included benchmarks under the paper’s evaluation protocol.\n\nWeaknesses:\n\nW1. Evidence for extension to open-ended autoregressive generation is limited; the NLP-facing experiment remains a classification proxy rather than a generative setting.\n\nW2. In few-shot classification, gains appear closely tied to noise injection, making it hard to separate “architecture benefit” from “regularization knob” without stronger controls.\n\nW3. The benchmark suite and baseline set remain relatively limited, leaving uncertainty about robustness on more complex tasks and against stronger alternatives.\n\nW4. The linear representation hypothesis is a meaningful but restrictive assumption; beyond being acknowledged, it limits the scope of the conclusions without additional validation.", "claims_and_evidence_yn": "No", "claims_explanation": "The paper provides clear evidence that CoQE improves the reported ICL/IWL metrics on the included regression and few-shot classification settings. However, broader implications e.g. extension to open-ended generation and attribution of improvements to the architectural separation rather than noise/regularization are not fully supported by the current experiments and controls.", "audience_yn": "Yes", "audience_explanation": "The ICL–IWL interaction is an active topic in ML community.", "requested_changes": "Suggestions:\n\nC1. Add an experiment in an autoregressive setting (even a small-scale next-token or conditional generation task) to substantiate claims about broader NLP applicability, or otherwise clearly explain the intended scope and why generation is out of scope.\n\nC2. Strengthen the causal story by adding controls to help disentangle architecture vs. noise/regularization effects (e.g., apply comparable noise/regularization to the Transformer baseline, or provide additional ablations), or explain why such controls are not feasible.\n\nC3. Broaden evaluation modestly (one more complex or longer-context task, or a more challenging synthetic/compositional benchmark) and/or add one or two stronger baselines; if not possible, clarify what conclusions should and should not be drawn from the current task set.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pLFXioksT2", "reviewer": "Reviewer_5PHU", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the conflicting dynamics between in-context learning (ICL) and in-weights learning (IWL), where good performance across both is hard-to-achieve. The paper identifies the shared entangled nature of the encoding in standard softmax-based self-attention architectures as a reason behind conflict. To this end, the paper proposes a novel architecture, called CoQE (Context Query Encoding) that achieves strong empirical performance across the synthetic tasks it considers. The main idea here is to separate IWL from ICL by separating context representations from the query, where the context representations effectively acts as a \"context-dependent function\" to process the query inputs. The key strengths of this paper is the principled understanding of the underlying ICL-IWL conflict, and the proposal of a novel architecture that provides demonstrable improvements in ICL / IWL capabilities. A key weakness of this paper is the focus on synthetic tasks as opposed to real-world tasks like language modelling, and on small transformer models as opposed large-scale ones used in practice.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The following are the main claims made by the paper:\n1. It is difficult to simultaneously achieve good performance in ICL and IWL: This is clearly demonstrated by experiments in Figure 1. I consider the UMAP visualizations to be anecdotal, but the accuracy curves in Figure 1 are more convincing regarding this fact.\n\n2. The conflict between ICL and IWL originates from the difficulty of simultaneously encoding sample information and contextual information within the representation space of Transformers: The paper provides two lines of evidence for this argument. First, there are the UMAP visualizations in Figure 2 which indicate a kind of bimodality in the nature of representations depending on the task parameters. Second, is Theorem 5.9, which demonstrates that under standard SoftMax self attention, there cannot be a task-representation decomposition as is present in Linear Self Attention (Proposition 5.8). Together, these lines of argument are convincing explanations regarding the underlying phenomena in the transformer models considered. It is important also to state that it is unclear whether this still holds for large-scale transformer models, and whether empirically they do learn some dual space representations implicitly, but no claim was made about this in the paper either way.\n\n3. CoQE enhances ICL performance, and successfully reconciles ICL and IWL capabilities: This is also demonstrated convincingly in the experiments in Section 7. The results in Figure 5 clearly show that CoQE is able to outperform the transformer-based model on a wide range of ICL tasks, and also does better out of distribution. Figure 6 and Table 2 also demonstrate its ability to simultaneously achieve strong performance across both ICL and IWL tasks. \n\nOne aspect where the paper can improve is that it misses the important baseline of Linear self attention (LSA) while demonstrating its results in Section 7. While proposition 5.8 clearly demonstrates that LSA does have the dual space structure, it is unclear whether SoftMax self-attention is strictly inferior to LSA on the synthetic tasks considered, and how it compares to the proposed CoQE architecture. Regardless of this, the CoQE architecture is still interesting in that it is fundamentally a non-linear architecture as opposed to LSA, which is bilinear.", "audience_yn": "Yes", "audience_explanation": "I believe the findings of this paper are likely to be of interest to researchers working on theory of ICL / transformers. The key idea behind this paper are the task-sample decomposition based analysis and the proposal of the CoQE architecture that explicitly maintains a dual space decomposition. Both ideas are interesting and have the potential to influence both (1) future theoretical work in this area, as well as (2) architectural improvements.", "requested_changes": "One aspect where the paper can improve is to strengthen its experiments, which can help strengthen the claims made, especially that standard transformer architecture have fundamental disadvantages with the ICL-IWL conflict that can be only resolved by novel architectures:\n\n1. As I mention in the claims section of the review, the main change I recommend in the paper is the inclusion of the linear self attention (LSA) baseline to experiments in Section 7. \n2. It would also be interesting to see the performance of both transformer models as well as CoQE upon scaling these models, e.g.: by increasing width and depth. It would be interesting to know whether for this task, standard transformer models have a fundamental performance ceiling, or whether at large enough scales, the performance gap might vanish.\n\nWith these additional experiments, I believe the underlying phenomena will be better understood (especially as it relates to model scale), and the claims made by the paper can be better extrapolated to large-scale settings.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1U1K3hsvf5", "reviewer": "Reviewer_Q6xq", "summary": "", "questions": "", "limitations": "", "rating": null, "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": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "**Summary**\n\nThis paper investigates the trade-off between In-Context Learning (ICL) and In-Weights Learning (IWL) within the Transformer architecture. The authors demonstrate that in a multi-task setting, by introducing a dual-space (referred to as task space), outputs can be decomposed into a scalar product between input representations and a task vector. They argue that standard Self-Attention (SA) modules entangle the representations of the context and the query, which complicates the separation of ICL and IWL mechanisms. To address this, the paper introduces CoQE, a novel architecture that independently encodes the context and the query. In this framework, the output is predicted via a scalar product between the query and a task vector derived from the context. Experimental results suggest that CoQE achieves strong performance for both ICL and IWL tasks.\n\n**Strengths:**\n\n- The \"dual-space modeling\" offers a theoretically sound and interesting explanation for the entangled structure of representations using standard Self-Attention.\n- The proposed CoQE architecture is well-motivated by the theoretical discussion.\n- Empirical results, while based on small models, demonstrate good performance across various settings for both ICL and IWL, which is the intended goal.\n\n**Weaknesses:**\n\n- The analysis in Section 4.2 relies only on qualitative visualizations (UMAP) and limited hyperparameter configurations.\n- The architectural assumption that queries are separable from context may not hold for all domains (e.g., in text).", "claims_and_evidence_yn": "No", "claims_explanation": "While the theoretical derivations are sounded and seem correct, the empirical evidence supporting the analysis in Section 4.2 is currently weak.\n\n1. The first observation relies heavily on qualitative analysis from UMAP plots (Figure 1, bottom). UMAP projections can be sensitive to hyperparameters and misleading regarding global structure. Relying solely on these plots without quantitative metrics (e.g., clustering separation scores) makes the evidence less convincing.\n2. The second observation regarding model size is derived from a limited set of experimental settings. Since there is a large combinatorial space of parameters that was not explored during the model size variation experiments (e.g., the $\\alpha$ parameter of the Zipfian distribution or $P_{bursty}$), the claims about scaling behavior may be over-extrapolated.", "audience_yn": "Yes", "audience_explanation": "The relationship and trade-off between In-Context Learning (ICL) and In-Weights Learning (IWL) is a central topic in current understanding of the behavior of Transformers. The \"dual-space modeling\" theoretical framework and the finding that separating context from query processing can improve this balance are valuable insights for researchers working on Transformer interpretability and efficient architecture design.", "requested_changes": "### Major changes\n \n**Strengthen Section 4.2 Analysis:**\n - Please provide quantitative evidence to support the claims derived from the UMAP plots (e.g., clustering metrics, cosine similarity analysis or measure of distance) rather than relying solely on visuals.  \n - Expand the experimental setups regarding model size to include variations in other parameters, and more granularity in current variations, to confirm the findings.\n\n**Discuss the Separability Assumption:** The CoQE architecture relies on a \"token-wise\" sample encoder that separates the query from the context. While this works for the tasks presented, it is a strong assumption that does not always hold in general, e.g. in NLP where the query is often an inseparable continuation of the context. Please add a paragraph discussing this limitation and how CoQE might (or might not) apply to standard autoregressive language modeling tasks.\n\n### Minor changes\n\nThe positioning of Figure 1 breaks the flow of the paper. It is first mentioned on page 1, but referenced multiple times later, requiring the reader to flip back and forth. Please consider moving it closer to the main analysis or splitting it to improve readability.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)