Discrete Adjoint Schrödinger Bridge Sampler
Abstract
Learning discrete neural samplers is challenging due to the lack of gradients and combinatorial complexity. While stochastic optimal control (SOC) and Schrödinger 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 state-space agnostic, and introduce discrete ASBS, a unified framework that extends AM and adjoint Schrödinger 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. Our code is available at https://github.com/AlexandreGUO2001/DASBS.
Lay Summary
Many problems in science and engineering require drawing examples that follow a known but complicated probability distribution. Traditional methods based on random walks (MCMC) can do this, but they are relatively slow to converge. A recent line of research replaces those slow-mixing random walks with neural network models that learn to generate samples directly, often orders of magnitude faster. These models have matured for problems on continuous state spaces, but the fastest among them have remained out of reach for problems with discrete states, where the underlying mathematics is fundamentally different. Our paper closes this gap. We identify the precise structural condition that lets the techniques on continuous spaces transfer to the discrete world, and we build an algorithm called DASBS that exploits it. On standard benchmarks, DASBS matches the quality of the best existing discrete samplers while training tens of times faster, opening a path to scaling these methods to harder discrete problems in chemistry, biology, and physics.