MPDiT：効率的なフローマッチングと拡散モデルのためのマルチパッチ・グローバルからローカルへのトランスフォーマー・アーキテクチャ

Abstract

Transformerアーキテクチャ、特にDiffusion Transformers（DiTs）は、畳み込みUNetと比べて強力な性能を発揮するため、拡散およびフローマッチングモデルで広く用いられるようになってきました。しかし、DiTの等方的な設計は、各ブロックで同じ数のパッチ化トークンを処理するため、学習プロセス中の計算量が比較的重くなります。本研究では、多段（マルチパッチ）Transformerの設計を導入します。ここでは、初期のブロックは粗いグローバルな文脈を捉えるためにより大きなパッチで動作し、後半のブロックは局所的な詳細を洗練するためにより小さなパッチを用います。この階層的な設計により、生成性能を良好に保ったまま、GFLOPsで最大50 %の計算コストを削減できる可能性があります。さらに、学習の収束を加速する時間埋め込みおよびクラス埋め込みの改良設計も提案します。ImageNetデータセットでの大規模な実験により、提案するアーキテクチャ選択の有効性が示されます。コードは url{https://github.com/quandao10/MPDiT} で公開されています