Categorical Schr""odinger Bridge Matching

2025-02-03

AI Summary

This work addresses generative modeling over discrete finite spaces—such as vector-quantized (VQ) codebooks, text tokens, and molecular atom types—and introduces the first extension of Schrödinger Bridge (SB) theory to the discrete domain. We propose Discrete-time Iterative Markovian Fitting (D-IMF), a provably convergent framework for discrete SB inference. Building upon D-IMF, we develop CSBM—the first scalable categorical Schrödinger Bridge method—integrating discrete optimal transport, gradient updates on the probability simplex, and vector-quantized representations. Evaluated on synthetic data and image VQ features, CSBM achieves significant improvements in reconstruction fidelity and distribution matching quality for unpaired domain translation tasks, consistently outperforming existing baselines.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Schrödinger Bridge Method
• Discrete Data Transformation
• Finite Categorical Data

Innovation

• Schrodinger Bridge Method
• Categorical Data
• Iterative Markov Fitting