SnapGen Generates 1024px Images in 1.4 s with Lightweight On‑Device Architecture

Problem

State‑of‑the‑art text‑to‑image diffusion models (e.g., SDXL, SD3) deliver high visual fidelity but are too large (multi‑billion parameters), have high inference latency, and require cloud execution, which limits mobile deployment, raises privacy concerns, and increases cost.

Efficient UNet redesign

Starting from the SDXL UNet, the authors thin and shorten the backbone:

Transformer block counts reduced from [0, 2, 10] to [0, 2, 4] across the three resolution stages.

Channel dimensions reduced from [320, 640, 1280] to [256, 512, 896].

Subsequent micro‑architectural tweaks are evaluated on ImageNet‑1K (256 px) using FID as the primary quality metric.

Self‑Attention removal at high resolution : keep SA only at the lowest‑resolution stage. FLOPs drop 17 %, latency 24 % (iPhone 15 Pro), and FID improves from 3.76 to 3.12.

Depthwise‑pointwise separable convolutions : replace standard Conv with DW‑PW layers. Parameters –24 %, latency –62 % (2.4× speed‑up) but FID rises to 3.38. Adding an intermediate channel expansion ratio of 2 restores quality, yielding a net 15 % parameter reduction, 27 % FLOP reduction and 2.4× overall speed‑up.

FFN trimming : shrink hidden‑channel expansion from 4 to 3. Parameters and FLOPs fall ≈12 % with comparable FID.

Multi‑Query Attention (MQA) : replace MHSA with MQA (shared keys/values across heads). Parameters –16 %, latency –9 % and negligible quality loss; the latency gain exceeds the modest FLOP reduction because memory traffic is reduced.

Early condition injection : add Cross‑Attention to the first UNet stage and replace the residual block with a transformer block (CA + FFN). This improves FID while keeping the model compact.

QK RMSNorm & 2‑D RoPE : incorporate Query‑Key RMSNorm and two‑dimensional rotary position embeddings. Training stabilises, softmax saturation is mitigated, and a marginal FID gain is observed with virtually no overhead.

Efficient Decoder

The SDXL/SD3 decoder cannot run on mobile neural engines for 1024 px output (OOM on iPhone 15 Pro). The tiny decoder makes five key changes:

Remove all attention layers.

Retain only minimal GroupNorm.

Thin channel widths and replace Conv with separable Conv.

Reduce the number of residual blocks at the highest resolution.

Eliminate Conv shortcuts; use up‑sampling layers for channel conversion.

Training uses MSE, LPIPS, and an adversarial loss (KL omitted), batch size 256, 1 M iterations on 256 px patches. The decoder attains PSNR comparable to the SDXL/SD3 baselines while delivering 35.9× (vs. SDXL) and 54.4× (vs. SD3) speed‑ups. On an iPhone 16 Pro‑Max, decoder latency is 119 ms and each UNet step costs 274 ms, giving a total of 1.2–2.3 s for 4–8 steps.

Training recipe (Rectified Flow)

Training follows a flow‑matching paradigm (Rectified Flows [18][19] ). The forward process linearly interpolates a latent image x_0 to a standard normal z: x_t = \sqrt{\alpha_t}\,x_0 + \sqrt{1-\alpha_t}\,z The denoising UNet predicts the velocity field v_θ(x_t, t) : v_θ(x_t, t) \approx \frac{d x_t}{d t} Logit‑norm sampling concentrates training samples on mid‑steps, improving stability. Inference uses a Flow‑Euler sampler: x_{t-Δt} = x_t + Δt\,v_θ(x_t, t) Multi‑level Knowledge Distillation Teacher: SD3.5‑Large‑Turbo (≈20 B parameters). Three complementary distillation streams are applied: Output distillation : L2 loss between student output and teacher output. Feature distillation : cross‑architecture (DiT → UNet) projection. The student’s final transformer layer is passed through a lightweight 2‑Conv projector to match the teacher’s final layer, then an L2 loss is applied. Time‑step‑aware scaling : instead of a fixed linear combination λ·L_task + μ·L_KD , per‑step weights w_{task}(t) and w_{KD}(t) are set proportional to the magnitude of the respective losses. Hard steps (t≈0 or 1) receive stronger teacher supervision; easy mid‑steps receive more data‑driven loss. The total loss is L = w_{task}(t)·L_{task} + w_{out}(t)·L_{outKD} + w_{feat}(t)·L_{featKD} Step Distillation Building on Latent Adversarial Diffusion Distillation (LADD [22] ), a few‑step teacher (SD3.5‑Large‑Turbo) guides the student through adversarial and output distillation. The student learns to generate high‑quality images in only 4 or 8 denoising steps, reducing inference time by >90 % while matching the 28‑step baseline on GenEval. Experimental setup Pre‑train UNet on ImageNet‑1K at 256 px for 120 epochs. Progressively fine‑tune to 512 px and then 1024 px. Use three text encoders (CLIP‑L, CLIP‑G, Gemma2‑2b) and fuse their embeddings as in SD3. Apply multi‑level KD with time‑step‑aware scaling using the large teacher. Perform step distillation with the few‑step teacher to obtain the final fast model. Results Quantitative evaluation on GenEval, DPG‑Bench, and COCO‑CLIP shows that the 0.38 B‑parameter SnapGen model outperforms much larger baselines: SnapGen beats SDXL (2.6 B), Playground (2.6 B) and IF‑XL (5.5 B) despite being 7–14× smaller. Knowledge distillation raises GenEval and DPG‑Bench scores, indicating better prompt‑following. Image‑Reward (aesthetic) scores are on par with Playground. After step distillation, 4‑step GenEval ≈ 28‑step baseline; 8‑step even closer, surpassing SDXL (50 steps) and PixArt‑α (100 steps). On‑device latency on an iPhone 16 Pro‑Max for 1024 px generation is 1.2–2.3 s for 4–8 steps, confirming real‑time feasibility. Key citations: [1][2][3][4][5][6][7][8][9][10][11][12][13][14][15][16][17][18][19][20][21][22] .