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

1. Problem Statement

The paper addresses the theoretical gap in understanding In-Context Learning (ICL) within multi-modal settings. While recent theoretical work has successfully explained ICL in unimodal (e.g., text-only) regression tasks using linearized attention models, these theories fail when applied to multi-modal data (e.g., text and images combined).

Core Challenge:
In standard unimodal ICL theory, it is assumed that the distribution of covariates (input features) remains invariant across different tasks/prompts. However, in multi-modal settings, the relationship between modalities often induces covariate shifts across prompts. Specifically, the paper models data where different modalities share a common latent factor but have distinct, task-specific noise and projection vectors. This coupling between the task vector (the rule to be learned) and the covariate distribution invalidates the assumptions required for single-layer linear self-attention (LSA) to succeed.

The central question is: Can attention-based neural networks learn prediction rules in-context when deployed on multi-modal datasets where covariate distributions shift across tasks?

2. Methodology

A. Statistical Model (Latent Factor Model)

The authors propose a mathematically tractable data generation model for multi-modal ICL:

• Latent Variable: A shared latent scalar $u^{(j)}_i \sim \mathcal{N}(0,1)$ drives the response $y$ and both modalities.
• Modalities: Two modalities, $\bar{z} \in \mathbb{R}^{d_1}$ and $\tilde{z} \in \mathbb{R}^{d_2}$, are generated as noisy linear projections of the latent variable:
  $$\bar{z} = u v + \xi, \quad \tilde{z} = u r + \eta$$
  where $v$ and $r$ are modality-specific vectors sampled independently for each prompt (task).
• Response: $y = \zeta u$, where $\zeta$ is a task-specific scalar.
• Key Feature: The covariate distribution depends on the prompt-specific vectors $m = [v; r]$. Consequently, the covariance matrix $\Lambda = I + mm^\top$ and the optimal regression vector $w$ vary across prompts.

B. Baseline Analysis: Failure of Single-Layer LSA

The authors first analyze a Single-Layer Linear Self-Attention (LSA) model, which has been proven optimal for unimodal ICL.

• Result: They prove that a single-layer LSA fails to achieve Bayes-optimal performance in this multi-modal setting.
• Reasoning: LSA relies on learning a fixed global inverse covariance matrix to whiten the inputs. In the multi-modal setting, the covariance $\Lambda$ is random and prompt-dependent. A fixed set of weights cannot invert a random, varying covariance matrix uniformly across the task distribution.

C. Proposed Solution: Multi-layer Linearized Cross-Attention (LCA)

To overcome the failure of LSA, the authors introduce a novel architecture:

• Architecture: A deep stack of Linearized Cross-Attention (LCA) layers followed by a single Linear Self-Attention (LSA) readout layer.
• Mechanism:
  • The model maintains an evolving embedding $F^t$ initialized at zero.
  • At each layer $t$, the embedding is updated via a recurrence: $F^t = F^{t-1} + S^{t-1} + A^{t-1}$.
  • Cross-Attention ($A^{t-1}$): Computes attention between the evolving embedding $F^{t-1}$ (queries) and the raw input data $X$ (keys/values). This allows the model to "read" the raw data at every step.
  • Skip Connection ($S^{t-1}$): A learnable linear injection of the raw data $X$ directly into the layer ($S^{t-1} = W^S X$). This is a novel component inspired by unrolled neural networks, crucial for theoretical tractability and performance.
• Linearization: The softmax is replaced with an identity map (linear attention) to allow for rigorous mathematical analysis of gradient flow dynamics.

3. Key Contributions

1. Theoretical Framework for Multi-modal ICL: Introduction of a latent factor model that captures the essential difficulty of multi-modal ICL: the coupling of task parameters with prompt-dependent covariate distributions.
2. Negative Result for LSA: A formal proof (Theorem 4.1) that single-layer linear self-attention cannot achieve Bayes-optimal ICL when covariate distributions shift across prompts, contrasting sharply with unimodal results.
3. Positive Result for Multi-layer CA: Proof that a deep Linear Cross-Attention (LCA) architecture, trained via gradient flow, converges to the Bayes-optimal predictor.
   • The model effectively learns to "whiten" the prompt-specific covariates implicitly, adapting to the random covariance $\Lambda$ of the current prompt.
4. Optimality of Depth: Demonstration that as the depth $T \to \infty$, the model parameters converge to values that minimize the worst-case error, achieving minimax optimality.

4. Main Results

The paper analyzes two simplified parameterizations of the proposed architecture trained via gradient flow on the population loss:

• One-Parameter Model: Weights are tied such that $W^S = -W^V = \alpha I$.

  • Theorem 6.2: Gradient flow converges to a unique limit $\alpha^*$. As depth $T \to \infty$, the predictor converges almost surely to the Bayes-optimal predictor $\langle w, x_q \rangle$.
  • Limiting Parameter: $\alpha^* = \frac{2}{2 + \underline{m} + \overline{m}}$, where $\underline{m}, \overline{m}$ are the essential infimum and supremum of the norm of the modality vector $m$.

• Two-Parameter Model: Weights are $W^S = \alpha I$ and $W^V = \beta I$.

  • Theorem 6.3: Under specific initialization conditions, gradient flow converges. As $T \to \infty$, the parameters converge such that $\beta^* \to -\alpha^*$ and $\alpha^* \to \frac{2}{2 + \underline{m} + \overline{m}}$.
  • Result: The resulting predictor is also Bayes-optimal.

Mechanism of Success:
The deep LCA stack acts as an iterative whitening filter. The recurrence relation allows the model to compute an approximation of $\Lambda^{-1} X$ (the inverse covariance times the input) implicitly from the prompt itself, rather than relying on a fixed global inverse. This enables the model to adapt to the specific covariate shift of each prompt.

5. Significance and Implications

• Bridging Theory and Practice: This work provides the first rigorous theoretical justification for why depth and cross-attention are necessary for multi-modal ICL, explaining empirical observations where shallow models fail on complex multi-modal tasks.
• Covariate Shift Handling: It identifies covariate shift across prompts as a fundamental barrier for shallow attention models and demonstrates how deep cross-attention mechanisms naturally overcome this by learning prompt-specific transformations.
• Architectural Insight: The inclusion of the $S^{t-1}$ skip connection (injecting raw data at every layer) is shown to be theoretically critical. Ablation studies (Appendix F) confirm that removing this component causes the model to fail, reverting to the performance of a simple sample mean.
• Future Directions: The paper suggests that these theoretical insights can guide the design of more robust multi-modal foundation models and points toward future research in random matrix theory for infinite token dimensions and sample-level generalization.

In summary, the paper proves that while shallow attention models are insufficient for multi-modal ICL due to prompt-dependent covariate shifts, deep cross-attention architectures are provably capable of learning the optimal prediction rule in-context by implicitly adapting to the statistical structure of each specific prompt.