TOC-SR: Task-Optimal Compact diffusion for Image Super Resolution

Technical Summary of TOC-SR: Task-Optimal Compact Diffusion for Image Super Resolution

Problem Statement
While diffusion models have demonstrated superior performance in image restoration tasks like Single-Image Super-Resolution (SR) by modeling natural image statistics and providing strong generative priors, they face significant barriers to practical deployment. Standard diffusion pipelines require large backbone networks and iterative sampling procedures (often tens of denoising steps), resulting in high latency, memory usage, and computational cost. Existing attempts to create compact, one-step SR models often suffer from unstable optimization and degraded quality because the model must simultaneously learn the generative prior and the restoration mapping within a limited capacity.

Methodology
The authors propose TOC-SR, a framework designed to build efficient one-step SR models by first discovering a compact diffusion backbone that preserves generative fidelity. The methodology proceeds in three distinct stages:

1. Latent Capacity Expansion (4 to 16 Channels):
   Recognizing that the standard four-channel latent representation in Stable Diffusion creates an information bottleneck for high-frequency details required in SR, the authors expand the latent dimensionality to sixteen channels while maintaining the same spatial compression factor. Since no standard 16-channel Stable Diffusion model exists, they construct a new backbone by distilling the representational capacity of a high-capacity VAE (from the FLUX model family) into a modified Stable Diffusion VAE. This involves training a student encoder to match the teacher's 16-channel latent distribution while maintaining reconstruction fidelity. The diffusion U-Net is subsequently modified to accept and produce 16-channel tensors, followed by joint fine-tuning to adapt the diffusion prior to this expanded space.

2. Compact Backbone Discovery via Surrogate-based $\epsilon$-Constrained Bayesian Optimization:
   To reduce the computational cost of the expanded 16-channel model, the authors employ a search strategy rather than direct training of compact models.

   • Surrogate Construction: The teacher U-Net is decomposed into sequential stages. For each stage, a library of parameter-efficient surrogate blocks is constructed through structured pruning and module reduction (e.g., removing attention modules or simplifying residual blocks) while strictly preserving input-output tensor dimensionality.
   • Feature-wise Generative Distillation (FwKD): Each candidate surrogate block is calibrated independently using a feature-wise distillation loss against the teacher's features on a calibration subset. This avoids the need for full diffusion training during the search phase.
   • $\epsilon$-Constrained Bayesian Optimization (BO): The architecture search is formulated as a constrained optimization problem. The objective is to minimize the parameter count ($f_{param}$) subject to a constraint on generative fidelity ($f_{acc} \leq \epsilon$), where $f_{acc}$ measures the discrepancy in noise predictions between the candidate and the teacher. A Gaussian Process surrogate model guides the BO process to efficiently explore the architecture space across a grid of $\epsilon$ values, identifying architectures that offer the best trade-off between compression and fidelity.

3. Task Specialization and One-Step Distillation:
   Once the optimal compact backbone is identified, it is adapted for the SR task through conditional diffusion training, where the model learns to denoise a latent representation conditioned on a low-resolution (LR) input. Finally, this iterative diffusion SR model is distilled into a single-step generator. The student network is trained to predict the final clean high-resolution latent directly from a noisy latent and the LR condition in a single forward pass, eliminating the need for iterative denoising.

Key Contributions

• Latent Expansion for SR: The paper demonstrates that expanding the latent space from 4 to 16 channels alleviates the bottleneck for modeling high-frequency details, providing a stronger starting point for compression than standard architectures.
• Surrogate-Based Architecture Search: The authors introduce a novel framework combining feature-wise generative distillation with $\epsilon$-constrained Bayesian Optimization to discover compact diffusion backbones. This approach systematically reduces parameters while ensuring the candidate architectures remain faithful to the teacher's generative dynamics without requiring task-specific training during the search.
• Efficient One-Step SR: The resulting framework, TOC-SR, successfully integrates a compact diffusion prior with a one-step inference mechanism, achieving high-quality reconstruction without iterative sampling.

Experimental Results
Experiments were conducted on standard benchmarks (DIV2K, RealSR, DRealSR) comparing TOC-SR against state-of-the-art diffusion-based SR methods (e.g., StableSR, DiffBIR, SeeSR) and existing compact/one-step models (e.g., SinSR, OSEDiff, PocketSR).

• Efficiency: Compared to the expanded 16-channel diffusion model, the discovered TOC-SR backbone achieves a 6.6× reduction in parameters and a 2.8× reduction in GMACs.
• Performance: On the DIV2K validation set, TOC-SR achieves a PSNR of 24.26 dB and the best LPIPS score (0.242) among the listed methods, while operating in a single step.
• Real-World Generalization: On real-world datasets (RealSR and DRealSR), TOC-SR outperforms other compact baselines in perceptual metrics (MUSIQ) and maintains competitive PSNR, demonstrating robustness to real-world degradations.
• Qualitative Analysis: Visual comparisons show that TOC-SR preserves fine textures and thin structures better than other compact models, avoiding common artifacts like over-smoothing and unnatural micro-textures.

Significance and Claims
The paper claims that TOC-SR provides a stable and effective pathway for building high-quality, efficient restoration systems. By decoupling the discovery of a compact generative prior from the task-specific adaptation, the framework avoids the optimization instability often seen in directly training compact SR models. The authors assert that their approach successfully balances the generative strength of diffusion models with the efficiency required for practical deployment, offering a significant reduction in computational cost (from tens of steps to one) while maintaining strong reconstruction quality. The work suggests that discovering compact diffusion priors via constrained optimization is a viable strategy for edge-friendly image restoration.