Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Entering the Era of Discrete Diffusion Models: A Benchmark for Schrödinger Bridges and Entropic Optimal Transport

Xavier Aramayo Carrasco, Grigoriy Ksenofontov, Aleksei Leonov, Iaroslav Sergeevich Koshelev, Alexander Korotin

TL;DR

This work introduces a benchmark for SB on discrete spaces, and obtains two new SB algorithms, DLightSB and DLightSB-M, and additionally extends prior related work to construct the $\alpha$-CSBM algorithm.

Abstract

The Entropic Optimal Transport (EOT) problem and its dynamic counterpart, the Schrödinger bridge (SB) problem, play an important role in modern machine learning, linking generative modeling with optimal transport theory. While recent advances in discrete diffusion and flow models have sparked growing interest in applying SB methods to discrete domains, there remains no reliable way to assess how well these methods actually solve the underlying problem. We address this challenge by introducing a benchmark for SB on discrete spaces. Our construction yields pairs of probability distributions with analytically known SB solutions, enabling rigorous evaluation. As a byproduct of building this benchmark, we obtain two new SB algorithms, DLightSB and DLightSB-M, and additionally extend prior related work to construct the $α$-CSBM algorithm. We demonstrate the utility of our benchmark by evaluating both existing and new solvers in high-dimensional discrete settings. This work provides the first step toward proper evaluation of SB methods on discrete spaces, paving the way for more reproducible future studies. The code for the benchmark and all associated experiments is available at https://github.com/gregkseno/catsbench.

Entering the Era of Discrete Diffusion Models: A Benchmark for Schrödinger Bridges and Entropic Optimal Transport

TL;DR

This work introduces a benchmark for SB on discrete spaces, and obtains two new SB algorithms, DLightSB and DLightSB-M, and additionally extends prior related work to construct the -CSBM algorithm.

Abstract

The Entropic Optimal Transport (EOT) problem and its dynamic counterpart, the Schrödinger bridge (SB) problem, play an important role in modern machine learning, linking generative modeling with optimal transport theory. While recent advances in discrete diffusion and flow models have sparked growing interest in applying SB methods to discrete domains, there remains no reliable way to assess how well these methods actually solve the underlying problem. We address this challenge by introducing a benchmark for SB on discrete spaces. Our construction yields pairs of probability distributions with analytically known SB solutions, enabling rigorous evaluation. As a byproduct of building this benchmark, we obtain two new SB algorithms, DLightSB and DLightSB-M, and additionally extend prior related work to construct the -CSBM algorithm. We demonstrate the utility of our benchmark by evaluating both existing and new solvers in high-dimensional discrete settings. This work provides the first step toward proper evaluation of SB methods on discrete spaces, paving the way for more reproducible future studies. The code for the benchmark and all associated experiments is available at https://github.com/gregkseno/catsbench.

Paper Structure

This paper contains 31 sections, 5 theorems, 49 equations, 5 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Dynamic and Static Schrödinger Bridges on Discrete Spaces
  4. Examples of suitable Reference Processes
  5. Uniform Reference Process ($q^{\text{unif}}$).
  6. Gaussian Reference Process ($q^{\text{gauss}}$).
  7. Entropic Optimal Transport on Discrete Spaces
  8. Discrete Generative Problem Setup and Evaluation Protocol
  9. Benchmark
  10. Main Theorem for Benchmark Construction
  11. Practical Parameterization
  12. High-dimensional Gaussian Mixtures Benchmark Construction
  13. Solvers for Evaluation
  14. Categorical Schrödinger Bridge Matching (CSBM)
  15. $\alpha$-Categorical Schrödinger Bridge Matching ($\alpha$-CSBM)
  16. ...and 16 more sections

Key Result

Theorem 3.1

Let $p_0\in \mathcal{P}(\mathcal{X})$ be a given initial distribution on a discrete space $\mathcal{X}$ and $v^*:\mathcal{X} \rightarrow \mathbb{R}$ be a given scalar-valued function. Consider a joint distribution $q^* \in \mathcal{P}(\mathcal{X}^{2})$ such that $q^*(x_0) = p_0(x_0)$ and $q^*(x_1| x

Figures (5)

  • Figure 1: Samples from all considered methods on our high-dimensional Gaussian mixture benchmark. Top row: $q^{\text{unif}}$ ($\gamma = 0.005$). Bottom row: Gaussian benchmark ($\gamma = 0.02$). CSBM, $\alpha$-CSBM, and DLightSB-M were trained with KL loss ($N\!+\!1=64$).
  • Figure 2: Samples from all methods on the high-dimensional Gaussian mixture benchmark using the Gaussian reference process $q^{\text{gauss}}$ with $\gamma = 0.02$.
  • Figure 3: Samples from all methods on the high-dimensional Gaussian mixture benchmark using the Gaussian reference process $q^{\text{gauss}}$ with $\gamma = 0.05$.
  • Figure 4: Samples from all methods on the high-dimensional Gaussian mixture benchmark using the uniform reference process $q^{\text{unif}}$ with $\gamma = 0.005$.
  • Figure 5: Samples from all methods on the high-dimensional Gaussian mixture benchmark using the uniform reference process $q^{\text{unif}}$ with $\gamma = 0.01$.

Theorems & Definitions (11)

  • Theorem 3.1: Benchmark Pair Construction for Discrete-Space EOT/SB
  • Proposition 3.1: Tractable Parameterization of Conditional Distributions
  • Proposition 3.2: Tractable Parameterization of Conditional SB Transition Distributions
  • Proposition 4.1: Feasible Discrete Reformulation of the Direct KL Objective
  • Proposition 4.2: Discrete-Space Optimal Projection with an Arbitrary Reference Process
  • proof : Proof of Theorem \ref{['theorem:benchmark_construction']}
  • proof : Proof of Proposition \ref{['prop:tractable_parameterization']}
  • proof : Proof of Proposition \ref{['proposition:dlight_sb_loss']}
  • proof : Proof of Proposition \ref{['proposition:optimal_projection']}
  • proof : Proof of Proposition \ref{['proposition:sb_trainsition_distribution']}
  • ...and 1 more