这项由谷歌DeepMind研究团队完成的研究以预印本形式发布于2026年4月,论文编号为arXiv:2604.09168,有兴趣深入了解的读者可以通过该编号查询完整论文。
深度学习模型变得越来越强大,但代价是越来越"重"——动辄数百亿个参数,需要占用大量内存,运行在昂贵的服务器上。普通人手机里的图片生成应用,能流畅运行的往往是功能大打折扣的"轻量版"。这种"要么强大要么轻便"的两难困境,困扰着整个AI领域。而这篇论文想做的,就是打破这个困境。
研究团队提出了一种叫做"弹性循环变换器"(Elastic Looped Transformers,缩写ELT)的新架构。用最简单的话来说:传统的AI模型就像一栋有很多楼层的大楼,每一层都有独立的装修和家具(也就是独立的参数)。而ELT则像一部电梯——它只有一套装置,但可以反复上下运行很多次,每运行一次,对图像的理解就更深一层。这样一来,你不需要建很多层楼,只需要让电梯多跑几趟,就能达到同样的效果。最终结果相当引人注目:在参数数量只有传统模型四分之一的情况下,ELT在标准图像生成测试中达到了相同的质量水平,并且在视频生成任务中还超越了对标模型。
一、为什么现有的AI模型会越来越"臃肿"?
要理解这项研究解决了什么问题,得先聊聊AI图像生成模型的工作方式。以目前主流的扩散模型(Diffusion Model)为例,它的工作原理可以类比为一位画家从一张布满噪点的画布开始,一遍又一遍地擦拭和修改,直到一幅清晰的图像浮现出来。这个"擦拭和修改"的过程,需要反复调用一个叫做"变换器"(Transformer)的神经网络模块来预测如何改进图像。
传统的做法是:把很多层变换器叠在一起,每一层都有自己独立的参数(可以理解为每层都有各自独立的"画笔工具箱")。层数越多,模型越强大,但占用的内存也越多。DiT-XL这类主流扩散变换器模型的参数量高达6.75亿个,而且在每一个去噪步骤中都需要把这些参数完整地调用一遍。
另一个思路是"循环"——让同一层参数被反复调用多次,就像那部反复上下的电梯。这种思路其实早就有人提出来了,最具代表性的是2018年提出的"通用变换器"(Universal Transformers)。然而,在图像和视频的高质量生成领域,循环架构一直没能真正站稳脚跟。原因在于一个棘手的问题:如果你训练模型时规定它循环8次,那它在循环3次或10次时生成的图像就会一团糟,因为中间状态的表示对模型来说毫无意义,只有跑完全部8圈,输出才是有效的。
这就像一个工人被训练成"必须拧满8圈螺丝才能停下"——如果你让他只拧3圈就停,螺丝根本没有固定住;让他拧10圈,螺丝可能就滑丝了。这种"只能在固定深度工作"的问题,让普通循环变换器失去了灵活性,也限制了它的实际应用价值。ELT要解决的,正是这个问题。
二、"弹性电梯"是怎么炼成的?
ELT的核心机制可以分成两个部分:循环机制本身,以及让循环变得"有弹性"的训练策略。
先说循环机制。研究团队把若干个变换器层组合成一个"复合块",用希腊字母Θ(theta)表示这个块里所有的参数。假设这个块由N层变换器组成,训练时让它循环L次,那么总的有效计算深度就是N×L。比如,一个有8层的块循环4次,等效于一个有32层的深层网络,但实际存储的参数只有8层的量。这样,模型的参数量由N决定,而计算能力由L决定,两者被解耦开来了。
然而,光有循环机制是不够的。如前所述,用传统方式训练出来的循环模型,中间状态(也就是电梯还没到顶层时的状态)是混乱无序的。为了解决这个问题,研究团队提出了"内循环自蒸馏"训练策略,英文缩写为ILSD。
"蒸馏"在AI领域是一个常用概念,原意是让一个小模型去模仿一个大模型的行为,从而让小模型获得大模型的能力。ILSD的创意在于:它不需要两个独立的模型,而是让同一个模型自己教自己。具体来说,在每一次训练时,模型被当作一个"双轨系统"来运行:一条轨道是"教师轨道",让模型跑满最大循环次数(比如8次),得到最高质量的输出;另一条轨道是"学生轨道",随机选一个中间循环次数(比如3次),只跑到这里就停下,得到一个中间状态的输出。
接着,训练目标包含三个部分。第一部分是让教师轨道的输出尽可能接近真实图像,这是主干任务。第二部分是让学生轨道的输出也尽可能接近真实图像,这给中间状态加上了直接的"生成质量"约束。第三部分是让学生轨道的输出尽可能接近教师轨道的输出,也就是让"跑了3圈的结果"去模仿"跑了8圈的结果"——这就是"自蒸馏"的含义。
这个设计有一个非常聪明的地方:学生轨道的计算过程,本身就是教师轨道计算过程的一个子集。教师要跑8圈,而学生是其中的前3圈。所以在一次前向传播中,两条轨道共享计算,不需要额外的计算开销。这与传统蒸馏方法形成了鲜明对比——传统方法需要分别跑教师模型和学生模型两遍,计算量翻倍。
在训练过程中,研究团队还引入了一个叫做"学生循环随机采样"的机制:每次训练时,学生轨道的中间停止点是从一个范围内随机选取的,而不是固定的某个数字。这使得模型学会了在任意中间深度都能产生有意义的输出。此外,第二部分损失(学生对真实图像的损失)和第三部分损失(学生对教师的蒸馏损失)之间的权重会随训练进程线性变化:训练初期,学生主要跟着真实答案学,因为此时教师本身还没训练好;随着训练推进,学生逐渐更多地去模仿教师,因为教师的输出质量越来越高。
通过这套训练策略,模型的每一个中间循环状态都被推向"解空间"(也就是有意义的图像表示空间)。用论文中的形象比喻:普通循环变换器就像一个只有终点才有意义的旅程,而ELT经过ILSD训练后,旅程中的每一站都变得有意义,随时可以下车。
三、从图像到视频,ELT的表现如何?
研究团队在两个主流任务上验证了ELT的效果:类别条件图像生成(给定一个类别标签,如"猫",生成对应图像)和类别条件视频生成。
在图像生成任务上,测试数据集是业界标准的ImageNet 256×256(包含超过100万张、1000个类别的图像)。衡量生成质量的指标是FID分数(Fréchet Inception Distance),这个分数越低代表生成质量越好。
研究团队实现了两种不同框架下的ELT:基于"掩码生成变换器"(Masked Generative Transformer,类似于MaskGIT框架)的版本,以及基于"扩散变换器"(Diffusion Transformer,类似于DiT框架)的版本。
在掩码生成框架下,最强的ELT-XL模型(以7层块循环4次的配置运行)实现了FID 2.0的成绩,与拥有4.46亿参数的MaskGIT-XL基线模型完全持平。而ELT-XL自身只有1.11亿参数,参数量减少了约75%。对应的ELT-L模型(12层块循环2次,或8层块循环3次)则分别实现了FID 2.1和2.2,参数量同样大幅低于基线。在扩散变换器框架下,以8层块循环4次(等效深度32层)的配置,ELT的FID为3.16,优于参数量为其4倍的32层标准DiT模型(FID 3.43)。以16层块循环2次的配置(参数量是标准32层DiT的约一半),FID更是达到2.83,进一步超越基线。
在视频生成任务上,测试数据集是UCF-101,这是一个只有约1370万个训练token的小规模数据集,对模型的过拟合防御能力有较高要求。衡量视频生成质量的指标是FVD(Fréchet Video Distance),同样是越低越好。ELT的7600万参数版本(6层块循环4次)在等算力设置下实现了FVD 72.8,优于拥有3.06亿参数的MAGVIT-L基线(FVD 76),参数量减少了约75%。如果允许更多的计算量(6层块循环6次,采样步数加倍),FVD可以进一步降低到60.8,取得更加显著的优势。值得关注的是,研究团队认为ELT在UCF-101这样数据量有限的场景中表现良好,与其循环架构的正则化效果有关——参数共享本身就起到了防止过拟合的作用。
四、"任意时刻推理":一个模型,多种计算预算
ELT最吸引人的特性之一,是研究论文中称为"任意时刻推理"(Any-Time Inference)的能力。这个词来自控制论领域,指的是一个算法能够在任何时刻被打断并输出当前最佳结果的能力。
对于普通循环变换器来说,这种能力根本不存在。如果你在它跑完8圈之前强行停下来,得到的只是一张混乱的噪点图,毫无用处。但经过ILSD训练的ELT不同:无论你在哪个循环次数停下来,输出都是有意义的图像,只不过循环次数越多质量越好。这意味着一个训练好的ELT模型可以同时服务于不同的应用场景:在高端云端服务器上,让它循环更多次,追求极致的生成质量;在手机或边缘设备上,让它循环较少次数,以换取更快的响应速度。用户可以根据自己的硬件情况和质量需求,动态调整循环次数,而完全不需要重新训练模型。
研究团队通过Pareto图(一种展示效率边界的图表)来展示这个特性:横轴是推理时需要的计算量(以GFLOPs衡量),纵轴是FID分数。对于一个固定的ELT模型,改变推理时的循环次数,就能在这张图上画出一条曲线。随着循环次数增加,计算量线性增大,同时FID稳步降低(质量提升),直到收益递减。研究团队发现,这条效率边界的拟合函数大约是FID = 1922.5 × G的负0.95次方 + 1.48(其中G代表GFLOPs),展现出非常规律的幂律缩放特性。
五、参数少了,速度反而更快?
ELT的参数效率不仅体现在模型文件更小,还带来了实际运行速度的提升。这背后有一个硬件层面的原因。
现代AI加速器(如GPU或TPU)的工作原理大致如下:它有一个非常快但容量有限的"片上内存"(SRAM),以及一个容量很大但读写速度慢得多的"外部高带宽内存"(HBM)。如果模型参数太多,就必须频繁地在HBM和SRAM之间来回搬运数据,这种"内存墙"瓶颈会严重拖慢速度。
由于ELT的参数量大幅减少,在很多情况下,整个模型的参数可以完整地放进片上内存,完全避免反复的数据搬运。研究团队在谷歌TPU v6e上测试了这一效果:对于L规模的ELT模型,吞吐量是同等计算量基线模型的2.9倍;XL规模的是3.3倍;H规模(最大测试规模)的是3.5倍。唯一例外的是B规模(最小的ELT版本),因为即使是B规模的基线模型本身也已经足够小,能完整放入内存,所以ELT在这个规模上没有速度优势。
此外,研究团队还发现ELT在训练收敛速度上也有明显优势。在等算力设置下,16层块循环2次的ELT比32层基线DiT快2倍达到相同的FID水平;8层块循环4次的ELT则快约1.4倍。这意味着训练成本也随之降低。
六、深入挖掘:缩放规律和失效场景
研究团队不只是展示了几个成功案例,还系统性地探索了ELT的缩放规律,以及它在哪些情况下会失效。
在缩放规律方面,研究发现:增加模型宽度(也就是每一层变换器的维度d)是提升质量最有效的手段;在固定宽度的前提下,增加循环次数可以稳步提升质量,但收益会递减。当循环次数增加到一定程度后,切换到更宽的模型比继续增加循环次数更有效。一个极端情况揭示了循环架构的本质限制:以1层块循环32次的配置(等效深度32),FID高达10.30,远差于任何有意义的对比基线。这说明单个变换器层的表达能力不足以支撑高质量生成,即使循环再多次也无济于事。模型需要在每次循环中处理一定复杂度的变换,才能有效地推进表示质量。
在失效场景方面,研究团队坦诚地指出了两个主要问题。其一,如前所述,当唯一层数N过少时,无论循环多少次都难以生成高质量内容。其二,当推理时的循环次数L远超训练时的最大循环次数L_max时,质量也会下降——模型的共享参数在超出训练范围的迭代次数下会"过度迭代",脱离已训练的收敛区域。有趣的是,研究团队发现ILSD的训练策略让模型具备了一定程度的"超范围外推"能力:在UCF-101视频生成实验中,一个以最大4圈训练的模型,在6圈时仍然能获得比4圈更好的FVD(69.2 vs 72.9),说明ILSD确实对迭代过程进行了有效的规律化约束。不过研究团队也指出,这种外推能力的边界和规律还需要更多研究。
归根结底,ELT这项研究做的事情,可以用一句话概括:它证明了AI生成模型不需要靠堆砌独立参数来变强,让同一组参数反复"深思熟虑"同样能达到高质量的结果,而且在合适的训练策略下,这种反复思考的过程可以在任意时刻输出当下最好的答案。这对于在资源有限的设备上部署高质量生成AI,以及动态调配算力资源,都有相当直接的实际意义。对于扩散模型的进一步加速,研究团队也提到了一个有前景的方向:现有的扩散模型在每个去噪步骤上都分配相同的算力,而ELT可以通过调整不同步骤的循环次数,把更多算力集中在"最需要精细调整"的步骤上,实现更智能的算力分配。
Q&A
Q1:弹性循环变换器和普通AI图像生成模型有什么根本区别?
A:普通AI图像生成模型里,每一层神经网络都有自己独立的参数,层数越多参数越多,占内存越大。弹性循环变换器则是把一组神经网络层反复循环使用,参数只存一份,但可以多次运行。训练时用一种叫"内循环自蒸馏"的策略,让模型在任意循环次数下都能输出有意义的图像。这样一来,用传统模型四分之一的参数,就能达到同等的生成质量,而且可以根据设备性能灵活调整循环次数。
Q2:内循环自蒸馏训练是怎么运作的?
A:内循环自蒸馏(ILSD)的核心思路是:在一次训练中,让模型既以完整循环次数运行(教师轨道),也随机选一个中间循环次数运行(学生轨道)。学生轨道的输出要同时接近真实图像和教师轨道的输出。因为学生的计算过程本身就是教师计算的一个子集,不需要额外的计算开销。随着训练推进,学生越来越多地去模仿教师,最终让每个中间循环状态都能输出有质量的结果。
Q3:弹性循环变换器在手机等低算力设备上真的能用吗?
A:从原理上看有很大潜力。弹性循环变换器参数量只有对标模型的四分之一,更容易放进设备内存,而且可以通过减少循环次数来降低计算量。实验中在TPU上已经测到最高3.5倍的吞吐量提升。具体到手机端部署,还需要结合量化、编译优化等工程手段,但ELT提供的参数效率优势确实为低算力高质量生成提供了更可行的技术路径。
热门跟贴