Reconciling In-Context and In-Weight Learning via Dual Representation Space Encoding

Guanyu Chen; Ruichen Wang; Tianren Zhang; Feng Chen

cs.LGMar 13, 2026

Reconciling In-Context and In-Weight Learning via Dual Representation Space Encoding

Authors: Guanyu Chen, Ruichen Wang, Tianren Zhang, Feng Chen

Organizations: Department of Automation, Tsinghua University

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 task representation space and a 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. Code: https://github.com/McGuinnessChen/dual-representation-space-encoding

Figures & tables

Figure 1: (Top left) Overview of our synthetic task setup. (Top right) ICL and IWL performances under different training settings ( $E$ , $L$ , $P_{\text{bursty}}$ , $\alpha$ ). The Transformers fluctuate between ICL and IWL capabilities, whereas our 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 simultaneously. Detailed discussions are presented in Section 4 .

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 $E$ 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 3 .

Figure 3: Comparison of Transformer and CoQE architectures. Unlike the Transformer, which encodes both context-level and sample-level information into the same representation space, CoQE implements dual representation spaces encoding that explicitly distinguishes between context and samples.

Figure 4: Construction and training of the task Representation space for few-shot classification.

Figure 5: Results of ICL regression. We provide optimal baselines for test settings except for combination functions. CoQE consistently achieves lower ICL error than the Transformer in both ID and OOD scenarios.

Figure 6: Learning curves of synthetic few-shot classification under different data distribution factors.

Table 7

Figure 7: (Left) Overview of our conditional pseudo-arithmetic task. (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 pre-training. However, after fine-tuning, all methods achieve nearly $100\%$ IWL performance, and CoQE achieves substantially higher ICL accuracy.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: The uncropped visualization results of training dynamics.

Figure 9: The $5$ -dimensional sample representation space learned by CoQE for the combination functions.

Model	Config	Metric	Settings $(P_{\text{bursty}},\alpha)$					Avg.
Model	Config	Metric	$(0.9,1)$	$(0.9,0)$	$(0.9,2)$	$(0.5,0)$	$(0.1,0)$	Avg.
Transformer	$E=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
	$E=64,L=8$	ICL	74.4	67.2	50.2	57.5	50.3	59.9
	#Param: 0.6M	IWL	65.2	8.1	97.9	21.4	57.9	50.1
	$E=64,L=4$	ICL	50.0	50.5	50.1	50.2	50.4	50.2

Appendix

Table 3: Performance comparison of different Transformer variants and baselines. All values are presented as percentages. The maximum and second-largest values are indicated by boldface and underline, respectively.

Model	Config	ICL	IWL
Trans. ( $E=64,H=64$ )	$L=12$	84.2	65.1
	$L=14$	86.5	62.9
	$L=16$	73.0	62.9
CoQE ( $E=256,H=64$ )	$L=4$	90.3	82.0
	$L=6$	89.3	84.7
	$L=8$	89.7	85.2

Appendix

Table 4: The effect of appropriately scaling both the Transformer and CoQE.

Panel of Figure 10 — Figure 10: Learning curves under different noise levels

In-Context Learning

In-context learning studies how models adapt to new tasks from examples or instructions supplied at inference time, without changing their weights.

See all 402 papers

Capturing In-Context Learning Dynamics with Task Operators
Oct 1, 2026Guangzhi Xiong, Zhenghao He, Bohan Liu +3Transformer Interpretability In-Context Learning
University of Virginia
Convergence and Emergence of In-Context Reinforcement Learning with Chain of Thought
May 8, 2026Zixuan Xie, Xinyu Liu, Rohan Chandra +1Reinforcement Learning Transformer
University of Virginia
Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning
Sep 24, 2026Zhongjie Shi, Rongjie Lai, Alexander Cloninger +1Transformer Multi-Task Learning
School of Mathematics, Georgia Institute of Technology, Atlanta, GA 30332, United States · Department of Mathematics, Purdue University, West Lafayette, IN 47907, United States · Department of Mathematics and Halıcıoğlu Data Science Institute, University of California, San Diego, La Jolla, CA 92093, United States

Oct 1, 2026J/K move · Enter open · S save

Capturing In-Context Learning Dynamics with Task Operators

Guangzhi Xiong, Zhenghao He, Bohan Liu, Sanchit Sinha+2

University of Virginia

Transformer Interpretability In-Context Learning

Open paper AskPDF arXiv

In-context learning (ICL) enables language models to perform new tasks from demonstrations without weight updates. However, every ICL inference requires processing the full set of examples, resulting in inefficient deployments, and how ICL works mechanistically is not fully understood. Prior work compresses ICL into fixed activation vectors extracted from specific layers or positions, but these input-independent interventions fail on complex tasks where the output depends on fine-grained interactions with the input. By analyzing the ICL forward pass, we show that each attention head's output is an affine transformation of its context-masked counterpart, and that the parameters of this transformation are empirically stable across samples for a given task. Building on this, we introduce Task Operator (TO), which replays this transformation as an analytically derived update to the attention output projection. Across lexical, algorithmic, and reasoning tasks, TO achieves the best overall performance among prior methods and substantially narrows the gap between zero-shot inference and ICL. We further show that the extracted knowledge concentrates in a task-specific sparse circuit across layers and positions, and that averaging operators from disjoint demonstration batches enables effective many-shot scaling without expanding the context window. Our code is available at https://github.com/gzxiong/task_operator.

Abstract

In-Context Learning

Explore similar work

Capturing In-Context Learning Dynamics with Task Operators

Convergence and Emergence of In-Context Reinforcement Learning with Chain of Thought

Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning