arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2501.04765cs.CVcs.AI

TREAD:用于高效架构无关扩散训练的令牌路由

TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

Felix Krause, Timy Phan, Ming Gui, Stefan Andreas Baumann, Vincent Tao Hu, Björn Ommer

AI总结:

本文提出TREAD,一种架构无关的令牌路由机制,通过将随机选择的令牌从早期层传输到深层,在不修改架构或增加参数的情况下,同时提升了扩散模型的训练效率和生成性能,在ImageNet-256基准上实现了最高37倍的收敛加速和2.09的FID。

摘要:

扩散模型已成为视觉生成的主流方法。然而,这些模型通常存在样本效率低和训练成本高的问题。因此,高效微调、推理和个性化的方法很快被社区采用。然而,从头训练这些模型仍然非常昂贵。虽然最近提出了几种方法——包括掩码、蒸馏和架构修改——来提高训练效率,但这些方法都存在权衡:它们以增加计算成本为代价实现性能提升,反之亦然。相比之下,这项工作旨在通过路由同时提高训练效率和生成性能,这些路由作为传输机制,将随机选择的令牌从模型的早期层传输到更深层。我们的方法不限于常见的基于Transformer的模型——它也可以应用于状态空间模型,并且在无需架构修改或额外参数的情况下实现这一点。最后,我们展示了TREAD在标准ImageNet-256基准测试的类条件合成中降低了计算成本,同时提升了模型性能。在400K训练迭代次数下,与DiT相比,这两项优势叠加实现了14倍的收敛加速;在7M训练迭代次数下,与DiT的最佳基准性能相比,实现了37倍的收敛加速。此外,我们在有引导和无引导设置下分别实现了2.09和3.93的具有竞争力的FID,这在无需架构更改的情况下改进了DiT。

英文摘要:

Diffusion models have emerged as the mainstream approach for visual generation. However, these models typically suffer from sample inefficiency and high training costs. Consequently, methods for efficient finetuning, inference and personalization were quickly adopted by the community. However, training these models in the first place remains very costly. While several recent approaches - including masking, distillation, and architectural modifications - have been proposed to improve training efficiency, each of these methods comes with a tradeoff: they achieve enhanced performance at the expense of increased computational cost or vice versa. In contrast, this work aims to improve training efficiency as well as generative performance at the same time through routes that act as a transport mechanism for randomly selected tokens from early layers to deeper layers of the model. Our method is not limited to the common transformer-based model - it can also be applied to state-space models and achieves this without architectural modifications or additional parameters. Finally, we show that TREAD reduces computational cost and simultaneously boosts model performance on the standard ImageNet-256 benchmark in class-conditional synthesis. Both of these benefits multiply to a convergence speedup of 14x at 400K training iterations compared to DiT and 37x compared to the best benchmark performance of DiT at 7M training iterations. Furthermore, we achieve a competitive FID of 2.09 in a guided and 3.93 in an unguided setting, which improves upon the DiT, without architectural changes.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑