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

arXiv:2602.04872v2 Announce Type: replace-cross 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, linear self-attention fails 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.

Related

• Gaussian Equivalence for Self-Attention: Asymptotic Spectral Analysis of Attention Matrix
• Distinct mechanisms underlying in-context learning in transformers
• Attention Sink in Transformers: A Survey on Utilization, Interpretation, and Mitigation

Source: arXiv cs.LG | 2026-04-29