Table of Contents
Fetching ...

From Compression to Expression: A Layerwise Analysis of In-Context Learning

Jiachen Jiang, Yuxin Dong, Jinxin Zhou, Zhihui Zhu

TL;DR

The paper investigates how in-context learning encodes task information across layers in large language models, uncovering a Layerwise Compression-Expression dynamic where early layers compress demonstration-derived task signals and later layers express them to generate outputs. It introduces the TDNV metric to quantify within-task compression versus between-task separation and demonstrates this phenomenon universally across architectures and tasks, with larger models and more demonstrations yielding stronger compression and better ICL performance. A bias-variance decomposition of task vectors shows that both bias and variance decay on the order of $\mathcal{O}(1/K)$ as the number of demonstrations $K$ increases, and analyses reveal how attention contributes to this reduction. The work also provides practical insights, such as using minimum-TDNV to identify optimal task-vector layers and applying contrastive fine-tuning to further compress task representations, which improves task-vector accuracy by about 20% on symbolic tasks. Overall, the findings offer a principled view of internal representation dynamics in ICL, with implications for interpretability, robustness to noisy demonstrations, and efficiency in large-scale models.

Abstract

In-context learning (ICL) enables large language models (LLMs) to adapt to new tasks without weight updates by learning from demonstration sequences. While ICL shows strong empirical performance, its internal representational mechanisms are not yet well understood. In this work, we conduct a statistical geometric analysis of ICL representations to investigate how task-specific information is captured across layers. Our analysis reveals an intriguing phenomenon, which we term *Layerwise Compression-Expression*: early layers progressively produce compact and discriminative representations that encode task information from the input demonstrations, while later layers express these representations to incorporate the query and generate the prediction. This phenomenon is observed consistently across diverse tasks and a range of contemporary LLM architectures. We demonstrate that it has important implications for ICL performance -- improving with model size and the number of demonstrations -- and for robustness in the presence of noisy examples. To further understand the effect of the compact task representation, we propose a bias-variance decomposition and provide a theoretical analysis showing how attention mechanisms contribute to reducing both variance and bias, thereby enhancing performance as the number of demonstrations increases. Our findings reveal an intriguing layerwise dynamic in ICL, highlight how structured representations emerge within LLMs, and showcase that analyzing internal representations can facilitate a deeper understanding of model behavior.

From Compression to Expression: A Layerwise Analysis of In-Context Learning

TL;DR

The paper investigates how in-context learning encodes task information across layers in large language models, uncovering a Layerwise Compression-Expression dynamic where early layers compress demonstration-derived task signals and later layers express them to generate outputs. It introduces the TDNV metric to quantify within-task compression versus between-task separation and demonstrates this phenomenon universally across architectures and tasks, with larger models and more demonstrations yielding stronger compression and better ICL performance. A bias-variance decomposition of task vectors shows that both bias and variance decay on the order of as the number of demonstrations increases, and analyses reveal how attention contributes to this reduction. The work also provides practical insights, such as using minimum-TDNV to identify optimal task-vector layers and applying contrastive fine-tuning to further compress task representations, which improves task-vector accuracy by about 20% on symbolic tasks. Overall, the findings offer a principled view of internal representation dynamics in ICL, with implications for interpretability, robustness to noisy demonstrations, and efficiency in large-scale models.

Abstract

In-context learning (ICL) enables large language models (LLMs) to adapt to new tasks without weight updates by learning from demonstration sequences. While ICL shows strong empirical performance, its internal representational mechanisms are not yet well understood. In this work, we conduct a statistical geometric analysis of ICL representations to investigate how task-specific information is captured across layers. Our analysis reveals an intriguing phenomenon, which we term *Layerwise Compression-Expression*: early layers progressively produce compact and discriminative representations that encode task information from the input demonstrations, while later layers express these representations to incorporate the query and generate the prediction. This phenomenon is observed consistently across diverse tasks and a range of contemporary LLM architectures. We demonstrate that it has important implications for ICL performance -- improving with model size and the number of demonstrations -- and for robustness in the presence of noisy examples. To further understand the effect of the compact task representation, we propose a bias-variance decomposition and provide a theoretical analysis showing how attention mechanisms contribute to reducing both variance and bias, thereby enhancing performance as the number of demonstrations increases. Our findings reveal an intriguing layerwise dynamic in ICL, highlight how structured representations emerge within LLMs, and showcase that analyzing internal representations can facilitate a deeper understanding of model behavior.

Paper Structure

This paper contains 31 sections, 1 theorem, 18 equations, 27 figures, 5 tables.

Key Result

Theorem 5.1

Suppose that each demonstration $\boldsymbol{h}_i,i =1,\ldots,K$ is i.i.d. randomly generated from a distribution $\mathcal{H}$ on $\mathbb{R}^d$. Then the output of the query token, $\boldsymbol{h}'_{q}(K) = [\mathrm{Attn}\left(\boldsymbol{h}_1,\ldots,\boldsymbol{h}_K, \boldsymbol{h}_{q})\right)]_{

Figures (27)

  • Figure 1: Layer-wise compression to expression in ICL representations. TDNV first decreases then increases from shallow to deep layers, splitting the model into compression and expression stages. During the compression stage, task vector accuracy increases as task information is progressively extracted from demonstration pairs. During the expression stage, early-exit accuracy increases as output information is progressively decoded based on the input query. Refer to \ref{['app:tv-early-acc']} for detailed explanation of task vector and early-exit accuracy.
  • Figure 2: Layerwise TDNV of different model architectures, including decoder-only transformers and state-space models.
  • Figure 3: Layerwise TDNV during training process. The phenomenon emerges and intensifies with training progress.
  • Figure 4: Symbolic ICL.
  • Figure 5: Language Understanding ICL.
  • ...and 22 more figures

Theorems & Definitions (3)

  • Theorem 5.1: Bias-variance decomposition with respect to $K$
  • proof
  • proof