---
title: 'FoundIR-v2: Diffusion-Based Image Restoration'
url: https://www.emergentmind.com/topics/foundir-v2
type: topic
---

# FoundIR-v2: Diffusion-Based Image Restoration

FoundIR-v2 is a high-capacity diffusion-based image restoration foundation model that leverages dynamic pre-training data mixture optimization and a Mixture-of-Experts (MoE) scheduler to address over 50 sub-tasks such as deblurring, dehazing, denoising, super-resolution, deraining, desnowing, and low-light enhancement within a unified framework. It is predicated on the observation that the proportions of task-specific datasets in the pre-training mixture directly affect multi-task performance, motivating the development of a generalizable architecture that couples diffusion modeling with data equilibrium scheduling for large-scale restoration [2512.09282].

## 1. Architectural Design and Objectives

FoundIR-v2 is constructed to serve as an all-in-one restoration foundation model, supporting over 50 diverse sub-tasks. Its design jointly optimizes (i) the mixture ratios of task-specific data via "data equilibrium scheduling" (DES) to prevent task imbalance, and (ii) an MoE-driven diffusion scheduler that provides task-adaptive generative priors in latent space. The architecture integrates a Stable Diffusion XL (SDXL) backbone for latent-space denoising, a frozen VAE encoder-decoder for mapping between image and latent domains, and a learned MoE scheduler applied at each diffusion timestep.

Key components include:
- **VAE Encoder (\(E_{\mathrm{VAE}}\))**: Transforms low-quality input \(I_{\mathrm{LQ}}\) to latent codes \(f^{\mathrm{LQ}}\), with stochastic resolution alignment for super-resolution tasks.
- **SDXL Denoiser (\(D_\theta\))**: Pre-trained backbone that produces clean latent codes \(f^{\mathrm{HQ}}\).
- **MoE Scheduler**: Injected at each diffusion step, selects among \(n\) expert blocks to condition the denoiser on both \(f^{\mathrm{LQ}}\) and the noisy latent \(x_t^{\mathrm{HQ}}\).
- **Data Equilibrium Scheduler**: Adjusts task sampling weights \(\{\lambda_i\}\) every \(T\) steps to maintain balanced learning signals across tasks.
- **VAE Decoder (\(D_{\mathrm{VAE}}\))**: Maps the final denoised latent code back to RGB space.

This configuration, together with multi-modal (including text) prompts, is intended to enhance the model's ability to generalize across heterogeneous and previously unseen image degradations.

## 2. Data Equilibrium Scheduling Paradigm

Central to FoundIR-v2 is the DES paradigm, which seeks optimal proportions in the mixture of \(k\) task-specific datasets \(\{\mathcal{D}_i\}_{i=1}^k\), where the overall sampling distribution is:
\[
P_\lambda = \sum_{i=1}^k \lambda_i\,\mathrm{Unif}(\mathcal{D}_i)\,, \quad \sum_{i=1}^k \lambda_i=1\,.
\]
The model \(\mathcal{M}_\theta\) is trained to minimize an \(\ell_1\) reconstruction loss:
\[
\mathcal{L}_\mathrm{rec}(\theta) = \mathbb{E}_{(I_{\mathrm{LQ}},I_{\mathrm{HQ}})\sim P_\lambda} \left\|I_{\mathrm{HQ}} - \mathcal{M}_\theta(I_{\mathrm{LQ}})\right\|_1\,.
\]

Every \(T\) training steps, the scheduler evaluates held-out task reference sets to compute score differentials \(\Delta s_i^{(t)}\) and updates mixing weights via softmax re-weighting:
\[
\lambda_i^{(t+1)}
  = \frac{ \lambda_i^{(t)} \exp(-\alpha\,\Delta s_i^{(t)}) }
         {\sum_{j=1}^k \lambda_j^{(t)} \exp(-\alpha\,\Delta s_j^{(t)}) }
\]
where \(\alpha>0\) is a tunable coefficient.

**Pseudocode for DES:**
```python
Input: initial λ⁽⁰⁾, interval T, max steps I
for t = 1 to I do
  Sample a mini-batch from each task i according to λ⁽ᵗ⁾
  Take T steps of diffusion-based training on those samples
  if t mod T == 0 then
    Evaluate s_i⁽ᵗ⁾ on D_ref for each task i
    Compute Δs_i = s_i⁽ᵗ⁾ - s_i⁽ᵗ⁻ᵀ⁾
    Update λ_i⁽ᵗ⁺¹⁾ by softmax re-weighting
  end if
end for
```
This iterative dynamic re-weighting enforces the "Data Mixing Law," which underlies FoundIR-v2's balanced multi-task convergence.

## 3. MoE-Driven Diffusion Scheduler

At each diffusion timestep \(t\) and for sub-task \(k\), the MoE scheduler fuses latent codes by concatenation:
\[
\mathbf{z}_t^{(k)} = \phi(f^{\mathrm{LQ}_k}, x_{t,k}^{\mathrm{HQ}})
\]
where each of \(n\) expert blocks \(E_i(\cdot)\) implements specialized attention (spatial, channel, sparse, etc.). The scheduler computes soft-gate weights:
\[
w_i^{(k)} = \frac{ \exp(\mathbf{g}_i^\top\,\mathbf{z}_t^{(k)}) }
                 { \sum_{j=1}^n \exp(\mathbf{g}_j^\top\,\mathbf{z}_t^{(k)}) }, \quad \sum_i w_i^{(k)}=1
\]
and forms the scheduled feature
\[
\mathcal{F}^{(k)}(z_t) = \sum_{i=1}^n w_i^{(k)} E_i(z_t^{(k)})
\]
which is passed to the SDXL noise predictor.

Diffusion training follows the denoising-score matching objective:
\[
q(x_t \mid x_0) = \mathcal{N}(x_t; \sqrt{\bar\alpha_t}x_0, (1-\bar\alpha_t)\mathbf{I}),\\
\mathcal{L}_\mathrm{diff} = \mathbb{E}_{t,\epsilon,x_0} \left\| \epsilon - \epsilon_\theta(x_t, t, \mathcal{F}^{(k)}) \right\|_2^2
\]
Training proceeds by first isolating MoE head pre-training (with frozen SDXL) and later joint end-to-end fine-tuning.

## 4. Training Protocol and Implementation

FoundIR-v2 is trained on a combination of publicly available datasets encompassing 50+ real-world sub-tasks, including but not limited to 4KRD (motion deblur), LSDIR (defocus deblur), PolyU (denoise), Dense-HAZE/NH-HAZE (dehaze), CSTNet HQ-NightRain (derain), UAV-Rain1k (raindrop removal), WeatherBench (desnow), UHD-LL (low-light), DIV2K/Flickr2K/DIV8K (super-resolution), FFHQ (faces), RealPhoto60 (real SR), and RealDeg (old-photo/face restoration). High-quality ground truth filtering is performed using deep multi-modal IQA metrics such as DA-CLIP and DepictQA.

The key hyperparameters are:
- **Hardware**: 2× NVIDIA H20 GPUs (96 GB each)
- **Batch size**: 32 (random \(512\times512\) crops)
- **Optimizer**: AdamW with default weight decay
- **Learning rate**: VAE encoder \(5\times10^{-6}\), others \(5\times10^{-5}\) (cosine annealing)
- **Total iterations**: 150k; evaluation interval \(T=1\)k, 10 reference samples per task
- **Diffusion inference**: Euler sampler, 20 steps, classifier-free guidance scale = 5; AdaIN color fix for SR tasks

This protocol supports scalable, balanced exposure to the full spectrum of restoration phenomena.

## 5. Empirical Results, Ablations, and Analysis

FoundIR-v2 achieves leading or near-leading performance on 80% or more of evaluated tasks, using the following metrics:
- **Full-reference**: PSNR↑, SSIM↑, LPIPS↓, MUSIQ↑, CLIPIQA+↑
- **No-reference**: PIQE↓, MANIQA↑, PaQ-2-PiQ↑

Table: Representative performance (DES vs. static mixing)

| Task        | Staticmix PSNR (dB) | DES PSNR (dB) | ∆ (dB)   |
|-------------|---------------------|---------------|----------|
| Deblurring  | 18.91               | 20.41         | +1.50    |
| Dehazing    | 18.69               | 19.93         | +1.24    |
| Low-light   | 19.93               | 20.41         | +0.48    |
| SR          | 18.91               | 20.09         | +1.18    |

Ablation studies indicate:
- DES provides +1.2–1.5 dB gain vs. static mixing across tasks.
- Soft MoE scheduling yields +0.3–0.7 dB vs. single-prior or hard MoE variants.
- Removing low-quality ground truth increases PSNR by ~0.2–0.4 dB.

Qualitative results demonstrate restoration of sharp edges in motion blur, detail retention in SR, superior handling of heterogeneous murals (outperforming GPT-5 and HYPIR baselines), and effective simultaneous resolution of cascaded tasks such as deraining plus SR. FoundIR-v2 also outperforms pipeline architectures (FoundIR + SUPIR) in joint restoration settings.

Generalization extends to medical imaging domains; limited fine-tuning enables superior recovery of diagnostic structures in laparoscopy and microscopy compared to the prior FoundIR.

## 6. Significance, Limitations, and Open Directions

FoundIR-v2 establishes the critical importance of dynamic data mixture balancing—formalized as the "Data Mixing Law"—in achieving robust all-in-one restoration. Coupling with an MoE-driven diffusion scheduler enables task-adaptive prior generation in latent space, promoting strong generalization across a diverse sub-task landscape and favorable zero-shot transfer.

Limitations persist with respect to some extreme degradations where task-specialized models outperform the all-in-one approach. Scheduling overhead incurs modest added complexity, and scope remains for lighter-weight MoE variants.

Open directions include:
- Developing adaptation signals beyond PSNR or MUSIQ for DES updates.
- Scaling to temporal (video) or multi-modal (e.g., depth, semantics) restoration contexts.
- Enabling continual learning for incremental addition of new tasks.
- Investigating compact MoE schedulers for resource-constrained deployment.

FoundIR-v2 offers a scalable foundation for multi-task, real-world restoration, and highlights the role of adaptive data mixing for foundation models in image processing [2512.09282].

Source: https://www.emergentmind.com/topics/foundir-v2