0. TL; DR

本文提出了一种名为Laughing Hyena Distillery的方法，用于从长卷积序列模型（LCSMs）中提取紧凑的递归模型，以实现高效的自回归生成。该方法通过将卷积层转换为低维线性状态空间模型（SSMs），显著降低了模型在生成任务中的计算和内存成本，同时保持了与原始模型相当的性能。Laughing Hyena Distillery通过多头设计和模态插值技术，实现了高吞吐量和低延迟的自回归生成，为长序列建模提供了一种高效且实用的解决方案。

1. 背景介绍

近年来，基于卷积的序列模型（如Hyena和H3）在许多任务中取得了与Transformer相当的性能，但这些模型在自回归生成任务中面临着效率问题。传统的长卷积模型在生成每个新标记时需要处理整个输入序列，这导致了计算和内存成本的显著增加。为了解决这一问题，本文提出了一种新的方法——Laughing Hyena Distillery，它能够将长卷积模型转换为递归模式，从而实现$O(1)$的计算和内存成本。

2. Laughing Hyena Distillery 模型

2.1 卷积模型

卷积操作是深度学习中的一个基本操作，它通过将输入信号与一个卷积核进行逐点相乘并求和来生成输出。数学上，卷积操作可以表示为：

\[(h * u)_t = \sum_{j=0}^{t} h_{t-j} u_j\]

其中，$h$ 是卷积核，$u$ 是输入信号，$t$ 是时间步。在深度学习中，卷积操作通常用于处理序列数据，例如时间序列或文本序列。

2.2 状态空间模型

状态空间模型（SSM）是一种用于描述线性系统的方法，它通过状态变量来表示系统的内部状态，并通过状态转移方程和输出方程来描述系统的动态行为。状态空间模型的一般形式为：

\[\begin{aligned} x_{t+1} &= A x_t + B u_t, \\ y_t &= C x_t + D u_t, \end{aligned}\]

其中$x_t$ 是状态变量，表示系统在时间步 $t$ 的内部状态。$u_t$ 是输入信号。$y_t$ 是输出信号。$A$ 是状态转移矩阵，描述状态变量如何从一个时间步转移到下一个时间步。$B$ 是输入矩阵，描述输入信号如何影响状态变量。$C$ 是输出矩阵，描述状态变量如何生成输出信号。$D$ 是直接传递项，描述输入信号如何直接影响输出信号。

为了简化计算，本文采用了模态形式的SSM，其中状态转移矩阵 $A$ 是对角矩阵，输入矩阵 $B$ 和输出矩阵 $C$ 是向量。这种形式的SSM可以高效地进行递归计算，并且可以通过优化模态参数来最小化与原始卷积滤波器的误差。具体来说，模态形式的SSM可以表示为：

\[\begin{aligned} x_{t+1} &= \text{diag}(\lambda) x_t + \mathbf{1} u_t, \\ y_t &= \mathbf{R} x_t + h_0 u_t, \end{aligned}\]

其中，$\lambda$ 是状态转移矩阵 $A$ 的对角元素，$\mathbf{1}$ 是全1向量，$\mathbf{R}$ 是输出矩阵 $C$ 的向量形式，$h_0$ 是直接传递项。

2.3 Laughing Hyena Distillery

Laughing Hyena Distillery的核心思想是将长卷积模型中的每个卷积层转换为一个低维的线性状态空间模型。这一过程包括以下几个关键步骤：

分析Hankel算子的谱：通过分析预训练模型的Hankel算子的谱，确定目标状态维度。Hankel算子的谱可以提供关于卷积滤波器的有效维度的信息。具体来说，Hankel矩阵的奇异值可以用来估计卷积滤波器的低维近似所需的维度。
模态插值：使用模态插值方法将卷积滤波器近似为一个低维的SSM。模态插值方法通过最小化卷积滤波器与SSM之间的误差来优化SSM的参数。具体来说，模态插值方法将卷积滤波器表示为一个有理函数的形式：

\[H(z) = h_0+ \frac{b_1 z^{-1} + b_2 z^{-2} + \cdots + b_d z^{-d}}{1 + a_1 z^{-1} + \cdots + a_d z^{-d}}\]

其中，$a_i$ 和 $b_i$ 是SSM的参数。通过优化这些参数，可以最小化卷积滤波器与SSM之间的误差。

为了提高预训练质量和减少蒸馏过程中的滤波器数量，本文提出了多头Hyena模型。在多头设计中，将输入序列的通道分成多个头，每个头独立地进行长卷积操作。这种设计不仅提高了模型的性能，还减少了蒸馏过程中的计算量。具体来说，多头Hyena模型的结构如下：

输入投影：将输入序列 $u$ 投影到多个头，每个头的投影结果分别为 $q_m, k_m, v_m$。
外积操作：对每个头的投影结果进行外积操作，生成中间张量 $z_m$。
长卷积操作：对每个头的中间张量 $z_m$ 进行长卷积操作，生成输出 $y_m$。
输出平均：将所有头的输出进行平均，得到最终的输出 $y$。

多头设计的具体公式如下：

\[\begin{aligned} q_m, k_m, v_m &= \text{Projection}(u), \quad \text{for } m \in [M], \\ z_m &= k_m \otimes v_m, \\ y_m &= T_h(z_m) q_m, \\ y &= \frac{1}{M} \sum_{m=1}^M y_m, \end{aligned}\]

其中，$M$ 是头的数量，$T_h$ 是长卷积操作。

通过将卷积层转换为SSM，模型可以在自回归生成任务中以递归模式运行。在递归模式下，模型只需要维护一个固定维度的状态向量，从而实现了$O(1)$的计算和内存成本。通过这种方式，模型可以在生成每个新标记时只处理当前的状态向量，而不需要重新处理整个输入序列，从而显著提高了生成任务的效率。具体来说，递归模式的实现如下：

状态初始化：在生成任务开始时，使用输入序列 $u$ 初始化状态向量 $x_T$。
递归更新：在每个时间步 $t$，使用状态转移方程更新状态向量：$x_{t+1} = A x_t + B u_t$
输出生成：使用输出方程生成输出信号：$y_t = C x_t + D u_t$

3. 实验分析

3.1 预训练

为了验证多头设计的有效性，本文在The Pile数据集上对MultiHyena模型进行了预训练。实验结果表明，多头Hyena模型在预训练阶段的困惑度（Perplexity）低于标准Hyena模型，表明多头设计能够提高模型的性能。

3.2 蒸馏分析

本文通过分析Hankel算子的谱来确定目标状态维度，并使用模态插值方法将卷积滤波器近似为低维的SSM。

实验结果表明，通过模态插值方法蒸馏的模型在保持与原始模型相当的性能的同时，显著降低了计算和内存成本。例如，在1.3B参数的模型中，Laughing Hyena Distillery能够实现比Transformer高10倍的吞吐量，同时在生成512个标记时所需的内存减少了3倍。

3.3 下游任务评估

为了验证蒸馏模型在下游任务中的性能，本文在LM-Eval-Harness和HELM基准测试中对蒸馏后的MultiHyena模型进行了评估。实验结果表明，蒸馏后的模型在这些任务中的表现与原始模型相当，表明Laughing Hyena Distillery能够在不损失性能的前提下显著提高模型的效率。

3.4 性能基准测试

本文对Laughing Hyena Distillery在自回归生成任务中的性能进行了基准测试，包括吞吐量、延迟和内存使用。实验结果表明，Laughing Hyena Distillery在这些方面均优于Transformer和Hyena模型。

例如，在生成256个标记的任务中，Laughing Hyena Distillery的吞吐量是Transformer的10倍，且在处理长提示时的延迟更低。此外，Laughing Hyena Distillery在生成任意数量的标记时所需的内存量是固定的，而Transformer的内存量则随着生成标记数的增加而线性增长。