Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learning

Nicholas Barnfield; Subhabrata Sen; Pragya Sur

stat.MLFeb 4, 2026

Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learning

Authors: Nicholas Barnfield, Subhabrata Sen, Pragya Sur

Organizations: Department of Statistics, Harvard University

Abstract

Recent progress has rapidly advanced our understanding of the mechanisms underlying in-context learning in modern attention-based neural networks. However, existing results focus exclusively on unimodal data; in contrast, the theoretical underpinnings of in-context learning for multi-modal data remain poorly understood. We introduce a mathematically tractable framework for studying multi-modal learning and explore when transformer-like architectures can recover Bayes-optimal performance in-context. To model multi-modal problems, we assume the observed data arises from a latent factor model. Our first result comprises a negative take on expressibility: we prove that single-layer self-attention models fail to recover the Bayes-optimal predictor uniformly over the task distribution. To address this limitation, we introduce a novel, linearized cross-attention mechanism, which we study in the regime where both the number of cross-attention layers and the context length are large. We show that this cross-attention mechanism is provably Bayes optimal when optimized using gradient flow. Our results underscore the benefits of depth for in-context learning and establish the provable utility of cross-attention for multi-modal distributions.

Figures & tables

Figure 1: Re-injection of covariates $\boldsymbol{X}$ throughout the $\mathsf{CA}_{T}$ ( $\mathsf{LCA}_{T}$ ) embeddings. The propagation of the raw data $\boldsymbol{X}$ occurs both through $\boldsymbol{S}_{t}$ as well as through the CA block $\boldsymbol{A}_{t}$ .

Figure 2: Prediction error for the one- and two-parameter LCA-based models and the LSA model. Left: ICL error at fixed depth for varying

L_{\mathrm{te}}

. Right: ICL error at fixed

L_{\mathrm{te}}

for varying depth

T

Figure 3: Diagnostics for less-restricted LCA parametrizations (setup same as in Figure 2 ; see Appendix I ). Left : layerwise scalar projections; Right : whitening error across layers.

Model	Params	Best test acc. (%)	Best epoch
Full paper-inspired model (Inj. & cross-attn)	$4.36\times 10^{4}$	63.85	64
No injection	$3.82\times 10^{4}$	62.72	65
No cross-attn	$4.36\times 10^{4}$	62.10	79
No inj. & no cross-attn	$3.82\times 10^{4}$	60.79	90
Standard bi-CA transformer	$6.07\times 10^{4}$	62.72	62

Table 1: Classification results on avMNIST for the full paper-inspired model, three ablations thereof, and a standard bi-directional cross-attention transformer. All models use (pre) layer-norm,

4

-headed multi-head attention,

10

attention layers, embedding dim.

d=16

, hidden dim.

d_{h}=32

, dropout rate of

0.1

, GELU MLP blocks, and an appended

\mathsf{CLS}

token. Training: cross-entropy loss over 100 epochs using AdamW optimizer with learning rate

3\times 10^{-4}

and weight decay

10^{-4}

. Model parameter counts and best test accuracy with corresponding epoch are shown.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Log-scale heat maps of the two-parameter excess loss $\ell(\alpha,\beta)$ with $d_{1}=d_{2}=16$ and $\|\boldsymbol{m}\|\sim{\rm Unif}(0,2)$ . ( Left ) depth $T=10$ . ( Middle ) depth $T=5$ . ( Right ) depth $T=5$ with pre-LayerNorm added between the CA layers. The dashed line marks $\alpha=-\beta$ over $\beta\in(-2/(\overline{m}+1),0)$ . Each panel carries its own colour scale, as the three surfaces span different magnitudes.

Figure 5: In-context performance at various

L_{\rm te}

(

T=10

d_{1}=d_{2}=16

). ( Left ) one- and two-parameter LCA-based models, the two ablations without

S_{t}

, and the sample mean

\bar{y}_{L_{\rm te}}

. ( Right ) one- and two-parameter LCA-based models, the two-parameter deep LSA model, the LSA model from ( 4.2 ), and the sample mean. Models are optimized on

\ell_{N,L_{\mathrm{tr}}}

(

L_{\rm tr}=100,N=2000

) using gradient descent. Performance is averaged over

1000

test-prompts where error bars represent standard deviation over

10

separate training runs. The sweep starts at

L_{\rm te}=128

, since for

L_{\rm te}

comparable to

d

the sample covariance fluctuates enough that the depth-

T

recurrence leaves its stable region on a small fraction of prompts.

Symbol	Meaning
$d=d_{1}+d_{2}$	total covariate dimension
$L$	context length (we use $L_{\mathrm{tr}}$ or $L_{\mathrm{te}}$ when needed)
$\boldsymbol{x}_{i}\in\mathbb{R}^{d}$	stacked covariate for context token $i\in[L]$
$\boldsymbol{x}_{q}\in\mathbb{R}^{d}$	query covariate
$y_{i}\in\mathbb{R}$	context response for token $i\in[L]$
$y_{q}\in\mathbb{R}$	query response

Appendix

Table 7

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 426 papers

Why Multimodal In-Context Learning Lags Behind? Unveiling the Inner Mechanisms and Bottlenecks
Apr 15, 2026Yu Wang, Sharon LiMultimodal ICL In-Context Learning
Department of Computer Sciences, University of Wisconsin-Madison
Sequential Correlations Change In-Context Learning: Effective Context Length and Architectural Mismatch
Jul 4, 2026Mary Letey, Yue M. Lu, Cengiz Pehlevan +1In-Context Learning Attention Mechanisms
John A. Paulson School of Engineering and Applied Sciences, Harvard University · Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University · Society of Fellows and Center for Brain Science, Harvard University
UniICL: Systematizing Unified Multimodal In-context Learning through a Capability-Oriented Taxonomy
Mar 25, 2026Yicheng Xu, Jiangning Zhang, Zhucun Xue +5Multimodal ICL Unified Multimodal Models
1Zhejiang University, China. · 2Shanghai Jiaotong University, China. · 3National University of Singapore, Singapore. +1

Apr 15, 2026J/K move · Enter open · S save

Why Multimodal In-Context Learning Lags Behind? Unveiling the Inner Mechanisms and Bottlenecks

Yu Wang, Sharon Li

Department of Computer Sciences, University of Wisconsin-Madison

Multimodal ICL In-Context Learning Multimodal Large Language Models Modality Gap in VLMs

Open paper AskPDF arXiv

In-context learning (ICL) enables models to adapt to new tasks via inference-time demonstrations. Despite its success in large language models, the extension of ICL to multimodal settings remains poorly understood in terms of its internal mechanisms and how it differs from text-only ICL. In this work, we conduct a systematic analysis of ICL in multimodal large language models. Using identical task formulations across modalities, we show that multimodal ICL performs comparably to text-only ICL in zero-shot settings but degrades significantly under few-shot demonstrations. To understand this gap, we decompose multimodal ICL into task mapping construction and task mapping transfer, and analyze how models establish cross-modal task mappings, and transfer them to query samples across layers. Our analysis reveals that current models lack reasoning-level alignment between visual and textual representations, and fail to reliably transfer learned task mappings to queries. Guided by these findings, we further propose a simple inference-stage enhancement method that reinforces task mapping transfer. Our results provide new insights into the mechanisms and limitations of multimodal ICL and suggest directions for more effective multimodal adaptation. Our code is available here.

Abstract

In-Context Learning

Explore similar work

Why Multimodal In-Context Learning Lags Behind? Unveiling the Inner Mechanisms and Bottlenecks

Sequential Correlations Change In-Context Learning: Effective Context Length and Architectural Mismatch

UniICL: Systematizing Unified Multimodal In-context Learning through a Capability-Oriented Taxonomy