MPDiT: Architettura Transformer Globale-to-Locale Multi-Patch per un Modello di Flow Matching e Diffusione Efficiente
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
March 27, 2026
Autori: Quan Dao, Dimitris Metaxas
cs.AI
Abstract
Le architetture Transformer, in particolare i Diffusion Transformer (DiT), sono diventate ampiamente utilizzate nei modelli di diffusione e di flow-matching grazie alle loro prestazioni superiori rispetto alle UNet convoluzionali. Tuttavia, la progettazione isotropa dei DiT elabora lo stesso numero di token suddivisi in patch in ogni blocco, comportando un calcolo relativamente oneroso durante il processo di addestramento. In questo lavoro, introduciamo una progettazione transformer multi-patch in cui i blocchi iniziali operano su patch più grandi per catturare il contesto globale approssimativo, mentre i blocchi successivi utilizzano patch più piccole per rifinire i dettagli locali. Questa progettazione gerarchica può ridurre il costo computazionale fino al 50\% in termini di GFLOPs, raggiungendo al contempo buone prestazioni generative. Inoltre, proponiamo anche progetti migliorati per gli embedding temporali e di classe che accelerano la convergenza dell'addestramento. Esperimenti estesi sul dataset ImageNet dimostrano l'efficacia delle nostre scelte architetturali. Il codice è rilasciato all'indirizzo 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