MPDiT: Глобально-локальная трансформаторная архитектура с множественными патчами для эффективного согласования потоков и диффузионной модели
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
March 27, 2026
Авторы: Quan Dao, Dimitris Metaxas
cs.AI
Аннотация
Архитектуры трансформаторов, в частности диффузионные трансформаторы (DiT), получили широкое распространение в диффузионных моделях и моделях с согласованием потоков благодаря их высокой производительности по сравнению со сверточными U-Net. Однако изотропная конструкция DiT обрабатывает одинаковое количество патчей-токенов в каждом блоке, что приводит к относительно высоким вычислительным затратам в процессе обучения. В данной работе мы представляем многомасштабную трансформаторную архитектуру, в которой начальные блоки работают с крупными патчами для захвата глобального контекста, а последующие блоки используют мелкие патчи для детализации локальных особенностей. Такая иерархическая конструкция позволяет снизить вычислительные затраты до 50% в пересчете на GFLOPs при сохранении высокого качества генерации. Дополнительно мы предлагаем усовершенствованные схемы для временных и классовых эмбеддингов, ускоряющие сходимость обучения. Масштабные эксперименты на наборе данных ImageNet подтверждают эффективность наших архитектурных решений. Код доступен по адресу 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