跳到主要内容
返回时间线
arXiv来源发表:

LiFT 让循环 DiT 在训练深度之外继续提升:L/2 与 XL/2 以更少参数和推理算力超过更大的稠密 DiT

相关研究与后续进展

核心概要

LiFT 把循环 DiT 的每次循环监督到一条从初始速度估计指向流匹配目标的直线路径上的对应点,使训练好的检查点可在推理时使用远超训练深度的循环数;在 ImageNet 类别条件生成上,L/2 与 XL/2 检查点通过把推理算力重新分配给循环,在更少参数与更少推理 FLOPs 下取得优于更大稠密 DiT 的 FID,而 B/2 规模上稠密模型仍更强。

AI-generated editorial illustration: LiFT: Loop Flow Transformers

深度剖析

LiFT 为循环 DiT 的每次循环分配各自的回归目标:从模型初始速度估计到流匹配目标的直线路径上的一个点,由连续深度坐标索引,最后一次读出仍承担完整流匹配任务。 此前循环扩散模型把每个中间预测都训练向完整输出,或依赖自蒸馏的全深度教师;LiFT 改为沿参考路径分配“还差多少修正”,使各次循环承担不同阶段的修正任务。 方法在 ImageNet 类别条件生成上验证:L/2 R10 从训练深度两循环的 FID 20.30 降至八循环的 10.95,XL/2 R12 从 17.07 降至十六循环的 9.31,XL/2 R8 从三循环的 16.95 降至三十二循环的 9.50。

推理循环数可在训练后自由选择,且循环与积分步数构成两条互补的推理算力轴,最优配置取决于预算与起点。 以往循环模型在超过训练深度后 FID 变差,或与稠密模型等算力下无优势;LiFT 报告深度外推持续改善,并给出循环与步数的联合分配结果。 L/2 R10 在 25 步、四循环下 FID 13.16、7.397 TFLOPs,优于稠密 L/2 在 50 步下的 16.94、8.069 TFLOPs;但同一模型在十六循环、十步下 FID 15.09、11.027 TFLOPs,反而更差更贵。

在 L/2 与 XL/2 规模上,把推理算力转向循环可获得优于更大稠密 DiT 的生成质量,同时参数与推理算力更少。 相对同规模稠密参考,LiFT 在训练深度处因共享块替代独立块而落后;优势出现在测试时算力重新分配之后,而非训练工作点。 LiFT L/2 R10 八循环 FID 10.95、28.24 TFLOPs,对比稠密 XL/2 在 250 步下 FID 14.30、59.31 TFLOPs,FID 低 3.34 点,参数约少 60%,推理 FLOPs 少 52%,训练算力约少 32%(61.98 对 91.10 EFLOPs)。

深度外推的收益依赖核心容量:单块核心与 B/2 规模收益有限,稠密模型在 B/2 的每个评估算力点仍更强。 论文把规模依赖作为明确的范围条件报告,而非普遍结论。 B/2 R1 从八循环 45.70 到十六循环 45.58 几乎持平;最佳 LiFT B/2 结果 FID 33.87 未达到稠密 B/2 在 25 步、1.150 TFLOPs 下的 32.50,作者归因于 89M 参数对 32.8B 训练 token 而言容量不足。

启示与展望

该结果面向 256 分辨率 ImageNet 类别条件生成、无分类器引导的设定,适用于希望在参数固定时按推理预算调整计算的部署场景,尤其是参数存储比计算更受限的情形;作者也指出这为视频与生成动力学模型中的预算感知生成提供了动机。方法上,LiFT 把 DiT 拆成 prelude、共享循环核心与仅做预测的 coda,核心按深度坐标做自适应归一化与残差门控,推理时每个积分步重新初始化 prelude 状态并只使用最终读出。

每个结果来自单次训练与单组 50,000 张样本,接近平台期的 FID 差异未做统计显著性检验;推理深度是预先指定的,联合分配研究只覆盖一个检查点与有限网格,自适应深度选择与内外算力预算的联合优化仍待研究。部署含义需要硬件评估,因为论文报告的是参数量与解析 FLOPs 估计,而非设备级延迟、能耗或峰值内存。此外,深度外推是否延伸到其他模态或引导采样仍是开放问题。

来源