一个模型，多种预算：面向扩散变换器的弹性潜在接口 One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers

ELIT的工作流程可以分为三个阶段：（1）空间头处理阶段：输入图像首先通过patch化层转化为N个空间token，然后经过Bin个短空间Transformer块处理，产生精炼的空间表示，避免直接从原始噪声patch中读取信息。（2）潜在核心计算阶段：引入一个可学习的潜在接口，包含K个潜在token。Read层通过交叉注意力将空间信息拉入潜在接口：潜在token作为查询，空间token作为键和值，输出更新后的潜在表示。接下来，Bcore个核心Transformer块在潜在接口上执行主要计算。（3）写回和输出阶段：Write层将更新后的潜在信息广播回空间token：空间token作为查询，潜在token作为键和值。最后Bout个尾部空间块完成处理并投影到输出速度空间。为减少计算成本，空间token和潜在token被分组处理，交叉注意力仅在对应组内进行，将成本从O(NK)降至O(NK/G)。

ELIT的技术新颖性体现在多个层面。首先，架构设计上，它首次将可变长度潜在接口引入DiT类生成模型，仅需添加两个轻量级交叉注意力层（Read和Write），参数增加极少（约3%），却能实现完全的计算弹性。其次，训练策略上，tail-dropping机制确保了重要性排序：在每个训练迭代中，随机采样保留的潜在token数量J_tilde，只保留每组的前J_tilde个token，丢弃后续token。这迫使前面的token学习全局信息，后面的token学习细节信息。第三，引导策略上，多预算训练天然产生了可用的弱模型（使用较少token的版本），使得autoguidance可以零成本实现。作者进一步提出Cheap CFG（CCFG），将autoguidance与CFG结合，在降低33%计算成本的同时保持甚至提升生成质量。最后，分组交叉注意力设计使方法可以扩展到高分辨率和长序列场景。

独立分析ELIT存在几个值得改进的弱点。首先，分组策略目前采用固定的2D/3D网格划分，未考虑语义边界——如果一个重要对象跨越多个组，其信息可能被分散处理，影响生成质量。改进方向可以是基于语义分割或注意力模式动态调整分组。其次，Read和Write层的设计相对简单（单层交叉注意力），论文消融实验显示增加Write层的容量（双层或更大FFN）能带来提升，但作者为保持轻量性选择了最简设计。一个可能的改进是根据模型规模自适应调整Read/Write层的复杂度。第三，当前的多预算训练使用均匀分布采样J_tilde，未探索更智能的采样策略，例如根据训练进度或噪声水平调整分布。最后，方法在视频生成上的验证相对有限（仅Kinetics-700 256px），在更长视频或更高分辨率场景下的表现有待进一步验证。

作者在讨论部分提出了几个有前景的未来研究方向。首先是噪声水平感知的预算调度——初步实验表明高噪声步骤可能需要更少的token（50% token在高噪声步骤，100%在其他步骤，性能相当但节省18%计算），一个原则性的噪声级别调度策略可能进一步提升效率。其次是训练和推理时的动态预算分配——根据每个采样步骤的实际需求分配不同的token预算。基于本文的多预算框架，还可以探索：（1）内容自适应的预算分配，为不同空间区域分配不同的token数量；（2）与蒸馏方法（如一致性蒸馏）的深度结合，进一步减少采样步骤；（3）扩展到更长的视频生成和3D内容生成场景；（4）探索ELIT在条件生成（如文本到图像、图像编辑）中的应用，利用潜在接口作为条件注入的接口。

从复现评估来看，ELIT具有较好的可复现性。代码和预训练模型已在项目主页（https://snap-research.github.io/elit）发布。方法实现相对简单，核心改动仅涉及添加Read/Write层和训练时的tail-dropping逻辑，不需要特殊的优化器或训练技巧。数据方面，实验使用公开数据集（ImageNet-1K、Kinetics-700），无需额外数据收集。算力方面，主要实验基于DiT-XL/2（675M参数），训练500k步，需要中等规模的GPU集群；Qwen-Image实验（20B参数）需要更大的计算资源。论文提供了详细的消融实验和超参数设置，降低了复现难度。不过，多预算训练需要仔细调整J_min、J_max和batch size以匹配基线的训练FLOPs，这部分可能需要一些调参经验。