Pre-trained Large Language Models Use Fourier Features to Compute Addition

📄 arXiv: 2406.03445v1 📥 PDF

作者: Tianyi Zhou, Deqing Fu, Vatsal Sharan, Robin Jia

分类: cs.LG, cs.CL

发布日期: 2024-06-05

💡 一句话要点

揭示预训练大语言模型使用傅里叶特征进行加法运算的机制

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 傅里叶特征 加法运算 预训练 可解释性

📋 核心要点

  1. 大型语言模型展现出强大的数学能力，但其内部算术运算机制尚不明确，特别是加法运算。
  2. 该研究发现LLM使用傅里叶特征表示数字，MLP层和注意力层分别利用低频和高频特征进行加法运算。
  3. 实验表明，预训练对于模型学习加法至关重要，从头训练的模型性能较差，引入预训练嵌入可以显著提升性能。

📝 摘要（中文）

预训练大语言模型（LLMs）展现出令人印象深刻的数学推理能力，但它们如何执行基本的算术运算（如加法）仍然不清楚。本文表明，预训练的LLMs使用傅里叶特征来加数字——隐藏状态中的维度通过一组在频域中稀疏的特征来表示数字。在模型内部，MLP和注意力层以互补的方式使用傅里叶特征：MLP层主要使用低频特征来近似答案的大小，而注意力层主要使用高频特征来执行模加法（例如，计算答案是偶数还是奇数）。预训练对于这种机制至关重要：从头开始训练以加数字的模型仅利用低频特征，导致较低的准确率。将预训练的token嵌入引入到随机初始化的模型中可以挽救其性能。总的来说，我们的分析表明，适当的预训练表示（例如，傅里叶特征）可以解锁Transformer学习算法任务精确机制的能力。

🔬 方法详解

问题定义：论文旨在探究预训练大语言模型（LLMs）如何执行基本的加法运算。现有方法缺乏对LLM内部算术运算机制的深入理解，无法解释其数学推理能力的来源。因此，理解LLM如何进行加法运算，有助于提升模型的可解释性和可控性。

核心思路：论文的核心思路是，LLM使用傅里叶特征来表示数字，并通过MLP层和注意力层协同完成加法运算。这种设计基于傅里叶分析的原理，将数字映射到频域，从而利用神经网络学习频域特征之间的关系。

技术框架：该研究主要通过分析预训练LLM的内部表示来揭示其加法运算机制。具体而言，研究人员首先识别出隐藏状态中代表数字的傅里叶特征。然后，他们分析MLP层和注意力层如何利用这些傅里叶特征来计算加法。最后，他们通过实验验证了预训练对于模型学习加法运算的重要性。

关键创新：该研究最重要的技术创新点在于发现了LLM使用傅里叶特征进行加法运算的机制。这与传统的神经网络学习方式不同，揭示了LLM内部的一种新的计算模式。此外，该研究还发现MLP层和注意力层在加法运算中扮演不同的角色，为理解LLM的内部工作原理提供了新的视角。

关键设计：论文的关键设计包括：1) 使用傅里叶变换将数字映射到频域；2) 分析隐藏状态中代表数字的傅里叶特征；3) 研究MLP层和注意力层如何利用这些特征进行加法运算；4) 通过从头训练模型和引入预训练嵌入来验证预训练的重要性。具体的参数设置和网络结构沿用了现有LLM的常用配置，重点在于分析和解释模型内部的计算过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

研究发现，预训练LLM使用傅里叶特征进行加法运算，其中MLP层利用低频特征近似答案大小，注意力层利用高频特征执行模加法。从头训练的模型性能显著低于预训练模型，而引入预训练嵌入可以有效提升性能，表明预训练对于学习精确的算法机制至关重要。

🎯 应用场景

该研究成果有助于提升大语言模型的可解释性和可控性，并为开发更高效、更可靠的数学推理模型提供理论指导。此外，该研究揭示的傅里叶特征表示方法，可能应用于其他算法任务，例如信号处理、图像识别等领域，具有潜在的实际应用价值。

📄 摘要（原文）

Pre-trained large language models (LLMs) exhibit impressive mathematical reasoning capabilities, yet how they compute basic arithmetic, such as addition, remains unclear. This paper shows that pre-trained LLMs add numbers using Fourier features -- dimensions in the hidden state that represent numbers via a set of features sparse in the frequency domain. Within the model, MLP and attention layers use Fourier features in complementary ways: MLP layers primarily approximate the magnitude of the answer using low-frequency features, while attention layers primarily perform modular addition (e.g., computing whether the answer is even or odd) using high-frequency features. Pre-training is crucial for this mechanism: models trained from scratch to add numbers only exploit low-frequency features, leading to lower accuracy. Introducing pre-trained token embeddings to a randomly initialized model rescues its performance. Overall, our analysis demonstrates that appropriate pre-trained representations (e.g., Fourier features) can unlock the ability of Transformers to learn precise mechanisms for algorithmic tasks.