Sakana AI 推出通用记忆模型:让模型学会遗忘,支持零样本跨模态
Sakana AI 推出神经网络注意力记忆模型 NAMMs,模仿人类大脑机制让模型主动剔除冗余 token。该系统通过演化算法优化注意力矩阵,在压缩显存的同时提升了长文本推理表现,还能零样本直接迁移到视觉和强化学习模型中。

人在记住重要事情的同时,大脑其实一直在拼命遗忘无聊的细节。但现在的 AI 并不是这样:主流 Transformer 模型就像一个收纳强迫症,不管读了几万行代码还是几个小时的视频,所有 token 都会原封不动塞在内存里,上下文越长,推理成本越吓人。
获得日本经济产业省 GENIAC 超算项目资助的 Sakana AI,决定教大模型学会「主动遗忘」。他们公布了论文《An Evolved Universal Transformer Memory》,推出了一套名为 NAMMs(神经注意力记忆模型)的注意力修剪系统。
用进化算法,教神经网络断舍离
在此之前,学术界减少长文本显存占用的常见做法是 H2O 或 L2 这类手工规则,比如死板地保留最近几个词或注意力最高的前几个词。问题是,硬性裁剪往往会伤到模型原本的理解能力,属于拿性能换速度。
Sakana AI 的思路不同,他们把「留还是扔」做成了一个小巧的神经网络分类器。但这里有个很棘手的问题:每个 token 要么留下来,要么彻底丢掉,这种非黑即白的离散决策是不可导的,没法用反向传播来梯度训练。
Sakana AI 拿出了他们最擅长的进化算法。通过模拟自然界的突变和试错筛选,系统不断演化这套小模型的网络参数,在没有梯度指导的情况下,找到了最会挑重点的淘汰策略。
具体执行时,NAMMs 会抓取注意力矩阵的数值,先把它转换成音频和地震学中常见的语谱图,再用指数移动平均压缩成紧凑特征,最后打出一个保留分数,低于门槛的直接丢弃。
读文本演化出来的经验,直接能给视频挑刺
更有意思的是这套系统的迁移能力。因为所有 Transformer 每一层都有注意力矩阵,NAMMs 盯的是注意力分布本身,而不是具体文字。
研究团队只在 Llama 3 8B 语言模型上演化了 NAMMs,然后直接做零样本迁移。把它挂到 70B 的大尺寸 Llama 上,或者给视频模型 Llava Next Video 剔除冗余帧,甚至给强化学习模型 Decision Transformer 扔掉次优动作,全部无需微调就能直接生效。
团队拆解发现,NAMMs 在不同层和不同任务上的行为完全不同。在模型浅层,它会保留任务前言和全局关键词;到了深层,这些信息已经被消化,它就会果断忘掉前言,专心处理局部细节。
面对代码时,被它修剪掉的通常是整块的空格、注释和样板代码;换成普通英文文本,它丢掉的反而是语法中冗余的虚词,把核心人名和概念留下来。为了测试长文本表现,他们还开源了日语长文本评测基准 ChouBun,涵盖 36 项长文本与编程任务。
为什么重要
长上下文现在的最大痛点不是模型读不懂,而是显存放不下、算力跑不起。对于开发者和部署端来说,NAMMs 这种插件式的修剪机制提供了一个不用重新预训练大模型就能给推理瘦身的新路子,而且官方开源了全部代码,值得长文本场景的工程师跟进验证。
对整个行业而言,这打破了过去大家只能在全量 KV 缓存硬抗、或者换成 Mamba 这种线性架构之间二选一的僵局。它证明了 Transformer 现有的注意力权重本身就藏着筛选特征的通用语言,不同模态之间竟然能直接共享这套剪枝逻辑。
我们认为,Sakana AI 这一步走得相当狡黠。他们巧妙绕开了不可导的数学死胡同,把进化计算的优势发挥得淋漓尽致。不过看起来它的局限也很实在:进化算法虽然训练出来的结果好用,但训练过程本身极其消耗算力,需要超算集群长时间试错。普通团队很难低成本复现演化过程,未来更现实的路径,或许是直接白嫖他们演化成熟的通用分类权重。
来源 (1)
- sakana.ai(在新标签页打开)原文sakana.ai