让谈话模型「轻装上阵」。
一直以来，矩阵乘法（MatMul）稳居神经网络操纵的主导地位，其中很大原因归结为 GPU 专门针对 MatMul 操纵进行了优化。这种优化使得 AlexNet 在 ILSVRC2012 挑战赛中一举胜出，成为深度进修崛起的历史性标志。
在这当中，有个值得注意的点是，AlexNet 利用 GPU 来提高训练速度，超越了 CPU 的能力，至此，GPU 的加入使得深度进修仿佛赢得了「硬件彩票」。
尽管 MatMul 在深度进修中很流行，但不得不承认的是它占据了计算开销的主要部分，主要表现为 MatMul 在训练和推理阶段消耗大部分执行时间和内存访问。
目前为止，钻研者已经开始借助其他更简单的操纵替代 MatMul，主要有两种。
第一种策略是利用初等运算代替 MatMul，例如，在卷积神经网络 (CNN) 中，用有符号加法代替乘法；
第二种方法是利用二值或三值化量化，将 MatMul 值在累加之前要么翻转要么清零。比如脉冲神经网络 (SNN) 利用二值激活，而二值化网络 BNN 利用量化权重。
在谈话建模方面，BitNet 等技术的出现表明量化的可扩张性， 但这种方式仍然保留了昂贵的矩阵 – 矩阵相乘（MMM）的自注意力机制。钻研者尝试过多种努力，但 MatMul 操纵在 GPU 上仍然是资源麋集型的。
既然 MatMul 占据了 LLM 整体计算成本，且随着 LLM 向更大的嵌入维度和上下文长度扩张时，这种成本只会增加。这引发了一个问题：是否有可能完全从 LLM 中消除 MatMul 操纵？
在这项工作中，来自加州大学圣克鲁兹分校等机构的钻研者证明了 MatMul 操纵可以完全从 LLM 中消除，同时在十亿参数尺度下保持强大的功能。
论文地址：https://arxiv.org/pdf/2406.02528
项目地址：https://github.com/ridgerchu/matmulfreellm
论文标题：Scalable MatMul-free Language Modeling
实验表明，该钻研提出的 MatMul-free 模型达到了与最先进的 Transformer 相当的功能，后者在推理期间需要更多的内存，规模至少为 2.7B 参数。
此外，论文还钻研了扩张定律，发现随着模型规模的增加，MatMul-free 模型与全精度 Transformer 之间的功能差距逐渐缩小。
钻研者还提供了一种高效的 GPU 模型实现方式，在训练期间相比未优化的基线模型减少了多达 61% 的内存利用。通过在推理时利用优化的内核，模型内存消耗可以比未优化的模型减少超过 10 倍。
最后，钻研者在 FPGA 上构建了一个自定义硬件解决方案，他们以 13W 的功耗处理了十亿参数规模的模型，超出了人类可读的吞吐量，使 LLM 更接近大脑般的效率。
网友看后不禁感叹道：看来有大事要发生了。
不过，受到计算资源的限制，钻研者还没有在非常大的模型（如参数大于100B的模型）上测试MatMul-free模型的有效性，因此其实战效果还有待观察。
方法介绍
该钻研构建了首个可扩张的MatMul-free谈话模型 (Matmul-free LM)，通过在麋集层中利用加法操纵以及在自注意力类函数中利用元素级 Hadamard 乘积完成。
具体来说，三值权重消除了麋集层中的 MatMul，类似于 BNN。为了从自注意力中移除 MatMul，该钻研优化了门控循环单元 (GRU) ，使其仅依赖于元素级乘积，并表明该模型可与最先进的 Transformer 相媲美，同时消除了所有 MatMul 运算。
论文详细介绍 MatMul-free 谈话模型（LM）的各个组成部分。
具有三值权重的 MatMul-free 麋集层
在标准麋集层中，输入和权重矩阵之间的 MatMul 可以表示为：
为了避免利用基于 MatMul 的麋集层，该钻研采用 BitNet 来替换包含 MatMul 的麋集层，即利用 BitLinear 模块，这些模块利用三值权重将 MatMul 操纵转换为纯加法操纵。当利用三值权重时，权重矩阵 W 中的元素被限制在集合 {-1, 0, +1} 中。带有三值权重的 MatMul 可以表达为：
由于三值化权重
只能从 {−1, 0, +1} 中取值，因而 MatMul 中的乘法运算可以用简单的加法或减法运算代替：
因此，三值化 MatMul 可以写成如下：
算法 1 为量化融合 RMSNorm 和 BitLinear 算法的流程图：
MatMul-free 谈话模型架构
钻研职员采用了 Metaformer 的观点，该观点认为 Transformer 由两部分组成：token mixer（用于混合时间信息，即自注意力机制）和 channel mixer（用于混合嵌入 / 空间信息，即前馈网络，门控线性单元 GLU ）。该架构的高级概览如图 2 所示。
自注意力机制是现代谈话模型中最常用的 token mixer，它依赖于三个矩阵 Q、K 和 V 之间的矩阵乘法。为了将这些操纵转换为加法，钻研职员至少对两个矩阵进行二值化或三值化处理。
假设所有麋集层的权重都是三值的，他们将 Q 和 K 量化，得到一个三值的注意力图，从而消除自注意力中的乘法操纵。但是，以这种方式训练的模型无法成功收敛。一个可能的解释是，激活值包含对功能至关重要但难以有效量化的异常值。
为了解决这一挑战，钻研职员探索了不依赖于矩阵乘法的替代方法来混合 token。
通过采用结合了逐元素操纵和累积的三值 RNN，可以构建一个MatMul-free 的 token mixer。在各种 RNN 架构中，GRU 因其简单高效而著称，它在比长短期记忆网络（LSTM）利用更少的门控单元和结构更简单的情况下，实现了类似的功能。因此，钻研职员选择 GRU 作为构建 MatMul-free token mixer 的基础。
为了实现 MatMul-free 的通道混合，钻研职员采用了门控线性单元（GLU），它在许多现代 LLM 中得到了广泛应用，包括 Llama 、Mistral 和 RWKV。一个适应了 BitLinear 的 GLU 版本可以表达如下：
这里的通道混合器仅由麋集层组成，这些层已被三值化累积操纵所替代。通过在 BitLinear 模块中利用三值权重，钻研职员可以消除对昂贵 MatMul 的需求，这样不仅使通道混合器在计算上更加高效，同时还保持了其在跨通道混合信息方面的有效性。
实验
该钻研的重点是在中等规模的谈话建模任务上测试 MatMul-free 的谈话模型。钻研职员将两种变体的 MatMul-free 谈话模型与复现的高级 Transformer 架构（基于 Llama-2 的 Transformer++）在三个模型巨细上进行比较：370M、13 亿和 27 亿参数。
为了公平比较，所有模型都在 SlimPajama 数据集上进行了预训练，其中 370M 模型训练了 150 亿个 token，而 13 亿和 27 亿模型各训练了 1000 亿个 token。
模型训练利用了 8 个 NVIDIA H100 GPU。370M 模型的训练时间大约为 5 小时，13 亿模型为 84 小时，27 亿模型为 173 小时。
MatMul-free 谈话模型的扩张规律
钻研团队评估了MatMul-free 谈话模型和 Transformer++ 在 370M、13 亿和 27 亿参数模型上的扩张规律，如图 3 所示。
为便于比较，MatMul-free LM 和 Transformer++ 中的每个操纵都被同等对待。但请注意，Transformer++ 中的所有权重和激活都是 BF16 格式，而MatMul-free 谈话模型中的 BitLinear 层利用三元参数，激活为 BF16。因此，MatMul-free 谈话模型的平均运算成本要低于 Transformer++。
有意思的是，与 Transformer++ 相比，MatMul-free 谈话模型的扩张投影显示出更陡峭的下降趋势，这表明MatMul-free谈话模型在利用额外计算资源提高功能方面更为高效。
因此，MatMul-free 谈话模型的扩张曲线预计将在大约 10^23 次浮点运算（FLOPs）处与 Transformer++ 的扩张曲线相交。这个计算规模大致相当于训练 Llama-3 8B（利用 1.5 万亿个 token 训练）和 Llama-2 70B（利用 2 万亿个 token 训练）所需的 FLOPs，这表明MatMul-free 谈话模型不仅在效率上胜出，而且扩张时在损失方面也可能表现更好。
下游任务
进修率是谈话模型训练中一个关键的超参数，当模型处于三元 / 二元权重状态时，对进修率变得更加敏感。为了确定最优进修率，钻研职员利用 370M 模型，批量巨细为 50k 个 token，在 1.5e−3 到 3e−2 的范围内进行了搜索。这次搜索的结果如图 4 (c) 所示。
结果显示，当进修率从 1.5e−3 增加到 1e−2 时，最终训练损失单调递减。只有当进修率超过 2e−2 时，模型才表现出不稳定。这一发现表明，以前利用三元权重的作品，如利用 1.5e−3 进修率的 BitNet，可能不是最优的，更高的进修率有可能带来更好的功能。
这些发现与 Deepseek LLM 的观察结果一致，后者发现传统大型谈话模型（LLMs）的最佳进修率实际上比大多数 LLM 训练设置中通常报告的值要大。
有趣的是，钻研职员还观察到，与利用较小进修率训练的模型相比，训练初期利用较大进修率训练的模型，在训练后期的阶段训练损失下降得更快。
钻研职员根据训练时间和内存利用情况评估了他们提出的融合型 BitLinear 和传统型 BitLinear 实现，如图 4 (a-b) 所示。
实验表明，他们的融合操纵器在更大的批量巨细下，能够带来更快的训练速度，并减少内存消耗。当批量巨细为 2 的 8 次方时，1.3B 参数模型的训练速度从每次迭代 1.52 秒提高到 1.21 秒，比 Vanilla 实现快了 25.6%。
此外，内存消耗从 82GB 减少到 32GB，内存利用减少了 61.0%。随着批量巨细的增加，融合实现的功能显著提高，允许同时处理更多的样本，并减少总迭代次数。
图 4 (d) 展示了不同模型巨细下，所提出的MatMul-free谈话模型与 Transformer++ 在 GPU 推理内存消耗和延缓方面的比较。
在MatMul-free谈话模型中，钻研职员采用 BitBLAS 进行加速，以进一步提高效率。评估是在批量巨细为 1，序列长度为 2048 的情况下进行的。
MatMul-free谈话模型在所有模型巨细上都显示出比 Transformer++ 更低的内存利用和延缓。
对于单层，MatMul-free 谈话模型只需要 0.12 GB 的 GPU 内存，并且实现了 3.79 毫秒的延缓，而 Transformer++ 消耗了 0.21 GB 的内存，并且有 13.87 毫秒的延缓。随着模型巨细的增加，MatMul-free 谈话模型的内存和延缓优势变得更加明显。
值得注意的是，对于大于 2.7B 的模型巨细，结果是利用随机初始化的权重进行模拟的。对于最大的 13B 参数模型，MatMul-free 谈话模型仅利用 4.19 GB 的 GPU 内存，并且有 695.48 毫秒的延缓，而 Transformer++ 需要 48.50 GB 的内存，并表现出 3183.10 毫秒的延缓。
这些结果突出了MatMul-free谈话模型所实现的效率增益，使其成为大规模谈话建模任务的有希望的方法，特别是在推理期间。
为了测试MatMul-free谈话模型功耗和有效性，钻研职员利用 SystemVerilog 创建了一个 FPGA 加速器。概览如图 5 所示。
表 2 显示了单块（N = 1）单核实现的资源利用、功耗和功能。
通过利用完整的 512 位 DDR4 接口，并对占据核心处理时间 99% 的 TMATMUL 功能单元进行并行化，预计在保持相同时钟速率的情况下，无需额外优化或流水线化，可以实现大约 64 倍的进一步加速，如表 3 所示。
