PaperlayerView exact revision on arXiv

Source-linked AI summary

Unpaired Image-to-Image Translation via Neural Schrödinger Bridge

Beomsu Kim, Gihyun Kwon, Kwanyoung Kim, Jong Chul Ye

arXiv:2305.15086v3cs.CVcs.AIcs.LGstat.ML

TL;DR

Unpaired image-to-image translation is constrained by diffusion models’ Gaussian prior, and prior SB methods had not successfully handled high-resolution images. UNSB reformulates SB as adversarial learning with regularization and demonstrates scalable translation between arbitrary distributions across unpaired tasks.

  • Problem

    Diffusion models’ Gaussian prior limits unpaired image-to-image translation, while prior SB methods had not successfully translated high-resolution images in the unpaired setting.

  • Method

    UNSB formulates the SB problem as a sequence of adversarial learning problems under a KL-divergence constraint, incorporating discriminators and regularization.

  • Results

    UNSB demonstrated scalable and effective unpaired image-to-image translation across various data-to-data and image-to-image tasks, often surpassing one-step models.

  • Takeaways & Limitations

    UNSB opens a research direction for applying diffusion models to unpaired image translation between arbitrary distributions.

  • Takeaways & Limitations

    UNSB’s latent diffusion model results are limited to Horse2Zebra and Summer2Winter because the model is not trained on Label2Cityscape or Map2Satellite images.

Abstract

from arXiv · show

Diffusion models are a powerful class of generative models which simulate stochastic differential equations (SDEs) to generate data from noise. While diffusion models have achieved remarkable progress, they have limitations in unpaired image-to-image (I2I) translation tasks due to the Gaussian prior assumption. Schrödinger Bridge (SB), which learns an SDE to translate between two arbitrary distributions, have risen as an attractive solution to this problem. Yet, to our best knowledge, none of SB models so far have been successful at unpaired translation between high-resolution images. In this work, we propose Unpaired Neural Schrödinger Bridge (UNSB), which expresses the SB problem as a sequence of adversarial learning problems. This allows us to incorporate advanced discriminators and regularization to learn a SB between unpaired data. We show that UNSB is scalable and successfully solves various unpaired I2I translation tasks. Code: \url{https://github.com/cyclomon/UNSB}

1 INTRODUCTION

Diffusion models’ Gaussian prior limits unpaired image-to-image translation, while prior Schrödinger Bridge methods had not successfully handled high-resolution unpaired translation. UNSB identifies dimensionality as the failure source and reformulates SB for adversarial learning between arbitrary distributions.

  • Diffusion models commonly assume Gaussian noise, limiting unpaired domain transfer, style transfer, and image restoration.
  • No previous SB method had successfully trained direct translation between high-resolution images in the unpaired setting.Existing methods were computationally demanding or produced poor results, and representative methods failed on concentric-sphere translation as dimension increased.
  • The paper identifies the curse of dimensionality as the cause of previous SB failures because sparse samples poorly capture image manifolds and bias transport maps.
  • UNSB formulates SB under a KL-divergence constraint and expresses it as a composition of adversarially learned generators.
  • UNSB enables translation between arbitrary distributions and is presented as a generalization of existing unpaired translation methods.

2 RELATED WORKS

Schrödinger Bridges support transport between arbitrary endpoint distributions, but earlier image-translation approaches often relied on simple or Gaussian intermediates. UNSB targets efficient SB learning for higher-resolution unpaired images.

  • Schrödinger Bridges learn stochastic processes transitioning between arbitrary initial and terminal distributions under a reference measure.
  • Recent methods achieved unpaired translation at resolutions up to 128 × 128 but were computationally intensive, often requiring several days of training.
  • UNSB is positioned as an efficient approach for learning SBs between higher-resolution unpaired images.

3 SCHRÖDINGER BRIDGES AND THE CURSE OF DIMENSIONALITY

The SB problem connects stochastic control and entropy-regularized optimal transport, but finite high-dimensional samples can distort the learned transport. This distortion produces poor correspondences as dimension increases.

  • The Schrödinger Bridge seeks the most likely stochastic process interpolating distributions π0 and π1 over time.
  • Static formulation: The static formulation interprets the SB as entropy-regularized optimal transport between its endpoint marginals and enables simple conditional sampling when the endpoint joint distribution is known.
  • Stochastic control formulation: SB sampling can be viewed through stochastic control, where the bridge drift minimizes energy among SDEs with the prescribed endpoint distributions.The stochastic-control formulation also gives Markov structure and connects to optimal-transport trajectories as τ approaches zero.
  • Static formulation: The Sinkhorn-Knopp algorithm can estimate the endpoint SB coupling for discrete distributions, motivating approximation algorithms.
  • Curse of dimensionality: With finite samples, SB training transports empirical distributions whose high-dimensional sparsity fails to represent image manifolds accurately.
  • Curse of dimensionality: High-dimensional sampling error biases the transport map, yielding image pairs without meaningful correspondence and nearly orthogonal interpolations.

4 UNPAIRED NEURAL SCHRÖDINGER BRIDGE (UNSB)

UNSB learns Schrödinger Bridges as a sequence of adversarially trained generator steps, using self-similarity, advanced discriminators, and regularization. On toy distributions, it remains robust to dimension and accurately recovers Gaussian-bridge statistics.

  • UNSB formulation: UNSB expresses SB as a composition of generators learned through adversarial learning over successive time steps.The method learns transition probabilities recursively from the known initial distribution, using intermediate samples generated by the learned chain.
  • UNSB formulation: Self-similarity makes the restricted SB on each time interval another SB, enabling conditional generators to recover the correct intermediate and target distributions.Theorem 1 states that matching the conditional target distribution yields matching transition and marginal distributions.
  • Training: A shared time-conditional DNN replaces separate per-time-step parameters, taking (x_ti, ti) as input and optimizing the summed objective across steps.This avoids sequentially maintaining separate networks for every time step.
  • Combating the curse of dimensionality: Advanced discriminators replace KL divergence with flexible distribution discrepancies; patch-level Markovian discrimination can capture high-frequency target-domain characteristics.The formulation permits adversarial metrics beyond a standard instance-level discriminator.
  • Combating the curse of dimensionality: Regularization adds an application-specific similarity bias between the predicted target image and the initial point within the final UNSB objective.The differentiable function R quantifies the desired similarity and is weighted alongside adversarial and SB losses.
  • Sanity checks on toy data: On two-shell data, all baselines either miss the target manifold or lose input-output cosine similarity, whereas UNSB remains robust as dimension increases.The experiment uses 1k samples from each distribution and selects a discriminator based on input norms with negative cosine similarity as regularization.
  • Sanity checks on toy data: In 50 dimensions, UNSB recovers the mean and covariance of the exact Gaussian Schrödinger Bridge relatively accurately.The test uses the bridge between N(−1, I) and N(1, I).

5 RESULTS FOR UNPAIRED IMAGE-TO-IMAGE TRANSLATION

UNSB is evaluated on large-scale unpaired image-to-image translation using multiple datasets, quality metrics, qualitative comparisons, NFE analysis, and ablations. It generally outperforms baseline methods, while quality depends on NFE and can degrade when artifacts emerge at large NFE.

  • Evaluation: UNSB is evaluated on four 256×256 datasets using FID and KID, with NFE and generation time also reported.The datasets are Horse2Zebra, Map2Cityscape, Summer2Winter, and Map2Satellite.
  • Quantitative comparison: UNSB outperforms baseline methods on all evaluated datasets, including recent models such as CUT.The comparison also reports degraded performance for Neural Optimal Transport.
  • Qualitative comparison: UNSB produces target-domain images while preserving source structure better than the compared one-step baselines.The qualitative comparison reports more realistic domain-changed outputs and fewer structural inconsistencies.
  • NFE analysis: Best FIDs occur at NFE values between 3 and 5 across all datasets, whereas NFE = 1 performs relatively poorly.Artifacts at large NFE can increase FID on datasets such as Map2Satellite.
  • Ablation study: The ablation study examines multi-step generation, advanced discriminators, and regularization as factors affecting Horse2Zebra performance.The configuration without advanced discriminators or regularization corresponds to the previous SB method SBCFM and gives poor results.
  • Stochasticity and transport cost: UNSB generates meaningful variation in outputs, indicating that it learns a stochastic map.Its input-output distances are smaller than those of SK pairs computed from actual dataset images.

6 CONCLUSION

The paper presents UNSB as an adversarial-learning solution to the Schrödinger Bridge problem for unpaired image translation. Experiments demonstrate scalability and effectiveness across translation tasks, opening a research direction for diffusion-based unpaired translation.

  • Conclusion: UNSB solves the Schrödinger Bridge problem via adversarial learning for unpaired image-to-image translation.The formulation combines Schrödinger Bridges with GAN training techniques.
  • Conclusion: UNSB demonstrates scalability and effectiveness across various data-to-data and image-to-image translation tasks.The paper reports that previous SB or OT methods fail on the considered tasks while UNSB often surpasses one-step models.

ETHICS AND REPRODUCIBILITY STATEMENTS

The paper identifies beneficial and malicious potential uses of UNSB and states that malicious content creation should be prevented through proper regulation. It reports that pseudocodes and hyper-parameters are provided in the paper and Appendix.

  • Ethics statement: UNSB may support applications such as medical image restoration but may also enable malicious content such as fake news.The ethics statement calls for proper regulation to prevent malicious use.
  • Reproducibility statement: Pseudocodes and hyper-parameters are described in the main paper and the Appendix.This is stated in the reproducibility statement.

A PROOFS

The proofs establish self-similarity of Schrödinger Bridges: restricting an SB to a time sub-interval yields another SB after accounting for the rescaled reference variance. The static restricted formulation and theorem proof connect this property to the conditional model used by UNSB.

  • Self-similarity: Restricting an SB trajectory to any sub-interval of [0, 1] produces another Schrödinger Bridge after a time change.The proof sketch states this self-similarity property and uses a rescaled time variable.
  • Proof argument: A competing drift on the restricted interval can be extended to the full interval, contradicting optimality of the original SB drift.This extension argument establishes uniqueness of the restricted solution within the proof.
  • Restricted formulation: The restricted interval uses a Wiener reference measure with variance τ(tb − ta).This reduced variance is incorporated when translating between stochastic-control and static formulations.
  • Restricted formulation: The static formulation of restricted SBs uses the normalized time coordinate s(t) = (t − ta)/(tb − ta).The proofs relate this formulation to the corresponding conditional transition and marginal equations.
  • Theorem proof: The theorem proof sets ta = ti and tb = 1 and shows that the conditional model qϕi(x1|xti) solves the restricted optimization problem under the stated distributional capacity assumption.The resulting solution is identified with the conditional distribution p(x1|xti).

B OMITTED EXPERIMENT DETAILS

The experiments use standardized training and evaluation protocols, with quantitative comparisons and visualizations spanning high-resolution translation results.

  • All experiments train UNSB for 400 epochs on a single RTX3090 GPU using batch size 1 and Adam optimization.Images are resized to 256 × 256, normalized to [−1, 1], and the unit interval is discretized into 5 uniformly spaced intervals.
  • Evaluation follows standard procedures, using metric calculation code and the protocol described by Park et al. (2020).
  • Baselines are trained from official code or evaluated using samples from official repositories, with some benchmark values taken from prior tables.Horse2Zebra and Label2Cityscape use values from Park et al. (2020), while Summer2Winter and Map2Satellite baselines are trained using official code.
  • Table 4 compares UNSB with another Schrödinger bridge method on Horse2Zebra, while Table 5 reports Male2Female results.
  • Figures 9 and 10 visualize generated results, conditional-flow-matching comparisons, and output pixel-wise standard deviation given an input.

C.1 OTHER SB METHODS

UNSB is evaluated against other Schrödinger bridge and unpaired translation methods, including high-resolution Male2Female results and stochasticity analysis.

  • SBCFM produces outputs unrelated to the input on Horse2Zebra, indicating failure to overcome the curse of dimensionality.
  • UNSB satisfies the necessary stochasticity condition for a Schrödinger bridge by producing one-to-many outputs.Pixel-wise standard deviation is used to visualize variation across generated outputs for the same input.
  • Greater foreground variation in UNSB outputs corresponds to meaningful one-to-many generation in target-domain regions rather than indiscriminate noise.
  • Male2Female translation: On CelebA-HQ-256 Male2Female translation, UNSB outputs are closer to the target domain according to FID and KID.The comparison includes EGSDE, StarGAN v2, and NOT under the CUT evaluation protocol.
  • Male2Female translation: The compared baselines often fail to preserve input structure, including eye direction, hairstyle, and mouth shape.

C.4 ENTROPIC OT BENCHMARK

The entropic OT benchmark shows UNSB outperforming an unconditional map but remaining below ENOT, while additional experiments examine reverse translation, ablations, and higher-resolution samples.

  • Entropic OT benchmark: UNSB performs much better than the unconditional p(x|y) = p(x) map but not quite as well as ENOT on the 64 × 64 entropic OT benchmark.The authors describe this as a preliminary result with ad-hoc hyper-parameter choices.
  • Entropic OT benchmark: Using only 100 x-samples per y-sample may reduce cFID relative to the benchmark’s provision of 5k x-samples per y-sample.
  • Entropic OT benchmark: The benchmark uses 64 × 64 images, whereas UNSB’s architecture is specialized to ≥256 × 256 images and is adapted through resizing.
  • Entropic OT benchmark: CUT regularization is specialized for clean-image translation, while this benchmark translates between noisy and clean images.
  • Reverse translation: UNSB is capable of reverse translation and beats CycleGAN on Zebra2Horse in both FID and KID while preserving input structure.
  • Ablation: Adding a Markovian discriminator, regularization, and multi-step sampling monotonically improves Summer2Winter performance.The ablation attributes orthogonal roles to these UNSB components.
  • Additional samples: UNSB produces preliminary Cat2Dog samples at 512 × 512 resolution with NFE = 5 after training epoch 100.

D DISCUSSION ON SCENE-LEVEL VS. OBJECT-LEVEL TRANSFER

The discussion distinguishes scene-level transfer, where structure and style must both be handled, from object-level transfer, where centered subjects make structural preservation easier.

  • Scene-level tasks mainly alter color and texture, whereas object-level tasks alter object identities.
  • Scene-level transfer: 35.7 FID: UNSB outperforms SDEdit at 97.3 FID and P2P at 60.9 FID on Horse2Zebra.
  • Scene-level transfer: P2P reverses horse direction and SDEdit fails to translate properly to the zebra domain, while UNSB avoids these problems using a Markovian discriminator and CUT regularization.
  • Object-level transfer: Diffusion-based methods perform reasonably on object-level tasks, plausibly because cropped, centered subjects make structural preservation easier.
  • Object-level transfer: Sinkhorn-Knopp can obtain semantically meaningful pairs on some object-level datasets despite not relying on deep learning.
Loading 2305.15086v3…