SFBD-OMNI: Bridge models for lossy measurement restoration with limited clean samples

Official implementation of "SFBD-OMNI: Bridge models for lossy measurement restoration with limited clean samples" (ICLR 2026).

• Paper: OpenReview | arXiv
• Authors: Haoye Lu, Yaoliang Yu, Darren Lo

Overview

We address the problem of restoring a clean distribution from abundant corrupted observations with only limited clean samples. The key challenge — distribution restoration — is generally unrecoverable without any clean data, but we show that even a small number of clean samples makes the problem tractable.

We show that this task can be formulated as a one-sided entropic optimal transport problem and solved using an EM-like algorithm, in which the posterior distribution is modeled with bridge models. In this implementation, we adopt flow matching. Specifically, we provide:

• Algorithm 1 (SFBD-OMNI): Iterative training that alternates between fitting a bridge model and regenerating pseudo-clean samples from corrupted observations.
• Algorithm 2 (Online SFBD-OMNI): An end-to-end variant that avoids optimizer resets by maintaining an exponentially-decaying mixture of pseudo-clean samples, refreshed incrementally at each training step.

Installation

# Create a conda environment
conda create -n sfbd python=3.10
conda activate sfbd

# Install PyTorch (adjust cuda version as needed)
conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia

# Install remaining dependencies
pip install torchdiffeq "torchmetrics[image]" pillow tqdm

# Install the flow_matching library (from repo root)
pip install -e .

Data Preparation

Experiments use CIFAR-10 and CelebA organized as ImageFolder datasets where each subfolder is a shard of images.

The expected structure is:

data_root/
  00000/   ← clean images (shard 0, e.g. first 1000 CIFAR-10 images)
  00001/   ← clean images (shard 1)
  00002/   ← noisy/unlabeled images (shard 2)
  ...
  00049/   ← noisy/unlabeled images (shard 49)

Use --clean_folders_to_include 00000,00001 to specify the clean shards and --noisy_folders_to_include 00002..00049 for the noisy shards. The clean set typically contains ~1000–2000 images (≈2 classes of CIFAR-10), while the noisy set contains the full 50,000-image training split.

To convert CIFAR-10 into this sharded ImageFolder layout, you can use the dataset_tool.py from NVlabs/eg3d, which packs datasets into evenly-sized subfolders.

Note on output_dir

train.py appends a descriptive suffix to --output_dir based on the corruption type and key hyperparameters:

<output_dir>/<corruption_type>[_aug<p>]_tau<tau>_noisy2update_rate<rate>_min_clean_floor<floor>_start_clean<start>/

For example, with --output_dir ./output/iter0 --corruption_type grayscale --augment 0.12:

./output/iter0/grayscale_aug0.12_tau100.0_noisy2update_rate0.0_min_clean_floor0.1_start_clean0.99/

Checkpoints are saved as checkpoint.pth inside this directory.

Running the Experiments

All training is launched from sfbd_omni/:

cd sfbd_omni

Algorithm 1 — SFBD-OMNI (Iterative)

Step 1 — Pretrain on clean samples only

python train.py \
    --dataset cifar10-corruption \
    --corruption_type grayscale \
    --data_path /path/to/cifar10-imagefolder \
    --noise_data_path /path/to/cifar10-imagefolder \
    --clean_folders_to_include 00000,00001 \
    --noisy_folders_to_include 00002..00049 \
    --arch ddpmpp \
    --batch_size 64 \
    --epochs 3000 \
    --lr 1e-4 \
    --use_ema \
    --skewed_timesteps \
    --edm_schedule \
    --ode_method adaptive_heun \
    --ode_options '{"nfe": 50}' \
    --cfg_scale 0.0 \
    --augment 0.12 \
    --percent_noisy2update_per_epoch 0.0 \
    --eval_frequency 50 \
    --fid_samples 3000 \
    --compute_fid \
    --output_dir ./output/iter0

This produces a checkpoint and a suffixed directory. With the above flags:

./output/iter0/grayscale_aug0.12_tau100.0_noisy2update_rate0.0_min_clean_floor0.1_start_clean0.99/

Step 2 — Generate pseudo-clean samples with the pretrained model

Pass the same flags as Step 1 (so output_dir resolves to the same path) plus --eval_only --save_fid_samples:

ITER0_DIR="./output/iter0/grayscale_aug0.12_tau100.0_noisy2update_rate0.0_min_clean_floor0.1_start_clean0.99"

python train.py \
    --dataset cifar10-corruption \
    --corruption_type grayscale \
    --data_path /path/to/cifar10-imagefolder \
    --noise_data_path /path/to/cifar10-imagefolder \
    --clean_folders_to_include 00000,00001 \
    --noisy_folders_to_include 00002..00049 \
    --arch ddpmpp \
    --batch_size 64 \
    --use_ema \
    --skewed_timesteps \
    --edm_schedule \
    --ode_method adaptive_heun \
    --ode_options '{"nfe": 50}' \
    --cfg_scale 0.0 \
    --augment 0.12 \
    --percent_noisy2update_per_epoch 0.0 \
    --fid_samples 50000 \
    --compute_fid \
    --save_fid_samples \
    --eval_only \
    --new_iteration \
    --resume ${ITER0_DIR}/checkpoint.pth \
    --output_dir ./output/iter0

Pseudo-clean images are written to ${ITER0_DIR}/fid_samples/0/ as individual .png files.

Before retraining, copy the original clean images into a clean/ subfolder alongside the generated samples so both are used during training:

mkdir -p ${ITER0_DIR}/fid_samples/clean
cp /path/to/cifar10-imagefolder/00000/* ${ITER0_DIR}/fid_samples/clean/
cp /path/to/cifar10-imagefolder/00001/* ${ITER0_DIR}/fid_samples/clean/

Step 3 — Retrain on pseudo-clean + original clean samples

python train.py \
    --dataset cifar10-corruption \
    --corruption_type grayscale \
    --data_path ${ITER0_DIR}/fid_samples \
    --noise_data_path /path/to/cifar10-imagefolder \
    --clean_folders_to_include 0,clean \
    --noisy_folders_to_include 00002..00049 \
    --arch ddpmpp \
    --batch_size 64 \
    --epochs 3000 \
    --lr 1e-4 \
    --use_ema \
    --skewed_timesteps \
    --edm_schedule \
    --ode_method adaptive_heun \
    --ode_options '{"nfe": 50}' \
    --cfg_scale 0.0 \
    --augment 0.12 \
    --percent_noisy2update_per_epoch 0.0 \
    --eval_frequency 50 \
    --fid_samples 3000 \
    --compute_fid \
    --output_dir ./output/iter1

The bridge model trains on pairs (corrupt(x), x) where x is drawn from both the pseudo-clean samples (fid_samples/0/) and the original clean images (fid_samples/clean/). Repeat Steps 2–3 for K total iterations.

---

Algorithm 2 — Online SFBD-OMNI (End-to-End)

Algorithm 2 unifies training and pseudo-sample generation into a single continuous loop. A fraction gamma (--percent_noisy2update_per_epoch) of the pseudo-clean buffer is refreshed each epoch, eliminating optimizer resets.

python train.py \
    --dataset cifar10-corruption \
    --corruption_type grayscale \
    --data_path /path/to/cifar10-imagefolder \
    --noise_data_path /path/to/cifar10-imagefolder \
    --clean_folders_to_include 00000,00001 \
    --noisy_folders_to_include 00002..00049 \
    --arch ddpmpp \
    --batch_size 64 \
    --epochs 5000 \
    --lr 1e-4 \
    --use_ema \
    --skewed_timesteps \
    --edm_schedule \
    --ode_method adaptive_heun \
    --ode_options '{"nfe": 50}' \
    --cfg_scale 0.0 \
    --augment 0.12 \
    --eval_frequency 50 \
    --fid_samples 3000 \
    --compute_fid \
    --denoised_data_path denoised_data \
    --percent_noisy2update_per_epoch 0.005 \
    --tau 100.0 \
    --min_clean_floor 0.5 \
    --start_clean 0.5 \
    --output_dir ./output/online

When --denoised_data_path is set, the pseudo-clean buffer is initialised before training begins (epoch −1) and then refreshed incrementally. The buffer is stored at <output_dir_auto_suffixed>/denoised_data/.

Key Algorithm 2 hyperparameters:

Argument	Default	Description
--denoised_data_path	""	Path (relative to output_dir) to store pseudo-clean samples. Setting this enables Algorithm 2.
--percent_noisy2update_per_epoch	0.1	Fraction γ of noisy samples refreshed each epoch. Use 0.005–0.01 for stable training.
--tau	100.0	Time constant for the exponential decay of the clean-data sampling weight.
--min_clean_floor	0.1 (code default); 0.5 (paper)	Minimum fraction of clean data kept in the training mixture.
--start_clean	0.99 (code default); 0.5 (paper)	Initial fraction of clean data in the mixture (epoch 0).
--start_epoch_to_update	0	Epoch at which pseudo-sample refreshing begins.

---

Corruption Types

--corruption_type	Description	Additional args
grayscale	Convert to grayscale	—
gauss_blur	Gaussian blur	--gauss_blur_kernel_size 9 --gauss_blur_sigma 2.0
pixelmask	Random pixel masking	--pixelmask_p 0.6 (paper value; default is 0.5)
gaussian	Additive Gaussian noise	--gaussian_corruption_level 0.2
watermark	Diagonal text watermark	—
randompatch	Random patch replacement	—

Note on --ode_options: When using --edm_schedule, pass '{"nfe": N}' to control the number of ODE steps via EDM time discretization. Without --edm_schedule, pass '{"step_size": 0.01}' (the default) for the midpoint/heun solvers, or use an adaptive solver (dopri5, adaptive_heun) with no step_size.

---

Multi-GPU Training

The code supports torch.distributed via torchrun:

torchrun --nproc_per_node=4 train.py [... same args ...]

---

Repository Structure

flow_matching/               ← flow matching library (CondOTProbPath, ODESolver, …)
sfbd_omni/
  train.py                   ← main entry point
  train_arg_parser.py        ← all CLI arguments
  train_loop.py              ← CDM loss, Algorithm 1 & 2 training loop
  eval_loop.py               ← ODE sampling, FID evaluation, pseudo-sample generation
  utils.py                   ← weighted sampler, clean/denoised probability schedule
  datasets.py                ← ImageFolder with on-the-fly corruption
  data_transform.py          ← image transforms
  edm_time_discretization.py
  load_and_save.py           ← checkpoint utilities
  distributed_mode.py        ← torch.distributed helpers
  grad_scaler.py             ← mixed precision gradient scaler
  models/                    ← UNet and EDM (ddpmpp / ncsnpp / adm) architectures
  corruption_functions/      ← grayscale, pixelmask, watermark, randompatch, …
  dnnlib/                    ← EDM utility library (EasyDict, construct_class_by_name, …)
  torch_utils/               ← EDM torch helpers

Acknowledgements

The model and infrastructure code partially uses code from:

• Guided Diffusion (MIT License) — models/unet.py and models/nn.py
• EDM by NVIDIA (CC-BY-NC-SA 4.0) — models/EDM_net.py, dnnlib/, and torch_utils/
• flow_matching by Meta (CC-BY-NC 4.0) — the flow_matching/ library bundled in this repo is a trimmed fork, retaining CondOTProbPath, ODESolver, and ModelWrapper

Citation

@inproceedings{lu2026sfbd,
  title     = {SFBD-OMNI: Bridge models for lossy measurement restoration with limited clean samples},
  author    = {Lu, Haoye and Yu, Yaoliang and Lo, Darren},
  booktitle = {International Conference on Learning Representations},
  year      = {2026},
  url       = {https://openreview.net/forum?id=28IuGdneJQ}
}

License

This project is licensed under the CC-BY-NC 4.0 License — see the LICENSE file for details.