SnapGen++: Liberando Transformadores de Difusão para Geração Eficiente de Imagens de Alta Fidelidade em Dispositivos de Borda
SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices
January 13, 2026
Autores: Dongting Hu, Aarush Gupta, Magzhan Gabidolla, Arpit Sahni, Huseyin Coskun, Yanyu Li, Yerlan Idelbayev, Ahsan Mahmood, Aleksei Lebedev, Dishani Lahiri, Anujraaj Goyal, Ju Hu, Mingming Gong, Sergey Tulyakov, Anil Kag
cs.AI
Resumo
Avanços recentes em transformadores de difusão (DiTs) estabeleceram novos padrões na geração de imagens, mas permanecem impraticáveis para implantação em dispositivos devido aos seus altos custos computacionais e de memória. Neste trabalho, apresentamos uma estrutura eficiente de DiT adaptada para dispositivos móveis e de borda que alcança qualidade de geração em nível de transformer sob rigorosas restrições de recursos. Nosso projeto combina três componentes principais. Primeiro, propomos uma arquitetura compacta de DiT com um mecanismo de atenção esparsa global-local adaptativa que equilibra a modelagem de contexto global e a preservação de detalhes locais. Segundo, propomos uma estrutura de treinamento elástica que otimiza conjuntamente sub-DiTs de capacidades variadas dentro de uma super-rede unificada, permitindo que um único modelo se ajuste dinamicamente para inferência eficiente em diferentes hardwares. Finalmente, desenvolvemos a Destilação por Correspondência de Distribuição Guiada por Conhecimento, um pipeline de destilação por etapas que integra o objetivo DMD com a transferência de conhecimento de modelos professores de poucas etapas, produzindo geração de alta fidelidade e baixa latência (por exemplo, 4 etapas) adequada para uso em tempo real no dispositivo. Juntas, essas contribuições permitem modelos de difusão escaláveis, eficientes e de alta qualidade para implantação em hardwares diversos.
English
Recent advances in diffusion transformers (DiTs) have set new standards in image generation, yet remain impractical for on-device deployment due to their high computational and memory costs. In this work, we present an efficient DiT framework tailored for mobile and edge devices that achieves transformer-level generation quality under strict resource constraints. Our design combines three key components. First, we propose a compact DiT architecture with an adaptive global-local sparse attention mechanism that balances global context modeling and local detail preservation. Second, we propose an elastic training framework that jointly optimizes sub-DiTs of varying capacities within a unified supernetwork, allowing a single model to dynamically adjust for efficient inference across different hardware. Finally, we develop Knowledge-Guided Distribution Matching Distillation, a step-distillation pipeline that integrates the DMD objective with knowledge transfer from few-step teacher models, producing high-fidelity and low-latency generation (e.g., 4-step) suitable for real-time on-device use. Together, these contributions enable scalable, efficient, and high-quality diffusion models for deployment on diverse hardware.