跳转至

ELT: Elastic Looped Transformers for Visual Generation

会议: ECCV 2026
论文: ECCV 官方页面
领域: 图像生成
关键词: 循环 Transformer、参数共享、循环内自蒸馏、弹性推理、视觉生成

一句话总结

ELT 让一组共享 Transformer 层在每个生成步骤内反复计算,并用完整循环输出自蒸馏中间出口,使同一模型能按预算调整循环次数;在等推理计算量下,ImageNet 上以 111M 参数达到 FID 2.0,而对应 MaskGIT-XL 基线需要 446M 参数。

研究背景与动机

MaskGIT、MAGVIT 和 DiT 虽然分别处理离散视觉 token 或连续潜变量,但都需要反复调用一个生成网络。通常,增加这个网络的计算深度意味着堆叠更多独立 Transformer 层,参数存储与权重搬运也随之增加。循环网络提供了另一种方式:保留一小组层,多次使用同一组权重,得到更深的有效计算路径。它减少的是独立参数,而不是让这些重复计算凭空消失。

困难在于,普通循环模型只要求最后一次循环的输出正确,中间状态完全可以是不可直接解码的过渡表示。一个训练时循环 8 次的模型,不一定循环 4 次就能生成稍差但可用的图像;提前退出可能直接破坏生成结果。因此,仅做权重共享可以得到小模型,却不能自然得到可随预算伸缩的模型。本文要解决的是如何让同一组共享参数在多个出口都能完成生成预测,而非为每个深度分别训练一套模型。

完整循环与提前退出并不是两个独立网络,后者本来就是前者计算轨迹的前缀。ELT 利用这一关系,把完整循环的预测作为在线教师,用真实标签和教师预测共同约束随机选中的中间出口。核心 idea:把循环块训练成每重复一次都能继续改善预测的迭代细化器,以循环内自蒸馏将参数效率转化为可用的弹性推理能力。

方法详解

整体框架

输入是当前生成步骤的掩码 token 或带噪潜变量,以及类别等条件;输出仍是原生成算法需要的 token 分布或去噪预测。ELT 不替换外层采样器,而是把每一次网络调用内部的独立深层堆叠,替换为共享块的多次循环。训练时跑到最大循环次数,同时监督一个中间出口;推理时根据预算选择出口。

整个过程依次涉及共享块循环、随机前缀出口和课程式循环内自蒸馏。下图实线表示网络计算及推理输出,虚线表示训练时的出口选择与监督关系;教师不是另一个独立参数模型。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前生成步骤输入<br/>掩码或带噪表示"] --> B["共享块循环"]
    B -->|重复使用同一组权重| B
    B -->|推理时选择循环次数| E["共享预测头<br/>更新采样状态"]
    B -.->|训练时选中间状态| C["随机前缀出口"]
    C -.->|学生预测| D["课程式循环内自蒸馏"]
    B -.->|完整循环教师预测| D
    Y["真实监督目标"] -.-> D
    D -.->|训练梯度| B

这里必须区分两种迭代:外层是生成算法的采样步骤,内层是单次网络调用中共享块的循环次数。减少内层循环会降低每一步的网络计算量,但不会自动减少外层采样步数。论文将这条加速轴与少步扩散区分开,也把两者结合视为未来可能性。

关键设计

1. 共享块循环:用重复计算替代重复存储权重

设一个共享块包含 \(N\) 个不同的 Transformer 层,单个生成步骤重复该块 \(L\) 次,其有效深度为 \(D=N\times L\)。块内各层有不同权重,但不同循环之间复用同一个块。因此,\(8\times4\) 配置走过 32 层等价的计算路径,却只保存 8 个独立层的权重;这与直接堆叠 32 个独立层不同,也不等于只执行 8 层计算。总模型还包含输入输出等组件,实际参数量应以实验表为准。

这种结构把物理参数规模与计算深度分开,却不能无限用循环取代容量。共享块太小,所有变换都必须由极少数参数反复完成,可能缺少足够丰富的表示能力;表 2 中单层循环 32 次的 FID 明显较差。另一方面,减少权重集合也可能减少片外存储访问,因此等 FLOPs 不一定等延迟,但是否获益取决于模型大小、硬件缓存和实现,而非循环结构的无条件保证。

2. 随机前缀出口:从一次完整前向中抽出学生

训练固定最大循环次数 \(L_{\max}\),每次迭代在 \(L_{\min}\leq L_{\mathrm{int}}<L_{\max}\) 的整数范围内均匀抽取一个中间出口。完整循环是教师,抽中的中间状态经共享预测头产生学生输出;学生没有另一套参数,也不需要从输入重新跑一遍。随着训练进行,不同出口轮流被监督,模型不能只把全部生成能力集中到固定终点。

学生轨迹是教师轨迹的严格前缀,所以复用的是已经计算的隐状态,而不是单独训练一个小模型再部署。论文强调额外计算很小,应理解为不需要额外的学生骨干前向,而不是宣称中间预测头、损失与梯度管理完全免费。这个随机抽样过程只发生在训练中;推理时由使用者选择循环预算,论文没有提供基于图像难度自动决定退出时刻的控制器。

3. 课程式循环内自蒸馏:先信真实标签,再信成熟教师

Intra-Loop Self Distillation(ILSD)同时保留完整循环的真实目标损失,并让学生接受真实目标与教师输出两种监督。教师预测作为蒸馏目标时停止梯度,避免它仅为了迎合学生而变化;但教师路径自身仍通过真实目标损失训练,且两条路径最终更新同一组共享参数。停止的是蒸馏目标这条梯度通路,不是冻结一个独立教师模型。

学生的真实目标项权重为 \(\lambda\),教师蒸馏项权重为 \(1-\lambda\),训练中将 \(\lambda\) 从 1 线性降低到 0。开始时教师尚不可靠,先用真实标签稳定中间出口;后期教师形成较成熟预测,再让学生学习其分布或去噪输出。这样既促使早期循环尽快完成有效变换,又保留后续循环继续细化的压力,而不是要求每次循环输出完全相同。

一个完整示例

以 \(N=4\)、\(L_{\max}=8\) 的配置说明计算流程,并假设某次训练抽到 \(L_{\mathrm{int}}=3\)。这只是机制示例,不是新增实验结果。当前带噪图像进入共享块后,循环 3 次得到经过 12 层等价计算的学生状态,继续使用同一组 4 层权重,直到循环 8 次得到经过 32 层等价计算的教师状态。

共享预测头读取两个出口。完整出口对真实去噪目标负责,中间出口同时参考真实目标和停止梯度后的教师预测;一次联合反向传播更新这组共享权重。下一次训练可能抽到另一个出口,因而逐渐覆盖不同计算预算。

部署时,若选择循环 4 次,就在每个外层采样步骤内执行 16 层等价计算,再把预测交还原采样器;选择循环 8 次则执行 32 层等价计算。它们使用同一套参数,但前者不是跳过外层的一半采样步骤,也不能保证所有图像质量不变。任意时刻推理在这里是对已评估预算范围的经验能力,而不是每个正整数循环次数都可靠的理论保证。

损失函数 / 训练策略

离散掩码生成对被掩码位置计算交叉熵:真实目标项使用原 token 标签,蒸馏项使用完整循环教师在整个词表上的软分布。这样学生不仅知道正确 token,还能学习教师对其他候选 token 的相对判断。图像使用 MaskGIT tokenizer,将 \(256\times256\) 图像编码为 \(16\times16\) 个 token;视频使用 MAGVIT tokenizer,将 \(16\times128\times128\) 序列编码为 \(4\times16\times16\) 个 token,词表大小均为 1024。

连续扩散采用随时间变化的 sigmoid 加权均方误差:真实目标项比较预测与干净潜变量,蒸馏项比较学生与教师去噪预测。缓存中的损失公式存在 PDF 提取损坏,此处依据相邻正文说明机制,不把猜补的排版当作作者的精确公式;复现时应核对论文原版。扩散模型用 Stable Diffusion v1.4 VAE 将图像压缩到 \(32\times32\times4\) 潜空间,采用 shifted cosine 噪声调度。

掩码模型默认训练 270 个 epoch,推理使用余弦解掩码调度;图像训练有 10% 的批次丢弃类别条件,以支持无分类器引导。扩散模型使用 Adam、batch size 512,训练 500K 步,默认采用 512 步 DDPM 和引导强度 3.0。循环预算与采样预算应分别记录,不能把训练循环数直接当作表中的生成步数。

实验关键数据

主实验

下表摘自原文表 1,均为 ImageNet \(256\times256\) 类别条件图像生成,保留同框架的直接比较。FID 越低越好,IS 越高越好;GFLOPs 与采样步数按原表报告。

模型 FID IS 参数 采样步数 GFLOPs
MaskGIT-L 2.1 270.1 303M 24 3.7k
ELT-L,8 层 × 3 循环 2.2 254.3 101M 24 3.7k
ELT-L,12 层 × 2 循环 2.1 281.8 152M 24 3.7k
MaskGIT-XL 2.0 294.8 446M 24 3.9k
ELT-XL,7 层 × 4 循环 2.0 266.1 111M 24 3.9k

ELT-XL 用约四分之一参数匹配了 FID,但 IS 从 294.8 降至 266.1,不能写成所有指标均持平。3.9k GFLOPs 相同也说明这里的核心结果是参数效率,而不是四倍 FLOPs 加速。

消融实验

原文表 2 在扩散框架下比较独立层数与循环次数。以下各配置有效深度均为 32,可用于观察固定计算深度下的容量分配;它不是去掉 ILSD 的单因素消融。

配置 FID 参数 有效深度
DiT,32 个独立层 3.43 2.1B 32
ELT,1 层 × 32 循环 10.30 69M 32
ELT,4 层 × 8 循环 3.96 271M 32
ELT,8 层 × 4 循环 3.16 539M 32
ELT,16 层 × 2 循环 2.83 1.1B 32

过度共享会损害容量:单层重复并没有换来与深层网络等价的质量。8 层与 16 层配置则在减少参数的同时优于 32 层 DiT,表明关键是合适的共享块容量,而不是循环次数越多越好。

原文表 4 给出 UCF-101 类别条件视频结果,不使用引导;最后一行同时增加循环次数和采样步数,不属于等计算量比较。

模型 FVD IS 参数 采样步数 GFLOPs
MAGVIT-L 76 ± 2 89.27 ± 0.15 306M 12 约 4.3k
ELT,6 层 × 4 循环 72.8 ± 2.5 88.27 ± 0.33 76M 12 约 4.3k
ELT,6 层 × 6 循环 60.8 ± 2.7 87.88 ± 0.39 76M 24 约 13k

关键发现

  • 图 8 比较有无 ILSD:普通循环模型离开训练深度后质量明显退化,ILSD 改善多个出口的可用性。缓存图中文字交错,此处不从模糊曲线读出额外精确数值。
  • 图 7 及第 5.4 节报告:视频模型训练最大循环次数为 4,却在推理循环 6 次时达到 FVD 69.20。该值属于深度外推分析,不与表 4 的 60.8 混为同一配置。
  • 表 3 在 TPU v6e、1 × 1 拓扑、batch size 8 下测得 XL 配置吞吐比为 3.3,H 为 3.5,B 为 1.0;收益明显依赖模型规模及内存访问条件。
  • 图 6 的 2.0 倍与 1.4 倍描述达到相近 FID 所需训练步数的收敛优势,不应直接解读为相同比例的端到端墙钟训练加速。

亮点与洞察

  • 让共享参数同时承担“尽早可用”和“继续改进”两个目标。 只监督末端允许中间状态任意漂移,随机出口监督则把可解码性提前。这比单纯减少独立层数更能解释弹性能力的来源。
  • 把蒸馏成本压到既有轨迹里。 学生不另跑骨干,而是读取教师前缀状态。这一思路适合其他重复细化网络,但需要确认浅出口与深出口的预测空间一致。
  • 区分模型容量、每步深度和采样步数。 三者分别影响参数存储、单次调用计算和生成轨迹长度。ELT 提供的是可独立调节的内层深度轴,不能把所有效率改善都归结为少步采样。

局限与展望

  • 任务范围有限。 实验主要是 ImageNet 类别条件图像与 UCF-101 短视频,没有证明文本到图像、长视频或更高分辨率部署同样有效。
  • 循环外推不是普遍保证。 作者明确认为超出训练深度后的行为仍需跨数据集和规模研究;过小共享块也已暴露容量瓶颈。
  • 硬件证据有明确边界。 吞吐测量来自特定 TPU 和批大小,参数减少不代表激活、反向传播存储或所有 GPU 延迟同比下降。
  • 自适应调度仍是展望。 按去噪步骤分配不同循环预算、结合一步或少步生成值得探索,但本文没有完成通用自动退出策略与这些组合的系统验证。

相关工作与启发

  • 对比 Universal Transformers 与普通循环网络: 权重共享本身不是新概念;ELT 的重点是用 ILSD 使视觉生成中的多个循环出口可用,而非仅验证递归计算能够工作。
  • 对比 MaskGIT、MAGVIT 与 DiT: ELT 沿用各自的离散解码或扩散生成范式,改变单次网络调用内部的参数组织与训练监督,因此属于兼容式结构改造。
  • 对比 E-DiT 与 MaGNeTS: 前者通过块跳过和宽度缩减调节预算,后者使用嵌套生成模型;ELT 则固定共享权重集合,通过重复次数调节有效深度。
  • 对比一致性模型与渐进蒸馏: 这些方法主要减少外层采样步骤,ELT 调整内层计算深度。二者概念上互补,但不能把这种互补性当作已经测得的组合收益。

评分

  • 新颖性: 4/5。循环共享已有基础,但前缀自蒸馏与课程监督针对视觉生成弹性问题形成了清晰方法。
  • 实验充分度: 4/5。覆盖离散图像、连续扩散和视频,并分析容量与吞吐;开放域任务和跨硬件验证仍不足。
  • 写作质量: 4/5。核心机制与实验组织清楚,但“无额外开销”和“任意时刻”的表述需要结合实现与评估范围理解。
  • 价值: 4/5。为参数受限的生成模型提供实用计算预算旋钮,真实部署价值取决于目标硬件和质量要求。