MPDiT: Multi-Patch Globaal-naar-Lokaal Transformer Architectuur voor Efficiënte Flow Matching en Diffusion Model
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
March 27, 2026
Auteurs: Quan Dao, Dimitris Metaxas
cs.AI
Samenvatting
Transformer-architecturen, met name Diffusion Transformers (DiTs), worden vanwege hun sterke prestaties in vergelijking met convolutionele UNets veelvuldig gebruikt in diffusie- en flow-matchingmodellen. Het isotrope ontwerp van DiTs verwerkt echter in elk blok hetzelfde aantal gepatchte tokens, wat leidt tot relatief zware rekenkosten tijdens het trainen. In dit werk introduceren we een multi-patch transformer-ontwerp waarin vroege blokken werken op grotere patches om grove globale context vast te leggen, terwijl latere blokken kleinere patches gebruiken om lokale details te verfijnen. Dit hiërarchische ontwerp kan de rekenkosten met tot 50% in GFLOPs verminderen en tegelijkertijd een goede generatieve prestatie behalen. Daarnaast stellen we ook verbeterde ontwerpen voor tijd- en klasse-embeddingen voor die de trainingsconvergentie versnellen. Uitgebreide experimenten op de ImageNet-dataset demonstreren de effectiviteit van onze architecturale keuzes. Code is vrijgegeven op 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