MPDiT: Multi-Patch Globale-zu-Lokale Transformer-Architektur für effizientes Flow Matching und Diffusionsmodell
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
March 27, 2026
Autoren: Quan Dao, Dimitris Metaxas
cs.AI
Zusammenfassung
Transformer-Architekturen, insbesondere Diffusion Transformer (DiTs), haben aufgrund ihrer überlegenen Leistung im Vergleich zu konvolutionalen UNets weite Verbreitung in Diffusions- und Flow-Matching-Modellen gefunden. Das isotrope Design von DiTs verarbeitet jedoch in jedem Block die gleiche Anzahl patchifizierter Tokens, was zu einem relativ hohen Rechenaufwand während des Trainings führt. In dieser Arbeit stellen wir ein Multi-Patch-Transformer-Design vor, bei dem frühe Blöcke mit größeren Patches arbeiten, um groben globalen Kontext zu erfassen, während spätere Blöcke kleinere Patches zur Verfeinerung lokaler Details nutzen. Dieses hierarchische Design kann die Rechenkosten um bis zu 50 % in GFLOPs reduzieren und dennoch eine gute generative Leistung erzielen. Zusätzlich schlagen wir verbesserte Designs für Zeit- und Klassen-Einbettungen vor, die die Trainingskonvergenz beschleunigen. Umfangreiche Experimente auf dem ImageNet-Datensatz demonstrieren die Wirksamkeit unserer architektonischen Entscheidungen. Der Code ist unter https://github.com/quandao10/MPDiT verfügbar.
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