No. 117·精读

两周13篇论文,显存省85%:Mamba 诞生90天,火从语言烧到了图像和 DNA

Mamba-Palooza: 90 Days of Mamba-Inspired Research with Jason Meaux: Part 2

Nathan 与 AI 侦察员 Jason Meaux 回顾 Mamba 后续研究:计算机视觉应用、上下文长度扩展实验、内部状态腐烂问题与生物领域的混合 SSM。

节目音频00:00 / 1:20:47
中文精读由 AI 基于原始访谈改写,非官方翻译,可能存在偏差或错误;表述以原节目为准。

2024年2月,Cognitive Revolution 主持人 Nathan Labenz 跟 AI 侦察员 Jason Meaux 打了个赌:两周之内,以 Mamba 为主题的新论文会出几篇?Nathan 开的盘口是5.5篇,Jason 买大。结果盘口被迫一路改到14,两周结束,最终数字停在13,只差一篇。

更值得玩味的是选题分布。原始 Mamba 论文2023年12月发布时,通篇没碰图像。而这90天里冒出来的后续论文,过半都在做视觉。

先给没听过 Mamba 的读者补一句背景:它是选择性状态空间模型(SSM),跟 Transformer 的注意力机制打对台。卖点是计算量随序列长度线性增长,注意力是平方增长,理论上下文可以无限长。这正是 Nathan 和 Jason 兴奋的起点。

这一期是两人 Mamba 系列精读的下篇,覆盖医学影像分割、多向扫描设计、超长上下文实验,还有一个可能改变药物研发的模型 Evo。Jason 维护着一个论文索引站 statespace.info,这期节目基本就是他90天刷论文的成果汇报。

先看那个赌局:指数曲线开始发力了

Nathan 复盘自己定盘口的心态:当时在想这条增长曲线走到哪了,弯得有多快。5.5明显低了。第一次改盘口到14时,他还觉得这个数很难够到,结果两周13篇。

两人顺势开了新盘口:未来90天,Mamba 新论文55篇。Jason 又买了大,理由很直接,55篇还不到前90天翻一倍的水平,按这条曲线大概率会超。6月5日可以验证。

两周13篇,这个架构已经从「一篇有意思的论文」变成了一个自循环的研究方向。没人发奖金,研究者就是自己想玩。

图像没有天然顺序,第一步是把它变成序列

语言、音乐、DNA 都有天然顺序,token 接着 token 读就行。图像是二维的,像素没有唯一读法,这是 Mamba 进视觉领域要过的第一道坎。Transformer 的老办法是切块,把小图块当 token,效果已经很好。但 Mamba 是循环结构,一次只能看一步,核心问题就变成了:怎么把二维图像摊平成一维序列。

打头阵的是多伦多大学的 UMamba,1月初发布。做法是让卷积块和 Mamba 块交错:卷积负责提取局部特征,盯着眼前这一小块;Mamba 负责长程依赖,把远处的关系串起来。整个网络是 U 形结构,在器官、CT、MRI 的基准上打平甚至超过了当时的最好水平。这篇论文点燃了分割论文的连环爆发,随后的 Swin-UMamba 直接在论文里引用并超过了 UMamba。Mamba 变体之间已经开始互相比较了,不再只盯着 Transformer 打。

Nathan 对分割任务有个提纲挈领的解释:它本质上更接近图像生成。灰度图进去,彩色标注图出来,不同颜色对应不同器官和组织。中间的网络结构叫 UNET,先从像素逐层抽象出高层概念,再逐层生成回一张新图,中间用跳跃连接保住原始细节。医生监控病情要翻几百张灰度图,手动把器官、肿瘤、组织分开,模型越准,这个反馈循环越快。

传统做分割靠卷积,卷积天生是局部的,一块一块往外扩,早期没法理解图像远处的部分。Mamba 补的就是这一块:完整扫一遍,或者双向、多向扫,拿到全局视野,再跟卷积的局部能力结合,最好的成绩就是这么来的。

节目里还插了一段虚拟染色的题外话。传统病理染色要切下组织、切片、滴化学试剂、人工看显微镜,全程几个小时,手术台上根本等不起。Tanishq Matthew Abraham 参与的项目用新型探针直接对体内组织拍二维图,再用模型生成染色效果,从取图到出结果可以压到秒级。这项工作还没进临床,但方向已经验证。他的另一个项目用视觉基础模型微调,从 fMRI 信号重建人当时看到的画面,数据少得可怜,微调路线再次奏效。

双向、四向、串行:扫描方式的军备竞赛

1月中旬,华中科技大学和地平线机器人发布了 Vim(Vision Mamba),第一个公开的多状态方案:图像序列正着扫一遍、反着扫一遍,两个选择性状态空间模型并行处理,结果直接相加。

数据很硬:比对比的 Transformer 快2.8倍,批量推理省下85%以上的 GPU 显存,测试分辨率1248×1248。同样的图,Transformer 要80GB 以上显存,Vim 只要11GB。11GB 是什么概念,一台用了几年的 MacBook Air 都装得下。论文原话:

「Vim 克服了 Transformer 式高分辨率图像理解在计算和内存上的限制,有潜力成为下一代视觉基础模型的骨干网络。」

地平线机器人的参与是个信号。机器人要在边缘设备上跑,80GB 显存没法随身带,11GB 就是能用和不能用的区别。

紧接着 vMamba 搞了四向交叉扫描:从四个角分别条带式扫过图像,四份表示并行处理后叠加。参数量在2000万到7000万之间,小模型跑四份没压力。Nathan 提出的疑问是,到70亿、700亿参数,四份状态还能不能塞进 SRAM,这是早期论文普遍绕开的问题。

2月8日的 Mamba-ND 把接线方式系统测了一遍:并行合并、全串行、混合接法都试了,还扩展到视频和气象数据。结果是串行最好:一个扫描处理完,把结果连同下一个扫描一起交给下一个状态空间模块。Nathan 打了个比方,与其克隆四个自己各看一遍再合并意见,不如自己先正着看一遍,带着第一遍的印象再换个方向看。Jason 补了原理:Mamba 的选择机制依赖输入,第一遍扫到的东西,只有串行才能影响第二遍的处理。论文末尾的有效感受野对比图也支持这一点,多向扫描下单个位置能「看到」的图像范围明显最大。

后面的几篇可以快速过。SegMamba 上三维,64×64×64 的立方体展开是26万长度的序列,三向扫描,只用了4张 A100。VIVIM 上视频,时间轴正扫反扫,第三向扫单帧内部。MambaMorph 做配准:MRI 和 CT 是两台机器拍的同一具身体,形状对不上,模型输出一个形变场,把两张图对齐到一起看。图结构那边,Graph Mamba 把图摊平成序列再加一个反向序列,1400个节点只花不到200MB 显存,对比的 Transformer 到800个节点左右就爆内存了。

Nathan 给这套玩法起了个名字,seeing like a mamba:任何数据,找到一种摊平成序列的表示,再不济就正反各扫一遍,新架构就能处理它。这是整个视觉部分最大的收获。

16K 训练,40K 才开始糊:超长上下文的第一份成绩单

LongMamba 用原版 Mamba 的训练数据子集,截到16000 token 训练。困惑度曲线在16000之后没有上升,到40000 token 才开始变差。训练窗口之外,模型还在泛化。大海捞针测试(Greg Kamradt 在 Twitter 上带火,拿 Claude 和 GPT-4 测过)的结果是16K 窗口内接近全对,测试内容是在一堆「草是绿的,天是蓝的」噪音里插入一句「密码是12362,记住它」,跑几百次画热力图。

有个细节跟 Transformer 的已知毛病正好相反:16K 窗口内的失败点出现在开头位置。Transformer 是「迷失在中间」,状态空间模型看起来是「迷失在开头」。超出训练窗口后还有点余量,20000 token 之后迅速衰减。

训练侧还有个便宜改进:作者在 GitHub 上延续 Tri Dao 等人的讨论,用 Transformer-XL 式的状态传递,一批训完把 SSM 隐状态存下来,下一批加载回去接着训。架构上没有任何东西阻止这么做,8张 A100 就能训到16K。Nathan 去年12月测过聊天版 Mamba,3000 token 左右就散架,后来发现训练数据里全是短对话。数据到位,能力就在。

更远的部分,经济学家罗宾·汉森有个观点:AI 系统会像生物和传统软件一样腐烂。困惑度在40K后回升,Nathan 的读法是状态开始被糊住,没有灾难性失败,但越走越黏。汉森还补了一刀:进化本可以造出永生并自我修复的生物,实际却选择了世代更替;软件也一样,没人无限打补丁,攒够了就重写。状态空间模型想无限跑下去,得发明自己的清理机制。Nathan 顺势下了个预测:未来90天会出现某种状态正则化,类似权重衰减的状态衰减,让旧信息逐步退场。粗糙版本是对状态里所有数值的总和加惩罚,不留不干活的信息。泰格马克组里 Ziming Liu 的「Seeing is Believing」论文在权重上玩过类似的花样,靠惩罚权重总量把网络压到最稀疏。

人脑的类比是:记不住昨天早饭吃了什么,但记得自己平时早饭吃什么。模糊记忆依然有用,问题是精度丢了,捞针就捞不到了。

Evo:拿 DNA 当训练数据,长出一个「细胞模型」

这一节被两人称为全场分量最重的部分。背景先说 ARC 研究所:Patrick Collison 参与创办,逻辑是很多人想做科学,不想当实验室主任,不想写基金、管学生、修仪器。ARC 发工资、配设备,你就专心做科学。

他们拿出的 Evo 用了 Together AI 的 StripedHyena 混合架构,在 DNA 序列上训练,一个碱基对一个 token,最长131000 token,主要用原核生物基因组。模型能生成很长的基因组序列,距离很远的片段之间表现出协调性。

Nathan 的判断是,语言模型靠预测下一个 token 长出世界模型,Evo 可能在做生物学版本的事,长出一个细胞模型。DNA 的长程依赖比语言狠得多:文本里隔几十万 token 才相关的情况是少数,细胞里基因组两端的片段各自表达蛋白,一个坏了整个细胞可能死,相互作用跟序列上的距离基本无关。

节目里还讨论了一个问题:光有 DNA 够不够?Jason 的看法是,细胞里真正的因果发生在 RNA 和蛋白层面的相互作用网里,那是一张乱糟糟的图。DNA 虽然是序列,但可以看成这张图天然摊平后的版本。人工处理图结构都得先摊平再扫描,大自然可能早就替我们做完了这一步。

一个耐人寻味的比例:LongMamba 训练16K,能力延到40K上下;Evo 训练131K,论文声称可用到650K。都是训练长度的数倍,比例量级接近。是巧合还是规律,等后续论文验证。

Nathan 现场发了招募帖,想找一位 AI 加生物学的侦察员搭档做下一期。他的理由是,语言这条线已经很重要,生物学这条路如果走通,分量一点不轻,而他自己读生物论文太吃力。Patrick Collison 的观察佐证了难度:没有记录显示有二十几岁的人颠覆过生物学,这个领域比机器学习复杂得多,没有设计,全是随机堆出来的相互作用,不熬过多年的细节就没法领航。

药物研发的战舰游戏,和 Transformer 的反击

AlphaFold 已经把门槛跨过去了:Alphabet 内部有专门部门做延伸和商业化,已有药物处在临床试验中或临近临床。Nathan 用一个游戏类比药物筛选:战舰。历史上大部分筛选在盲打,往细胞上扔东西,看有没有东西杀死癌细胞,命中率很低。传统方法里也有从自然找线索的路数,比如某种青蛙不得癌症,就去研究那只青蛙。计算模型有假阳性假阴性,但只要能把搜索空间缩小几个数量级,先超算打分、只把前1%送去湿实验,研发的投入产出比就换了量级。

一款药上市成本10亿美元,大头是临床试验,这部分短期不会变,但前端的命中环节会变。命中率提高百倍,意味着这件事本身值得做百倍的量。

Transformer 会被取代吗?Nathan 之前被问过:两年内状态空间模型会进前沿系统吗?他的答案是概率不小,理由很朴素:全世界都在抢算力,Mamba 省算力,光这一条就够它进主流。反过来问,如果进不了,原因会是什么?唯一说得通的答案是注意力真的够用:ring attention 把视野推到千万 token,Gemini 1.5 展示了超长上下文。把整个代码库扔进 Gemini 1.5 已经帮到很多人,而这还没做产品化优化。

注意力阵营还有一手叫 memory token。GPT-4V 里一张低分辨率图片只花85个 token,一图千言。照这个思路,把一大段文本压缩成一个 token 呢?10比1、100比1、1000比1,深层历史用压缩表示,下游照常跑全注意力,命中了再内置 RAG 式地把原文拉回来。这条路走顺的话,状态空间革命可能被 Transformer 的持续进化掐灭在半路。

还有个野路子数据点:magic.dev 宣布了约1亿美元融资,出资人包括 Nat Friedman 和 Daniel Gross,官网明说:

「Transformer 不是最终架构。我们手里有数百万 token 上下文窗口的东西。」

Jason 的收尾判断:90天的论文潮摆在那,没人需要被激励,大家抢着玩。但 Transformer 这边也没躺平,两边都可能赢。

明天就能用的四件事

  • 想跟踪 Mamba 研究,直接收藏 Jason 维护的 statespace.info,比刷社交媒体系统得多。
  • 做小数据领域模型的人,学 Swin-UMamba:先在通用数据上预训练,让模型学会边缘检测这类视觉原语,再用稀缺的领域数据微调。医学影像和脑扫描重建都靠这招拿到了几个点的提升。
  • 做多向扫描的,参考 Mamba-ND 的结论:串行接线,让第二遍扫描带着第一遍的结果跑,目前成绩好于并行合并。
  • Nathan 的招募帖是给所有人的:做这类精读不需要机器学习博士学位,需要的是极强的好奇心和搞明白的执念。选一个自己懂的交叉领域,照着这两期的样子做文献综述,这样的位置现在空着很多。

In this second part of a two episode series, Nathan and AI scout Jason Meaux provide a sweeping overview of the first 90 days of Mamba-inspired research. They discusstMamba's application to computer vision, experiments in extending effective context length, the potential problem of internal rotting states, and the use of hybrid SSMs in biology.Try the Brave search API for free for up to 2000 queries per month at https://brave.com/api

LINKS:
Show Notes and Paper Links: https://docs.google.com/document/d/1NK_a3deVL_aczORmSRw8LyujNPotpO7Kd90sIRj9Qx0/edit?usp=sharing
Nathan's original Mamba Deep Dive

Spotify: https://open.spotify.com/episode/3b3zeDVU6fE2yV9vdNVLqX?si=dc7a885d601d4a58

Apple: https://podcasts.apple.com/us/podcast/emergency-pod-mamba-memory-and-the-ssm-moment/id1669813431?i=1000639397566

X/SOCIAL:
@labenz (Nathan)

SPONSORS: ORACLE | OMNEKY | BRAVE | ON DECK | SQUAD
Oracle Cloud Infrastructure (OCI) is a single platform for your infrastructure, database, application development, and AI needs. OCI has four to eight times the bandwidth of other clouds; offers one consistent price, and nobody does data better than Oracle. If you want to do more and spend less, take a free test drive of OCI at https://oracle.com/cognitive

Omneky is an omnichannel creative generation platform that lets you launch hundreds of thousands of ad iterations that actually work customized across all platforms, with a click of a button. Omneky combines generative AI and real-time advertising data. Mention "Cog Rev" for 10% off www.omneky.com

The Brave search API can be used to assemble a data set to train your AI models and help with retrieval augmentation at the time of inference. All while remaining affordable with developer first pricing, integrating the Brave search API into your workflow translates to more ethical data sourcing and more human representative data sets. Try the Brave search API for free for up to 2000 queries per month at https://brave.com/api

ODF is where top founders get their start. Apply to join the next cohort and go from idea to conviction-fast. ODF has helped over 1000 companies like Traba, Levels and Finch get their start. Is it your turn? Go to http://beondeck.com/revolution to learn more.

💥 Access global engineering without the headache and at a fraction of the cost: head to choosesquad.com and mention “Turpentine” to skip the waitlist.

TIMESTAMPS:
(00:00) - Episode Start
(00:01:14) - Nathan and Jason’s Bet
(00:04:51) - U-Mamba
(00:15:10) - Sponsors: Oracle | Omneky
(00:06:08) - Swin U-Mamba
(00:18:59) - Vision Mamba
(00:24:30 - VMamba
(00:27:26) - VM-UNet
(00:30:37) - Sponsors: Brave | On Deck
(00:35:15) - SegMamba
(00:36:59) - Vivim
(00:37:31) - Mamba Morph
(00:39:01) - Graph Mamba 
(00:43:15) - Long Mamba
(00:59:54) - Evo and StripedHyena

The Cognitive Revolution is produced by Turpentine: a media network covering technology, business, and culture.
Producer: Vivian Meng
Editor: Graham Bessellieu
For sponsor or guest inquiries, email: vivian@turpentine.co

本集资源与出处