gregkseno/catsbench

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

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

This repository contains the benchmark checkpoints associated with the paper "Entering the Era of Discrete Diffusion Models: A Benchmark for Schrödinger Bridges and Entropic Optimal Transport", accepted at ICLR 2026.

📦 CatSBench (Package)

Benchmark usage is provided via catsbench, a standalone package that includes benchmark definitions, evaluation metrics, and reusable utilities, including a Triton-optimized log-sum-exp (LSE) matmul kernel.

📥 Installation

Install the benchmark package via pip:

pip install catsbench

🚀 Quickstart

Load a benchmark definition and its assets from a pretrained repository:

from catsbench import BenchmarkHD

bench = BenchmarkHD.from_pretrained(
    "gregkseno/catsbench",
    "hd_d2_s50_gaussian_a0.02_gaussian",
    init_benchmark=False,  # skip heavy initialization at load time
)

To sample marginals $p_0$ and $p_1$:

x_start, x_end = bench.sample_input_target(32) # ([B=32, D=2], [B=32, D=2])

Or sample them separately:

x_start = bench.sample_input(32) # [B=32, D=2]
x_end = bench.sample_target(32)  # [B=32, D=2]

Both examples above sample independently, i.e., $(x_0, x_1) \sim p_0(x_0),p_1(x_1)$.

To sample from the ground-truth EOT/SB coupling, i.e., $(x_0, x_1) \sim p_0(x_0),q^*(x_1 | x_0)$, use:

x_start = bench.sample_input(32) # [B=32, D=2]
x_end = bench.sample(x_start)    # [B=32, D=2]

See the end-to-end benchmark workflow (initialization, evaluation, metrics, plotting) in notebooks/benchmark_usage.ipynb.

🎓 Citation

@misc{
  carrasco2025enteringeradiscretediffusion,
  title={Entering the Era of Discrete Diffusion Models: A Benchmark for {Schr\"odinger} Bridges and Entropic Optimal Transport}, 
  author={Xavier Aramayo Carrasco and Grigoriy Ksenofontov and Aleksei Leonov and Iaroslav Sergeevich Koshelev and Alexander Korotin},
  year={2025},
  eprint={2509.23348},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2509.23348}, 
}

🙏 Credits

• Comet ML — experiment-tracking and visualization toolkit;
• Inkscape — an excellent open-source editor for vector graphics;
• Hydra/Lightning template — project template used as a starting point.