|
|
* n9 B) u& i# D$ l+ h. O2 G
( _) X- q, d- b! ~! B这项由上海交通大学、西安交通大学与华为联合完成的研究,于2026年6月30日以预印本形式发布,论文编号为arXiv:2606.29215v2,感兴趣的读者可通过该编号查阅完整原文。0 j; h4 \% p" P9 R% H$ m
" U: B( M% j; t5 E7 F& A, E
当你在手机上向AI助手提问,或者让大模型帮你写一段代码、解一道数学题时,你有没有想过,它生成那些文字的时候,究竟是怎么工作的?大多数人熟悉的那种AI,比如ChatGPT,是像打字员一样,一个字一个字地往外蹦,每输出一个词都要等前一个词定稿之后才能开始。这种方式虽然直观,但速度上有天花板——毕竟,你总不能让打字员同时打两个字吧。
6 C# \ n3 F% T5 t& h! {正是为了突破这个瓶颈,一类叫做"扩散语言模型"的AI生成方式应运而生。这类模型的思路截然不同:它不是逐字输出,而是先把整段文字打满问号,然后一步一步把问号"洗掉",最终得到完整的答案——有点像先画出模糊草图,再一遍遍擦掉不清晰的地方,直到图画变得清晰为止。这种方式天生适合"批量处理",可以同时推敲一大段文字,理论上速度快很多。9 T. x; F: B, H1 C' D7 ?0 G, F! w
然而,理论归理论。现实中的扩散语言模型在实际部署时遇到了一个尴尬问题:它们虽然能同时处理一段文字,但段与段之间依然是排队等候的关系——第一段画完了,才能开始画第二段,第二段画完了,才能开始存档,然后才能开始第三段。这就像一条流水线上,明明有多台机器,却只能让一台在工作,其余都在等待,浪费了大量时间。
% p) A8 i- }0 T) s1 q' t这篇论文的核心贡献,就是打通了这条流水线,让多个"段落"可以同时处理,真正实现了批量并行。研究团队将这套新方法称为"多块扩散语言模型"(MBD-LMs),并配套提出了训练策略和推理引擎两大核心技术。实验数据表明,在数学推理和代码生成等典型任务上,新方法在保持生成质量基本不变的前提下,每步前向计算处理的有效词数提升了约78%,整体吞吐量也显著提升。, j* D4 }2 |. s* P4 f
一、流水线上的"等待气泡"——问题究竟出在哪里
" I2 J2 p" _# [( y. I/ Z要理解这项研究解决了什么,先得搞清楚原来的系统卡在哪里。4 b' r" o1 D6 f5 C
现有的"块扩散语言模型"(BD-LM)把输出文本切成若干"块",每一块相当于生产线上的一个工件。模型每次只负责加工当前这一块,加工完了,把它"存档"(写入KV缓存,这是一种加速计算的内存机制),然后才能开始下一块。这个流程本身是清晰的,但问题就出在"存档"这个步骤上。! u" [% R9 q: R! e: }
具体来说,存档是需要时间的。在存档这段时间里,模型什么都不能做——既不能继续打磨刚完成的这块,也不能提前预热下一块。这段空白时间,研究团队称之为"存储气泡",就像工厂流水线上的停工等待。如果能在存档的同时,让模型提前开始处理下一块的初步加工,不就把这段空白时间利用起来了吗?
" v$ K7 I5 y/ F$ a% C) _这正是"多块扩散"(MultiBD)的基本思路:维护一个"运行窗口",窗口里同时容纳多个正在加工中的块。前一块加工完成、正在存档时,后一块已经在炉子上预热了。存档完成的块悄悄退出窗口,新的块悄悄进来。整条流水线始终保持忙碌,没有等待的空隙。
: H" n2 |7 Q) _( {% E这个思路听起来简单,但实现起来有一道难关:原有的模型在训练时,根本没见过"窗口里同时有多个噪声块"这种情况。它只见过一种场景——前面的块都是干净清晰的,当前只有一块是模糊的。突然让它同时面对多块模糊的文字,它会手足无措,生成质量会大幅下滑。这就是论文着力解决的"训练与推理的错位"问题,也是整个研究的核心矛盾所在。0 @! J& i q. S0 t+ |0 b
二、现有方案的局限——为什么"差不多"并不够用( S& w( t$ t2 A$ R! P7 @8 t
在这篇论文之前,学界已经有人注意到了这个问题。一种叫做"离散扩散强迫"(D2F)的训练策略,尝试让模型在训练时也见到多个噪声块共存的场景。它的做法是:取一段完整文字,把后半段按块切开,给每块分配一个递增的噪声比例——越靠前的块越清晰,越靠后的块越模糊,然后让模型同时面对这整段带噪声的后缀来学习。
7 n: k9 |8 X7 I6 M$ f这个方向是对的,但研究团队发现它的训练场景和真正的多块推理场景之间,存在两个关键的错位。( T8 t& ?# P1 [6 a% s1 m2 V) M
第一个错位是"窗口大小"。在真实的多块推理过程中,模型每次处理的窗口是有限的,通常只有两到三个块同时活跃。但D2F训练时一下子把整个后缀都塞进来,窗口可以非常大,这和实际推理时的受限窗口不匹配。
3 X R# {3 |/ R4 y o; h第二个错位是"噪声分布的形状"。在真实推理中,窗口里相邻两块之间的噪声差异往往很大——前一块已经几乎清晰了,后一块还是一团模糊。但D2F采用的是单调递增的噪声调度,相邻块之间的噪声比例差距很小,彼此之间的区分度不够。打个比方,真实推理就像流水线上,前一道工序刚刚收尾、后一道工序还是毛坯,两者差距显著;而D2F的训练场景,更像所有工序都处于"差不多完成一半"的状态,彼此之间看不出谁先谁后。 S0 Z* F- f ^" e4 c
这两个错位叠加在一起,导致用D2F训练出来的模型虽然能应付多块推理,但生成质量下降明显。研究团队通过统计数据清晰地展示了这一点:D2F诱导的各块噪声比例分布高度重叠,相互混淆,而真实推理中各块的噪声比例分布则泾渭分明,差异显著。' u1 P* K8 l, S, q4 V" V
三、核心训练方法——给模型"排练真实场景"
% \, Q6 Y/ U- W研究团队提出的解决方案叫做"多块教师强迫"(MultiTF)。这是一种后训练方法,意思是先用原来的方式训练好一个基础模型,然后用MultiTF对它进行"针对性补课",让它学会在真实的多块推理场景下也能表现良好。
" Z! p+ E( j6 @0 }+ dMultiTF的核心思想是:既然训练场景和推理场景不匹配,那就在训练时精心构造那些和推理场景高度相似的状态,让模型充分见识到这类场景,从而学会如何应对。
* _7 A: ^6 Z& r, b+ L+ J# Q具体怎么构造呢?首先,要把训练文本切分成"噪声组",每个噪声组就是一小批相邻的块,对应推理时窗口里同时活跃的那几个块。切分的方式分两类:一类是"系统化布局",按照固定间隔和偏移量枚举所有可能的切分方式,确保每个块都能出现在组内的每个相对位置上;另一类是"随机布局",随机决定每个组的大小和边界,增加多样性。这两类布局结合使用,能全面覆盖各种可能的推理窗口状态。9 ^+ e( n7 v6 |# Q% W+ E" ?
其次,最关键的是噪声分配方式。研究团队专门设计了一种叫做"链式均匀噪声调度器"的机制。对每个噪声组,先随机抽一个基础噪声水平作为"地板",然后第一块在地板到上限之间随机取一个噪声比例,这个比例成为第二块的新"地板",第二块再在新地板和上限之间随机取值,如此链式递进。这样做的结果是,组内相邻两块之间的噪声差距往往很大,而且是随机的——有时候第一块八九成清晰、第二块还有六七成模糊,有时候差距更大。这种"大且随机的噪声落差",正好模拟了真实推理中相邻活跃块之间的状态差异。
6 B; e( ~* f+ a3 x3 t在注意力机制的设计上,研究团队还引入了"组感知双流注意力掩码"。组内的噪声块可以互相看见对方(就像同一个生产批次的工件可以参照彼此),每个噪声组只能看见它之前已经清晰存档的前缀块(就像当前批次的工人能看到已完成的历史档案),而干净的前缀块之间保持原有的因果关系,不会被噪声块影响。这套掩码设计,使得模型在训练时能够精准模拟推理时的信息可见范围。
/ k _( {: W" v& L4 [$ \训练目标是经典的掩码词元交叉熵损失,只在噪声序列中被掩盖的位置计算损失,确保模型专注于学习如何在多块噪声场景下恢复正确的内容。对于那些本身就有额外训练目标的模型(比如带有DMax自去噪机制的版本),MultiTF也能无缝兼容,在相同的输入结构上叠加对应的模型专属损失。% k# l0 H8 @* d9 s7 M, }
研究团队还从理论层面为MultiTF提供了支撑。他们证明,通过系统化的偏移布局,每一种可能出现在推理窗口中的有界连续块组合,都会在训练时至少出现一次。训练分布与推理分布之间的差距,可以被分解为"窗口组合分布的差距"和"噪声比例分布的差距"两部分,两者都随着MultiTF的设计改进而缩小。这给了该方法严密的理论保证,而不仅仅是经验上的有效。% w* ~/ j4 M* N7 A* ~
四、推理引擎设计——让并行计算真正落地
4 T6 J" K$ h5 n训练好了模型,还需要一套与之配套的推理引擎,才能把多块并行的理论优势转化为实际的速度提升。这里的挑战同样不小。- x. q4 b9 |4 B B, }" _
多块推理的本质是维护一个动态变化的运行窗口:前面的块加工完就退出,新的块进来激活。如果直接按照窗口的实际内容来构造每次前向计算的输入,那么每次前向计算的输入大小都会变化——有时窗口里有两个活跃块,有时有三个,有时还有一个已经完成但还没存档。这种动态变化的输入形状,对GPU的高效执行是个噩梦。现代GPU依赖一种叫做"CUDA图"的机制来加速重复计算——它的前提是每次计算的形状完全固定,一旦形状变化,这个加速机制就失效了。
9 r+ }; h: ?+ |研究团队为此设计了"块缓冲区"机制。核心思路是引入一个固定大小的"缓冲区",缓冲区里永远有固定数量的槽位。每个槽位有四种状态,按时间顺序依次经历:空闲(占位但不参与计算)、活跃(正在被推敲加工)、待存档(加工完毕、等待写入缓存)、已存档(已写入缓存,以后通过缓存访问)。( T- t& X0 S( D8 J# @, Y# U _1 A$ Z
关键在于"空闲"槽位的设计。当一个新的块需要进入窗口时,它不是去延长输入序列,而是激活缓冲区里已经存在的空闲槽位。当前面的块完成存档后,缓冲区向前滑动,在末尾补充一个新的空闲槽位。整个物理输入序列的长度始终是"缓冲区大小 × 块大小",永远不变。CUDA图因此能够一次捕获、反复回放,极大降低了运行时开销。
' J& O5 O% G" V5 h7 B, J8 A% y) a除了静态形状,这套推理引擎还仔细处理了前缀缓存的问题。已经完成并存档的块,其KV状态会作为稳定的前缀缓存保留,后续所有前向计算都能直接复用,不需要重新计算。活跃块每次推敲时只需处理缓冲区内的内容,再加上对前缀缓存的查询。这套机制和传统BD-LM的前缀缓存接口完全兼容,也就是说,之前为单块推理积累的系统优化成果,在多块推理中可以直接延续。, H7 U/ C; s' I
推理过程中还有几个阈值参数控制窗口的推进节奏。比如,只有当前最后一个活跃块的"完成进度"超过设定阈值时,才会激活下一个空闲槽位让新块进来,避免过早引入新块导致质量不稳定。另有"半完成阈值",允许后一块在前一块还没完全存档时,先参考前一块的最优预测状态,进一步减少等待。* t# r! b7 E5 Q# p
五、实验结果——数字背后的真实提升; J# f# [- W, [) f @
研究团队在数学推理(GSM8K、MATH500)和代码生成(MBPP+、HumanEval+)四个基准测试上进行了系统评测,对象涵盖了LLaDA2系列和SDAR系列多个模型变体。, `. Z T9 w( c3 e( c
以LLaDA2-Mini为基础的实验最具代表性。原始的单块推理模式下,平均每步前向计算能贡献约3.47个有效词。经过MultiTF后训练、采用多块推理后,这个数字提升到了6.19,增幅约78%。与此同时,四个任务上的平均准确率不降反升,从79.95%提升到了81.03%。也就是说,速度快了将近八成,质量还略有改善,这种"两全其美"的结果在系统优化领域并不常见。
9 F$ o' B0 V( p/ b' K! w与"不训练直接多块推理"的对照相比,差异更加直观。直接把多块推理施加在未经训练的基础模型上,虽然每步有效词数从3.47提升到了4.41,但准确率从79.95%跌到了78.59%。而MultiTF后训练不仅把有效词数进一步推高到6.19,还把准确率拉回到81.03%,完整体现了"训练与推理对齐"的价值。
; o0 `1 T1 ]# H+ G- T当MultiTF与另一种加速技术DMax(一种提升单步有效词数的token级加速方法)结合使用时,效果进一步叠加。MBD-LLaDA2-Mini-DMax的平均每步有效词数达到9.34,相比对应的单块推理版本提升了47%,准确率仅下降约1个百分点。吞吐量方面,平均每秒处理词数从779.50提升到926.67,增幅约19%。. X* Q' I0 ~: h$ I
从实际延迟数据来看,多块推理的代价是可控的。单块推理时每步耗时约7.07毫秒,多块推理后上升到8.78毫秒,约增加了24%的时间成本。而有效词数提升了约78%。两者相除,实际吞吐量提升约44%,与实测数据高度吻合。这也从侧面说明,多块推理的大部分计算额外开销,来自缓冲区内空闲槽位和已完成但未存档的块占用的计算资源,属于为保持静态形状而必须付出的代价;在这个代价之上,有效词数的大幅增长最终带来了可观的实际速度提升。- r5 G3 Z5 V: h! }
在SDAR系列模型上,多块推理同样表现出良好的迁移性。SDAR-8B-Chat使用块大小为32的配置时,每步有效词数从2.54提升到4.46,增幅超过75%,准确率也略有提升。即便是块大小仅为4的极小粒度配置,多块推理依然能将有效词数从1.25提升到2.42,准确率几乎持平。
5 T/ W4 K& O1 i( ^1 g v: ]* `& F六、消融实验——逐个验证每个设计选择的必要性
V/ V7 w$ k+ t, P# E0 S; x为了检验每个设计要素的具体贡献,研究团队对MultiTF的各个组件做了系统的消融分析,以LLaDA2-Mini-DMax为对象,在GSM8K和HumanEval+上综合评估。
) V; d6 v) k; e$ Q p9 i在组布局的构造方式上,单独使用系统化布局可以将每步有效词数从6.57推至9.71,准确率为83.22%;单独使用随机布局可以推至9.42,准确率82.72%。两者都比原始单块推理快了很多,但准确率都有一定下滑。把两类布局结合起来,每步有效词数进一步提升到9.87,准确率恢复到84.59%,几乎与原始单块推理的84.67%持平,同时速度提升了50%以上。这说明系统化布局和随机布局并不是互相替代的关系,而是互补的:前者保证了对所有有界窗口组合的完整覆盖,后者提供了结构之外的分布多样性,两者缺一不可。
4 c2 n' U0 L$ ? ~$ b在噪声调度器的对比中,差异更加显著。用D2F风格的单调调度器替代链式均匀调度器,每步有效词数提升到8.76,但准确率暴跌到79.34%,下降了超过5个百分点。这印证了前面的分析:仅仅让模型见到多个噪声块共存还远远不够,噪声分布的"形状"不对,模型照样会出错。
2 E* T' F6 o" X. O5 V2 J随机调度器(在整个范围内独立随机采样,不保证单调性)和排序均匀调度器(独立采样后排序保证单调性)的表现介于中间:两者的每步有效词数都超过了9.7,但准确率分别为83.14%和81.28%,均不如链式均匀调度器的84.59%。排序均匀调度器的问题在于,它虽然产生了单调的噪声序列,但相邻块之间的噪声差距是由均匀分布的次序统计量决定的,差距通常较小;而链式均匀调度器通过"前一块的噪声成为下一块的地板"这一机制,主动强迫相邻块之间的差距扩大,更贴近推理时相邻活跃块之间的实际状态落差。( L6 ]: D9 X7 N" q6 B
七、从吞吐量到墙钟速度——效率分析的完整逻辑4 M' ^; Y2 k# |! U* I
研究团队还对推理效率做了详细的理论分析,把"每步有效词数"与"实际每秒处理词数"之间的关系梳理得非常清楚。
) f, X; n N* C& u4 p) t4 i& f每步前向计算的工作量由两部分决定:一是处理缓冲区内所有词元(包括活跃块、待存档块和空闲块)时的计算量,这部分与缓冲区大小成正比;二是在前缀缓存上的注意力查询,这部分与已存档的前缀长度成正比。当缓冲区变大时,每次前向计算需要更多的算力,但查询前缀缓存的开销被分摊到了更多的词元上,每个词元平均承担的计算强度实际上更高,这使得GPU的计算单元利用率更高,整体效率更好——这就像一次性处理多件衣服的洗衣机比逐件手洗更高效,虽然每批次时间稍长,但平均到每件衣服的时间反而更短。! }2 I" D9 N; g" z5 a% J
真正的吞吐量增益,等于"有效词数增益"除以"单步延迟增益"。以LLaDA2-Mini为例,有效词数增益约1.78倍,延迟增益约1.24倍,理论吞吐量增益为1.78/1.24≈1.44倍,与实测的1.44倍完全吻合。这种高度一致性,说明研究团队对系统行为的分析和建模是准确的,也给后续在不同硬件和模型规模上的估算提供了可靠的分析框架。/ [' G: ^/ G3 i( g7 N% `% @
说到底,这篇研究解决的是一个很实际的工程问题:同样的模型,通过更聪明的训练方式和推理调度,让它在生成文字时不再"一块一块地等",而是"多块流水线地跑"。它没有改变模型的基本结构,没有引入大量新的参数,只是通过后训练和推理引擎的优化,把原本浪费的"存储气泡"时间转化成了实际产出。* {7 D3 i7 X- ~# u' W" a, N# ^- E
对普通用户来说,这意味着同样的AI模型在做数学题、写代码时,响应速度可以更快,而质量不会打折。对AI开发者和部署方来说,这意味着同样的算力可以支撑更多的并发用户,降低服务成本。
; j% X. y$ T9 |% @: c* f' P当然,这项研究目前的实验对象还集中在特定的扩散语言模型系列上,多块窗口大小在实验中也保持在2到4块的范围内。随着模型规模增大、窗口进一步扩展,如何保持质量稳定和系统效率,还有进一步探索的空间。有兴趣深入了解的读者,可以通过arXiv编号2606.29215查阅完整论文和实验细节。
+ M& R( E1 q! V6 z& a" D5 jQ&A' E7 Y1 V8 M% G$ u% C
Q1:多块扩散语言模型和普通扩散语言模型有什么区别?7 L$ s7 f9 s- ?6 ^
A:普通扩散语言模型每次只能对一块文字进行推敲,上一块完成并存档后才能开始下一块,中间存在等待空隙。多块扩散语言模型通过维护一个滑动窗口,允许多块同时处于不同的加工阶段,前一块存档的同时后一块已经开始处理,消除了等待气泡,从而提升整体吞吐量。
" ^1 b% N* w6 Z6 _6 H% qQ2:MultiTF后训练为什么不直接用D2F训练方法替代?
3 h z/ d; A! P7 i) KA:D2F训练时把整段文字的所有后缀块一次性引入,噪声比例单调递增但差距很小,与真实多块推理中相邻块之间噪声落差大、窗口有限的特点不匹配。MultiTF专门设计了有界噪声组和链式均匀噪声调度,使训练场景与实际推理状态高度对齐,因此在质量上明显优于直接用D2F。7 t; d/ a# d/ ^
Q3:块缓冲区机制是如何保持输入形状固定的?. o |' f' ~0 G2 q+ o
A:块缓冲区预先分配固定数量的槽位,新块进入时激活已有的空闲槽而不是延长输入序列,完成的块转入缓存后在末尾补充新空闲槽,整个缓冲区的物理大小始终不变,因此每次前向计算的输入形状固定,可以被CUDA图一次捕获后反复高效回放。 |
|