arXiv:2506.23679cs.LGcs.AI2025-06中稿 · NeurIPS

Transformer学会模幂运算,揭示其内部数学推理机制。

Learning Modular Exponentiation with Transformers

  • 通过反向输入训练,模型突然掌握跨模数的模幂计算。
  • 最终层部分注意力头即可实现完整计算性能。
  • 发现模型内存在专用算术电路,利于解释与优化。

模幂运算是数论与密码学的核心,但其机制可解释性仍不明确。我们训练了一个4层编码器-解码器Transformer模型完成该任务,并研究训练过程中数值推理的涌现。通过合理的采样策略、基于PCA的嵌入分析和激活修补,我们考察了数论性质在模型中的编码方式。发现反向操作符训练带来显著性能提升,且在相关模数间出现突现泛化。这种同步精度跃升反映了类似grokking的动力学特征,表明模型内化了共享的算术结构。此外,我们发现仅由最后一层全部注意力头构成的子图即可实现完整模幂运算性能。结果表明,Transformer通过专用计算回路学习模算术,为更可解释、高效的神经模幂方法铺平道路。

原文摘要 · Abstract (English)

Modular exponentiation is crucial to number theory and cryptography, yet remains largely unexplored from a mechanistic interpretability standpoint. We train a 4-layer encoder-decoder Transformer model to perform this operation and investigate the emergence of numerical reasoning during training. Utilizing principled sampling strategies, PCA-based embedding analysis, and activation patching, we examine how number-theoretic properties are encoded within the model. We find that reciprocal operand training leads to strong performance gains, with sudden generalization across related moduli. These synchronized accuracy surges reflect grokking-like dynamics, suggesting the model internalizes shared arithmetic structure. We also find a subgraph consisting entirely of attention heads in the final layer sufficient to achieve full performance on the task of regular exponentiation. These results suggest that transformer models learn modular arithmetic through specialized computational circuits, paving the way for more interpretable and efficient neural approaches to modular exponentiation.

Transformer模幂可解释性算术

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