给4096个词打4个补丁,也能训好一个大模型? Original en chino

先说一个可能让你意外的事实。
现在主流的大模型,处理一句话的时候,理论上要把这句话里每一个词,跟前面出现过的所有词都两两比对一遍。一句4096个词的话,就要做几百万次这种比对。这就是Transformer架构里最核心也最烧钱的部分:注意力机制。
注意力机制*:模型判断"这个词该重点参考前面哪些词"的计算方式,是Transformer大模型的核心组件。
正因为这个计算量随句子长度的平方增长,业内一直在想办法给它"减负"。一种做法是滑动窗口注意力,简单说就是规定"只看前面固定的N个词,别的一律不看"。这招确实省了不少计算,但代价是:不管这句话讲的是什么内容,模型永远只能盯着那几个固定位置看,内容说什么跟看哪儿完全脱钩。
这篇论文的作者觉得,这不是唯一的出路。他提出了一个叫Graph Machine(图机器,简称GM)的架构,核心思路是:既然滑动窗口这种"死看固定位置"的省钱方法丢了灵活性,那能不能设计一种"看得少但看得准"的机制,而且这个"准"是由内容动态决定的,不是提前写死的?
结果相当有意思。在一个4096个词的句子里,如果每个注意力头只允许看2个词,损失几乎不受影响;如果允许看4个词,效果甚至比原模型还略好一点点。你没看错,4096个词里挑4个,比全都看一遍的效果还要好。
排队的比喻:四种"记忆方式"
要理解GM在解决什么问题,得先把几种常见的序列建模方式摆在一起比一比。
第一类是RNN和SSM这类循环模型。它们的特点是,不管你输入多长的文本,脑子里始终只留一块巴掌大的"记忆板",新信息进来,旧信息就要给挤走一部分。
RNN*:一种按顺序逐步处理序列的神经网络,靠一个固定大小的隐藏状态传递历史信息。
SSM*:结构化状态空间模型,一种用状态方程来建模长序列的架构,Mamba是其代表。
这就像你参加一个只能记住5件事的会议,第6件事说出来,你必须忘掉前面某一件才能记住它。信息注定要被压缩,压缩就意味着丢失。
第二类是标准Transformer,也就是原始的注意力机制。它反过来,把所有历史信息原样保留,每个新词进来都要跟所有历史词逐一核对。这就像开会时你把每个人说的每句话都完整记在笔记本上,想查什么随时能查,但如果会开了三天,你的笔记本翻页速度会越来越慢,因为每次查阅都要把之前所有内容重新扫一遍。
第三类是滑动窗口这种"静态稀疏"方法。它不再全记,只记最近的一段,而且不管这次讨论的主题是什么,规则永远是"只看最近N句"。如果讨论内容其实跟很早之前提过的一个关键点有关,滑动窗口是看不到的,因为那超出了它规定的范围。
第四类,就是GM试图填补的位置:既想保留全部历史(状态是O(n)大小,n是句子长度),又想每次只看很少几个位置(访问量是O(1)),但这几个位置具体是哪几个,要由当前内容动态决定,不能提前写死。
论文里给这四类做了个信息论上的推算:如果你想在n个位置里动态选出固定几个来看,那你至少需要log(n)那么多"比特"的信息来标明"我要看哪儿"。这听起来抽象,但换算到实际场景里很简单:一个能存下几十亿数字的普通整数变量,就够用了,存储成本几乎可以忽略。
边:除了记住"是什么",还要记住"该看哪儿"
GM的做法,是把传统神经网络里只有的"节点特征"这一种状态,拆成三份状态来存。
节点特征*:每个词对应的浮点数向量,记录这个词当前"理解到什么程度",跟普通Transformer里的隐藏状态是一回事。
除了节点特征,GM还额外存了两样东西:边索引和边权重。边索引记录的是"这个词现在指向谁",是一个整数;边权重记录的是"指向的程度有多强",是一个0到1之间的小数,所有权重加起来等于1。
边索引 / 边权重*:边索引是整数,表示当前词关联的目标词位置;边权重是浮点数,表示这种关联的强度,类似于给每条"关系"打分。
这里有一个关键设计,叫referral,中文可以理解成"转介"或者"指针传递"。
referral*:一种更新边的机制,让每个词的"邻居"通过邻居的邻居,递归地找到新的、更相关的邻居,类似于顺着指针一路追下去。
具体怎么运作?想象一个电话转接的场景。你打电话找客服,客服说"这个问题我处理不了,转给你另一个同事",那个同事又说"这块归另一个部门",于是电话被转了两次,你最终接通的人跟你最初拨的号码之间隔了两跳。这就是referral里的"两跳组合":第一条边告诉你"从A能到B",第二条边告诉你"从B能到C",两条边拼起来就变成一条新边"A直接连到C",权重是两段转接权重的乘积。
如果每个词只保留少量"邻居"(论文里通常是16到32个),两跳组合出来的候选路径数量会指数级增长。如果要保留固定数量的邻居,就必须做一次"硬选择",从一大堆候选路径里挑出最有价值的几条。
问题是,"硬选择"这个动作本身没法直接求导,梯度传不过去,模型学不会该怎么选。这就好比让一个新员工去筛选简历,你只告诉他"选5份最好的",但从不告诉他"好"在哪个维度,他没法通过试错慢慢调整判断标准,因为每次筛选结果都是非黑即白的0或1,没有中间地带可以微调。
论文的解法是,不直接学"怎么选",而是学"怎么在权重上做加减"。先构造一个软性的、带权重的连接矩阵,把它做两次幂运算(对应两跳组合),再把结果稀疏化,只保留权重最高的那几条。整个过程里,被保留下来的那些边,权重是可以求导的,梯度能顺着这些权重一路传回去调整模型参数。没被选中的边虽然被丢弃了,但保留下来的边权重的相对大小,已经足够传递足够的学习信号。
稀疏注意力:2个词也能撑起一个大模型?
有了边索引和边权重之后,接下来就是怎么用它们做注意力计算。
论文把这个环节叫做Sparse Edge Attention,简称SEA。
SEA*:稀疏边注意力,GM架构里负责根据边索引取出对应的词、再计算注意力权重的模块。
普通注意力计算的是"查询向量"跟"键向量"的匹配程度,匹配分数高说明这个位置的信息更值得参考。SEA的创新点在于,它不只看这个匹配分数,还把边权重的对数值加进去,两者相加后再做softmax归一化。
这个设计背后的逻辑,叫做product of experts,专家乘积。
product of experts*:一种融合多个"专家意见"的方式,数学上表现为把各专家给出的概率取对数后相加,再统一归一化,效果上相当于让每个专家都对最终结果有"否决权"。
打个比方,你去看病,一个医生根据检查报告判断"这个位置有问题",另一个医生根据经验判断"这类人这个部位容易出问题",两人的判断如果方向一致,会互相增强你对诊断结果的信任;如果方向冲突,哪怕一个医生特别自信,另一个的强烈反对也能把结论拉回来。SEA里的"查询键匹配分数"和"边权重"就是这两位医生,谁都不能单独说了算,必须两者都认可,这个位置才会被高度重视。
如果没有这层边权重的参与会怎样?那SEA就退化成了普通的稀疏注意力,只能凭内容临时决定看哪儿,却没法利用referral阶段已经积累下来的"关系线索"。相当于每次看病都从零开始检查,完全不参考病人过去的病历,诊断效率会大打折扣。
实测效果:用四分之三的层做减法,损失几乎没变
接下来看实际的训练结果。
作者拿阿里的Qwen3-0.6B模型当底座,这是一个6亿参数规模的语言模型,作为对照基准。他们把这个模型里75%的标准Transformer层换成了GM的稀疏层,搭配25%保留原样的稠密层,整体排布是"稀疏、稀疏、稠密、稀疏"这样的模式循环7次。然后从零开始,用157亿个词元(token)重新训练。
这里出现了两套配置,论文分别给它们起名叫Theia和Hyperion,都是希腊神话里的名字。Theia系列每个注意力头只检索2个位置,Hyperion系列检索4个位置。放到4096个词的句子里换算成比例,分别只占了0.098%和0.195%的信息量,相当于全量注意力所需计算量的千分之一左右。
结果表格显示,原版Qwen3训练到底最终损失是2.587。Theia系列里最好的配置(代号Theia-K24-R3,K代表边的数量,R代表转介轮数)最终损失比原版高0.014,可以理解成"性能略微下滑,但下滑幅度很小"。而Hyperion系列里最好的配置(Hyperion-K16-R3-S,S代表启用了额外的"稠密层刷新"机制)最终损失反而比原版低了0.003,也就是说,用四分之一的位置检索量,效果比全量注意力还略好。
更值得一提的是训练成本。这个表现最好的Hyperion-K16-R3-S模型,参数量只比原版多了11%,但"转介加注意力"这部分的估算计算量反而比原版少了19%。换句话说,它用更少的核心计算,拿到了更好的效果。
这组数字放在一起看,冲击力就出来了。原本大家默认"注意力看得越全,模型越聪明",这个假设在这里被撬动了一角:至少在这个规模、这个训练量下,把绝大多数词的关注范围砍到几乎可以忽略不计的程度,模型照样学得不错,甚至学得更好一点。
当然论文作者自己也提到,这几个模型的训练规模都不算大(6亿参数,157亿词元),而且只测了训练损失这一个指标,没有测下游任务表现,所以这个结论目前只能算是"在这个尺度下站得住",还需要更大规模的验证。
温度、刷新与重新对齐:几个让机制更稳的补丁
论文里还提到几个辅助设计,虽然听起来偏工程细节,但对效果的稳定性有实打实的贡献。
一个是温度缩放。在把多条边混合成新边的过程中,论文没有简单做加权平均,而是先给权重做一次指数运算再归一化,相当于人为拉大或缩小权重之间的差距。这一步的用意是防止边权重在多轮转介之后变得"温吞水",谁都差不太多,失去区分度。
另一个是refresh机制,可以翻译成"边刷新"。
refresh*:允许referral过程重新引入最初的边,或者引入最近一次稠密注意力层算出的边,防止多轮转介之后信息越传越偏。
这个设计解决的问题是,转介进行的轮数越多,边的指向可能会越来越偏离原始信息,类似于传话游戏传到第五个人的时候,内容已经跟第一个人说的不太一样了。定期让"最初的话"重新插进来校正一下方向,能有效防止跑偏。实验里,启用了这个机制的Hyperion-K16-R3-S,比不启用的版本损失又低了0.004。
还有一个叫realignment,重新对齐。
realignment*:在做完注意力计算之后,用注意力算出来的最终权重反过来修正一部分存储的边,让边的指向更贴合当前内容实际计算出的关联程度。
这个机制的作用,是给referral这种"预测式"的路径搭建打个补丁。referral是提前猜"这两个词大概有关系",但猜测终究是猜测,真正算出注意力权重之后,才是基于内容的实际验证。realignment就是拿这份"事后验证"的结果,去纠正之前"事先猜测"里猜错的部分。
写在后面
读完这篇论文,最让我意外的一点,其实是它对"注意力到底在做什么"这件事的重新拆解。
我们通常把注意力当成一个整体动作:模型看一眼所有词,算出一堆权重,加权求和。但这篇论文把它拆成了两件事:一件是"决定该看哪儿"(靠索引这种离散的、逻辑性的信息),另一件是"看了之后该信多少"(靠权重这种连续的、可微分的信息)。这两件事原本被揉在一次矩阵乘法里,现在被拆开,分别用整数和浮点数两种完全不同的数据类型去承载。
这个拆分带来一个挺有意思的副产品:模型的"决策过程"变得可读了。你可以直接打印出某个词此刻的边索引,看它现在"指"向哪几个词,这跟传统Transformer里那一整片密密麻麻、很难解释的注意力权重矩阵完全是两种体验。论文里也提到,这种"高度自解释的关系状态"是留给未来做机制可解释性研究的一个方向,虽然这次没有深入展开。
另外一个细节值得单独说一说:论文里提到不同配置的表现并不是单调的,比如在Theia这种更稀疏的设定下,转介轮数从3轮加到4轮,反而让损失变差了一点。这说明"多转几轮总是更好"这个直觉在这里并不成立,转介轮数、边的数量、稀疏预算这几个参数之间存在某种没被完全摸清楚的耦合关系。这大概也是这篇论文自己承认的,规模还太小,很多现象背后的原因还说不清楚。
如果一个模型可以只靠几个整数指针,就完成原本需要海量浮点数计算才能做到的事,那我们对"理解语言"这件事所需要的计算量,是不是一直被高估了?
Q&A
Q1:Graph Machine是什么?
A:Graph Machine(简称GM)是一种新型序列建模架构,它用O(n)大小的状态配合动态稀疏路由的方式访问历史信息,通过整数边索引和浮点边权重记录词与词之间的关系,让模型只检索少量位置就能完成注意力计算。
Q2:GM相比普通Transformer注意力有什么优势?
A:普通Transformer需要每个词跟历史上所有词逐一比对,计算量随句子长度平方增长。GM每个注意力头只检索2到4个位置就能完成计算,实验显示检索4个位置时损失甚至比全量注意力略低,同时估算计算量减少了约19%。
Q3:referral机制是做什么的?
A:referral(转介)是GM更新边关系的核心机制,让每个词的邻居通过邻居的邻居递归寻找新连接,类似电话转接层层传递,最终把多跳关系压缩成一条新的直接连接,帮助模型建立更准确的关注路径。