NeurIPS 2024 | 中科院自动化所提出MetaLA！线性模型架构的大一统
https://csdnimg.cn/release/phoenix/outside_default.png©PaperWeekly 原创 · 作者 | 李国齐课题组
单位 | 中国科学院自动化所
目前，各种线性复杂度模型来取代 Transformer 布局中的传统 Softmax 留意力被提出，比方线性 Transformer（LinFormer），状态空间模型（SSM） 和线性 RNN（LinRNN）。
然而，这些线性模型的最佳设计仍然是一个悬而未决的题目。在这项工作中，本研究试图从理论角度找到 Softmax 留意力的最佳线性近似往返答这个题目。
我们起首将现有的线性复杂度模型同一为线性留意力形式，然后确定最佳线性留意力设计的三个条件：i）动态记忆能力；ii）静态近似能力；iii）最小参数近似。
我们发现当前的线性模型都不能满意所有三个条件，导致性能不佳。相反，我们提出了元线性留意力（MetaLA）作为满意这些条件的解决方案。我们在多查询联想回忆 (MQAR) 任务、语言建模、图像分类和长距离依靠（LRA）基准上的实行表明，MetaLA 比现有的线性模型更有效。
https://i-blog.csdnimg.cn/img_convert/1ea5b45f3710b352a6a8ad9f11bf7194.png
论文链接：
https://arxiv.org/abs/2411.10741
代码链接：
https://github.com/BICLab/MetaLA
https://i-blog.csdnimg.cn/img_convert/7d018b6356629ed91b93eef37c7fb898.png
背景
Transformer 模型凭借高效的并行训练能力和杰出的性能，在深度学习应用中体现出色。然而传统的 Softmax 留意力机制在训练时，盘算复杂度随输入长度呈二次增长；在推理时每个时间步和隐状态交互，时间和空间复杂度成线性增长。因此，Transformer 面对着盘算本钱过高的题目。
为此，当前研究主要致力于开辟线性留意力模型，如 LinFormer、SSM（状态空间模型）和 LinRNN，试图到达训练时以线性复杂度替换 Softmax 留意力、推理时到达常数级别的时间和空间复杂度。然而这些模型在当前流行的功能和体现上仍与 Softmax 留意力存在差距。
https://i-blog.csdnimg.cn/img_convert/83f413123b3b3e15a647fceb09d1f521.png
主要贡献
克日，研究者提出了一种新型线性留意力模块——MetaLA，能够实现对 Softmax 留意力映射的最优线性逼近。MetaLA 的设计突破了传统线性模型的限定，同一了现有线性模型的布局，具有以下明显贡献：
a）同一框架下的线性模型解读
本课题组团队首次将 LinFormer、SSM 和 LinRNN 等线性模型抽象为同一的线性留意力形式，从模型的隐蔽状态巨细、隐蔽状态维护方式以及参数映射策略等角度剖析其关键设计。这种同一视角不但帮助理解现有模型的功能差别，还从理论上展现了它们在实现 Softmax 留意力功能方面的潜力和不足。
b）定义线性逼近的必要条件
为评估线性模型能否逼近 Softmax 留意力映射，研究者提出了两个必要条件：动态记忆和静态逼近。动态记忆要求线性留意力模型能存储最重要的信息并忘记无关信息，而静态逼近则要求模型能够拟合任意 Softmax 留意力映射。
基于这一理论分析，研究者指出，现有诸如 TransNormer、RetNet、RWKV-4、LRU、HGRN 等模型未能满意必要条件，而某些模型（如 Mamba 和 GLA）因利用多余的 Key 矩阵而非最优参数化方案。
c）最佳性能
实行表明，基于 MetaLA 的 Transformer 在关联记忆、语言建模、长序列建模和图像分类等任务上均取得了明显性能提拔。同时，研究团队通过消融实行验证了 MetaLA 中各改进的有效性，并进一步探究了怎样提拔线性留意力的逼近能力以及线性留意力的容量上限题目。
https://i-blog.csdnimg.cn/img_convert/0047d90562b4cde5316e4adc289601c8.png
方法
3.1 一种通用的线性模型形式
观察现有的 LinFormer、SSM 和 LinRNN 模型，研究发现它们的推理过程可以同一为维护隐蔽状态的递归形式。Softmax 留意力通过 KV 缓存实现无穷隐蔽状态，而线性模型通过限定隐蔽状态实现对 Softmax 功能的逼近。
模型具有如下的串行形式，此中，差别线性模型的差别点主要在于 q，k，v，alpha 等向量信息等天生运算过程差别：
https://i-blog.csdnimg.cn/img_convert/7662734c104aca5a32fc14e3f37bf2b1.png
针对线性模型，其不但有同一的串行形式，也存在同一的并行形式：
https://i-blog.csdnimg.cn/img_convert/b08e53d888eb057858ab650af30412ad.png
基于如许的串行形式和并行形式，我们可以绘制出如下的通用线性模型流程图：
https://i-blog.csdnimg.cn/img_convert/791b4289cab24fda06a6fca88f3d3db9.png
▲ 图1：线性模型的通用形式的信息处理流程图。该图可视化了上述的同一公式在并行、串行视角下的信息处理流程
上述公式和流程图同一了 LinFormer 、SSM 和 LinRNN 等模型的递归形式，为理论分析奠定了底子。
在上文给出了同一的线性模型框架的底子上，我们指出，差别的线性模型，比方线性 Transformer （LinFormer），线性 RNN（LinRNN），状态空间模型（SSM）都是通用框架下的特例。而差别的具体模型，其差别体现在 q，k，v，alpha 等向量的天生过程，以及所维护的状态维度上。
https://i-blog.csdnimg.cn/img_convert/7ce17d13b62bdbb5becdcd925ea1f3ea.png
▲ 表1：LinFormer ，LinRNN ，SSM 均为线性模型的特例
3.2 逼近softmax的留意力图的最优线性模型的必要条件
研究者从理论上定义了实现 Softmax 留意力最优线性逼近的必要条件：
[*]线性复杂度：训练的时间和空间复杂度需为 O（n），推理复杂度为 O（1）。
[*]动态记忆能力：通过有限隐蔽状态动态存储重要信息，忘记不重要信息。
[*]静态逼近能力：能逼近任意 Softmax 留意力映射。
[*]最优参数化：在满意上述条件的条件下，利用最少的参数。
https://i-blog.csdnimg.cn/img_convert/2cc84f552ff51beb105423d058856cdb.png
▲ 表2：差别模型对最优理论的满意性分析，可以看，GLA ，RWKV ，Transformers ，等模型都不能满意分析中的全部条件
理论分析表明，动态衰减和 Query 矩阵是实现上述条件的关键。而 Key 矩阵在理论上并非必要，可通过优化动态衰减机制替换。
3.3.基于通用形式的MetaLA架构设计
去除 Key 矩阵，用动态衰减 alpha 替换 Key 矩阵，淘汰参数冗余并加强动态记忆能力。一方面，这一机制有利于更好的参数调配。另一方面，这一机制保证了我们对最优逼近分析得来的必要条件（动态记忆能力和静态逼近能力）。
引入自加强机制和短卷积加强 Token 对自身的留意力，避免留意力稀释题目，提高当前 Token 的信息表达能力，强化局部特征建模能力。
https://i-blog.csdnimg.cn/img_convert/8dfb542be28713b14513f6155d85725a.png
▲ 图2：基于最优理论，为线性模型找到了最佳的设计方案
https://i-blog.csdnimg.cn/img_convert/60e0edbd8139eb5cac6a48334ea7f9c8.png
实行
我们在多查询联想回忆 (MQAR) 任务、语言建模、图像分类和长距离依靠 (LRA) 基准上的实行表明，MetaLA 比现有的线性模型更有效。
多查询联想回忆（MQAR）旨在测试模型在多查询场景下的联想记忆能力和信息检索服从。该实行的核心任务是让模型记住一系列键值对（Key-Value Pair），并在稍后根据给定的查询键返回精确的值。
通过这一实行，我们能够相识模型在处理动态记忆和高效查询中的体现，以及其是否能够成功应对多次查询的累积负担。
https://i-blog.csdnimg.cn/img_convert/f235c3b5ba6e238967c06e6b4a8f2ab6.png
▲ 图3：MetaLA 模型中在多查询联想回忆任务中的性能，反映模型记忆能力
语言建模（CommonsenseReasoning）评估模型对一样平常生活中常识性知识的掌握程度以及基于常识举行推理的能力。
实行利用了如 Winograd Schema Challenge、HellaSwag 等常用基准数据集，要求模型推理隐含信息或基于有限背景知识做出决策。我们重点分析了模型在处理常识性推理能力。该实行有助于衡量模型在广泛现实应用场景中的泛化和推理能力。
https://i-blog.csdnimg.cn/img_convert/207a0e461bc357e07374b9d0c6699e83.png
▲ 表3：模型在 CommonSense Reasoning 上的性能对比，反映语言建模能力，模型性能明显高于 Pythia ，Gated Linear Attention ，Mamba 等主流模型
图像分类（ImageNet-1k）实行 是经典的视觉分类评测任务，旨在验证模型对图像内容的识别能力。我们利用了包含 1000 个类别的大规模 ImageNet 数据集，测试模型的 Top-1 分类准确率。
https://i-blog.csdnimg.cn/img_convert/164501e95eb2ed98d79e17bf4010445c.png
▲ 表4：模型在ImageNet上的性能对比，反映图像建模能力
长序列任务 Long Range Arena（LRA）实行 旨在评估模型在处理长距离依靠关系和复杂布局数据方面的性能。LRA 基准任务包罗文天职类、布局预测和图形匹配等，挑战模型在长文本或大规模图形数据上的捕捉能力。我们特殊关注模型在长序列中的局部信息整合和全局依靠建模能力。
https://i-blog.csdnimg.cn/img_convert/33295e060c829f6d9b9b1c40aa0f694e.png
▲ 表5：模型在长距离依靠任务上的建模能力，反映模型对长序列关系的捕捉能力
https://i-blog.csdnimg.cn/img_convert/531e8d5c1876824233effe3b0ce255ff.png
总结
MetaLA 模块通过去除冗余的 Key 矩阵、引入自加强机制以及加强局部交互的短卷积设计，成功实现了对 Softmax 留意力的最优线性逼近。其创新性地同一了现有线性留意力模型的通用形式，并满意动态记忆和静态逼近的必要条件，同时有效降低了参数复杂度。
这一设计为线性留意力模型在长序列建模任务中的应用提供了全新思绪，并明显提拔了盘算服从和模型性能。
https://csdnimg.cn/release/phoenix/outside_default.png
参考文献
https://csdnimg.cn/release/phoenix/outside_default.png
Yang S, Wang B, Shen Y, et al. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. Qin Z, Li D, Sun W, et al. Scaling transnormer to 175 billion parameters. arXiv preprint arXiv:2307.14995, 2023.
Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Smith J T H, Warrington A, Linderman S W. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933, 2022.
Qin Z, Yang S, Sun W, et al. Hgrn2: Gated linear rnns with state expansion. arXiv preprint arXiv:2404.07904, 2024.
Peng B, Alcaide E, Anthony Q, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048, 2023.
Katharopoulos A, Vyas A, Pappas N, et al. Transformers are rnns: Fast autoregressive transformers with linear attention//International conference on machine learning. PMLR, 2020: 5156-5165.
Biderman S, Schoelkopf H, Anthony Q G, et al. Pythia: A suite for analyzing large language models across training and scaling//International Conference on Machine Learning. PMLR, 2023: 2397-2430.
更多阅读
https://i-blog.csdnimg.cn/img_convert/3569626587b2d20ee63cdc63c66de767.png
https://i-blog.csdnimg.cn/img_convert/fc7a082cddccd3c8be781f9335133c94.png
https://i-blog.csdnimg.cn/img_convert/c6a3eae458ea662c758a8c15003eb387.png
https://i-blog.csdnimg.cn/img_convert/17cbbf5e4daf64acd38374e7013751e7.gif
#投 稿 通 道#
让你的文字被更多人看到
怎样才能让更多的优质内容以更短路径到达读者群体，缩短读者探求优质内容的本钱呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使差别背景、差别方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实行室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或比赛履历解说等。我们的目的只有一个，让知识真正活动起来。
页:
[1]