NanoSD: Edge Efficient Foundation Model for Real Time Image Restoration

AI summary

Overview

Research area: Computer vision — efficient latent diffusion models for image restoration, with a focus on on-device (edge/mobile NPU) deployment.

Technical level: Advanced. The paper assumes familiarity with latent diffusion models, knowledge distillation, Fréchet Inception Distance, and multi-objective optimization.

Scope: The paper introduces NanoSD, a family of lightweight diffusion backbones distilled from Stable Diffusion 1.5 through hardware-aware network surgery, block-level feature distillation, and Bayesian multi-objective search, then validates them as drop-in backbones across super-resolution, face restoration, deblurring, dehazing, deraining, desnowing, and monocular depth estimation.

What This Paper Is About

Diffusion models like Stable Diffusion 1.5 carry strong generative priors that are valuable for restoring degraded images, but their full pipelines are too heavy to run on phones and other edge chips. Prior lightweight variants mostly shrink the denoising U-Net or shorten the diffusion trajectory, which the authors argue disrupts the latent manifold and limits generalization to more than one task. The goal of NanoSD is to compress the entire pipeline — U-Net plus VAE encoder and decoder — in a way that is guided by measured hardware latency rather than theoretical FLOPs, while keeping the model usable as a general-purpose backbone for many restoration tasks.

Key Contributions

1. Hardware-aware reformulation of the SD 1.5 U-Net. The U-Net is decomposed into stage-wise dimensions, and shape-preserving block variants (residual-only, reduced-attention, hybrid) are constructed for each stage and profiled on the target edge accelerator. The design space contains 32,768 (4 × 4 × 4 × 8 × 8 × 8) candidate U-Net architectures.

2. Block-level generative distillation. Each candidate block variant is distilled independently against its corresponding SD 1.5 teacher block using an ℓ2 feature-matching loss, producing 30 (3+3+3+7+7+7) distilled surrogate blocks and avoiding full-model retraining for every candidate.

3. Multi-objective Pareto search. Backbone selection is framed as a bi-objective problem over teacher-aligned FID (taFID) together with either on-device latency or parameter count, solved with Bayesian optimization using Gaussian Process surrogates and Expected Hypervolume Improvement (EHVI), yielding seven Pareto-optimal architectures — the NanoSD family.

4. Full-pipeline VAE distillation and cross-task validation. From the Pareto set, a balanced U-Net (Model 2, "NanoSD-Prime") is frozen and the corresponding VAE encoder and decoder are distilled from SD 1.5 to complete a fully lightweight latent diffusion pipeline, which is then evaluated across super-resolution, face restoration, deblurring, dehazing, deraining, desnowing, and monocular depth estimation.

Main Findings

• The NanoSD family spans 130M–315M parameters. Seven Pareto-optimal architectures were obtained: five from the latency–taFID front and three from the parameter–taFID front, with one model appearing in both.

• Pareto-optimal operating points (Table 1, Qualcomm NPU, 8-bit weights / 16-bit activations). NanoSD 1: 41 ms, 309M, taFID 10. NanoSD 2: 27 ms, 315M, taFID 10. NanoSD 3: 24 ms, 306M, taFID 10.5. NanoSD 4: 20 ms, 297M, taFID 11.1. NanoSD 5: 12 ms, 170M, taFID 17.3. NanoSD 6: 28 ms, 160M, taFID 18.2. NanoSD 7: 27 ms, 130M, taFID 22.

• Model naming and selection. Model 5 is the lowest-latency point (NanoSD-Latency); Model 7 has the smallest parameter footprint (NanoSD-Parameters); Models 1 and 2 achieve the best taFID (both 10), but Model 1 has substantially higher latency, so Model 2 is chosen as NanoSD-Prime for all downstream experiments.

• Baselines fall outside the frontier. The hand-tuned SD 1.5 variant records 53 ms, 276M, taFID 20.6, and Segmind TinySD records 74 ms, 323M, taFID 13.8 — both worse trade-offs than the NanoSD points.

• Latency does not track parameter count. The paper reports that parameter reduction alone does not correlate with hardware efficiency: shape-compatible block alternatives show 3–8× lower latency despite comparable parameter counts, and several variants diverge sharply (for example, within Decoder 1, the R variant measures 1.7 ms while the RARA variant measures 39 ms).

• Cross-platform generalization. On the Apple A17 Pro Neural Engine, the same U-Nets — originally optimized using Qualcomm SM8750 data — reproduce the relative latency ordering: NanoSD 2 measures 27 ms on Qualcomm and 38 ms on Apple; NanoSD 5 measures 12 ms on Qualcomm and 20 ms on Apple; TinySD measures 74 ms on Qualcomm and 192 ms on Apple.

• Super-resolution results (DIV-2K Val). Nano-OSEDiff attains the best PSNR (24.29) and SSIM (0.628) among the listed methods, with 340G MACs and 448M parameters at 1 step. Nano-S3Diff attains the best NIQE (4.09) and MUSIQ (70.44) and the second-best FID (22.34), with 285G MACs and 318M parameters at 1 step. For comparison, TinySR reports FID 22.94 with 427G MACs and 341M parameters, and PocketSR reports 225G MACs and 146M parameters.

• Face restoration results (synthetic CelebA-Test). Nano-OSDFace achieves 0.341 LPIPS, 0.182 DISTS, 3.913 NIQE, 76.01 MUSIQ, 45.92 FID, and 5.172 LMD at 1 step, with 479G MACs and 415M parameters, close to OSDFace (0.336 LPIPS, 0.177 DISTS, 3.884 NIQE, 75.64 MUSIQ, 45.41 FID, 5.286 LMD) but far cheaper than OSDFace's 2465G MACs and 1887M parameters. Nano-DiffBIR reaches 70.89 FID and 5.231 LMD at 50 steps with 4734G MACs and 726M parameters, versus DiffBIR's 71.77 FID and 5.104 LMD at 24234G MACs and 1717M parameters.

• Multi-task restoration (Table 4). Nano-Diff-Plugin records FID/KID of 34.83/5.63 (desnowing), 35.23/15.03 (dehazing), 52.41/15.48 (deblurring), and 50.78/14.21 (deraining) at 20 steps, with 17120G MACs and 712M parameters, compared with Diff-Plugin's 34.30/5.20, 34.68/14.38, 51.81/14.63, and 50.55/13.84 at 30400G MACs and 1256M parameters.

• Monocular depth estimation (Table 5). Nano-Marigold reports AbsRel 7.2 and δ1 94.6 on NYU V2, and AbsRel 11.8 and δ1 86.3 on KITTI. The strongest listed baselines are Marigold (5.5, 96.4 on NYU V2; 9.9, 91.6 on KITTI) and HDN (6.9, 94.8 on NYU V2; 11.5, 86.7 on KITTI).

• Generative prior is preserved. Latent-space interpolation between two random seeds produces smooth, semantically consistent transitions for both SD 1.5 and NanoSD. Quantitatively, NanoSD vs SD 1.5 measures 0.57 LPIPS and 0.84 embedding cosine similarity, close to SD 1.5's own run-to-run variation (0.48 LPIPS, 0.89 cosine similarity), while a regression-based U-Net baseline measures 1.92 LPIPS and 0.41 cosine similarity.

• SD 1.5 is not deployable as-is. SD 1.5 contains approximately 829 million parameters, roughly 3.3 GB of raw weights in full precision and about 0.8 GB even with INT8 quantization. The authors report that deployment attempts failed to generate executable binaries for Qualcomm NPUs (Android) and the Apple Neural Engine (iOS) at FP32, INT16, or INT8. With aggressive INT4 quantization, the average latency was 116 ms per tile on the Qualcomm Snapdragon SM8750 NPU, implying about 10.2 seconds end-to-end (88 tiles × 116 ms per tile) for a 1000 × 750 input at a single diffusion step.

• Tiled inference for high resolution. A 1000 × 750 input is split into 128 × 128 overlapping tiles with 25% overlap, producing 88 tiles that are processed independently and reassembled into a final 4K output.

Methodology in Plain English

The researchers start from Stable Diffusion 1.5 and treat its U-Net as a set of separable stages. Following prior lightweight-diffusion findings (SnapFusion, Segmind TinySD, Koala) that the deepest encoder, the middle block, and the deepest decoder contribute minimally to fidelity, they remove those three stages, leaving three encoders (E1–E3) and three decoders (D3–D1) as independent design axes.

For each retained stage, they build alternative block configurations — written as strings like R, RA, RR, RRA, RAR, RRRA, RARR, RARA, where R is a residual module and A an attention module — that keep the exact same input and output tensor shapes. Each alternative is then measured on a Qualcomm NPU (SM8750) to capture real latency rather than theoretical FLOPs.

Because training all 32,768 combinations is infeasible, they distill each candidate block separately against its SD 1.5 teacher block using a squared-error feature-matching loss. This is massively parallel and cheap, and it means any composition of distilled blocks still produces a valid U-Net that behaves locally like SD 1.5.

A complete architecture is then described by a six-element vector selecting one surrogate per stage. Two bi-objective searches are run: (FID, latency) and (FID, parameters), where the FID is a "teacher-aligned" taFID computed by generating images from identical prompts and seeds as SD 1.5 and comparing distributions. Bayesian optimization with a continuous relaxation of the discrete vector, Gaussian Process models, and Expected Hypervolume Improvement selects candidates efficiently.

From the resulting Pareto fronts, one balanced U-Net is frozen, and the VAE encoder and decoder are distilled from SD 1.5 using feature-matching losses. Finally, each NanoSD candidate is fine-tuned end-to-end with the standard diffusion denoising objective on image-text pairs from LAION-style web-scale data to correct accumulated discrepancies from block-level distillation.

Why This Matters

Impact on research. The work argues that parameter count and FLOPs are poor proxies for on-device latency, and it shows a concrete methodology — profiling shape-preserving block alternatives directly on the target accelerator and embedding that cost into the search objective — that generalizes across chips (the Qualcomm-derived architecture ordering holds on Apple's ANE). It also shows that a distilled backbone can be swapped into several existing frameworks (OSEDiff, S3Diff, OSDFace, Diff-Plugin, DiffBIR, Marigold) rather than being tied to one task.

Real-world applications.

• On-device photo enhancement in smartphone camera pipelines, where a 4K capture is restored in real time.
• Face restoration in mobile photo galleries and video calling, where identity preservation (measured by LMD) matters.
• Real-time depth estimation for AR and computational photography on phones and headsets.
• Adverse-weather image cleanup (dehazing, deraining, desnowing) for autonomous driving and surveillance feeds.

Industry relevance. The authors are from Samsung Research India, Bangalore, and the entire framing targets commercial mobile NPUs (Qualcomm Snapdragon SM8750, Apple A17 Pro) and the memory and binary-compilation constraints of real devices. The paper explicitly notes that SD 1.5 could not be compiled into an executable binary on either Qualcomm NPUs or the Apple Neural Engine, which is a deployment blocker that published accuracy numbers alone do not capture.

Future Directions

• Beyond latency and parameters. The paper optimizes taFID against latency and against parameter count; additional objectives such as peak memory, energy per inference, or thermal budget are not explored.
• Hardware-specific versus hardware-agnostic search. The search space was profiled on a single accelerator (Qualcomm SM8750). The authors demonstrate cross-platform transfer to Apple's ANE without modification, but whether a joint multi-device objective would produce a better shared frontier remains open.
• Task-specific Pareto adaptation. Models 1, 3, 4, and 6 are described as candidates for task-specific or resource-specific constraints (referred to as NanoSD-S and NanoSD-XS), but the paper does not systematically report which task benefits from which operating point.
• Depth and depth-of-quality gap. Nano-Marigold trails the strongest listed depth baselines (Marigold at 5.5 AbsRel / 96.4 δ1 on NYU V2 versus Nano-Marigold at 7.2 / 94.6), leaving room to study whether a different Pareto point or a depth-specific conditioning path closes that gap.

Target Audience

Researchers and engineers working on efficient diffusion models, on-device or edge machine learning, and mobile ISP/Camera pipelines. It is most useful to readers who already understand latent diffusion, knowledge distillation, and hardware-aware neural architecture search, and to practitioners who need to choose an accuracy–latency–size operating point for a specific mobile chip. Readers looking for a beginner-friendly introduction to diffusion restoration should look elsewhere first, given the density of architecture notation, distillation losses, and multi-objective optimization detail.

Authors’ abstract

Latent diffusion models such as Stable Diffusion 1.5 offer strong generative priors that are highly valuable for image restoration, yet their full pipelines remain too computationally heavy for deployment on edge devices. Existing lightweight variants predominantly compress the denoising U-Net or reduce the diffusion trajectory, which disrupts the underlying latent manifold and limits generalization beyond a single task. We introduce NanoSD, a family of Pareto-optimal diffusion foundation models distilled from Stable Diffusion 1.5 through network surgery, feature-wise generative distillation, and structured architectural scaling jointly applied to the U-Net and the VAE encoder-decoder. This full-pipeline co-design preserves the generative prior while producing models that occupy distinct operating points along the accuracy-latency-size frontier (e.g., 130M-315M parameters, achieving real-time inference down to 20ms on mobile-class NPUs). We show that parameter reduction alone does not correlate with hardware efficiency, and we provide an analysis revealing how architectural balance, feature routing, and latent-space preservation jointly shape true on-device latency. When used as a drop-in backbone, NanoSD enables state-of-the-art performance across image super-resolution, image deblurring, face restoration, and monocular depth estimation, outperforming prior lightweight diffusion models in both perceptual quality and practical deployability. NanoSD establishes a general-purpose diffusion foundation model family suitable for real-time visual generation and restoration on edge devices.