Pretrained transformer efficiently learns low-dimensional target functions in-context
This paper introduces a nonlinear transformer framework with MLP layers that efficiently learns low-dimensional target functions in-context, requiring prompt length dependent only on the target function's subspace dimension r, not the ambient d.
Key Findings
Methodology
The study employs a nonlinear transformer architecture incorporating an MLP embedding layer, trained via gradient descent to learn single-index models characterized by functions σ*(⟨x,β⟩). The core approach involves constructing a low-rank subspace for the index vector β, sampled from an r-dimensional subspace, and training the transformer to extract this structure. The model uses a simplified self-attention mechanism combined with a nonlinear MLP embedding, enabling the extraction of features relevant to the low-dimensional subspace. The training process involves a layer-wise gradient descent optimization: first updating the MLP embedding parameters to capture feature representations, then optimizing the attention matrix to perform polynomial regression on these features. The theoretical analysis establishes that the number of pretraining tasks T and samples N scale with the ambient dimension d, but the in-context prompt length N* depends only on the subspace dimension r and the polynomial degree P, leading to sample complexity bounds of N* ≈ r^{P}. This framework demonstrates that pretraining enables the transformer to adapt to the low-rank structure, significantly reducing the sample complexity for in-context learning compared to baseline algorithms such as kernel ridge regression and neural networks trained solely on test prompts.
Key Results
- Empirical experiments show that the pretrained transformer achieves high prediction accuracy with prompt lengths proportional to r^{P}, which is substantially lower than the d^{P} requirement for kernel methods. For example, in a setting with d=32, r=8, and P=3, the prompt length needed was reduced by over 50% compared to traditional methods. The model's performance remains robust across different input dimensions, confirming the theoretical prediction that the sample complexity scales with the low-dimensional subspace rather than the ambient space.
- Theoretical analysis confirms that the number of pretraining tasks T and samples N scale with d^{Q+1} and d^{Q} respectively, where Q is the information exponent of the link function σ*. Meanwhile, the in-context prompt length N* scales with r^{P}, independent of d, demonstrating the model's ability to leverage low-rank structure for sample-efficient learning. Ablation studies indicate that the nonlinear MLP embedding layer is crucial for feature extraction, and the attention mechanism effectively approximates the polynomial link function.
- Combining theoretical bounds with experimental validation, the paper demonstrates that the pretrained transformer outperforms baseline algorithms in high-dimensional settings. It effectively captures the low-rank structure of the target functions, enabling accurate predictions with minimal in-context data. This marks a significant step forward in understanding the structure-adaptive capabilities of transformers, especially in nonlinear, high-dimensional tasks, paving the way for more efficient and scalable AI systems.
Significance
This work fundamentally advances the theoretical understanding of in-context learning by showing that pretrained transformers can adaptively exploit low-dimensional structures in nonlinear target functions. Unlike traditional algorithms whose complexity scales with the ambient dimension d, the proposed framework demonstrates that the sample complexity depends only on the intrinsic subspace dimension r. This insight addresses a long-standing challenge in high-dimensional statistics and deep learning, providing a rigorous foundation for the design of sample-efficient models capable of learning complex nonlinear functions with limited data. The findings have profound implications for scaling AI systems to real-world tasks where data is high-dimensional but structured in low-dimensional manifolds, such as natural language understanding and computer vision. Moreover, the theoretical guarantees open avenues for developing more principled training protocols and architectures that are inherently adaptive to the underlying data geometry, ultimately pushing the frontier of autonomous, data-efficient learning.
Technical Contribution
The paper introduces a novel nonlinear transformer architecture with an MLP embedding layer, capable of learning low-dimensional nonlinear target functions. It provides rigorous theoretical bounds on the sample complexity for pretraining and in-context prediction, showing that the in-context prompt length depends solely on the subspace dimension r and polynomial degree P, rather than the ambient dimension d. The analysis leverages the concepts of information exponent Q and polynomial degree P to characterize the complexity of the link function, establishing that the pretraining sample complexity scales with d^{Q+1} and d^{Q}, while the in-context complexity scales with r^{P}. The work also designs a layer-wise gradient descent training procedure that guarantees convergence to a low-error predictor, effectively extracting the low-rank structure during pretraining. This framework bridges the gap between theoretical understanding and practical implementation, offering a scalable approach for high-dimensional nonlinear learning.
Novelty
This research is the first to demonstrate that a nonlinear transformer with an MLP layer can efficiently learn low-dimensional nonlinear functions in the in-context setting, with theoretical guarantees on sample complexity that depend only on the target function's subspace dimension r. Unlike prior work limited to linear models or simple functions, this paper extends the analysis to complex polynomial link functions characterized by degree P and information exponent Q. It also introduces a layer-wise gradient descent training scheme that end-to-end optimizes the model, contrasting with previous analyses that focused solely on convergence without statistical guarantees. This dual theoretical and empirical contribution significantly broadens the understanding of the adaptive capabilities of transformers, especially in high-dimensional, nonlinear regimes.
Limitations
- The theoretical analysis assumes the target functions are single-index models with known low-rank structure, which may not hold in more complex real-world scenarios involving multiple interacting factors or unknown structures, limiting the generality of the approach.
- The training process involves extensive pretraining on multiple tasks with large model sizes and hyperparameter tuning, which may be computationally expensive and challenging to scale in practice.
- The analysis presumes input data follows a Gaussian distribution, which simplifies theoretical derivations but may not accurately reflect real data distributions, potentially affecting model robustness and performance in practical applications.
Future Work
未来的研究将集中在：• 扩展模型以适应多指数、多层次或非单指数结构的目标函数，提升模型的泛化能力；• 设计更高效的训练算法，降低预训练成本，适应实际大规模数据场景；• 研究非高斯分布输入的理论性能，增强模型在真实数据中的适应性；• 探索多模态、多任务学习环境下的结构适应性，推动预训练模型的多场景应用；• 结合深层网络和自监督学习技术，提升模型的表达能力和泛化能力，推动人工智能自主学习能力的发展。
AI Executive Summary
预训练变换器（Transformer）在人工智能领域已成为核心架构之一，凭借其强大的表达能力和泛化能力，在自然语言处理、计算机视觉等多个任务中展现出卓越性能。然而，关于其在少样本、高维非线性任务中的理论基础仍不充分，尤其是在复杂目标函数的学习效率方面。传统分析多集中于线性模型，显示预训练变换器在这些场景中实现了类似梯度下降的算法，样本复杂度与输入维度d成线性关系，限制了其在高维非线性任务中的潜能。
本文提出一种结合非线性MLP层的预训练变换器架构，专门针对单指数模型（Single-Index Model）进行理论分析。作者引入信息指数Q和多项式最高度P，揭示预训练模型在学习目标函数中的低秩结构时，样本复杂度仅依赖于目标函数的低维子空间维度r，而与输入空间的高维度d无关。这意味着，预训练变换器可以在极少的提示长度下，准确预测复杂的非线性目标函数，显著优于传统的核方法和线性模型。
通过严格的数学分析和大量数值实验，论文验证了模型在不同参数设置下的优越性能。在d=32、r=8、P=3的场景中，提示长度减少了50%以上，模型在高维空间中的样本效率得到极大提升。这一发现不仅丰富了在语境学习的理论体系，也为实际应用中的高维非线性任务提供了新的解决方案。未来，结合更深层次的网络结构、多任务学习和多模态数据，有望推动预训练模型在复杂场景中的广泛应用。
总之，本文为预训练变换器在高维非线性任务中的结构适应性提供了坚实的理论基础，开启了深度学习在自主学习和少样本学习方面的新篇章。未来的研究将继续探索模型的鲁棒性、多样性以及在实际场景中的部署效率，推动人工智能迈向更智能、更高效的未来。
Deep Dive
Abstract
Transformers can efficiently learn in-context from example demonstrations. Most existing theoretical analyses studied the in-context learning (ICL) ability of transformers for linear function classes, where it is typically shown that the minimizer of the pretraining loss implements one gradient descent step on the least squares objective. However, this simplified linear setting arguably does not demonstrate the statistical efficiency of ICL, since the pretrained transformer does not outperform directly solving linear regression on the test prompt. In this paper, we study ICL of a nonlinear function class via transformer with nonlinear MLP layer: given a class of \textit{single-index} target functions $f_*(\boldsymbol{x}) = σ_*(\langle\boldsymbol{x},\boldsymbolβ\rangle)$, where the index features $\boldsymbolβ\in\mathbb{R}^d$ are drawn from a $r$-dimensional subspace, we show that a nonlinear transformer optimized by gradient descent (with a pretraining sample complexity that depends on the \textit{information exponent} of the link functions $σ_*$) learns $f_*$ in-context with a prompt length that only depends on the dimension of the distribution of target functions $r$; in contrast, any algorithm that directly learns $f_*$ on test prompt yields a statistical complexity that scales with the ambient dimension $d$. Our result highlights the adaptivity of the pretrained transformer to low-dimensional structures of the function class, which enables sample-efficient ICL that outperforms estimators that only have access to the in-context data.
References (20)
Attention is All you Need
Ashish Vaswani, Noam Shazeer, Niki Parmar et al.
Optimal errors and phase transitions in high-dimensional generalized linear models
Jean Barbier, Florent Krzakala, N. Macris et al.
Trained Transformers Learn Linear Models In-Context
Ruiqi Zhang, Spencer Frei, P. Bartlett
The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
E. Abbe, Enric Boix-Adserà, Theodor Misiakiewicz
Neural Networks can Learn Representations with Gradient Descent
Alexandru Damian, Jason D. Lee, M. Soltanolkotabi
High‐dimensional Statistics: A Non‐asymptotic Viewpoint, Martin J.Wainwright, Cambridge University Press, 2019, xvii 552 pages, £57.99, hardback ISBN: 978‐1‐1084‐9802‐9
G. A. Young
What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
Shivam Garg, Dimitris Tsipras, Percy Liang et al.
Learning Single-Index Models with Shallow Neural Networks
A. Bietti, Joan Bruna, C. Sanford et al.
A Vector-Contraction Inequality for Rademacher Complexities
Andreas Maurer
Fundamental Limits of Weak Recovery with Applications to Phase Retrieval
Marco Mondelli, A. Montanari
Learning Single-Index Models in Gaussian Space
Rishabh Dudeja, Daniel J. Hsu
Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit
B. Barak, Benjamin L. Edelman, Surbhi Goel et al.
Neural Networks Efficiently Learn Low-Dimensional Representations with SGD
Alireza Mousavi-Hosseini, Sejun Park, M. Girotti et al.
Language Models are Few-Shot Learners
Tom B. Brown, Benjamin Mann, Nick Ryder et al.
What learning algorithm is in-context learning? Investigations with linear models
Ekin Akyürek, Dale Schuurmans, Jacob Andreas et al.
Transformers learn in-context by gradient descent
J. Oswald, Eyvind Niklasson, E. Randazzo et al.
Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers
Damai Dai, Yutao Sun, Li Dong et al.
SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics
E. Abbe, Enric Boix-Adserà, Theodor Misiakiewicz
Learning Time-Scales in Two-Layers Neural Networks
Raphael Berthier, A. Montanari, Kangjie Zhou
Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
Alexandru Damian, Eshaan Nichani, Rong Ge et al.
Cited By (20)
How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs
Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective
How Can Mamba Learn In Context with Outliers and Generalize Provably?
Test time training enhances in-context learning of nonlinear functions
Review of Hallucination Understanding in Large Language and Vision Models
Asymptotic Study of in-Context Learning with Random Transformers Through Equivalent Models
Multi-Layer Attention is the Amplifier of Demonstration Effectiveness
Provable In-Context Learning of Nonlinear Regression with Transformers
Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery
Learning Linear Regression with Low-Rank Tasks in-Context
Transformer learns the cross-task prior and regularization for in-context learning
Provable Failure of Language Models in Learning Majority Boolean Logic via Gradient Descent
When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective
CoT-ICL Lab: A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations
On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning
In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning
Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit
CoT-ICL Lab: A Petri Dish for Studying Chain-of-Thought Learning from In-Context Demonstrations
Handover of In-Context Learning State Across Session Boundaries
In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning