重新思考少步扩散 Transformer 中缓存什么：求解器感知的目标选择

扩散 Transformer（DiTs）能够生成高质量的图像和视频，但生成每个样本都需要多次高成本的 DiT 前向传播。加速 DiT 采样的两种常见方法是步骤蒸馏和缓存：前者减少采样步数，后者通过复用较早步骤计算得到的张量来跳过部分 DiT 计算。大多数缓存方法会预先决定要复用哪个张量。经过蒸馏后，相邻采样步骤之间的间隔更大。在这一更大的间隔内复用张量会引入更多误差，因此选择缓存对象变得尤为重要。为此，我们提出 AutoTarget，这是一种针对给定模型、求解器和复用计划选择缓存张量的方法。AutoTarget 使用少量不进行缓存复用的运行，测量复用每个候选张量所造成的误差，然后选择误差最低的候选张量。我们还分析了某个复用步骤中的误差如何影响最终样本。对于 Euler 采样，我们确定了能够产生相同轨迹的缓存目标，并说明为什么存储的求解器更新量可能无法实现这一点。在蒸馏图像和视频 DiT 上的实验表明，最佳缓存目标会随模型、图像分辨率和求解器而变化。AutoTarget 减少了 DiT 计算次数和保留的缓存存储量，同时生成质量仍接近相应的不使用缓存的运行结果。在测试的 PixArt-LCM 和 FLUX.1-schnell 设置上，其校准排序与留出缓存运行所得的排序一致。为帮助他人复现该方法，我们在 GitHub 上提供了其核心实现：https://github.com/wali1024-offical/AutoTarget。