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

2025-10-29

AI Summary

This work investigates the theoretical mechanisms underlying in-context learning (ICL) in pretrained Transformers for nonlinear tasks, focusing on performance limits under nonlinear MLP heads and heterogeneous multi-source data (with distinct input, task, and noise distributions). We propose an asymptotic equivalence framework: when a Transformer is coupled with a two-layer MLP—where the first layer undergoes single-step gradient updates and the second layer is fully optimized—its ICL behavior is asymptotically equivalent to a structured polynomial predictor. Leveraging high-dimensional asymptotics, Gaussian universality, and orthogonal polynomial theory, we establish the first sufficient conditions under which nonlinear heads enhance ICL capability and identify the feature-learning emergence threshold. Theoretically and empirically, we show that low-noise, structured-covariance data sources significantly improve generalization. Our framework is validated on real-world multilingual sentiment analysis, confirming practical efficacy.

Technology Category

• Machine Learning: Structured Learning
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Pretrained Transformers demonstrate remarkable in-context learning (ICL) capabilities, enabling them to adapt to new tasks from demonstrations without parameter updates. However, theoretical studies often rely on simplified architectures (e.g., omitting MLPs), data models (e.g., linear regression with isotropic inputs), and single-source training, limiting their relevance to realistic settings. In this work, we study ICL in pretrained Transformers with nonlinear MLP heads on nonlinear tasks drawn from multiple data sources with heterogeneous input, task, and noise distributions. We analyze a model where the MLP comprises two layers, with the first layer trained via a single gradient step and the second layer fully optimized. Under high-dimensional asymptotics, we prove that such models are equivalent in ICL error to structured polynomial predictors, leveraging results from the theory of Gaussian universality and orthogonal polynomials. This equivalence reveals that nonlinear MLPs meaningfully enhance ICL performance, particularly on nonlinear tasks, compared to linear baselines. It also enables a precise analysis of data mixing effects: we identify key properties of high-quality data sources (low noise, structured covariances) and show that feature learning emerges only when the task covariance exhibits sufficient structure. These results are validated empirically across various activation functions, model sizes, and data distributions. Finally, we experiment with a real-world scenario involving multilingual sentiment analysis where each language is treated as a different source. Our experimental results for this case exemplify how our findings extend to real-world cases. Overall, our work advances the theoretical foundations of ICL in Transformers and provides actionable insight into the role of architecture and data in ICL.

Problem

• Analyzing in-context learning in Transformers with nonlinear MLPs
• Establishing asymptotic equivalence to structured polynomial predictors
• Investigating data mixing effects across heterogeneous data sources

Innovation

• Transformer with nonlinear MLP heads for in-context learning
• Two-layer MLP trained via gradient step and optimization
• Equivalence to structured polynomial predictors under asymptotics