arXiv:2502.07553cs.LG2025-02被引 2

Transformer仅用极少量参数就能快速学会稀疏XOR,突破传统网络瓶颈。

Transformers Provably Learn Sparse XOR with Polylogarithmic Parameters

  • 用单层双头Transformer学习稀疏XOR,仅需polylog(d)参数
  • 一步梯度下降后所有输入损失趋近于0,发现关键特征
  • 依赖精确softmax注意力,适合研究模型特征提取机制

学习稀疏奇偶函数已成为检验神经网络特征学习能力的理论基准。现有分析多集中于前馈神经网络(FFNN),而对具备结构优势的Transformer在该任务中的理论理解仍有限。本文分析单层双头Transformer如何学习稀疏XOR问题:输入样本$(\mathbf{x}, y) \in \{\pm 1\}^d \times \{\pm 1\}$,标签定义为$y = -x_{i^*} x_{j^*}$,其中$i^*, j^* \in [d]$未知。我们证明,仅需$O(\mathrm{polylog}(d))$可训练参数,Transformer即可在一步梯度下降中成功识别相关特征,并使所有输入的损失趋近于0。该结果表明,Transformer突破了FFNN在该问题上的$Ω(d)$参数下限。此外,实验显示这种快速特征发现完全依赖于精确softmax注意力,优于线性或逐分量注意力。最后,我们给出了有限数据下的理论样本复杂度界,验证了Transformer在此任务中的泛化能力。

原文摘要 · Abstract (English)

Learning sparse parity functions has become a theoretical testbed for studying feature learning in neural networks. However, existing analyses primarily focus on Feed-Forward Neural Networks (FFNNs). Meanwhile, theoretical understanding of Transformers in this setting remains limited, despite their empirical success and structural suitability for discovering sparse support over long sequences. To address this gap, we analyze how a single-layer, two-head Transformer learns the sparse XOR problem. Considering samples $(\mathbf{x}, y) \in \lbrace\pm 1\rbrace^d \times \lbrace\pm 1\rbrace$, where the label is defined by $y = -x_{i^*} x_{j^*}$ for some unknown $i^*, j^* \in [d]$, we prove that, with only $O(\mathrm{polylog}(d))$ trainable parameters, Transformers can successfully discover the relevant features and drive the loss for every input to nearly 0 with one gradient step. This result establishes that Transformers break the fundamental $Ω(d)$ parameter bottleneck inherent to FFNNs for this problem. Furthermore, we empirically show that this rapid feature discovery is uniquely driven by the exact softmax attention, outperforming common substitutes such as linear or component-wise attention. Finally, we provide a theoretical sample complexity bound for learning from finite data, demonstrating the generalization ability of Transformers in this task.

Transformer稀疏学习理论分析

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