Fourier Head: Helping Large Language Models Learn Complex Probability Distributions
作者: Nate Gillman, Daksh Aggarwal, Michael Freeman, Saurabh Singh, Chen Sun
分类: cs.LG, cs.AI, cs.CL, stat.ML
发布日期: 2024-10-29 (更新: 2025-03-10)
备注: Camera ready version (ICLR 2025). Code at https://nategillman.com/fourier-head
💡 一句话要点
提出傅里叶头，增强LLM建模复杂概率分布的能力，提升非语言token序列建模效果。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 傅里叶级数 大型语言模型 序列建模 决策Transformer 时间序列预测 连续动作空间 非语言Token
📋 核心要点
- 现有LLM在建模非语言token时，离散softmax可能无法有效捕获连续结构和复杂分布。
- 提出傅里叶头，利用傅里叶级数构建神经网络层，替代线性层，使输出具有更连续的结构。
- 实验表明，傅里叶头在决策制定和时间序列预测任务中显著提升了性能，最高提升达377%。
📝 摘要（中文）
随着大型语言模型质量的提高，人们对使用它们来建模非语言token的兴趣日益增加。例如，Decision Transformer将agent决策制定重新定义为一个序列建模问题，使用仅解码器的LLM来建模Atari agent离散动作空间上的分布。然而，当将LLM应用于非语言领域时，softmax在离散bins上的应用是否能捕获token的连续结构以及高质量token生成所需的潜在复杂分布，仍然不清楚。我们引入了一个使用傅里叶级数构建的神经网络层，如果希望输出具有更连续的结构，可以很容易地将其替换为任何线性层。我们对合成数据集以及大规模决策制定和时间序列预测任务进行了广泛的分析。我们还提供了理论证据，表明该层可以更好地从数据中学习信号，同时忽略高频噪声。所有结果都支持我们提出的傅里叶头在底层数据分布具有自然连续结构的情况下是有效的。例如，傅里叶头将Decision Transformer agent在四个基准Atari游戏中的回报提高了高达377%，并将最先进的时间序列基础模型在训练期间未见过的20个基准测试中的预测性能提高了3.5%。
🔬 方法详解
问题定义：现有的大型语言模型在处理非语言token时，通常采用softmax函数对离散的动作空间进行建模。这种方法的局限性在于，它无法有效地捕捉到非语言token中潜在的连续结构以及复杂的数据分布。尤其是在决策制定和时间序列预测等任务中，这种离散化的处理方式可能会导致性能瓶颈。
核心思路：论文的核心思路是利用傅里叶级数来构建一个新的神经网络层，称为“傅里叶头”。傅里叶级数能够将任意周期函数分解为一系列正弦和余弦函数的叠加，因此可以更好地表示连续的信号。通过将传统的线性层替换为傅里叶头，模型能够更好地学习和生成具有连续结构的token。
技术框架：该方法的核心在于将LLM中的线性层替换为傅里叶头。整体框架保持了LLM的结构不变，只是在输出层进行了修改。具体来说，傅里叶头接收来自LLM隐藏层的输出，然后通过傅里叶级数进行变换，最终输出具有连续结构的token。这种替换可以很容易地集成到现有的LLM架构中。
关键创新：最关键的创新点在于利用傅里叶级数来建模连续的概率分布。与传统的softmax函数相比，傅里叶头能够更好地捕捉到数据中的连续结构，从而提高模型的性能。此外，论文还提供了理论分析，证明傅里叶头能够更好地从数据中学习信号，同时抑制高频噪声。
关键设计：傅里叶头的关键设计包括傅里叶级数的阶数、频率的选择以及系数的初始化。论文中可能讨论了如何选择合适的傅里叶级数阶数，以平衡模型的复杂度和表达能力。此外，损失函数的设计也至关重要，可能采用了某种形式的距离度量来衡量生成token与真实token之间的差异。
🖼️ 关键图片
📊 实验亮点
实验结果表明，傅里叶头在多个任务上取得了显著的性能提升。在Atari游戏中，使用傅里叶头的Decision Transformer agent的回报提高了高达377%。在时间序列预测任务中，傅里叶头将最先进的时间序列基础模型的预测性能提高了3.5%，这些基准测试在训练期间是未见过的，表明了傅里叶头具有良好的泛化能力。
🎯 应用场景
该研究成果可广泛应用于需要对连续动作空间或复杂概率分布进行建模的领域，例如强化学习中的连续控制、时间序列预测、机器人控制、以及其他涉及非语言token序列建模的任务。傅里叶头有望提升模型在这些领域的性能和泛化能力，推动相关技术的发展。
📄 摘要（原文）
As the quality of large language models has improved, there has been increased interest in using them to model non-linguistic tokens. For example, the Decision Transformer recasts agentic decision making as a sequence modeling problem, using a decoder-only LLM to model the distribution over the discrete action space for an Atari agent. However, when adapting LLMs to non-linguistic domains, it remains unclear if softmax over discrete bins captures the continuous structure of the tokens and the potentially complex distributions needed for high quality token generation. We introduce a neural network layer, constructed using Fourier series, which we can easily substitute for any linear layer if we want the outputs to have a more continuous structure. We perform extensive analysis on synthetic datasets, as well as on large-scale decision making and time series forecasting tasks. We also provide theoretical evidence that this layer can better learn signal from data while ignoring high-frequency noise. All of our results support the effectiveness of our proposed Fourier head in scenarios where the underlying data distribution has a natural continuous structure. For example, the Fourier head improves a Decision Transformer agent's returns across four benchmark Atari games by as much as 377%, and increases a state-of-the-art times series foundation model's forecasting performance by 3.5% across 20 benchmarks unseen during training.