PaperlayerView exact revision on arXiv

Source-linked AI summary

I$^2$SB: Image-to-Image Schrödinger Bridge

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, Anima Anandkumar

arXiv:2302.05872v3cs.CVcs.LGstat.ML

TL;DR

Image restoration is ill-posed, and existing Schrödinger Bridge methods are difficult to train at scale for complex nonlinear diffusions. I2SB introduces tractable diffusion bridges with analytic marginals and simulation-free, scalable training, achieving strong restoration performance with interpretable generation.

  • Problem

    Image restoration is ill-posed, while existing Schrödinger Bridge computational frameworks are unfavorable for high-dimensional nonlinear diffusion learning because of discretization error, high variance, or divergence.

  • Method

    I2SB constructs tractable nonlinear diffusion bridges between clean and degraded image distributions, with analytically computable marginals and a simulation-free training framework using scalable diffusion-model techniques.

  • Results

    I2SB sets new records across super-resolution, deblurring, inpainting, and JPEG restoration, with more interpretable generation and little performance drop as the NFE decreases.

  • Takeaways & Limitations

    Diffusion bridges can directly leverage degraded images as structurally informative priors for image restoration while retaining scalable training and interpretable generation.

  • Takeaways & Limitations

    Existing Schrödinger Bridge methods can suffer from discretization error, high variance, or divergence, leaving practical large-scale nonlinear diffusion learning an open question.

Abstract

from arXiv · show

We propose Image-to-Image Schrödinger Bridge (I$^2$SB), a new class of conditional diffusion models that directly learn the nonlinear diffusion processes between two given distributions. These diffusion bridges are particularly useful for image restoration, as the degraded images are structurally informative priors for reconstructing the clean images. I$^2$SB belongs to a tractable class of Schrödinger bridge, the nonlinear extension to score-based models, whose marginal distributions can be computed analytically given boundary pairs. This results in a simulation-free framework for nonlinear diffusions, where the I$^2$SB training becomes scalable by adopting practical techniques used in standard diffusion models. We validate I$^2$SB in solving various image restoration tasks, including inpainting, super-resolution, deblurring, and JPEG restoration on ImageNet 256x256 and show that I$^2$SB surpasses standard conditional diffusion models with more interpretable generative processes. Moreover, I$^2$SB matches the performance of inverse methods that additionally require the knowledge of the corruption operators. Our work opens up new algorithmic opportunities for developing efficient nonlinear diffusion models on a large scale. scale. Project page and codes: https://i2sb.github.io/

1. Introduction

I2SB addresses the mismatch between image restoration and noise-initialized diffusion by learning tractable diffusion bridges between degraded and clean image distributions. It enables scalable, simulation-free nonlinear diffusion and performs strongly across restoration tasks.

  • Motivation: Image restoration is ill-posed, and degraded images provide more structural information about clean images than random noise.This motivates starting generation from degraded images rather than Gaussian white noise.
  • Prior limitations: Existing Schrödinger bridge methods are computationally unfavorable in high-dimensional settings because their distinct training frameworks incur discretization error, high variance, or divergence.On 256×256 resolution, SB is 6× slower and consumes 3× memory than SGM.
  • Results: I2SB provides more interpretable generation and exhibits little or no performance drop as the number of function evaluations decreases.Its generation progressively restores intact images from degradations.
  • I2SB: I2SB uses a tractable Schrödinger bridge class with analytically computable marginals and a simulation-free training framework.Its framework adopts practical techniques from standard score-based diffusion models.
  • Results: I2SB surpasses standard conditional diffusion models and matches diffusion-based inverse models across super-resolution, deblurring, inpainting, and JPEG restoration on ImageNet 256×256.Unlike inverse methods, I2SB does not exploit corruption operators.
  • I2SB: I2SB directly learns fully nonlinear diffusion bridges between clean and degraded distributions as a conditional diffusion model.It is designed for image-to-image restoration rather than generation from random noise.

2. Preliminaries

This section introduces score-based generative models and Schrödinger bridges through forward and backward stochastic processes, their marginal densities, and practical score-based training. Schrödinger bridges extend score-based models with nonlinear transport between boundary distributions.

  • Score-based Generative Models: Score-based generative models construct stochastic differential equations that transform data into an approximate Gaussian terminal distribution and reverse them for generation.The forward process starts from X0 sampled from pA, while the reverse process traverses from noise toward data using learned score functions.
  • Score-based Generative Models: The reverse SDE shares the forward process's path-wise measure and marginal densities, with its drift determined by the time-dependent score ∇log p.This equivalence makes the reversed process a valid generative dynamics corresponding to the forward diffusion.
  • Score-based Generative Models: Training samples tuples (X0, t, Xt), where Xt is sampled analytically, and regresses a neural network's output using a rescaled denoising score-matching objective.The variance of Xt|X0 supplies the rescaling that gives the regression target unit variance.
  • Schrödinger Bridges: Schrödinger bridges define entropy-regularized optimal transport through forward and backward SDEs connecting boundary distributions pA and pB.Time-varying energy potentials Ψ and bΨ solve coupled PDEs, and the two SDEs induce equivalent path measures and marginal densities.
  • Schrödinger Bridges: Schrödinger bridges generalize score-based models by adding a nonlinear forward drift ∇logΨ, enabling transport beyond Gaussian priors.The backward drift ∇logbΨ is related to the forward process through Nelson’s duality and reverses the nonlinear forward SDE.

3. Image-to-Image Schr¨odinger Bridge (I2SB)

I2SB constructs tractable Schrödinger bridges between clean and degraded image distributions and reformulates their drifts for score-based training. Analytic boundary-pair posteriors enable simulation-free training and generation, while the zero-noise limit yields an optimal-transport formulation.

  • 3. Image-to-Image Schrödinger Bridge: I2SB proposes a tractable Schrödinger bridge class for directly transporting between two image domains, including clean and degraded distributions.The construction targets image-to-image translation and image restoration.
  • 3. Image-to-Image Schrödinger Bridge: Theorem 3.1 reformulates nonlinear Schrödinger bridge drifts as score functions of auxiliary linear SDEs.The nonlinear component is absorbed into boundary distributions, allowing score-network parameterization.
  • 3. Image-to-Image Schrödinger Bridge: Corollary 3.2 removes one coupling by using a Dirac delta boundary, making the backward drift tractable and directed toward the clean sample.This boundary choice generalizes the denoising objective’s per-sample construction beyond Gaussian terminal distributions.
  • 3. Image-to-Image Schrödinger Bridge: The Dirac delta boundary may hinder generalization beyond training samples, although the authors report practical generalization to unseen samples from the same distributions.The optimality statement is with respect to minimum energy.
  • 3.2. Algorithmic Design: I2SB requires paired clean/degraded data during training and adopts a standard score-network parameterization without conditioning modules.Training samples X0, X1, and an intermediate Xt before applying the diffusion objective.
  • 3.2. Algorithmic Design: The analytic posterior given a boundary pair lets training sample intermediate states directly and lets generation start from degraded images using standard DDPM updates.This avoids solving the nonlinear forward diffusion and preserves the marginal densities of the Schrödinger bridge paths when the predicted clean image is accurate.
  • 3.3. Connection to Flow-based Optimal Transport (OT): When diffusion noise vanishes, the Schrödinger bridge becomes an ODE whose posterior mean follows optimal-transport paths, enabling a simulation-free OT variant.The posterior mean reduces to the OT displacement interpolation µt = (1−t)X0 + tX1.

4. Related Work

I2SB is positioned against conditional score-based restoration models and diffusion-based inverse models. It directly bridges clean and degraded image distributions without requiring corruption operators.

  • Conditional score-based models: Conditional score-based restoration models bias diffusion toward intact images by conditioning the network on degraded measurements.They have been applied to deblurring, super-resolution, and inpainting, but begin generation from noise.
  • I2SB: Unlike these approaches, I2SB directly bridges clean and degraded image distributions and does not require knowledge of the corruption operators.This makes it applicable without task-specific operator manipulations.
  • Diffusion-based inverse models: Diffusion-based inverse models restore clean signals through measurement-aligned projections during generation, but require the degradation operators.The operators are needed during both training and test time.

5. Experiment

I2SB is evaluated across diverse ImageNet restoration tasks and compared with conditional diffusion, inverse, and standard Schrödinger bridge baselines. It delivers strong task coverage, interpretable restoration trajectories, and improved sampling efficiency, while requiring paired training data.

  • Experimental Setup: I2SB is evaluated on inpainting, JPEG restoration, deblurring, and 4× super-resolution tasks on ImageNet 256×256.Each restoration problem includes multiple settings such as quality factors, filtering kernels, or mask types.
  • Quantitative Results: I2SB surpasses standard conditional diffusion models on six of nine tasks, including super-resolution, JPEG restoration, and all evaluated inpainting masks.Although ADM and Palette achieve higher CA on some settings, I2SB obtains lower FID there.
  • Quantitative Results: I2SB achieves state-of-the-art FID on seven of nine tasks and matches diffusion-based inverse models without corruption operators during training or generation.It also sets new CA records for both JPEG quality factors and Freeform 10–20% inpainting, while outperforming standard SB on all tasks.
  • Interpretability and Efficiency: I2SB progressively restores clean images from degradations, producing more interpretable generative processes than noise-starting diffusion models.The bridge begins from a structurally informative degraded image rather than random Gaussian noise.
  • Interpretability and Efficiency: I2SB maintains performance better than Palette as the number of function evaluations decreases.For Freeform 20–30% inpainting, I2SB reaches similar best performance with 2–10 NFEs, whereas Palette requires at least 100.
  • Sampling and Ablations: Sampling from q(Xt|X0, X1) is important: both metrics deteriorate as proposals move toward the distribution induced by the forward SDE.The posterior proposal better covers regions traversed by the generative process.
  • Sampling and Ablations: Replacing stochastic sampling with OT-ODE favors deterministic restoration such as deblurring but is biased against high-uncertainty tasks such as JPEG restoration.This result highlights a task-dependent role for stochasticity in the generative dynamics.
  • Limitations: The tractable I2SB framework requires paired clean and degraded data during training, limiting applications to unpaired translation settings.The paper identifies unpaired tasks such as CycleGAN-style translation as outside this current scope.

6. Conclusion

The paper concludes that I2SB transports between clean and degraded image distributions through a tractable Schrödinger bridge. It reports interpretable generation, sampling efficiency, and new image-restoration records.

  • Conclusion: I2SB is a conditional diffusion model that transports between clean and degraded image distributions using a tractable class of Schrödinger bridge.The conclusion presents this bridge construction as the basis of the method.
  • Conclusion: I2SB yields interpretable generation, sampling efficiency, and new records on image restoration.The conclusion also identifies combining I2SB with inverse-problem techniques as a future direction.

A. Proof

The appendix establishes analytical and limiting properties of the Schrödinger bridge construction, including its Fokker–Planck formulation, posterior marginals, and deterministic OT-ODE limit.

  • PDE and SDE Relations: The Schrödinger bridge construction is connected to Fokker–Planck equations through its coupled PDE formulation.The proof interprets the PDEs as density evolutions for the forward and reverse SDEs.
  • Boundary Constraints: The boundary constraints are satisfied by constructing the forward and backward potentials with the required endpoint distributions.The proof verifies the remaining constraint through the transformed Hamilton–Jacobi–Bellman solution.
  • Reduction to SGM: When the terminal forward distribution is Gaussian, the Schrödinger bridge reduces to the score-based generative-model setting.The proof sets the corresponding potential to one throughout time under the stated terminal condition.
  • Posterior Marginals: The posterior q(Xt|X0, X1) has an analytic Gaussian form and matches the corresponding DDPM posterior marginal.The equality is established inductively using Gaussian means and covariances.
  • Deterministic Limit: As βt approaches zero, the posterior variance vanishes while its mean remains, yielding the deterministic trajectory Xt = µt(X0, X1).The stochastic diffusion term disappears and the drift approaches a vector field, producing the OT-ODE.
  • Deterministic Limit: For constant small β, the resulting trajectory becomes µt = (1 − t)X0 + tX1 with constant velocity between the endpoints.This provides the linear interpolation interpretation of the OT-ODE limit.

B. Introduction to Schr¨odinger Bridge

Schrödinger bridges formulate transport between endpoint distributions as an optimally controlled diffusion. Their coupled PDEs and stochastic control interpretation provide the theoretical basis for I2SB’s bridge construction.

  • Bridge Formulation: The dynamic Schrödinger bridge connects path measures with prescribed marginal densities at the two temporal boundaries.Its formulation is related to optimal transport and stochastic optimal control.
  • Stochastic Control: The stochastic-control formulation minimizes accumulated control energy while satisfying the boundary distributions.The optimal control is represented through the gradient of a Schrödinger potential.
  • PDE Formulation: The coupled Schrödinger bridge PDEs arise from applying the Hopf–Cole transform to the optimal-control conditions.This yields the controlled SDE representation used by the bridge model.

C.1. Additional Experimental Setup

The additional setup specifies task-specific corruption procedures, inpainting handling, evaluation protocols, and a controlled Palette implementation for reproducible comparisons.

  • Task-specific setup: Deblurring and JPEG restoration use implementations from prior diffusion inverse methods, with FID computed on 10k ImageNet validation reconstructions.Deblurring follows Kawar et al. (2022a), while JPEG restoration follows Kawar et al. (2022b).
  • Task-specific setup: For 4× super-resolution, 64×64 images are upsampled to 256×256 before I2SB processes same-sized clean and degraded images.FID is evaluated on the entire ImageNet validation set against training-set reference statistics.
  • Ablation: OT-ODE shows severe degradation with noiseless inpainting masks but slightly better results after additional noise is injected into masked regions.This is summarized in the additional stochasticity ablation.
  • Task-specific setup: Inpainting restricts training and generation to masked regions and fills those regions with Gaussian noise as inputs.Evaluation uses 10,000 freeform masks for each 10%-20% and 20%-30% mask ratio.
  • Evaluation: Evaluation uses cleanfid for FID and a pretrained ResNet50 for classifier accuracy, while avoiding pixel-level metrics.Reference statistics are taken from ADM for ImageNet training data and computed for resized, center-cropped validation images.
  • Baselines: Palette comparisons use matched unconditional ADM initialization, conditional-input concatenation, and the same 1000-step quadratic-discretization SDE.The implementation is used for selected tables and figures, while official reported values are used elsewhere.

C.2. Additional Qualitative Results

Additional figures show I2SB generation across restoration tasks, comparisons with Palette at varying sampling effort, and diversity among inpainting outputs.

  • Palette comparisons: Figures 17–19 compare I2SB with Palette on inpainting and deblurring under various NFE sampling settings.The comparisons include center and freeform inpainting as well as uniform-kernel deblurring.
  • Sample diversity: Figure 20 demonstrates that I2SB can generate diverse samples for inpainting tasks.The figure focuses on diversity among I2SB outputs rather than a single reconstruction.

C.3. Additional Discussions

The additional discussions examine deterministic OT-ODE behavior, generalized score parameterization, and qualitative generation across restoration settings.

  • OT-ODE ablation: OT-ODE appears unsuitable for restoration tasks with large unobserved regions.The discussion attributes this behavior to severe information loss in degraded inputs hindering deterministic reconstruction.
  • Score parameterization: The training objective uses ϵ(X_t, t; θ), where X_t interpolates between clean and corrupted image pairs.This parameterization is described for the bridge between the two image distributions.
  • OT-ODE ablation: Injecting Gaussian noise into masked regions removes OT-ODE’s severe inpainting degradation and produces comparable results.The comparison contrasts noiseless masks with stochastic inputs.
  • Score parameterization: An alternative parameterization follows Karras et al. by choosing time-varying coefficients so ϵ inputs and outputs have unit variance and approximation error is minimized.The coefficients are re-derived for interpolations between clean and corrupted pairs rather than i.i.d.-noise perturbations.
  • Coefficient analysis: The generalized coefficients are selected by minimizing c2_out, and they recover the Karras et al. coefficients in the i.i.d.-noise special case.Their empirical differences from the standard coefficients are reported as negligible in practice.
  • Qualitative results: Qualitative figures cover inpainting masks, JPEG quality factors, deblurring kernels, super-resolution filters, Palette comparisons, and output diversity.The displayed settings include center and freeform masks, QF=5 and QF=10, uniform and Gaussian kernels, and pool and bicubic filters.
Loading 2302.05872v3…