Mamba一作最新工作:多亏Transformer,让Mamba更强了!

点击下方卡片,关注“CVer”公众号

AI/CV重磅干货,第一时间送达

点击进入—>【Mamba/多模态/扩散】交流群

添加微信号:CVer111,小助手会拉你进群!

扫描下方二维码,加入CVer学术星球!可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料,及最前沿应用!发论文/搞科研/涨薪,强烈推荐!

9d3600bf8fc25b55dc431ad34237f10c.png

明敏 发自 凹非寺
转载自:量子位(QbitAI)

Attention is all you need.

至少在矩阵这儿是。

Mamba架构最新进展:仅需1%计算量,新模型性能达SOTA

能做到这一点,还多亏了Transformer。

ec0c34a4eddc497864a4be8e4aec76f5.png

通过将Transformer模型中的知识有效迁移到Mamba等替代架构中,模型能在保持较低计算成本的同时,性能更好。

这就是由Mamba主创之一Albert Gu领衔的最新成果。

值得一提的是,这种方法还适用于Mamba以外的非Transformer架构。

6fa9066f4941f15a6c7f010129e2abd7.png

论文:https://arxiv.org/abs/2408.10189

从Transformer到SSMs

Transformer由于依赖二次自注意力机制,所需计算量很大。

二次自注意力机制能让模型在处理序列数据时有效捕捉序列内部的长距离依赖关系,但是由于二次时间复杂度(如果输入规模翻倍,模型计算所需时间增加4倍),导致处理长序列的计算成本很高。

为了解决这个问题,学界提出了很多新架构,比如Mamba、RWKV等,它们的微调和推理成本更低。

考虑到Transformer模型预训练已经投入了大量计算资源,研究人员想到,为什么不能在此基础上进行提升?

所以在本项研究中,他们提出了一种蒸馏方法MOHAWK,利用Transformer预训练模型来训练SSMs模型。

其核心在于注意力机制、线性注意力、Mamba的结构化掩码注意力SMA等,都是跨输入长度维度的序列转换。因此它们都有各自的矩阵混合器,比如softmax。

60a449007849619d71a8394cb2636fdd.jpeg

通过将注意力和SSMs视为通过应用不同类别的矩阵来混合不同token嵌入的序列变换,序列模型架构可以分解为独立序列混合和通道混合块。

比如Transformer由注意力(序列混合器)和MLP(通道混合器)块组成,使用这种分解可以蒸馏模型的每个元素。

具体蒸馏分为三个阶段

第一阶段:矩阵对齐(Matrix Orientation)。对齐序列变换矩阵本身。

第二阶段:隐藏状态对齐(Hidden-State Alignment)。对齐网络每个单独层的隐藏状态表示,且不牺牲预先学习的表示。

第三阶段:权重转移和知识蒸馏(Weight-Transfer and Knowledge Distillation)。通过一个端到端训练阶段,将权重转移,最终使用只有一小部分训练数据来蒸馏网络的最终输出。

利用这个方法来实际修改一个模型,比如Phi-Mamba。

579becbfd8f9c11164db9c76743c58f0.png

它结合了Mamba-2和Phi-1.5。

通过MOHAWK方法,该模型从预训练的Transformer模型中学习,同时作为状态空间模型,它在处理长序列上比传统Transformer架构更高效。

该模型仅使用3B token进行蒸馏,数据量为从头训练模型的1%,但是性能达到开源非Transformer架构中的SOTA。

7b4ac5d107b80ecb4dccfc49e67e72fe.png

实验发现,隐藏状态对齐更好,可以提高后续阶段的性能。

a564b1b84716208a9b999c96c497b742.png

研究团队也发布了混合Phi-Mamba-1.5B,通过5B token蒸馏,模型与类似混合模型表现相当,但是注意力层只用了4层

aa67c8d94d5f3e5472cc4830b5cc55a3.png

值得一提的是,这种蒸馏方法不止适用于Mamba。

954c97a6e7175e670fa033d7dad0773a.png

该研究由CUM助理教授、Cartesia AI联合创始人及首席科学家Albert Gu领衔。

去年,他和FlashAttention作者Tri Dao一起提出了Mamba,成为第一个真正实现匹配Transformer性能的线性时间序列模型。

在CVer公众号后台回复:Mamba,即可下载Mamba最全综述!

 
 

何恺明在MIT授课的课件PPT下载

 
 

在CVer公众号后台回复:何恺明,即可下载本课程的所有566页课件PPT!赶紧学起来!

ECCV 2024 论文和代码下载

在CVer公众号后台回复:ECCV2024,即可下载ECCV 2024论文和代码开源的论文合集

CVPR 2024 论文和代码下载

在CVer公众号后台回复:CVPR2024,即可下载CVPR 2024论文和代码开源的论文合集

Mamba、多模态和扩散模型交流群成立

 
 
扫描下方二维码,或者添加微信号:CVer111,即可添加CVer小助手微信,便可申请加入CVer-Mamba、多模态学习或者扩散模型微信交流群。另外其他垂直方向已涵盖:目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注:研究方向+地点+学校/公司+昵称(如Mamba、多模态学习或者扩散模型+上海+上交+卡卡),根据格式备注,可更快被通过且邀请进群

 
 
▲扫码或加微信号: CVer111,进交流群
CVer计算机视觉(知识星球)来了!想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料,欢迎扫描下方二维码,加入CVer计算机视觉(知识星球),已汇集上万人!

▲扫码加入星球学习
 
 
▲点击上方卡片,关注CVer公众号
整理不易,请赞和在看
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值