TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

要約

拡散モデルは、視覚生成の主流のアプローチとして浮上しています。
ただし、これらのモデルは通常、サンプルの非効率性と高いトレーニングコストに悩まされています。
その結果、効率的な微調整、推論、パーソナライズの方法は、コミュニティによって迅速に採用されました。
ただし、そもそもこれらのモデルのトレーニングは非常に費用がかかります。
マスキング、蒸留、建築の修正を含むいくつかの最近のアプローチは、トレーニング効率を改善するために提案されていますが、これらの各方法にはトレードオフが伴います。計算コストの増加を犠牲にしてパフォーマンスが向上します。
対照的に、この作業は、初期層からモデルのより深い層までランダムに選択されたトークンの輸送メカニズムとして機能するルートを通じて、トレーニング効率と生成パフォーマンスを同時に改善することを目的としています。
私たちの方法は、一般的な変圧器ベースのモデルに限定されません。また、状態空間モデルにも適用でき、建築の変更や追加のパラメーターなしでこれを達成することもできます。
最後に、トレッドが計算コストを削減し、同時にクラス条件合成における標準のImagenet-256ベンチマークのモデルパフォーマンスを向上させることを示します。
これらの利点はどちらも、7mトレーニングの繰り返しでのDITの最高のベンチマークパフォーマンスと比較して、DITと37Xと比較して、400Kトレーニングイテレーションで14倍の収束速度を増やします。
さらに、ガイド付きで2.09の競争力のあるFID、ガイド付き設定で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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google