无矩阵乘法语言模型:提升效率的突破
在大型语言模型(LLMs)中,矩阵乘法(MatMul)是计算压力最大的操作,尤其是在使用Transformer架构时。随着模型规模的扩大,MatMul操作的成本急剧上升,导致训练和推理时内存使用和延迟显著增加。
加州大学圣克鲁斯分校、苏州大学和加州大学戴维斯分校的研究人员开发了一种创新架构,成功消除了语言模型中的矩阵乘法,同时在大规模情况下依然保持强劲性能。
无矩阵乘法模型的介绍
在其开创性论文中,研究人员提出了无矩阵乘法的语言模型,这些模型在推理时所需内存大幅减少,但仍与最先进的Transformer性能相当。
深度学习中的矩阵乘法
矩阵乘法在深度学习中至关重要,负责将数据与神经网络中的权重结合,以转换输入数据并生成预测。由于其并行架构,GPU在执行多个MatMul操作时表现卓越,这对于高效训练和部署复杂模型至关重要。
然而,随着LLMs参数数量达到数百亿,MatMul操作逐渐成为瓶颈,训练和推理需要庞大的GPU集群。摆脱矩阵乘法有可能带来显著的内存和计算成本节约。然而,过去的替代尝试往往成果不佳,常常因为GPU性能不足而导致速度下降。
用三元权重革新计算方式
研究人员提出了一种令人振奋的替代方案:将传统的16位浮点权重替换为可以表示三种状态(-1、0和+1)的3位三元权重。他们引入了加法运算来替代MatMul,从而显著降低计算成本。这些模型使用“BitLinear层”来实现三元权重的运算。
“通过将权重限定在{-1, 0, +1}集并采用额外的量化技术,我们用加法和取反操作替代了MatMul。”研究人员表示。
创新的架构改变
该架构与传统Transformer有根本不同,传统Transformer由令牌和通道混合器组成。负责通过自注意机制整合序列令牌信息的令牌混合器,转变为无矩阵乘法线性门控递归单元(MLGRU)。MLGRU通过简单的三元运算更新隐状态,高效绕过昂贵的矩阵乘法。
除了令牌混合器外,通道混合器则对令牌表示的不同特征通道进行信息整合,采用修改过的门控线性单元(GLU),以适应三元权重。这一调整不仅降低了计算复杂度和内存使用,同时有效整合特征。
“通过结合MLGRU令牌混合器和使用三元权重的GLU通道混合器,我们的架构仅依赖加法和元素级乘法。”研究人员指出。
无矩阵乘法语言模型的性能评估
研究人员将无矩阵乘法LM与在Llama-2中使用的先进Transformer++架构进行了对比,涵盖了多种模型规模。结果表明,无矩阵乘法LM能够更有效地利用额外的计算资源,从而提升性能。
在语言任务的评估中,2.7B无矩阵乘法LM在ARC-Challenge和OpenbookQA等基准测试上的表现超越了其Transformer++对手,而在其他任务中也取得了相当的成绩。
“这些结果表明,无矩阵乘法的架构在问答和常识推理等多种语言任务中展现了强大的零-shot性能,”研究人员强调。
随着模型规模的增加,无矩阵乘法LM在内存消耗和延迟方面的优势愈发明显。例如,13B模型仅需4.19GB的GPU内存和695.48毫秒的延迟,而Transformer++则需要48.50GB内存和3183.10毫秒的延迟。
优化实施与未来方向
研究人员为无矩阵乘法语言模型开发了优化的GPU实现和自定义FPGA配置。通过这一优化,他们在训练中实现了25.6%的加速,并在内存使用上减少了多达61.0%。
“这一工作超越了仅依赖软件的轻量模型实现,展示了可扩展和高效的语言模型如何有效降低计算需求和能耗。”研究人员总结道。
尽管测试在超过1000亿参数的模型上受到限制,研究人员希望能鼓励机构投资于轻量模型,为更易获取的语言模型铺平道路。研究人员已向研究界提供了其代码和模型。
“通过优先发展无矩阵乘法架构,未来的LLMs将朝着更高的可获取性、效率和可持续性发展,”研究人员倡导道。