Gradient Discrete Schrodinger Bridge Matching for Peptide Generation

Source abstract

The optimal control of a discrete Schrödinger bridge is a discrete gradient, a difference of vertex potentials, but recent models such as MadSBM instead learn a free control on the edges by cross-entropy. We prove that this leaves the action uncontrolled: cross-entropy identifies the control only up to a state-dependent shift (a gauge), and the action is $\operatorname{KL}(r\|p)$ plus a gauge term the loss never constrains. We propose GDSBM, which writes the control as a difference of vertex potentials and adds the action to the loss; pinning the gauge costs no expressiveness and makes the action a property of the model, so the penalty acts on the model and not on the gauge. On peptide generation our MadSBM reimplementation has a free-gauge action of order $10^{8}$ for a gauge-fixed value of $6.5$, against $1.5$ for GDSBM, with no quality loss that three seeds can detect. The gradient form also fits guidance, which becomes the addition of a potential: guided toward a net charge of $+3$ at pH~7, GDSBM is closer than MadSBM, over three seeds, to its exact tempered target in total variation at every guidance strength (at $β=1$, $81\%$ of samples in the target window against $74\%$, for a target of $97\%$); we claim no statistical significance.