Pretrained transformer efficiently learns low-dimensional target functions in-context

TL;DR

本论文提出一种基于预训练变换器的非线性低维目标函数学习方法,利用梯度优化实现对单指数模型的高效在语境学习,样本复杂度仅依赖目标函数的低维结构。

cs.LG 🔴 高级 2024-11-05 34 引用 68 次浏览
Kazusato Oko Yujin Song Taiji Suzuki Denny Wu
深度学习 变换器 在语境学习 低维结构 非线性函数

核心发现

方法论

本文采用结合非线性MLP层的预训练变换器,通过梯度下降优化模型参数,实现对单指数目标函数的在语境学习。研究中引入信息指数Q和多项式最高度P,分析预训练任务的样本复杂度与目标函数的低维结构关系。具体方法包括:• 构建以σ*(⟨x,β⟩)为目标的单指数模型,β从r维子空间中采样;• 利用带有非线性激活的MLP嵌入层提取特征;• 通过自注意力机制实现高效的函数逼近;• 采用逐层梯度下降优化MLP和注意力参数,确保模型在预训练阶段学习到低维子空间信息;• 最终在有限的提示长度下实现对目标函数的高精度预测。该方法充分利用目标函数的低秩结构,通过预训练实现样本效率的提升。

关键结果

  • 实验结果显示,预训练变换器在单指数模型任务中,所需的提示长度仅依赖目标函数的子空间维度r,而与高维输入空间d无关。在特定设置下,样本复杂度达到N* ≈ r^{P},远优于传统的核方法和线性模型,其在d=32、r=8、P=3的场景中,提示长度减少了50%以上,显著优于仅在测试提示上学习的算法。
  • 在多任务预训练中,模型在不同任务间成功捕获低维子空间结构,展现出优越的适应性。与线性变换器仅实现一次梯度步不同,本文的非线性MLP层能有效逼近复杂的非线性目标函数,提升了在高维空间中的统计效率。
  • 通过理论分析和数值模拟,验证了预训练样本复杂度与目标函数的低秩结构密切相关,而与输入空间维度d无关。这一发现为大规模预训练模型在高维非线性任务中的应用提供了理论基础,表明预训练变换器具有强大的结构适应能力。

研究意义

本研究突破了以线性模型为代表的在语境学习理论的局限,首次系统性地证明了带有非线性MLP层的预训练变换器在学习低维非线性目标函数中的统计优势。该成果不仅丰富了深度学习的理论体系,也为实际中的高维非线性任务提供了高效的模型设计思路。通过揭示预训练模型在捕获目标函数低秩结构方面的潜能,推动了预训练技术在复杂任务中的应用前沿。未来,结合更深层次的网络结构和多任务学习,有望实现更广泛的泛化能力和样本效率,为人工智能的自主学习能力提供坚实的理论支撑。

技术贡献

本文提出了结合非线性MLP层的预训练变换器架构,首次系统性分析了其在学习单指数模型中的样本复杂度,证明了预训练可以有效捕获目标函数的低秩结构,显著降低在高维空间中的样本需求。通过引入信息指数Q与多项式最高度P,建立了模型参数规模与目标函数复杂度的关系,为深度学习中的结构适应性提供了理论依据。此外,论文设计了端到端的梯度优化算法,保证模型在有限预训练任务和提示长度下实现高精度预测,填补了非线性函数在语境学习中的理论空白。

新颖性

本研究的创新点在于:一是首次系统性地将非线性MLP层引入预训练变换器,用于学习低维非线性目标函数;二是提出基于信息指数Q和多项式度P的样本复杂度分析框架,揭示模型在高维空间中的结构适应性;三是通过理论和实验验证,证明预训练变换器在捕获低秩目标函数方面优于传统的核方法和线性模型。这些创新极大丰富了在语境学习的理论体系,为未来深度模型的结构设计提供了新思路。

局限性

  • 本研究假设目标函数为单指数模型,且目标函数的低秩结构已知,实际应用中可能存在更复杂的非线性关系和未知结构,限制了模型的泛化能力。
  • 模型训练依赖大量预训练任务和参数调优,计算成本较高,实际部署时需考虑效率与资源限制。
  • 理论分析主要集中在理想化的高维正态输入分布,实际数据中的非高斯分布可能影响模型性能,未来需扩展到更复杂的数据分布场景。

未来方向

未来的研究方向包括:• 扩展模型以适应更复杂的非线性目标函数和未知结构,提升实际应用中的鲁棒性;• 结合深层网络结构和多任务预训练,增强模型的泛化能力;• 研究非高斯输入分布下的理论性能,适应实际数据的多样性;• 优化训练算法以降低计算成本,实现更高效的预训练与推理;• 探索在自然语言处理、计算机视觉等多模态任务中的应用潜力,推动预训练模型的广泛落地。

AI 总览摘要

在人工智能领域,预训练变换器(Transformer)凭借其强大的表达能力和泛化能力,已成为深度学习的核心架构之一。然而,尽管在大量任务中表现出色,其在语境学习(In-Context Learning, ICL)能力的理论基础仍不充分,尤其是在非线性高维任务中。传统的理论分析多集中于线性模型,显示预训练变换器在这些场景中实现了类似梯度下降的算法,样本复杂度与输入维度d成线性关系,限制了其在高维非线性任务中的应用潜力。

本文突破性地提出一种结合非线性MLP层的预训练变换器架构,专门针对单指数模型(Single-Index Model)进行理论分析。作者引入信息指数Q和多项式最高度P,揭示预训练模型在学习目标函数中的低秩结构时,样本复杂度仅依赖于目标函数的低维子空间维度r,而与输入空间的高维度d无关。这意味着,预训练变换器可以在极少的提示长度下,准确预测复杂的非线性目标函数,显著优于传统的核方法和线性模型。

通过严格的数学分析和大量数值实验,论文验证了模型在不同参数设置下的优越性能。在d=32、r=8、P=3的场景中,提示长度减少了50%以上,模型在高维空间中的样本效率得到极大提升。这一发现不仅丰富了在语境学习的理论体系,也为实际应用中的高维非线性任务提供了新的解决方案。未来,结合更深层次的网络结构、多任务学习和多模态数据,有望推动预训练模型在复杂场景中的广泛应用。

总之,本文为预训练变换器在高维非线性任务中的结构适应性提供了坚实的理论基础,开启了深度学习在自主学习和少样本学习方面的新篇章。未来的研究将继续探索模型的鲁棒性、多样性以及在实际场景中的部署效率,推动人工智能迈向更智能、更高效的未来。

深度解读

原文摘要

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.

cs.LG stat.ML

参考文献 (20)

Attention is All you Need

Ashish Vaswani, Noam Shazeer, Niki Parmar 等

2017 189977 引用 ⭐ 高影响力 查看解读 →

Optimal errors and phase transitions in high-dimensional generalized linear models

Jean Barbier, Florent Krzakala, N. Macris 等

2017 315 引用 ⭐ 高影响力 查看解读 →

Trained Transformers Learn Linear Models In-Context

Ruiqi Zhang, Spencer Frei, P. Bartlett

2023 370 引用 ⭐ 高影响力 查看解读 →

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

2022 155 引用 ⭐ 高影响力 查看解读 →

Neural Networks can Learn Representations with Gradient Descent

Alexandru Damian, Jason D. Lee, M. Soltanolkotabi

2022 195 引用 ⭐ 高影响力 查看解读 →

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

2020 1643 引用 ⭐ 高影响力

What Can Transformers Learn In-Context? A Case Study of Simple Function Classes

Shivam Garg, Dimitris Tsipras, Percy Liang 等

2022 854 引用 ⭐ 高影响力 查看解读 →

Learning Single-Index Models with Shallow Neural Networks

A. Bietti, Joan Bruna, C. Sanford 等

2022 116 引用 查看解读 →

A Vector-Contraction Inequality for Rademacher Complexities

Andreas Maurer

2016 315 引用 查看解读 →

Fundamental Limits of Weak Recovery with Applications to Phase Retrieval

Marco Mondelli, A. Montanari

2017 144 引用 查看解读 →

Learning Single-Index Models in Gaussian Space

Rishabh Dudeja, Daniel J. Hsu

2018 65 引用

Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit

B. Barak, Benjamin L. Edelman, Surbhi Goel 等

2022 209 引用 查看解读 →

Neural Networks Efficiently Learn Low-Dimensional Representations with SGD

Alireza Mousavi-Hosseini, Sejun Park, M. Girotti 等

2022 70 引用 查看解读 →

Language Models are Few-Shot Learners

Tom B. Brown, Benjamin Mann, Nick Ryder 等

2020 62453 引用 查看解读 →

What learning algorithm is in-context learning? Investigations with linear models

Ekin Akyürek, Dale Schuurmans, Jacob Andreas 等

2022 770 引用 查看解读 →

Transformers learn in-context by gradient descent

J. Oswald, Eyvind Niklasson, E. Randazzo 等

2022 849 引用 查看解读 →

Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers

Damai Dai, Yutao Sun, Li Dong 等

2022 257 引用 查看解读 →

SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics

E. Abbe, Enric Boix-Adserà, Theodor Misiakiewicz

2023 156 引用 查看解读 →

Learning Time-Scales in Two-Layers Neural Networks

Raphael Berthier, A. Montanari, Kangjie Zhou

2023 62 引用 查看解读 →

Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models

Alexandru Damian, Eshaan Nichani, Rong Ge 等

2023 67 引用 查看解读 →

被引用 (20)

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

2025 1 引用 ⭐ 高影响力 查看解读 →

Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective

2025 6 引用 查看解读 →

How Can Mamba Learn In Context with Outliers and Generalize Provably?

2025 2 引用 查看解读 →

Test time training enhances in-context learning of nonlinear functions

2025 4 引用 查看解读 →

Review of Hallucination Understanding in Large Language and Vision Models

2025 2 引用 查看解读 →

Asymptotic Study of in-Context Learning with Random Transformers Through Equivalent Models

2025 3 引用 查看解读 →

Multi-Layer Attention is the Amplifier of Demonstration Effectiveness

2025 1 引用 查看解读 →

Provable In-Context Learning of Nonlinear Regression with Transformers

2025 5 引用 查看解读 →

Neural Interpretable PDEs: Harmonizing Fourier Insights with Attention for Scalable and Interpretable Physics Discovery

2025 3 引用 查看解读 →

Learning Linear Regression with Low-Rank Tasks in-Context

2025 2 引用 查看解读 →

Transformer learns the cross-task prior and regularization for in-context learning

2025 6 引用 查看解读 →

Provable Failure of Language Models in Learning Majority Boolean Logic via Gradient Descent

2025 10 引用 查看解读 →

When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective

2025 6 引用 查看解读 →

CoT-ICL Lab: A Synthetic Framework for Studying Chain-of-Thought Learning from In-Context Demonstrations

2025 3 引用 查看解读 →

On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning

2025 6 引用 查看解读 →

In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning

2024 4 引用 查看解读 →

Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit

2024 49 引用 查看解读 →

CoT-ICL Lab: A Petri Dish for Studying Chain-of-Thought Learning from In-Context Demonstrations

2025 1 引用

Handover of In-Context Learning State Across Session Boundaries

In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning

2025 12 引用 查看解读 →