MPDiT: Arquitetura Transformer Global-para-Local de Múltiplos Patches para Modelos Eficientes de Flow Matching e Difusão
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
Resumo
As arquiteturas Transformer, particularmente os Diffusion Transformers (DiTs), tornaram-se amplamente utilizadas em modelos de difusão e *flow-matching* devido ao seu forte desempenho em comparação com as UNets convolucionais. No entanto, o design isotrópico dos DiTs processa o mesmo número de *tokens* em *patches* em todos os blocos, resultando em um processamento computacionalmente pesado durante o treinamento. Neste trabalho, introduzimos um design de transformer multi-*patch* no qual os blocos iniciais operam em *patches* maiores para capturar o contexto global geral, enquanto os blocos posteriores usam *patches* menores para refinar os detalhes locais. Este design hierárquico pode reduzir o custo computacional em até 50\% em GFLOPs, ao mesmo tempo que alcança uma boa performance generativa. Além disso, também propomos designs aprimorados para as incorporações de tempo e classe que aceleram a convergência do treinamento. Experimentos extensivos no conjunto de dados ImageNet demonstram a eficácia das nossas escolhas arquiteturais. O código está disponível em 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