arXiv:2606.00605cs.LGstat.ML2026-06被引 1

带层归一化的循环Transformer能自动学习幂法,实现主成分预测。

Looped Transformers with Layer Normalization Provably Learn the Power Method

论文配图:Looped Transformers with Layer Normalization Provably Learn the Power Method
图 1 · 摘自论文原文
  • 用梯度下降训练循环线性Transformer,每层自注意力完成一次幂迭代。
  • 有层归一化时模型收敛到幂法解,无则无法精确学习该算法。
  • 揭示了层归一化在算法学习中的关键作用,适合研究模型内在机制的读者。

Transformer在众多应用中表现卓越,其优势部分源于学习和执行算法的能力。然而,对带有层归一化(LN)的Transformer如何学习算法的理解仍有限。本文以主成分预测为测试任务,证明:通过梯度下降训练的循环线性Transformer(带LN)会收敛至实现幂法的解,每层自注意力执行一次幂迭代。值得注意的是,模型仅被监督进行主成分预测,而非显式学习幂法。这揭示了带LN的循环Transformer存在‘算法隐式偏差’——尽管主成分预测可通过多种机制实现,梯度下降却选择幂法路径。进一步对比显示,即使有逐层指导,无LN的Transformer无法精确学习幂法,而带LN版本可以,导致主成分预测性能出现可证明差距。本工作首次提供了对带LN的循环与单层Transformer训练动态的理论分析,阐明了层归一化在算法学习中的作用。

原文摘要 · Abstract (English)

Transformers have achieved remarkable success across a wide range of applications, and a growing body of work suggests that part of their strength comes from their ability to learn and execute algorithmic procedures. However, our understanding of how transformers learn such algorithms remains limited, especially in the presence of layer normalization (LN). In this work, we study principal component prediction as a concrete testbed for understanding the training dynamics of transformers with LN. We prove that a looped linear transformer with LN, trained by gradient descent, converges to a solution that implements the power method, with each self-attention layer performing one power iteration. Notably, the model is trained only for principal component prediction, rather than being explicitly supervised to implement the power method. Our finding thus reveals an "algorithmic implicit bias" of looped transformers with LN: principal-component prediction can in principle be achieved by many mechanisms, yet gradient descent selects one that realizes the power method. We further provide a concrete comparison between transformers with and without LN: even with layerwise guidance from power iterations, a transformer without LN cannot exactly learn the power method, whereas the corresponding transformer with LN can, leading to a provable performance gap in principal component prediction. Our results provide, to our knowledge, the first theoretical analysis of the training dynamics of looped and single-layer transformers with LN, and shed light on the role of LN in transformer models.

Transformer幂法层归一化算法学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。