本文目录

1 DiJiang：频域角度降低 LLM 的成本
(来自北京华为诺亚方舟实验室)
1.1 背景：大语言模型需要极简注意力机制
1.2 频域的核化注意力机制
1.3 不同尺寸的结果测评
1.4 不同模型的结果测评
1.5 与线性 Transformer 的对比
1.6 推理时间对比
1.7 可视化

太长不看版

本文着眼于大语言模型的训练和使用代价，提出一种从频域角度降低 LLM 的成本的 [帝江] 大语言模型。减少 Transformer 架构使用成本的一种常见方法是基于线性注意力机制 (Linear Attention)，但是构建带有线性注意力机制的 Transformer 就需要重新训练整个模型，花费的计算代价太高了，对于巨量参数的 LLM 显然不切实际。

因此，本文提出 频域核化 (Frequency Domain Kernelization) 方法 ：使用离散余弦变换 (DCT) 有效且精准地将 Transformer 的 Query 和 Key 映射到频域。这种映射能够有效地消除 Self-Attention 机制中的 Softmax 操作，使注意力计算复杂度转化为线性。而且，作者从理论上证明，这种频域映射是与原始注意力机制等效近似，允许预训练的原始 Transformer 模型在训练成本很小的情况下，转换为线性复杂度模型。

本文提出的加权准蒙特卡罗方法提供了优越的逼近效率。为了进一步降低计算复杂度，核化方法基于离散余弦变换 (Discrete Cosine Transform, DCT) 操作。本文所提出的方法实现了与原始 Transformer 相当的性能，但大大降低了训练成本 (约 1/10) 和更快的推理速度 (最快约 10 倍) 。

1 DiJiang：频域角度降低 LLM 的成本

论文名称：DiJiang: Efficient Large Language Models through Compact Kernelization

论文地址： http://arxiv.org/pdf/2403.19928.pdf

代码地址： http://github.com/YuchuanTian/DiJiang

1.1 背景：大语言模型需要极简注意力机制

Transformer 彻底改变了自然语言处理的领域，也带来了大语言模型 (LLM) 所主导的时代。LLM 可以处理很复杂的任务，但同时也带来了大量的计算需求：显著的推理成本和能耗，使得在手机和机器人这类端侧设备的部署显著受阻。

在大量的模型压缩策略中，简化注意力机制 (simplifying the attention mechanism) 是一种极具前景的方法。比如 Linear Transformer[1]，Performer[2]。还有很多经典的改进注意力机制复杂度的技术路线，比如：

• RWKV：RWKV: Reinventing RNNs for the Transformer Era
• RetNet：Retentive Network: A Successor to Transformer for Large Language Models
• Mamba：Mamba: Linear-Time Sequence Modeling with Selective State Spaces

但是，大多数现有的优化 Transformer 的方法，通常需要对模型架构进行重大修改，且通常需要从头训练整个的模型以实现最佳性能。这样的重新训练过程对于 LLM 这种参数量巨大的模型而言，的确是个不小的挑战。比如，训练一个 LLaMA-7B[3] 量级的模型需要 82,432 GPU-hours，总功耗约为 36 MWh。对于这种量级的模型，再训练不仅会带来比较可观的经济问题，还会引发不小的环境问题。那么就需要更有效的方法来适应和优化这些大模型。简化的注意力机制建模方法在大语言模型上面还没有得到很完善的验证。

1.2 频域的核化注意力机制

作者首先重新审视了 Self-attention 的一般形式，如图1所示。为了简化符号的表征，只考虑 single head，并省略了 Layer Norm。Self-attention 的一般形式可以写成下式：

式中, 代表 tokens 的数量, 代表 Embedding dimension。当把 表示为 表示为 表示为 表示为 时, 上式1可以重写成:

可以看出, Transformer 模型中计算每个输出的计算和内存复杂度为 代表 tokens 的数量, 代表 Embedding dimension。因此, 处理长度为 的句子的时间和内存复杂度呈二次方, 变为 。这种二次方的计算复杂度带来了比较大的计算负担, 尤其是对于 比较大的长序列, 使得计算变得密集。

为了减轻这种复杂性, 一些工作 已经引入了 Kernel 的概念来降低注意力机制的计算需求。具体而言, 引入核函数 , 是正定核, 能够测量相似性。上式2可以重写成:

通过应用核技巧，可以线性分解注意力机制，上式3可以重写成：

式中, 为一个映射函数, 将 维的向量映射为一个 维的向量。这个分解可以借助 Key 和 Value 可以先合并计算的优势, 将计算复杂度从 降低为 。因为维度 通常都是远远小于 的, 因此, 注意力机制的这种线性化导致计算量大幅下降。

在大语言模型中，重新训练的成本非常高。在这种情况下，必须找到一个可以等价地替换香草注意机制的核，而不需要大量的重训练：

Positive Random Features (PRF)[2]是一个不错的候选：

式中, , 理论证明表明: 。这意味着当特征空间的维度 足够大时, PRF 映射可以作为原始注意力机制的等效。

这种等价性表明: 在理论上, 使用 PRF 映射将现有的注意力机制转化为线性注意力机制是可行的, 从而在不影响功能的前提下实现加速。但是, 这样做的一个问题是需要把 设定为一个很大的值, 才能够 减少近似带来的误差 。这一要求导致计算需求也会增加不少。比如, 在 Performer 中, 为了实现无损线性注意力, 通常需要设置为大于 , 从而减少了线性注意力降低计算复杂度的好处。

为了解决这个问题，作者首先对基于核的方法进行理论分析，以逼近注意力机制：

Theorem 1 (Bochner's Theorem )： 连续移位不变缩放核函数 是正定的,当且仅当它是 上唯一有限概率测度 的傅里叶变换。

其中, 代表 的复数共轭。

根据这个理论, 核函数 与概率密度 存在一一映射关系。取 , 则:

其中, 是蒙特卡洛方法的采样结果。 是显式的有限维特征图, 它取决于核 。作者使用准蒙特卡罗方法 。这个转变允许使用特定的均匀分布来估计积分, 而不是随机采样的分布。

利用 Bochner 定理允许转变 Transformer 模型中注意力机制。对于高斯核而言：

由于注意力机制中的 和 通常是归一化的, 因此高斯核可以看作是 , 这与 Query 和 Key 之间的计算相同。

Theorem 2： Positive Fixed Features (PFF) 可以表示为：

式中, 是渐近均匀分布, 。 是高斯核 的无偏估计。通过这种转换, 作者证明了基于 任何渐近均匀分布的近似 可以密切逼近 原始高斯核 。

Theorem 3： Weighted Positive Fixed Features (WPFF) 可以表示为：

其中, 是一个可学习的参数, 可以通过输入 进行优化。然后 WPFF (加权正固定特征) 方法目标函数积分估计误差的上界不大于 PFF (正固定特征) 方法目标函数积分估计误差的上界。

为了进一步加快训练速度，作者使用频域变换来减少所需的计算资源。快速傅里叶变换 (Fast Fourier Transform, FFT) 和离散余弦变换 (Discrete Cosine Transform, DCT) 是一种常用的这种转换方法。与普通正交变换相比，频域变换具有快速计算的算法，显着降低了映射的计算成本，可以把 的计算复杂度降低为 。而且，由于 DCT 只在实数域操作，需要更少的计算资源，更硬件友好。因此，作者选择 DCT 来执行核映射。

频域中的 DCT 系数 定义为:

式中, 当 时, 。其他情况 。使用了 DCT 之后的加权映射可以写成:

式中,