SANA: Efficient Image Generation with Adaptive Flow Scheduling and Model Pruning

An extension of the SANA architecture with two novel efficiency contributions: Adaptive Flow-DPM Scheduling and Sensitivity-Based Transformer Block Pruning.

Authors: Luv Valecha (B23CS1093) · Dheeraj Kumar (B23CS1016) · Neeraj Mansingh (B23CS1095)

---

Overview

This project reproduces and extends the SANA diffusion transformer framework targeting extreme inference acceleration for edge-device deployment. We introduce two orthogonal efficiency techniques on top of the SANA-0.6B backbone, validated by training from scratch on the CelebA face dataset at 256×256 resolution on a single T4 GPU.

---

Table of Contents

• Background
• Contributions
  • Training on CelebA
  • Adaptive Flow-DPM Scheduling
  • Sensitivity-Based Block Pruning
• Installation
• Results
• References

---

Background

SANA addresses the scalability bottlenecks of modern Diffusion Transformers (DiTs) — namely quadratic attention complexity and excessive token counts — through four key innovations:

• DC-AE (32× compression): Reduces latent token count by 16× compared to standard AE-F8 autoencoders.
• Linear Attention: Replaces O(N²) vanilla attention with O(N) linear attention, augmented by Mix-FFN with 3×3 depth-wise convolutions for local context.
• Decoder-only LLM Text Encoding: Uses Gemma-family models with Complex Human Instruction (CHI) prompting for richer text-image alignment.
• Flow-DPM-Solver: Reduces required sampling steps from 28–50 down to 14–20.

This project builds on the SANA-0.6B variant and further pushes inference efficiency through dynamic scheduling and structural pruning.

---

Contributions

1. Training from Scratch on CelebA

We trained a 0.6B-parameter SANA variant from scratch on the CelebA dataset to validate architectural efficiency on a specialized face generation domain.

• Dataset: 202,599 face images, pre-processed via DC-AE to 32×32 latent grids
• Hardware: Single NVIDIA T4 GPU
• Resolution: 256×256

export NCCL_P2P_DISABLE=1
export NCCL_IB_DISABLE=1
export PYTHONPATH=$PYTHONPATH:$(pwd)

torchrun --nproc_per_node=1 train_scripts/train.py \
    --config_path configs/sana_config/celeba_t4.yaml \
    --resume_from latest

---

2. Adaptive Flow-DPM Scheduling

Standard solvers (e.g., Flow-Euler) use fixed step sizes across the diffusion trajectory t ∈ [0, 1]. We observed that image formation velocity is non-linear: critical semantic features form during middle timesteps (t ∈ [0.4, 0.7]), while the denoising tail contains largely redundant function evaluations.

Method: A velocity-aware adaptive scheduler that dynamically adjusts step size based on the L₂-norm of the predicted velocity field:

$$\Delta t_i = \eta \cdot \frac{1}{||v(t_i)||_2 + \epsilon}$$

This allocates finer resolution to high-variance regions and aggressively skips near-converged steps, reducing NFE from 20 → 12 with no significant perceptual degradation.

python ./latency.py

---

3. Sensitivity-Based Block Pruning

Large diffusion transformers frequently contain blocks that perform near-identity transformations, contributing minimally to latent evolution. We quantify this via a Block Importance (BI) score.

Method:

1. Register forward hooks on all 28 transformer blocks of SANA-0.6B.
2. Run forward passes on a calibration set of 100 images across multiple diffusion timesteps.
3. Compute BI score per block as the cosine similarity deviation between input and output activations, averaged across the calibration set.
4. Remove the 8 lowest-importance blocks (predominantly in the middle of the stack).
5. Fine-tune the pruned model on CelebA to recover any performance loss.

This reduces model depth from 28 → 20 blocks and parameters from 590M → 421M.

# Block importance analysis
python phase2_analyze.py

# Load the model, delet the blocks identified in Phase 2, and train the remaining blocks to recover image quality
python phase3_4_train.py

# Quality benchmark
python phase6_evaluate.py

---

Installation

# Clone the repository
git clone https://github.com/Luv-valecha/SANA
cd SANA

# Set up environment (follows base SANA setup)
./environment_setup.sh sana

---

Results

Adaptive Flow-DPM Scheduling

~6.6% reduction in inference latency with NFE reduced from 20 to 12.

---

Sensitivity-Based Block Pruning

---

Qualitative Comparison

The pruned + adaptive model preserves key facial structures and visual consistency across a range of portrait prompts, demonstrating the SANA architecture's robustness to structured pruning.

Sample prompts evaluated:

• "A high-quality portrait of a young male with black wavy hair."
• "A high-quality portrait of a young female with wavy hair, arched eyebrows, high cheekbones, wearing heavy makeup, earrings, a necklace, lipstick."

---

References

1. E. Xie et al., "SANA: Efficient High-Resolution Image Generation with Linear Diffusion Transformers," arXiv 2024. [Paper]
2. Y. Liu et al., "Flow-DPM: A Faster Solver for Flow-based Diffusion," 2023.

---

Acknowledgements

This project is built upon the SANA codebase by NVIDIA Labs and MIT HAN Lab. We thank the authors for open-sourcing their work.