Source-linked AI summary
SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices
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
TL;DR
Large diffusion transformers offer high image quality but are difficult to deploy on mobile and edge devices because of their computational and memory costs. The paper combines efficient sparse-attention architecture, elastic training, and K-DMD to support scalable few-step generation, reporting 1.8s inference for 1024×1024 images on an iPhone 16 Pro Max.
Problem
Large DiT models provide strong generation quality but remain impractical for on-device deployment because of high computational and memory costs.
Method
The framework combines adaptive global–local sparse attention, an elastic DiT supernetwork, and K-DMD with few-step-teacher knowledge transfer.
Results
1.8s inference on an iPhone 16 Pro Max supports 1024×1024 generation, reaching 85.2% on DPG-Bench and 0.70 on GenEval.
Takeaways & Limitations
The unified framework supports high-fidelity, scalable diffusion-transformer generation across heterogeneous mobile and edge hardware.
Takeaways & Limitations
Standard DMD is sensitive to hyperparameter choices that vary across model capacities and can cause unstable convergence in compact models.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper targets the gap between high-quality but resource-intensive DiTs and efficient but less scalable U-Net-based on-device models. It combines sparse attention, elastic subnetworks, and knowledge-guided distillation for high-quality generation across devices.
- Motivation: Large DiT models achieve strong image quality but often require tens of billions of parameters, server-grade GPUs, and custom CUDA kernels.These requirements increase computational cost and reliance on cloud infrastructure.
- Efficient DiT Architecture: The proposed three-stage DiT uses adaptive global–local sparse attention to combine compressed global context with blockwise local spatial modeling.Content-dependent attention allocation addresses quadratic attention complexity at high resolutions such as 1K.
- Elastic Deployment: Elastic training jointly optimizes sub-DiTs with varying capacities in one supernetwork, allowing hardware-specific subnetworks to be selected without retraining.This design targets heterogeneous platforms while maintaining consistent generation quality.
- Distillation: Knowledge-guided DMD adds few-step-teacher supervision to standard DMD after cloud-scale knowledge distillation, stabilizing compact-model training for low-latency generation.The added supervision addresses instability and visual artifacts observed when standard DMD is applied directly to compact models.
- Results: 1.8s inference on an iPhone 16 Pro Max supports 1024×1024 generation, with 85.2% on DPG-Bench and 0.70 on GenEval.The reported result combines on-device latency, image resolution, and quality scores.
2 Related Work
Prior work improves DiT efficiency, on-device generation, elastic deployment, sparse attention, and step distillation, but these efforts address separate aspects of efficient generative modeling. The paper positions its approach as combining these directions for diffusion-transformer deployment across heterogeneous devices.
- T2I Diffusion Models: Diffusion Transformers improve scalability, quality, and generalization over earlier U-Net-based diffusion models but remain costly because of billion-scale parameter counts.The related-work discussion frames this cost as a barrier to on-device deployment.
- Efficient Diffusion Transformers: Efficient DiT research explores key–value compression, linear self-attention, Mamba-based attention, and hybrid convolutional-transformer designs.These methods target reduced attention or architectural cost for image synthesis.
- On-Device Generative Models: On-device generative models use quantization, pruning, and knowledge distillation to reduce model size and latency, primarily with U-Net architectures.Prior systems include 512^2 generation within seconds and compact-U-Net 1024^2 generation with quality and editing-flexibility trade-offs.
- Elastic Networks: Elastic-network research enables models to adapt to computational budgets, but elastic architectures remain underexplored for generative models.The paper extends this deployment idea to diffusion transformers across heterogeneous devices.
- Sparse Attention: Sparse-attention methods exploit block, neighborhood, local, spatiotemporal, or adaptive sparsity, often relying on selective activation or specialized kernels.The paper distinguishes its mobile-oriented design from prior hardware-specific sparse-attention approaches.
- Step Distillation: Step distillation compresses multi-step diffusion sampling into a few iterations through trajectory, consistency, adversarial, or distribution-matching objectives.The related work situates K-DMD within these approaches to few-step generation.
3 Method
The method combines an efficient three-stage DiT architecture with adaptive sparse attention, elastic training, and knowledge-guided distillation for high-fidelity edge generation. These components reduce attention and deployment costs while preserving generation quality and stabilizing compact-model distillation.
- The framework introduces an efficient DiT architecture, an elastic training framework, and a multi-stage distillation pipeline for edge generation.
- Efficient DiT Architecture: The three-stage design assigns Down and Up blocks to high-resolution latents with ASSA, while Middle blocks use downsampled latents and standard self-attention.Cross-attention provides text conditioning and feed-forward layers complete the blocks.
- Efficient DiT Architecture: ASSA combines compressed global key-value attention with blockwise local attention, reducing overhead while preserving global structure and fine-grained spatial detail.The global branch reduces key/value token length by a factor of four, while blockwise neighborhoods enable contiguous, hardware-efficient computation.
- Efficient DiT Architecture: 293 ms latency with validation loss 0.513 demonstrates that sparse attention can reduce attention cost without loss of generation quality.
- Efficient DiT Architecture: The combined efficient DiT reaches 360 ms latency and 0.509 validation loss, offering a strong efficiency–accuracy trade-off and outperforming SnapGen in validation loss.
- Elastic DiT Framework: Elastic training jointly optimizes subnetworks of different widths within a supernetwork, using shared supervision and diffusion loss across granularities.Parameter sharing maintains comparable validation loss and DINO-FID to standalone training while reducing model-state footprint.
- Distillation Pipelines: K-DMD addresses unstable compact-model DMD training by adding knowledge transfer from a few-step teacher to the distribution-matching objective.The combined supervision stabilizes training across capacities without additional hyperparameter tuning.
4 Experiments
Experiments evaluate generation quality, deployment scalability, and few-step distillation across model sizes, baselines, devices, and human preferences. The results show competitive quality at multiple scales, elastic hardware adaptation, and four-step performance comparable to 28-step models.
- Evaluation Setup: The evaluation covers generation quality, deployment scalability, and step-distillation effectiveness using benchmarks, device latency, human preferences, and few-step comparisons.The study includes DPG-Bench, GenEval, T2I-CompBench, CLIP, mobile-device measurements, Fig. 7 preference comparisons, and Figs. 8–9 distillation results.
- Quantitative Results: The full 1.6B variant achieves results comparable to Flux.1-dev and SD3.5-Large, while the 0.4B variant remains competitive at substantially smaller size.The models are evaluated across multiple generation benchmarks and qualitative comparisons.
- Quantitative Results: The 0.4B variant surpasses models up to 8B parameters on multiple benchmarks while retaining on-device efficiency comparable to SnapGen.The 0.3B tiny variant achieves the highest throughput.
- Deployment Scalability: The elastic design provides a controllable quality–efficiency trade-off within one unified model, adapting width configurations to diverse hardware without retraining.This supports deployment across heterogeneous devices using shared model parameters.
- Human Preference Study: The full variant surpasses SANA, SD3-M, and Flux.1-dev in human-rated fidelity, realism, and text–image alignment, while the small variant outperforms SANA and SD3-M on most attributes.The user study follows Parti prompts and evaluates realism, visual fidelity, and text alignment.
- Few-Step Generation: K-DMD produces four-step models with performance comparable to 28-step baselines despite substantially fewer sampling steps.The comparison uses DPG-Bench and GenEval scores; K-DMD also mitigates over-saturation and grainy artifacts seen with standard DMD.
5 Conclusion
The paper concludes that an efficient DiT framework can bring high-quality image generation to mobile and edge devices through architectural efficiency, elastic deployment, and knowledge-guided distillation.
- Conclusion: The framework combines a lightweight DiT, adaptive global–local sparse attention, elastic subnetworks, and K-DMD for scalable edge deployment.The components reduce computational and memory overhead, adapt inference to heterogeneous hardware, and support reliable few-step generation.
- Conclusion: The approach achieves competitive quality against large server-side systems while operating efficiently on-device.The conclusion presents the framework as a comprehensive approach for real-world edge deployment.
A Discussion of On-Device Latency
The latency discussion attributes mobile feasibility to Adaptive Sparse Self-Attention and reports efficient quantized full-model execution without out-of-memory failures.
- Latency and Memory: The quantized full model runs on mobile devices without out-of-memory issues while computing attention over 4096 tokens at dimension 1024.Adaptive Sparse Self-Attention is identified as the enabling design component.
- Latency and Memory: Latency reporting separates per-step denoising and total generation time, while VAE decoding contributes approximately 120 ms.Latent scaling, scheduler stepping, and CLIP embedding add negligible latency according to the discussion.
B Demo on Mobile Device
The mobile demonstration runs the compact and full variants on an iPhone 16 Pro Max, showing 1024×1024 generation with quantized models and a 1.8-second small-model runtime.
- On-Device Demonstration: The small 0.4B model generates high-quality 1024×1024 images in 1.8 seconds per image on an iPhone 16 Pro Max.The application uses the open-source Swift Core ML Diffusers framework.
- On-Device Demonstration: The demo compares the small 0.4B model with the full 1.6B variant, whose weights use 4-bit quantization in the shown application.Screenshots present outputs from both variants on the iPhone 16 Pro Max.
- Mobile Optimization: Mobile deployment reduces transpose and reshape operations, uses convolutional tensor layouts, and reimplements attention with split einsum operations.These implementation choices target on-device computational efficiency.
- Mobile Optimization: The 1.6B full model quantizes most layers to 4 bits and assigns 8 bits to more sensitive layers using layerwise output MSE.K-means clustering over weight values provides the quantization procedure.
D Training Implementation Details
Training uses staged distributed optimization with elastic training, followed by high-resolution knowledge distillation. Blockwise Neighborhood Attention exposes sparsity controls through block number and neighborhood radius.
- Training Setup: Training begins at 256^2 resolution across 32 nodes with 8 A100 GPUs each, using elastic training.The initial stage uses FSDP2 and a global batch size of 8192 for 400K iterations.
- Training Setup: The second training stage increases resolution to 1024^2, enables gradient checkpointing, and adds knowledge distillation under elastic training.This stage uses a global batch size of 2048 for an additional 100K iterations.
- Step Distillation: K-DMD uses a time shift of 3, classifier-free guidance 4, and LoRA on both student and critic networks.LoRA uses rank 64 and α = 128; the student is updated every 5 iterations.
- Sparse Attention: BNA adjusts block number b and neighborhood radius r to control attention sparsity and balance computational efficiency with representational fidelity.Specific configurations match the spatial coverage of standard neighborhood attention with three- or five-token neighbors.
F Detailed Results on T2I Benchmarks
The paper reports detailed comparisons on DPG-Bench, GenEval, and T2I-CompBench, alongside qualitative ImageNet-1K results. The supplied passages identify the benchmark tables but do not provide their numerical entries.
- Benchmark Results: Detailed comparison results are presented for DPG-Bench in Table 2.The supplied passage names the benchmark and its corresponding table but gives no numerical scores.
- Benchmark Results: Detailed comparison results are presented for GenEval in Table 3.The supplied passage identifies the benchmark table without reporting its metric values.
- Benchmark Results: Detailed comparison results are presented for T2I-CompBench in Table 4.The supplied passage identifies the benchmark table without reporting its metric values.
- Qualitative Results: The 0.4B small model records Validation Loss = 0.5090 versus 0.5131 for the 0.4B SnapGen U-Net.The accompanying qualitative comparison reports sharper textures, more consistent colors, and improved structural fidelity across diverse categories.
H Additional Qualitative Comparison on T2I
Additional T2I comparisons evaluate diverse prompts spanning objects, scenes, and artistic compositions. The reported outcomes emphasize competitive visual quality, prompt adherence, and visually consistent generation.
- Prompt Evaluation: The models are evaluated on diverse text-to-image prompts covering objects, scenes, and artistic compositions.The comparison targets image–text alignment, semantic accuracy, and visual consistency.
- ImageNet-1K Comparison: The 0.4B small model produces sharper textures, more consistent colors, and improved structural fidelity than the 0.4B SnapGen U-Net comparison.The reported Validation Loss values are 0.5090 for the small model and 0.5131 for SnapGen U-Net.
- Prompt Evaluation: Additional comparisons report competitive visual quality and superior prompt-following ability at 1024^2 resolution.The input prompts are displayed above the image grids in Figures 4 and 5.