机器之心原创

2024/06/11 00:42

从LLM中完全消除矩阵乘法，效果出奇得好，10亿参数跑在FPGA上接近大脑功耗

让语言模型「轻装上阵」。

一直以来，矩阵乘法（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 所示。

工程MatMul-freeLLMMatMul

相关数据

深度学习技术

深度学习（deep learning）是机器学习的分支，是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。深度学习是机器学习中一种基于对数据进行表征学习的算法，至今已有数种深度学习框架，如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

来源：LeCun, Y., Bengio, Y., & Hinton, G. (2015). Deep learning. nature, 521(7553), 436.

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

二值化技术

二值化是将像素图像转换为二进制图像的过程。

来源：Felix Niklas

迭代技术

模型的权重在训练期间的一次更新。迭代包含计算参数在单个批量数据上的梯度损失。

来源：Google ML glossary

自注意力技术

自注意力（Self-attention），有时也称为内部注意力，它是一种涉及单序列不同位置的注意力机制，并能计算序列的表征。自注意力在多种任务中都有非常成功的应用，例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系，这种内部关系对于翻译以及序列任务的性能非常重要。

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

脉冲神经网络技术

第三代神经网络，脉冲神经网络（Spiking Neural Network，SNN），旨在弥合神经科学和机器学习之间的差距，使用最拟合生物神经元机制的模型来进行计算。脉冲神经网络与目前流行的神经网络和机器学习方法有着根本上的不同。SNN 使用脉冲——这是一种发生在时间点上的离散事件——而非常见的连续值。每个峰值由代表生物过程的微分方程表示出来，其中最重要的是神经元的膜电位。本质上，一旦神经元达到了某一电位，脉冲就会出现，随后达到电位的神经元会被重置。对此，最常见的模型是 Integrate-And-Fire（LIF）模型。此外，SNN 通常是稀疏连接的，并会利用特殊的网络拓扑。

来源：机器之心

收敛技术

在数学，计算机科学和逻辑学中，收敛指的是不同的变换序列在有限的时间内达到一个结论（变换终止），并且得出的结论是独立于达到它的路径（他们是融合的）。通俗来说，收敛通常是指在训练期间达到的一种状态，即经过一定次数的迭代之后，训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说，如果采用当前数据进行额外的训练将无法改进模型，模型即达到收敛状态。在深度学习中，损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变，暂时形成收敛的假象。

来源：Wikipedia Google ML glossary

学习率技术

在使用不同优化器（例如随机梯度下降，Adam）神经网络相关训练中，学习速率作为一个超参数控制了权重更新的幅度，以及训练的速度和精度。学习速率太大容易导致目标（代价）函数波动较大从而难以找到最优，而弱学习速率设置太小，则会导致收敛过慢耗时太长

来源：Liu, T. Y. (2009). Learning to rank for information retrieval. Foundations and Trends® in Information Retrieval, 3(3), 225-331. Wikipedia

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集（或特征）的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的，这让「神经网络在执行预测任务时可以更多关注输入中的相关部分，更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时，源句子中仅有少部分是相关的；因此，可以应用一个基于内容的注意力机制来根据源句子动态地生成一个（加权的）语境向量（context vector）, 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

来源：机器之心

神经网络技术

（人工）神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型，那时候研究者构想了「感知器（perceptron）」的想法。这一领域的研究者通常被称为「联结主义者（Connectionist）」，因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型，它们都是前馈神经网络：卷积神经网络（CNN）和循环神经网络（RNN），其中 RNN 又包含长短期记忆（LSTM）、门控循环单元（GRU）等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习，但也有一些为无监督学习设计的变体，比如自动编码器和生成对抗网络（GAN）。

来源：机器之心

门控循环单元技术

门控循环单元（GRU）是循环神经网络（RNN）中的一种门控机制，与其他门控机制相似，其旨在解决标准RNN中的梯度消失／爆炸问题并同时保留序列的长期信息。GRU在许多诸如语音识别的序列任务上与LSTM同样出色，不过它的参数比LSTM少，仅包含一个重置门（reset gate）和一个更新门（update gate）。

来源：Understanding GRU networks, Medium

卷积神经网络技术

卷积神经网路（Convolutional Neural Network, CNN）是一种前馈神经网络，它的人工神经元可以响应一部分覆盖范围内的周围单元，对于大型图像处理有出色表现。卷积神经网路由一个或多个卷积层和顶端的全连通层（对应经典的神经网路）组成，同时也包括关联权重和池化层（pooling layer）。这一结构使得卷积神经网路能够利用输入数据的二维结构。与其他深度学习结构相比，卷积神经网路在图像和语音识别方面能够给出更好的结果。这一模型也可以使用反向传播算法进行训练。相比较其他深度、前馈神经网路，卷积神经网路需要考量的参数更少，使之成为一种颇具吸引力的深度学习结构。卷积网络是一种专门用于处理具有已知的、网格状拓扑的数据的神经网络。例如时间序列数据，它可以被认为是以一定时间间隔采样的一维网格，又如图像数据，其可以被认为是二维像素网格。

来源：Goodfellow, I.; Bengio Y.; Courville A. (2016). Deep Learning. MIT Press.维基百科

参数模型技术

在统计学中，参数模型是可以使用有限数量的参数来描述的分布类型。这些参数通常被收集在一起以形成单个k维参数矢量θ=（θ1，θ2，...，θk）。

来源：Wikipedia

长短期记忆网络技术

长短期记忆(Long Short-Term Memory) 是具有长期记忆能力的一种时间递归神经网络(Recurrent Neural Network)。其网络结构含有一个或多个具有可遗忘和记忆功能的单元组成。它在1997年被提出用于解决传统RNN(Recurrent Neural Network) 的随时间反向传播中权重消失的问题（vanishing gradient problem over backpropagation-through-time)，重要组成部分包括Forget Gate, Input Gate, 和 Output Gate, 分别负责决定当前输入是否被采纳，是否被长期记忆以及决定在记忆中的输入是否在当前被输出。Gated Recurrent Unit 是 LSTM 众多版本中典型的一个。因为它具有记忆性的功能，LSTM经常被用在具有时间序列特性的数据和场景中。

来源：Hochreiter, Sepp & Schmidhuber, Jürgen. (1997). Long Short-term Memory. Neural computation. 9. 1735-80. 10.1162/neco.1997.9.8.1735.

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。由于字词与句子都是任意组合的长度，因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题)，也使得在语料库中估算字串的机率变得很困难，这也是要使用近似的平滑n元语法(N-gram)模型之原因。

来源：维基百科

量化技术

深度学习中的量化是指，用低位宽数字的神经网络近似使用了浮点数的神经网络的过程。

来源：Medium