How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs

The Transformer architecture, introduced by Vaswani et al., has become the backbone of modern AI systems due to its scalability and effectiveness. Early theoretical work focused on simplified models, such as linear attention and linear tasks, to understand the mechanisms behind ICL. Recent empirical findings show that nonlinear MLP components are crucial for achieving state-of-the-art performance, especially on complex tasks. However, the theoretical understanding of how these nonlinearities contribute, particularly in multi-source heterogeneous data environments, remains incomplete. Existing studies often assume homogeneous data or ignore the impact of data quality, limiting their applicability. As real-world applications involve diverse data sources with varying noise levels and covariance structures, understanding the interplay between data properties and model architecture is vital for advancing AI capabilities.

Despite the empirical importance of nonlinear MLPs in Transformers, there is a lack of rigorous theoretical analysis explaining their role in ICL, especially under realistic multi-source data conditions. Current models often assume idealized, single-source, homogeneous data, which does not reflect practical scenarios. The core challenge lies in characterizing how data heterogeneity, quality, and structure influence feature learning and ICL performance. Additionally, understanding the asymptotic behavior of such models in high-dimensional regimes remains an open problem. Addressing these issues is critical for designing models that can effectively leverage complex, multi-source data in real-world tasks, such as multilingual NLP and multi-task learning.

This paper introduces several key innovations: 1) It establishes that a Transformer with a nonlinear MLP head is asymptotically equivalent to a finite-degree polynomial model, providing a new theoretical lens for understanding nonlinear ICL. 2) It incorporates multi-source heterogeneous data, analyzing how data covariance structures and noise levels influence feature learning and model performance. 3) It employs a novel training scheme combining a single gradient step and ridge regression, simplifying the analysis while capturing feature learning dynamics. These advances extend the theoretical framework of neural networks into complex, realistic data environments, bridging the gap between empirical success and mathematical understanding.

• �� Assume high-dimensional asymptotics: input dimension d, context length ℓ, sample size n, and hidden dimension k tend to infinity with fixed ratios. • Use Gaussian universality to replace complex random features with Gaussian equivalents, deriving the asymptotic distribution of F·vec(HZ). • Model training involves a single gradient step on the first layer’s weights, followed by ridge regression on the second layer, enabling tractable analysis. • Expand the nonlinear activation σ via Hermite polynomials, approximating it with a finite-degree polynomial in the limit. • Prove that the model’s output converges to a polynomial predictor, establishing the asymptotic equivalence. • Analyze the impact of data covariance structures, noise levels, and task distribution structure on feature learning and ICL error. • Validate the theoretical predictions through extensive simulations and real-world multilingual sentiment analysis experiments.

• �� Synthetic datasets simulate multiple sources with varying covariance structures and noise levels to test the theoretical predictions. • Compare linear and nonlinear MLP Transformer models across different sample sizes, context lengths, and model capacities, measuring ICL error. • Use ReLU and tanh activations to assess robustness of the polynomial approximation. • Real-world multilingual sentiment analysis dataset, treating each language as a source, to validate practical applicability. • Multiple runs ensure statistical significance, analyzing how data quality, structure, and model size influence performance. The experiments confirm the theoretical insights, showing nonlinear MLPs outperform linear counterparts, especially with structured data.

• �� The nonlinear MLP Transformer achieves over 15% lower ICL error compared to linear models in high-dimensional regimes, consistent with polynomial equivalence theory. • Data sources with structured covariance matrices and low noise levels significantly boost feature learning, leading to performance improvements exceeding 20%. • Richer task distribution structures promote deeper feature extraction, with models leveraging covariance complexity for better generalization. • Empirical results on multilingual sentiment analysis demonstrate that the theoretical findings extend to real-world, heterogeneous data scenarios, validating the model’s practical relevance.