SnapGen++: Liberando el Potencial de los Transformadores de Difusión para una Generación de Imágenes de Alta Fidelidad y Eficiente en Dispositivos de Borde
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
Resumen
Los recientes avances en transformadores de difusión (DiTs) han establecido nuevos estándares en generación de imágenes, pero siguen siendo poco prácticos para implementación en dispositivo debido a sus altos costos computacionales y de memoria. En este trabajo, presentamos un marco de DiT eficiente diseñado para dispositivos móviles y de edge que logra calidad de generación a nivel de transformador bajo estrictas restricciones de recursos. Nuestro diseño combina tres componentes clave. Primero, proponemos una arquitectura DiT compacta con un mecanismo de atención dispersa global-local adaptativa que equilibra el modelado de contexto global y la preservación de detalles locales. Segundo, proponemos un marco de entrenamiento elástico que optimiza conjuntamente sub-DiTs de distintas capacidades dentro de una superred unificada, permitiendo que un solo modelo se ajuste dinámicamente para inferencia eficiente en diferentes hardwares. Finalmente, desarrollamos Knowledge-Guided Distribution Matching Distillation, una canalización de destilación por pasos que integra el objetivo DMD con transferencia de conocimiento de modelos maestro de pocos pasos, produciendo generación de alta fidelidad y baja latencia (ej. 4 pasos) adecuada para uso en dispositivo en tiempo real. En conjunto, estas contribuciones permiten modelos de difusión escalables, eficientes y de alta calidad para implementación en diversos hardwares.
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.