MPDiT: Arquitectura Transformadora Global-a-Local de Múltiples Parches para Modelos Eficientes de Emparejamiento de Flujo y Difusión
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
March 27, 2026
Autores: Quan Dao, Dimitris Metaxas
cs.AI
Resumen
Las arquitecturas Transformer, particularmente los Transformadores de Difusión (DiTs), se han adoptado ampliamente en modelos de difusión y flow-matching debido a su rendimiento superior comparado con las UNets convolucionales. Sin embargo, el diseño isotrópico de los DiTs procesa la misma cantidad de tokens segmentados en cada bloque, lo que conlleva un cómputo relativamente intenso durante el entrenamiento. En este trabajo, presentamos un diseño de transformer multi-parche en el que los bloques iniciales operan con parches más grandes para capturar contexto global general, mientras que los bloques posteriores utilizan parches más pequeños para refinar detalles locales. Este diseño jerárquico puede reducir el coste computacional hasta un 50\% en GFLOPs manteniendo un buen rendimiento generativo. Adicionalmente, también proponemos mejoras en el diseño de los embeddings temporales y de clase que aceleran la convergencia del entrenamiento. Experimentos exhaustivos en el dataset ImageNet demuestran la efectividad de nuestras decisiones arquitectónicas. El código está disponible en 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