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

TL;DR

This paper establishes that multi-layer cross-attention mechanisms are provably optimal for multi-modal in-context learning, demonstrating the importance of depth and cross-attention in achieving Bayes-optimal performance.

Contribution

The paper introduces a mathematically tractable framework for multi-modal in-context learning and proves the optimality of multi-layer cross-attention mechanisms under this model.

Findings

Single-layer linear self-attention cannot recover Bayes-optimal predictor.

Linearized multi-layer cross-attention is provably Bayes optimal with gradient flow.

Depth enhances the in-context learning capability of attention-based models.

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.