MPDiT：効率的なフローマッチングと拡散モデルのためのマルチパッチ大域-局所トランスフォーマーアーキテクチャ

MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model

March 27, 2026

著者: Quan Dao, Dimitris Metaxas

cs.AI

要旨

Transformerアーキテクチャ、特にDiffusion Transformer（DiT）は、畳み込みUNetと比較して優れた性能を発揮することから、拡散モデルやフローマッチングモデルにおいて広く利用されるようになった。しかし、DiTの等方性設計では、すべてのブロックで同じ数のパッチ化トークンを処理するため、訓練プロセスにおける計算コストが比較的高いという課題がある。本研究では、大規模パッチを用いて大域的な文脈を粗く抽出する初期ブロックと、細かいパッチを用いて局所的な詳細を洗練する後期ブロックから構成される、マルチパッチTransformerを提案する。この階層的設計により、GFLOPsで最大50％の計算コスト削減を実現しつつ、良好な生成性能を達成する。さらに、時間埋め込みとクラス埋め込みの改良設計も提案し、訓練の収束を加速させる。ImageNetデータセットを用いた大規模な実験により、本アーキテクチャの有効性を実証する。コードはhttps://github.com/quandao10/MPDiTで公開している。

English

Transformer architectures, particularly Diffusion Transformers (DiTs), have become widely used in diffusion and flow-matching models due to their strong performance compared to convolutional UNets. However, the isotropic design of DiTs processes the same number of patchified tokens in every block, leading to relatively heavy computation during training process. In this work, we introduce a multi-patch transformer design in which early blocks operate on larger patches to capture coarse global context, while later blocks use smaller patches to refine local details. This hierarchical design could reduces computational cost by up to 50\% in GFLOPs while achieving good generative performance. In addition, we also propose improved designs for time and class embeddings that accelerate training convergence. Extensive experiments on the ImageNet dataset demonstrate the effectiveness of our architectural choices. Code is released at https://github.com/quandao10/MPDiT