EdgeDiT: Hardware-Aware Diffusion Transformers for Efficient On-Device Image Generation

1. Problem Statement

Diffusion Transformers (DiTs) have achieved state-of-the-art results in high-fidelity image synthesis, surpassing traditional convolutional U-Net architectures. However, their massive computational complexity, large memory footprints, and high inference latency make them impractical for deployment on resource-constrained edge devices (e.g., smartphones).

• The Gap: While cloud-based inference solves hardware limitations, it introduces privacy risks, network dependency, and high energy costs.
• The Challenge: Simply reducing parameters or FLOPs (Floating Point Operations) does not guarantee improved latency on mobile Neural Processing Units (NPUs) like Qualcomm Hexagon or Apple Neural Engine (ANE), as these chips are optimized for specific operations (e.g., GEMM) rather than raw arithmetic counts.

2. Methodology: The EdgeDiT Framework

The authors propose EdgeDiT, a hardware-aware framework that systematically optimizes DiT architectures for mobile NPUs. The approach combines Feature-wise Knowledge Distillation (FwKD) with Multi-Objective Bayesian Optimization (MOBO).

A. Hardware-Aware Search Space Construction

Instead of manual pruning, the authors define a structured search space of "surrogate" transformer blocks derived from the baseline DiT-XL/2. The search space is built in two stages:

1. Block Removal: Replacing pairs of consecutive DiT layers with a single layer to reduce depth.
2. Structural Simplification:
   • MLP Ratio Modification: Adjusting the expansion ratio of the Feed-Forward Network (FFN).
   • Hidden Dimension Reduction: Lowering the internal projection dimension ($d$) to create low-rank representations.

B. Feature-wise Knowledge Distillation (FwKD)

Training every candidate architecture in the search space ($2^{14} + 4^{28}$ combinations) from scratch is computationally infeasible. To solve this:

• Divide-and-Conquer Strategy: The authors train lightweight surrogate blocks independently to mimic the behavior of the corresponding blocks in a pre-trained "Teacher" model (DiT-XL/2).
• Loss Function: The distillation loss minimizes the L2 distance between the teacher's output features and the surrogate's output features ($L_{KD} = ||T_l(x) - S_l(x)||^2_2$).
• Benefit: This allows the assembly of candidate architectures from pre-trained surrogates without requiring full end-to-end retraining for every candidate.

C. Multi-Objective Bayesian Optimization (MOBO)

To identify the optimal architecture, the authors treat model selection as a multi-objective optimization problem:

• Objectives: Maximize generative quality (measured by Fréchet Inception Distance, FID) and minimize hardware cost (measured by on-device inference latency).
• Process: A Gaussian process model predicts the performance of untested architectures based on evaluated configurations. An acquisition function (Expected Hyper-volume Improvement) balances exploration and exploitation to find the Pareto frontier—the set of models where improving one objective (e.g., speed) would degrade the other (e.g., quality).

D. End-to-End Training

Once the Pareto-optimal candidates are selected, they undergo a final end-to-end training phase using standard diffusion objectives to fully adapt the network parameters to the generation task.

3. Key Contributions

1. EdgeDiT Architecture: A family of lightweight diffusion transformers specifically engineered for mobile NPUs (Qualcomm and Apple).
2. Surrogate-Based Search Framework: A novel method combining feature-wise knowledge distillation with Bayesian optimization to efficiently explore a massive architecture search space without exhaustive training.
3. Hardware-Aware Pruning: Identification and removal of structural redundancies that are particularly taxing for mobile data flows, rather than just reducing parameter counts.
4. Empirical Validation: Demonstration that transformer-based diffusion models can be effectively scaled down for real-world edge deployment without sacrificing the scaling advantages of the original architecture.

4. Experimental Results

The authors evaluated EdgeDiT on the ImageNet 256x256 dataset and benchmarked it against baseline DiT variants (S, B, L, XL) and optimized mobile U-Nets.

• Efficiency Gains:
  • Parameters: Reduced by 20–30%.
  • FLOPs: Decreased by 36–46%.
  • Latency: Achieved a 1.65× speedup on Samsung Galaxy S25 Ultra (Qualcomm NPU) and 1.45× speedup on iPhone 16 Pro Max (Apple ANE) compared to DiT-XL/2.
• Generative Quality:
  • EdgeDiT models achieved competitive or superior Fréchet Inception Distance (FID) scores compared to the baseline DiT-XL/2, despite having fewer parameters and significantly less training iterations.
  • Visual comparisons showed that EdgeDiT generated images with quality comparable to or better than the baseline in class-conditioned scenarios.
• Pareto Optimality: The selected EdgeDiT models (EdgeDiT-1 and EdgeDiT-6) lie on the Pareto frontier, offering a superior trade-off between image fidelity and inference latency compared to existing mobile U-Nets and vanilla DiT variants.

5. Significance and Impact

• Privacy and Offline Capability: EdgeDiT enables responsive, private, and offline generative AI directly on user devices, eliminating the need for cloud inference.
• Scalability Blueprint: The paper provides a scalable blueprint for transitioning large-scale foundation models from high-end GPUs to mobile devices.
• Hardware-Centric Design: By acknowledging that theoretical compute metrics (FLOPs) do not always correlate with real-world latency on NPUs, the paper shifts the paradigm toward hardware-aligned architecture search.
• Future Potential: The framework is model-agnostic and could be extended to other diffusion families (e.g., SiT, MDT) and specialized for other accelerators (TPUs, AR/VR devices).

In conclusion, EdgeDiT successfully bridges the gap between the high fidelity of modern Diffusion Transformers and the strict constraints of mobile hardware, proving that efficient, high-quality generative AI is feasible on the edge.