SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices

TL;DR

This work addresses the challenge of deploying high-fidelity diffusion models on edge devices by introducing an Efficient Diffusion Transformer (DiT) with adaptive global-local sparse attention. It combines three pillars: a compact three-stage DiT architecture with Adaptive Sparse Self-Attention, an Elastic DiT framework enabling subnetworks of 0.3B, 0.4B, and 1.6B to share a single supernetwork, and a Knowledge-Guided Distribution Matching Distillation (K-DMD) pipeline that leverages few-step teachers for fast, high-quality generation. The approach yields strong on-device performance, achieving near server-level image quality at substantial runtime gains, including 4-step generation with competitive fidelity and real-time deployment feasibility on mobile hardware. These contributions pave the way for scalable, on-device diffusion with cross-device adaptability and minimal cloud dependence.

Abstract

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.

Figures (12)

• Figure 1: Efficient DiT Overview.Left: Our model consists of three stages: Down, Middle and Up. Down and Up blocks operate on high-resolution latent while using our novel Adaptive Sparse Self-Attention (ASSA) layers. Middle blocks operate at latents downsampled by $2\times2$ window and use standard Self-Attention (SA) layers. Other layers in the blocks are Cross-Attention (CA) for modulating with input text conditioning and Feed-Forward (FFN) layer. Right: We delve deeper into our ASSA layer. It consists of two parallel attention processing branches: (i) coarse-grained key-value compression for overall structure, and (ii) fine-grained blockwise neighborhood attention features. Finally, the layers to weight these two features are adaptively per head through an input-dependent weighting function.
• Figure 1: On-device Image Generation Demo. Screenshots from our on-device application running on an iPhone 16 Pro Max. The left panel shows results from the small (0.4B) model, and the right panel shows results from the full variant with 4-bit quantization.
• Figure 2: Efficient DiT Ablations. We plot the performance (validation loss) and model footprint (parameters & latency on iPhone 16 Pro Max) for various stages in our ablations. Using a baseline DiT yields extremely high latency. Our multi-stage design with ASSA layers and additional enhancements results in an Efficient DiT with comparable latency and better performance than the state-of-the-art on-device model SnapGen hu2024snapgen.
• Figure 2: Illustration of Blockwise Neighborhood Attention (BNA). Visualization of BNA under different hyperparameter settings of block number ($b$) and neighborhood radius ($r$), showing the corresponding spatial neighbor coverage and attention sparsity.
• Figure 3: Illustration of Blockwise Neighborhood Attention.(a) Naive Neighborhood Attention, where each query attends to its local window of 3 neighboring tokens. (b) Corresponding self-attention mask showing the limited receptive field for each query. (c) Blockwise Neighborhood Attention extends this concept by grouping tokens into 8 local blocks, enabling efficient attention computation while preserving locality.