Discrete Adjoint Schr\"odinger Bridge Sampler

TL;DR

This paper introduces a novel discrete Schr"odinger bridge sampler that extends continuous adjoint methods to discrete spaces, improving training efficiency and scalability in neural sampling tasks.

Contribution

It reveals the state-space agnostic nature of adjoint matching and develops a unified framework for discrete Schr"odinger bridge sampling, bridging a gap in existing methods.

Findings

Achieves competitive sample quality

Offers significant training efficiency improvements

Demonstrates scalability in discrete neural sampling

Abstract

Learning discrete neural samplers is challenging due to the lack of gradients and combinatorial complexity. While stochastic optimal control (SOC) and Schr\"odinger bridge (SB) provide principled solutions, efficient SOC solvers like adjoint matching (AM), which excel in continuous domains, remain unexplored for discrete spaces. We bridge this gap by revealing that the core mechanism of AM is $\mathit{state}\text{-}\mathit{space~agnostic}$, and introduce $\mathbf{discrete~ASBS}$, a unified framework that extends AM and adjoint Schr\"odinger bridge sampler (ASBS) to discrete spaces. Theoretically, we analyze the optimality conditions of the discrete SB problem and its connection to SOC, identifying a necessary cyclic group structure on the state space to enable this extension. Empirically, discrete ASBS achieves competitive sample quality with significant advantages in training efficiency and scalability.