本文提出光谱字典生成模型（SDGM），通过学习全局傅里叶字典和 token 混合系数替换自注意力机制，实现 O(KL) 复杂度的高效语言建模，并在基准数据集上取得竞争性 perplexity 和显著的资源节省。
Spectral Modeling, Language Modeling, Efficiency, Interpretability, Generative Modeling, Transformer Alternative
Andrew Kiruluta
UC Berkeley
Generated by grok-3-mini-latest
Background Problem
变压器架构的自注意力机制 revolutionized 序列建模，但其 O(L²) 的计算和内存复杂度在处理长序列（如整个文档或长代码序列）时变得不可接受，导致效率低下。为解决这一问题，本文提出了一种基于光谱字典学习的替代方法，旨在减少计算复杂度、保持竞争性性能并提升模型的可解释性。背景包括现有注意力近似方法（如稀疏注意力、内核方法和低秩投影），以及光谱方法的初步探索（如 FNet），但这些方法要么适应性有限，要么未充分优化语言建模任务。
Method
- 核心思想： 本文提出 Spectral Dictionary Generative Model (SDGM)，通过学习一个全局时变傅里叶字典和每个 token 的混合系数来替换自注意力机制，实现高效的序列混合。
- 如何实现：
- 字典参数化： 学习 K 个光谱原子，每个原子由幅度矩阵 A、频率矩阵 F 和相位矩阵 Φ 参数化。对于每个原子 k 和特征维度 d，原子值定义为 S_k(t)_d = a_{k,d} \sin(2π f_{k,d} rac{t}{L} + \phi_{k,d})，其中 t 是归一化时间索引。
- 混合系数编码： 使用一维卷积 (Conv1D) 编码器处理输入嵌入序列 X，产生每个 token 的混合系数 C，例如 ，其中 σ_act 是激活函数。
- 重建解码器： 通过加权求和重建嵌入，公式为 ，复杂度为 O(KL)。
- 训练目标： 最小化复合损失函数 ，其中包括时域均方误差、频域 STFT 幅度损失、负对数似然损失和 GMM 先验损失。
- 生成过程： 训练后拟合 GMM 于混合系数，生成时从 GMM 采样系数 z，然后解码到嵌入并预测下一个 token。
- 关键步骤： 端到端训练，参数包括嵌入表、字典参数和卷积权重，生成时采用自回归采样。
Experiment
- 数据集和基线： 使用 WikiText-2 和 Penn Treebank 数据集，前者约 2M 训练 token，后者约 1M 训练 token。基线包括 Transformer-XL、GPT-2 Small 和 Linformer，所有在相同数据划分和分词方案下重新训练。
- 实验设置： 嵌入维度 D=512，字典大小 K=256，序列长度 L=128。使用 STFT 参数 (nfft=256, hop length=64, Hann 窗口)。优化器为 Adam，学习率 10^{-3}，权重衰减 10^{-5}，批量大小 32，训练最多 10 个 epoch，以验证集 perplexity 早停。评价指标包括 perplexity (PPL，下越好)、推理速度 (tok/s，上越好)、参数量 (百万)、内存占用 (GB，下越好) 和嵌入重建余弦相似度。
- 结果： SDGM 在 WikiText-2 上验证集 PPL 为 31.2，在 PTB 上为 57.1，与基线竞争性（Transformer-XL PPL 32.1 和 58.7），但参数量少 22.8M (基线 41.2-117M)，内存占用低 6.5GB (基线 8.7-12.5GB)，推理速度高 2100 tok/s (基线 1200-1800 tok/s)。消融实验显示，移除频域损失 (β=0) 使 PPL 上升到 33.5，移除语言模型损失 (γ=0) 使 PPL 上升到 35.0，确认损失组件重要性。重建余弦相似度为 0.92，优于无频域损失的 0.88。结果符合预期，证明了光谱监督和语言建模损失的协同作用。
Further Thoughts
这个方法可能启发其他领域，如时间序列预测或音频信号处理，因为光谱方法在这些领域已有应用；此外，学习的可解释性傅里叶原子可能有助于神经科学中理解脑波模式，或在 AI 安全中提升模型透明度；未来可以探索与注意力机制的混合架构，或将字典大小动态调整以优化容量-效率权衡，类似于稀疏编码在视觉任务中的扩展。