Hyena:面向大型卷积语言模型.
0. TL; DR
本文提出了 Hyena 架构,一种用于大规模语言建模的新型无注意力机制(attention-free)的卷积语言模型。Hyena 通过交错使用隐式参数化的长卷积和数据控制的门控机制,实现了与注意力机制相当的性能,同时显著降低了计算复杂度。在标准数据集(如 WikiText103 和 The Pile)上,Hyena 在不依赖注意力机制的情况下达到了与 Transformer 相当的困惑度,并且在训练计算量上减少了 20%。此外,Hyena 在长序列处理上表现出色,其效率比优化后的注意力机制高出数倍,为处理更长上下文提供了新的可能性。
1. 背景介绍
Transformer 架构因其强大的上下文建模能力和可扩展性,在自然语言处理(NLP)领域取得了巨大成功。然而,Transformer 的核心组件——注意力机制(self-attention)——在处理长序列时面临着计算复杂度呈二次方增长的问题,这限制了模型能够处理的上下文长度。
尽管已有研究通过低秩、稀疏等近似方法来降低注意力机制的计算成本,但这些方法通常需要与密集注意力层结合才能达到 Transformer 的性能水平,这表明在大规模语言建模中,仍存在性能与效率之间的差距。
为了解决这一问题,本文提出了 Hyena 架构,旨在通过高效的子二次复杂度操作替代传统的注意力机制,同时保持或提升模型的性能。作者指出自注意力的成功基于以下几个关键特性:
- 数据控制(Data Control):注意力机制通过数据控制实现了高效的线性变换,能够根据输入动态调整其行为。
- 参数解耦(Parameter Decoupling):注意力机制的参数数量与序列长度无关,使得模型可以在其他部分(如前馈网络)分配更多参数。
- 无限制上下文(Unrestricted Context):注意力机制能够处理序列中任意两个元素之间的依赖关系,不受局部性限制。
2. 方法介绍
(1)Hyena的整体结构
Hyena 架构由以下步骤组成:
- 计算输入$x_t$的$N+1$个线性投影$v_t,x_t^1,…,x_t^N$。其中$v_t$作为值向量($z_t^1=v_t$),$x_t^n$作为生成注意力矩阵的向量。
- 引入一系列可学习的滤波器$h^1,…,h^N$,递归地使用隐式参数化的长卷积(long convolution)和数据控制的门控机制(gating mechanism)来构建高效的子二次复杂度操作:
- 线性输入输出的结果可以表示为$y_t=z_t^{N+1}$。
(2)Hyena中的卷积运算
卷积运算$h^n * z_t^n$可以使用快速傅里叶变换(FFT)的卷积定理来快速计算:首先将输入序列的 FFT 相乘,然后应用逆 FFT 来有效地计算卷积的输出。要将此定理用于上述非循环卷积,需要用零填充等长的输入序列和卷积核,然后丢弃输出序列的填充部分。
\[\text{pad}(h^n * z_t^n) = \text{iFFT}(\text{FFT}(\text{pad}(h^n)) \cdot \text{FFT}(\text{pad}(z_t^n)))\]傅里叶域中执行卷积的复杂度为$O(L \log L)$,因此Hyena循环的总复杂度为$O(NL \log L)$。
Hyena操作的每次循环是由一次时域卷积($h^n * z_t^n$,相当于频域乘积)和一次时域乘积($x_t^n \cdot (h^n * z_t^n)$,相当于频域卷积)组成的,相当于交替地在时域和频域中应用卷积(或者在时域和频域中应用逐元素乘积)。时域中的卷积增加了上下文长度;时域中的逐元素乘法允许对信号的特定频率分量进行更细粒度的选择。
(3)Hyena定义的卷积滤波器
Hyena 使用隐式参数化来定义卷积滤波器,避免了显式存储长卷积滤波器带来的参数数量问题。具体来说,卷积滤波器 $h^n$ 可以通过一个前馈神经网络(FFN)来参数化:
\[h_t = \text{Window}(t) \cdot (\text{FFN}\odot \text{PositionalEncoding})(t)\]其中,$\text{Window}(t)$ 是一个衰减窗口函数,用于调制卷积滤波器;$\text{PositionalEncoding}(t)$ 是位置编码,用于引入序列中的位置信息;$\text{FFN}$ 是一个浅层前馈神经网络,用于生成卷积滤波器的值。
3. 实验分析
在关联记忆的合成任务中,Hyena 表现出色,尤其是在长序列和大词汇表的情况下。例如,在序列长度为 131k、词汇表大小为 30 的设置下,Hyena 的准确率达到了 97.2%,而其他无注意力机制模型(如 H3、AFT-conv 等)的准确率远低于此。这表明 Hyena 在处理长序列和复杂任务时具有显著优势。
在语言建模任务中,Hyena 在 WikiText103 和 The Pile 数据集上达到了与 Transformer 相当的困惑度。具体来说,在 The Pile 数据集上,Hyena 在 150 亿 tokens 的训练后,困惑度达到了 9.2,与标准 Transformer 模型相当,但训练计算量减少了 20%。这表明 Hyena 在保持性能的同时,显著降低了计算成本。
在 SuperGLUE 数据集上,Hyena 的零样本和少样本性能与标准 Transformer 模型相当。例如,在 WSC 任务中,Hyena 的零样本准确率为 21.2%,少样本准确率为 39.4%,与 GPTNeo 模型相当。这表明 Hyena 在处理下游任务时具有良好的泛化能力。
在效率评估中,Hyena 在长序列处理上表现出色。例如,在序列长度为 8k 时,Hyena 的运行时间比优化后的注意力机制(如 FlashAttention)快 2 倍;在序列长度为 64k 时,Hyena 的运行时间比 FlashAttention 快 100 倍。这表明 Hyena 在处理长序列时具有显著的效率优势。