深度与循环结构决定Transformer多任务学习能力,且循环结构更鲁棒。
On the Role of Depth and Looping for In-Context Learning with Task Diversity
- 通过分析不同任务的线性回归,揭示深度对多任务学习的关键作用。
- 深层Transformer需至少√κ层才能解决复杂任务,但易受分布偏移影响。
- 循环Transformer在保持表达力的同时,具备更强鲁棒性和深度单调性。
深度Transformer模型在上下文学习(ICL)中的惊人表现引人关注。现有研究多基于单一模态高斯数据,认为ICL源于Transformer模拟梯度下降的能力。然而,这些研究未能捕捉其在多种任务共存情境下的学习能力。本文研究具有多样化任务的线性回归问题,任务由条件数范围为[1, κ]的数据协方差矩阵表征,强调深度的重要性。理论上证明:在无限制(或受限)注意力设置下,需至少log(κ)(或√κ)层线性注意力。进一步表明,多层Transformer可实现该下界所需层数。但发现其表达力以牺牲鲁棒性为代价:即使在Wasserstein距离仅O(e⁻ᴸ)的微小分布偏移下也不稳定,其中L为网络深度。随后证明,具有权重共享的循环Transformer不仅具备相似表达力,且在温和假设下可证明鲁棒性。此外,循环Transformer是唯一展现出损失随深度单调下降的模型,显著提升泛化能力。
原文摘要 · Abstract (English)
The intriguing in-context learning (ICL) abilities of deep Transformer models have lately garnered significant attention. By studying in-context linear regression on unimodal Gaussian data, recent empirical and theoretical works have argued that ICL emerges from Transformers' abilities to simulate learning algorithms like gradient descent. However, these works fail to capture the remarkable ability of Transformers to learn multiple tasks in context. To this end, we study in-context learning for linear regression with diverse tasks, characterized by data covariance matrices with condition numbers ranging from $[1, κ]$, and highlight the importance of depth in this setting. More specifically, (a) we show theoretical lower bounds of $\log(κ)$ (or $\sqrtκ$) linear attention layers in the unrestricted (or restricted) attention setting and, (b) we show that multilayer Transformers can indeed solve such tasks with a number of layers that matches the lower bounds. However, we show that this expressivity of multilayer Transformer comes at the price of robustness. In particular, multilayer Transformers are not robust to even distributional shifts as small as $O(e^{-L})$ in Wasserstein distance, where $L$ is the depth of the network. We then demonstrate that Looped Transformers -- a special class of multilayer Transformers with weight-sharing -- not only exhibit similar expressive power but are also provably robust under mild assumptions. Besides out-of-distribution generalization, we also show that Looped Transformers are the only models that exhibit a monotonic behavior of loss with respect to depth.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。