PaperWeekly 22小时前
比Transformer还快?闭环递归Comba登场!标量低秩变换加速前向传播,CV/NLP齐夺SOTA
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

香港科技大学广州与上海 AI Lab 联合发布了Comba,一种创新的递归模型架构,旨在提升序列建模的效率。该模型基于Bilinear RNNs,引入标量加低秩(SPLR)状态变换和输出校正技术,实现了硬件友好的分块并行优化。Comba在速度和性能上均有显著提升,尤其是在CV和NLP等任务上。通过闭环控制策略,Comba在长序列任务中展现出优越性,并已开源代码。

💡 Comba 采用了标量加低秩(SPLR)状态变换,这种结构优势使得模型只需计算一次逆矩阵,相较于Gated-DeltaNet,前向传播速度提升了40%。

🔄 Comba 模型的核心在于闭环控制策略,通过两阶段反馈机制,先修正输入信息,再对输出进行校正,增强了系统的适应性,使其能处理更复杂的动态任务。

🚀 Comba 架构在硬件友好性方面进行了优化,通过分块并行计算,最大化矩阵乘法的效率,从而在序列维度上实现高效并行计算,减少了计算开销。

📊 Comba 在多项任务上取得了SOTA(State-of-the-art)效果,尤其是在长序列建模、图像分类和目标追踪等视觉任务中表现出色,并已开源代码供研究者使用。

原创 让你更懂AI的 2025-06-18 21:32 北京

闭环递归新突破!

近日由香港科技大学广州与上海 AI Lab 联合发布了一项研究成果:Improving Bilinear RNNs with Closed-loop Control,该文章首先对最近高效序列建模方法,包括 Transformer,线性注意力,状态空间模型的发展进行了全面的总结,将近期提出的 RWKV7,Gated-DeltaNet 等模型归结为 Bilinear RNN 并分析了其理论优势。


随后该文章提出了一种新型的递归模型架构 Comba,采用标量加低秩(SPLR)状态变换以及输出矫正技术,基于 Triton 实现了硬件友好的分块并行优化。


由于 SPLR 的结构优势,Comba 只需要计算一次逆矩阵,从而在前向传播相比  Gated-DeltaNet 提升了 40% 的速度,并在 CV,NLP 等多项任务上取得 SOTA 效果。代码已开源并即将整合到 FlashLinearAttention 库中。

论文标题:

Comba: Improving Bilinear RNNs with Closed-loop Control

论文链接:

https://arxiv.org/abs/2506.02475

代码链接:

https://github.com/AwesomeSeq/Comba-triton



背景知识和相关工作

Transformer 依赖于 Softmax 注意力,虽然可以依赖于为矩阵乘法所设计的现代 GPU 高效并行计算并且效果强大,但在训练中时间复杂度为二次并且在推理过程中 KV-Cache 会使得内存无限制的增长,使得其在长序列任务中受限。


一些工作包括滑动窗口注意力、稀疏注意力和更高效的 KV Cache 管理在一定程度上缓解了瓶颈,却仍未摆脱二次注意力的基本范式。


自 mamba 爆火之后,递归模型快速崛起。这些模型把长序列压缩进固定大小的状态 S,使得推理只需恒定大小的内存开销以及 (1)时间复杂度,逐渐成为了 Transformer 的有力替代品。

1.1 第一阶段:数据依赖的门控

递归模型有着很多发展路径,例如线性注意力系列,最初对 q 和 k 分别通过核函数来近似整体的 softmax 操作,由此可以将注意力计算重写为关于时间的线性复杂度:

例如状态空间模型系列,最初源自 Hippo 理论,提出了一类正交多项式投影算子,当选取闭合递归多项式族时可以表达为递归形式。这些模型在最初提出的时候效果并不如 transformer,因为它们缺少了像 transformer 中自注意力机制那样的自适应带权信息融合。


随后,门控线性注意力或者选择状态空间(Mamba)系列模型提出了数据依赖的门控机制,长序列任务上取得了不错的效果。并在 2024 年下半年的各类工作中被统一称呼为 Linear RNNs,这些模型可以看作一个关于 key-value 联想记忆的线性寄存器系统:

1.2 第二阶段:从赫布生物学习法则到Delta学习法则

尽管第一阶段中带有数据门控的线性 RNNs 在长序列任务上取得了不错的效果,但仍然存在着一些问题,例如为了能够在现代 GPU 上面快速计算,这些模型往往将门控简化为标量,导致状态更新的表达能力不足。


并且这类模型基于加法的更新机制,不断的将 key-value 所构成的外积记忆加到记忆状态中,导致模型缺少一个准则来判断哪些记忆是重要的,哪些记忆是不重要的,数据依赖的门控会使得状态中存储的记忆会平等的遗忘。


为了解决这些问题,研究者们开始探索更复杂的状态更新机制。最近 Gated-DeltaNet 通过 Delta 学习法则对模型引入了纠错能力,使得通过 key 从状态中查询出对应的历史记忆更接近于当前输入 value,这个机制也等价于基于 Delta 损失对记忆进行随机梯度下降。


这使得递归模型推动到如下形式:

一方面,系统的状态变换矩阵从原本的标量进化为了类似 Household 变换的矩阵形式,大大增强了模型的表达能力,同时相比于密集实矩阵易于处理。有趣的是这和 Hippo 理论不谋而合:任何正交多项式投影算子都可以被表示为对角加低秩分量的组合形式。


另一方面,从几何视角来看 Householder 定义了一个以当前时间步 key 为法线的镜面,根据强度的不同可以选择性地保留或抑制模型状态中储存的记忆,从而对有限的记忆容量进行更高效的记忆管理。某种程度上可以视为对递归模型 “KV Cache” 的施密特正交分解。

Bilinear RNNs:这类模型由于引入了状态 S 和 key 的乘积项,其动力学已经不再是线性的了,再叫 Linear RNNs 有些不合适了。由于这类模型对于状态和输入分别是线性的但整体不是,动力学形式上属于双线性系统,一种保持着非线性行为的可控系统,因此本篇文章将这类模型称为 Bilinear RNNs。


在此基础上衍生了各种改进工作(具体参考论文第二章节),同时像是 TTT(Test-Time-Learning)等一系列模型则从优化器的角度对模型进行设计,并引入了基于 MLP 的深度记忆,高阶优化器,以及动量等操作,这类模型又不属于 Blinear RNNs,可以被看作现代版本的 Nonlinear RNNs。


虽然这些非线性操作理论上比 Bilinear RNNs 有着更强的表现力,但是牺牲了硬件高效的分块并行计算能力,哪怕是采用小批次梯度下降来训练模型,速度也大幅度落后 Gated-DeltaNet 等模型。


并且最近研究表明对于结构化数据如视频(帧)等采用小批次梯度下降是有效的,但是对于语言建模任务,单步梯度下降(类似 Gated-DeltaNet)仍然是性能最优的。


综合考虑,Comba 沿用了 Bilinear RNNs 的形式,并基于 Triton 编写了高效分块并行计算算子。该文章在文章中 Table1 对这些模型进行了全面总结:


Comba架构:基于闭环控制理论

根据控制理论,Mamba,GLA 等模型属于开环系统(关于 key-value 联想记忆),其输出不会对控制行为产生反馈。相比之下,另一类系统,被称为闭环控制系统,采用负反馈来增强系统的适应性,并允许它处理更复杂的动态任务。


在该文章中,如下图所示,Comba 采用两阶段反馈策略:首先通过基于状态的反馈 P 对输入信息(value)进行修正,然后使用相同的反馈机制对输出进行类似的校正。


从这个角度来看,TTT、DeltaNet 和 RWKV-7 等模型只包含了第一阶段的状态反馈校正,缺乏输出校正,因此不是严格的闭环系统。

该文章选用 key 向量来作为反馈(关于反馈的选择文章中进行了细致讨论),最终该文章可以将 Comba 写为如下递归形式:

相比于之前的 Bilinear RNNs,Comba 主要有两点不同:


1. 状态变换:Comba 采用了标量加低秩(SPLR)状态变换,该文章会在随后部分表明这个形式只需要计算一次逆矩阵,从而相比 Gated-DeltaNet 速度上有了大幅度提升。


并且相比于 RWKV7 中的对角加低秩(DPLR)状态变换,SPLR 状态变换已经有着足够的表达能力并且相比于标准 Transformer 几乎不引入额外参数,更易于从 Transformer 蒸馏出该递归形式。


2. 输出校正:Comba 在输出端引入了输出矫正机制,从神经记忆的角度来看,key 保证记忆 value 被尽可能清晰地存储,从而使 query 能够精确地查询。输出校正直接促进了这一过程,并在实验结果中显著降低了模型的困惑度。



硬件高效的分块并行计算

由于现代 GPU 是为矩阵乘法所设计的,导致递归模型在较短序列的计算时通常比 Transformer 要慢,为了最大化矩阵乘法 Flops,类似于 Mamba2,GLA,DeltaNet 等模型,Comba 采取如下块间递归,块内并行的方式,在序列维度上高效并行计算。

其中在计算中间变量 W,U 时,Comba 采用了 SPLR 状态变换,只需要一次逆矩阵计算(逆矩阵计算是主要速度瓶颈),而 Gated-DeltaNet 中需要两次。对于头维度 128 的模型来说,Comba 相比于 Gated-DeltaNet 在前向传播速度上提升了 40%。


在反向传递中,对于中等规模的模型(例如,隐藏维度 1024 或 2048),由于重用缓存的逆矩阵 M 和块级状态(会产生额外显存占用),速度相比 Gated-Deltanet 提升不大。然而,在需要重新计算的大规模模型中,Comba 有望提供显著的性能优势。

算子速度对比如下:PyTorch 中的循环 Comba 产生了大量的计算开销,限制了其大规模预训练的可扩展性。


对于较短的序列(例如 1024),Flash-attention 达到最快的速度,但其二次复杂度导致随着序列长度的增加效率降低。由于使用对角门控矩阵,GLA 的操作速度较慢。虽然 RetNet 实现了最快的速度,但由于缺乏依赖数据的内存管理,它在性能上落后于其他模型。

最终该文章构建了 Comba 家族,分别就线性注意力形式,mamba 形式,以及混合形式架构进行了实验探究,其中 mamba 形式类似于多 value 注意力,并共享维度为 128 的 query 和 key。对于混合模型,该文章每 8 层混合一次 softmax 注意力,并基于 FlashAttention 实现。



实验结果分析

该文章在此仅列举部分语言建模实验,对于长序列建模,图像分类,目标追踪等视觉任务请阅读原论文。


可以发现大多数递归模型在常识性推理任务中优于 Transformer,因为它们的递归结构类似于一条思维链。SPLR 结构的 Comba 优于 IPLR 结构的 Gated-DeltaNet 以及 DPLR 结构的 RWKV7,并且保持着最快的计算速度。输出纠正可以显着减少困惑,从而实现精确记忆检索。


该文章发现 Mamba 架构设计是次优的。MLP 作为一种特殊的(非 hebbian)键值记忆,补充了状态下的键值关联记忆,对于推理等任务尤为重要。该文章发现对于模型尺度,初始化输出校正因子 d 的选择应该有所不同,例如对于 340M 模型,d=0.02,对于 1.3B 模型,d=1 比较好。


递归模型由于其有限的状态空间,召回能力通常弱于拥有着无限记忆的 Transformer,然而,Comba 的召回性能与变压器非常接近。


消融研究表明,输出校正对 Comba 的回忆性能有显著的促进作用,优化 qk 的相似度将增强模型检索精确记忆的能力,从而提高回忆性能。该文章还测试了 Comba 在 MQAR 任务上的综合回忆能力。


更多阅读



#投 稿 通 道#

 让你的文字被更多人看到 



如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。


总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 


PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。


📝 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算


📬 投稿通道:

• 投稿邮箱:hr@paperweekly.site 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿


△长按添加PaperWeekly小编



🔍


现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧


·


阅读原文

跳转微信打开

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

Comba 递归模型 Bilinear RNNs 序列建模 闭环控制
相关文章