Categorical Schr\"odinger Bridge Matching

TL;DR

This paper extends the Schr"odinger Bridge framework to discrete data spaces using a new algorithm, Categorical Schr"odinger Bridge Matching, enabling applications in areas like VQ representations and text tokens.

Contribution

It provides the first theoretical and algorithmic foundation for applying Schr"odinger Bridge methods to discrete data, including convergence guarantees and a practical algorithm.

Findings

Demonstrates convergence of the discrete-time IMF to SB in discrete spaces.

Develops the CSBM algorithm for practical discrete Schr"odinger Bridge matching.

Shows effectiveness of CSBM on synthetic data and VQ image representations.

Abstract

The Schr\"odinger Bridge (SB) is a powerful framework for solving generative modeling tasks such as unpaired domain translation. Most SB-related research focuses on continuous data space $\mathbb{R}^{D}$ and leaves open theoretical and algorithmic questions about applying SB methods to discrete data, e.g, on finite spaces $\mathbb{S}^{D}$. Notable examples of such sets $\mathbb{S}$ are codebooks of vector-quantized (VQ) representations of modern autoencoders, tokens in texts, categories of atoms in molecules, etc. In this paper, we provide a theoretical and algorithmic foundation for solving SB in discrete spaces using the recently introduced Iterative Markovian Fitting (IMF) procedure. Specifically, we theoretically justify the convergence of discrete-time IMF (D-IMF) to SB in discrete spaces. This enables us to develop a practical computational algorithm for SB, which we call Categorical Schr\"odinger Bridge Matching (CSBM). We show the performance of CSBM via a series of experiments with synthetic data and VQ representations of images. The code of CSBM is available at https://github.com/gregkseno/csbm.