PaperlayerView exact revision on arXiv

Source-linked AI summary

Tilted Schr\"odinger Bridge Matching

Sergei Kholkin, Evgeny Burnaev, Alexander Korotin

arXiv:2609.34642v1cs.LG

TL;DR

Reward finetuning for Schrödinger bridges is underexplored, despite its relevance to aligning pretrained models with preferences or constraints. TSBM adapts a pretrained bridge to a reward-tilted target through alternating optimization, and steers attributes in unpaired image translation while largely preserving source content.

  • Problem

    Reward finetuning is well studied for diffusion models but remains underexplored for Schrödinger bridges, where endpoint dependence can bias the desired terminal tilt.

    Source: §1 · p. 1

  • Method

    TSBM formulates reward tilting as stochastic optimal control relative to a pretrained bridge and learns an incremental control through alternating controller and corrector matching.

    Source: §1 · p. 1, §4 · p. 6

  • Results

    On Colored MNIST and CelebA, TSBM steers attributes in unpaired image translation while largely preserving source content.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The results suggest pretrained bridges can be adapted to new preferences without training a new bridge from scratch.

    Source: §7 · p. 9

    The method requires an accessible endpoint transition score and a differentiable reward.

    Source: §7 · p. 9

Abstract

from arXiv · show

Schr\"odinger bridges provide an entropy-regularized framework and a principled solution for unpaired domain translation. In practice, a pretrained bridge may need to be adapted to human preferences or physical constraints through a reward a problem closely related to reward tilting in diffusion models but underexplored for Schr\"odinger bridges. We introduce Tilted Schr\"odinger Bridge Matching (TSBM), a post-training method for fine-tuning a learned bridge $P$ between source $p_0$ and target $p_1$ toward a reward-tilted target $p_1^r\propto p_1e^r$, while preserving source $p_0$. We formulate this adaptation as alternating optimization initialized from $P$, provide theoretical justification, and derive a practical algorithm based on Adjoint Matching. We evaluate TSBM on unpaired image-to-image translation targeting digit properties in MNIST and facial attributes in CelebA.

1 INTRODUCTION

Schrödinger bridges offer an entropy-regularized framework for transport, but adapting pretrained bridges to rewards is underexplored. TSBM formulates this adaptation as optimal control relative to the pretrained bridge and learns an incremental control while preserving the source marginal.

1 INTRODUCTION

  • Schrödinger bridges minimize relative entropy to a reference process while matching prescribed endpoint distributions, connecting stochastic control, entropic optimal transport, and diffusion generative modeling.The paper notes applications including unpaired domain transfer, trajectory inference, and Boltzmann sampling.Source: §1 · p. 1
  • Reward finetuning for Schrödinger Bridges is underexplored, although it is used to align generative models with preferences, constraints, or image attributes.The desired target is p1^r(x) ∝ p1(x)e^r(x).Source: §1 · p. 1
  • TSBM casts reward-tilted bridge adaptation as stochastic optimal control relative to an exact pretrained bridge, preserving the source marginal while targeting p1^r.The method learns an incremental control around the pretrained bridge.Source: §1 · p. 1
  • Figure 1 depicts TSBM adding control ut to the pretrained drift bt, yielding bt + σtut to target the reward-tilted marginal.The pretrained bridge transports p0 to p1.Source: Figure 1 · p. 2

2 BACKGROUND

The background contrasts reward-based stochastic optimal control with Schrödinger bridges, which impose both endpoint marginals. It also explains why non-memoryless bridges complicate reward tilting and how adjoint matching supports controller and corrector updates.

2 BACKGROUND

  • A terminal reward yields the desired tilted marginal for memoryless references, but generally not for non-memoryless diffusion or Schrödinger bridges.Dependence between X0 and X1 can bias the terminal marginal away from the reward-tilted target.Source: §2 · p. 2
  • Adjoint Matching defines an optimal controller through a matching objective using a lean adjoint state and the reference-drift Jacobian.The SDE sampling and adjoint ODE simulations are still required and can be expensive.Source: §2 · pp. 2–3
  • Unlike regular stochastic optimal control with one marginal constraint, a Schrödinger bridge matches both source and target marginals while retaining a minimal-energy control objective.The bridge is the diffusion closest to the reference in KL divergence; with a Wiener reference, this is equivalent to minimizing expected control energy.Source: §2 · p. 3
  • Corrector Matching learns the gradient of the terminal Schrödinger potential, and alternates corrector updates with controller updates because samples from the optimal bridge are unavailable.The corrector update can use the target score when its density is known up to a normalizing constant.Source: §2 · pp. 3–4

3 TILTING A SCHR ¨ODINGER BRIDGE

Section 3 defines the reward-tilted bridge problem and develops an equivalent formulation that uses the pretrained bridge as its reference. The relative formulation removes the need for explicit access to the original target density in its terminal cost.

3 TILTING A SCHR ¨ODINGER BRIDGE

  • The section relates the base-process and pretrained-bridge references, their path-space solutions and Schrödinger potentials, and a relative stochastic-control formulation.The relative formulation is developed to avoid explicit knowledge of p1.Source: §3 · p. 4

3.1 TILTED SCHR ¨ODINGER BRIDGE PROBLEM

  • TiltedSBP seeks a Schrödinger bridge from p0 to the reward-tilted target p1^r(x) ∝ p1(x)e^r(x).The definition assumes the reward-tilted density is normalizable.Source: §3.1 · p. 4
  • The formulation notes a practical difficulty: constructing the tilted target can be complicated when p1 is given by empirical samples and r by a reward function.Source: §3.1 · p. 4

3.2 RELATIVE STOCHASTIC OPTIMAL CONTROL FORMULATION FOR TILTEDSBP

  • Changing the reference to the pretrained bridge preserves the desired tilted-bridge solution, allowing terminal retargeting through an incremental control around its drift.The stated result holds under the paper's assumptions.Source: §3.2 · p. 4
  • The relative terminal cost cancels explicit p1, requiring the reward and old or reward-tilted Schrödinger potentials instead.The formulation requires the pretrained bridge and its terminal potential, but no further access to p1 through samples or density evaluations.Source: §3.2 · p. 5

4 TILTED SCHR ¨ODINGER BRIDGE MATCHING

TSBM adapts a pretrained bridge through alternating controller and corrector updates, learning an incremental control around the pretrained drift. The procedure is initialized at the pretrained bridge and implemented with Adjoint Matching.

4 TILTED SCHR ¨ODINGER BRIDGE MATCHING

  • The controller update uses Controller Matching with a fixed corrector, while the corrector update uses Corrector Matching with a fixed controller.The controller is updated first, then the corrector, in each stage.Source: §4 · p. 5
  • The algorithm initializes the incremental control at zero and the corrector at the pretrained corrector, then iterates controller and corrector updates for K stages.It samples paths under the controlled drift for controller updates and endpoints from the controlled process for corrector updates.Source: §4 · p. 6
  • TSBM alternates Controller and Corrector Matching updates, each using the previous stage’s result to relax their interdependence.The paper states a convergence theorem for the exact population updates under Assumptions 1 and 2, but the supplied theorem passage does not include its conclusion.Source: §4 · p. 6
  • In practice, TSBM fine-tunes a pretrained bridge represented by forward and backward drift networks, using the forward drift for controller updates and the backward drift in both updates.The method uses the same parameterization as the pretrained bridge.Source: §4 · p. 6

5 RELATED WORK

Prior work develops Schrödinger Bridge solvers and reward-steering methods, but the supplied related-work discussion identifies reward tilting of a general pretrained bridge with source preservation as unestablished. The paper positions its approach against this gap.

5 RELATED WORK

  • Existing Schrödinger Bridge methods include entropic optimal transport, iterative fitting procedures, and Adjoint Matching, but the passage reports no current method supporting exact marginal reward tilting of pretrained SB models.Applications cited include unpaired domain translation, trajectory inference, Boltzmann sampling, mutual-information estimation, and data generation.Source: §5 · p. 6
  • Reward steering includes inference-time guidance, reward optimization, preference learning, and stochastic optimal control methods targeting exponential tilts, but much of this work concerns noise-to-data models.Inference-time methods can increase sampling costs or introduce artifacts under strong guidance.Source: §5 · p. 7
  • For bridges, source-target dependence means conditional tilting generally does not yield the desired terminal marginal when the source marginal is fixed.The discussed prior methods do not establish the joint guarantee of a prescribed terminal reward tilt and preservation of the source marginal for a general pretrained bridge.Source: §5 · pp. 7–8

6 EXPERIMENTS

The experiments test TSBM for reward-guided unpaired image translation on Colored MNIST and CelebA, against conditional-guidance and population-based baselines. Across both settings, results report improved reward alignment alongside measures of source faithfulness or output quality.

6.1 COLORED MNIST UNPAIRED TRANSLATION ATTRIBUTES MANIPULATION

  • Colored MNIST tests digit-2-to-digit-3 translation with thin-stroke, thick-stroke, and red-chroma rewards, using a 20-stage DSBM pretrain and five TSBM stages.The TSBM reward strength is λ = 1000.Source: §6.1 · p. 8
  • For Colored MNIST, TSBM changes the targeted digit property while largely preserving unrelated attributes, such as source color for stroke targets and source thickness for red-chroma targets.The qualitative comparisons include the source, pretrained DSBM, TSBM, DPS, and CondSMC.Source: §6.1 · p. 8
  • TSBM improves upon the non-TSBM reward-faithfulness Pareto frontier on Colored MNIST, achieving the most favorable trade-off among evaluated methods.DPS introduces visible artifacts, while CondSMC makes limited changes to targeted attributes.Source: §6.1 · p. 9

6.2 CELEBA UNPAIRED TRANSLATION ATTRIBUTES MANIPULATION

  • CelebA tests unpaired male-to-female translation at 128×128 for natural toothy smile, heavy makeup, and elderly-woman prompts, using frozen ImageReward for training and PickScore for evaluation.The pretrained DSBM and TSBM are trained for 20 and five stages, respectively; evaluation uses 100 NFE and reward strength λ = 30.Source: §6.2 · p. 9
  • On CelebA, TSBM obtains the highest evaluation PickScore in each prompt comparison, while its CLIP-IQA is highest for smile and elderly prompts but less favorable for heavy makeup.TSBM improves ImageReward over the pretrained bridge for all three prompts, though less than DPS; the authors interpret DPS’s large training-reward gains with modest evaluation gains as reward hacking.Source: §6.2 · p. 9, Table 1 · p. 9

7 DISCUSSION

The discussion presents TSBM as a way to steer pretrained bridges toward rewarded attributes while largely preserving source content, and notes practical assumptions and computational costs. The accompanying appendix relates the tilted problem to the pretrained bridge and analyzes alternating controller and corrector updates.

7 DISCUSSION

  • TSBM steers attributes in unpaired image translation on Colored MNIST and CelebA while largely preserving source content.The authors suggest pretrained bridges can be adapted to new preferences without training a bridge from scratch.Source: §7 · p. 9
  • The method assumes exact optimization and valid theoretical assumptions, requires an accessible endpoint transition score and differentiable reward, and incurs repeated simulation and adjoint-ODE costs.The authors note that these assumptions may fail in practice and that distillation may reduce computational cost.Source: §7 · pp. 9–10

A.1 PROOF OF PROPOSITION 1

  • The proof identifies the tilted bridge as the unique minimizer of both the original and transformed Schrödinger bridge problems.The change of reference shifts the KL objective by a constant for fixed marginals.Source: §A.1 · p. 15

A.2 PROOF OF PROPOSITION 2

  • The potential relations recover the reward-tilted terminal marginal and, under smoothness conditions, yield a differentiated identity.The text states that continuity extends the identity through the open region before differentiation.Source: §A.2 · p. 16

A.3 PROOF OF PROPOSITION 3

  • The tilted bridge can be formulated as stochastic optimal control relative to the pretrained bridge, with the same optimizer.The proof compares the objective with reference P and target p1^r.Source: §A.3 · p. 16

A.4 ASSUMPTIONS AND CONTROLLED DIFFUSIONS FOR CONVERGENCE

  • The convergence analysis assumes population matching, endpoint and reward conditions, regular transition scores, and admissible controlled diffusions with unique nonexplosive solutions.Additional regularity supports the adjoint ODE and derivatives used in matching.Source: §A.4 · pp. 16–17

A.5 CONTROLLER MATCHING AND THE FORWARD HALF BRIDGE

  • The exact Controller Update solves the forward half-bridge problem by replacing the initial marginal with p0 while preserving trajectories conditional on X0.Source: §A.5 · pp. 18–19

A.6 CORRECTOR MATCHING AND THE BACKWARD HALF BRIDGE

  • After the exact Controller Update, the backward control given by the Corrector Update identifies the backward half bridge.Source: §A.6 · p. 19

A.7 PROOF OF THEOREM 1

The proof establishes convergence of the alternating half-bridge updates to the tilted bridge under stated assumptions. The appendix also describes competing reward-guidance and sampling approaches and practical choices for TSBM training.

A.7 PROOF OF THEOREM 1

  • The controller sequence converges in integrated mean square under the optimal controlled diffusion, with its limit solving the tilted Schrödinger bridge problem.The proof first establishes endpoint entropy convergence and then lifts it to path-space convergence.Source: §A.7 · pp. 22–23

B COMPETITORS

  • The competitors share the frozen pretrained bridge as their sampler and use a terminal log reward for image experiments.The described sampler uses Euler steps, with noise omitted on the final image-sampling step.Source: B COMPETITORS · p. 23

B.1 REWARD TILTING WITH NON-MEMORYLESS PRIORS

  • Reward-only terminal tilting is unbiased for memoryless diffusions, but source-dependent endpoint effects mean diffusion bridges generally miss the desired tilted marginal.Source: §B.1 · p. 24

B.2 DPS-STYLE GUIDANCE

  • DPS-style guidance differentiates the reward at a predicted endpoint to guide the frozen bridge.The guidance scale is separate from reward strength, and the experiments use the full endpoint-prediction Jacobian.Source: §B.2 · p. 24

B.3 CONDITIONAL SELF-NORMALIZED IMPORTANCE SAMPLING (CONDSNIS)

  • CondSNIS samples old-bridge trajectories for each source and, as particle count grows, targets the reward-tilted conditional terminal distribution.It selects a final sample using log weights.Source: §B.3 · p. 24

B.4 CONDITIONAL SEQUENTIAL MONTE CARLO (CONDSMC)

  • CondSMC propagates old-bridge particles, resamples when effective sample size falls below N/2, and targets the conditional tilt as particle count grows.Source: §B.4 · pp. 24–25

C GENERAL DESIGN CHOICES FOR TSBM

  • TSBM freezes pretrained networks, initializes trainable drift and corrector networks from them, and carries their parameters forward between stages.Current-process simulation uses one trainable-drift evaluation per time step; replay buffers store trajectories and lean adjoints.Source: C GENERAL DESIGN CHOICES FOR TSBM · p. 25

D.1 TOY ILLUSTRATIVE EXPERIMENT

The appendix specifies a toy reward-tilting setup, Colored MNIST energy constructions, and CelebA prompts and scoring procedures. These passages describe experiment design rather than comparative outcome values.

D.1 TOY ILLUSTRATIVE EXPERIMENT

  • The toy experiment changes mixture weights from (0.25, 0.25, 0.25, 0.25) to (0, 0.2, 0.2, 0.6), leaving Gaussian means and covariances unchanged.Source: §D.1 · p. 26

D.2 COLORED MNIST UNPAIRED TRANSLATION

  • Colored MNIST uses red-chroma and stroke-thickness energies; the stroke-width estimate is based on a soft foreground mask and perimeter estimate.Table 2 pairs each source with the terminal TSBM generation from the corresponding comparison-grid column and lists terminal energies and desired outputs.Source: §D.2 · pp. 26–27, Table 2 · p. 26

D.3 CELEBA UNPAIRED TRANSLATION

  • CelebA male-to-female translations target a natural toothy smile, heavy makeup, or an elderly-woman appearance using prompt-conditioned rewards.The elderly-woman training reward also uses a negative prompt; smile and makeup use positive prompts only.Source: §D.3 · pp. 27–28
  • ImageReward defines the CelebA training preference score, while PickScore is used only for post-hoc evaluation.Both reported evaluation scores use the positive prompt; PickScore does not subtract the elderly-woman negative prompt.Source: §D.3 · p. 28

E EXPERIMENTAL DETAILS

The experimental details identify the TSBM and DSBM hyperparameter tables and describe image preprocessing and simulation procedures.

E EXPERIMENTAL DETAILS

  • Image experiments omit Gaussian noise on the final sampling step, and parameter counts exclude frozen pretrained counterparts.The SDE uses Euler–Maruyama simulation followed by lean-adjoint ODE integration.Source: E EXPERIMENTAL DETAILS · p. 28
  • Image data are normalized to [-1, 1], and the experiments use the official DSBM GitHub repository.Source: E EXPERIMENTAL DETAILS · p. 28
  • The appendix provides separate hyperparameter tables for TSBM and DSBM pretraining.Source: Table 3 · p. 29, Table 4 · p. 29

F ADDITIONAL QUALITATIVE RESULTS

The additional qualitative results collect Colored MNIST examples for three reward targets and CelebA examples for several facial attributes, alongside a toy experiment’s trajectories and terminal samples.

F ADDITIONAL QUALITATIVE RESULTS

  • Colored MNIST figures show full qualitative comparisons for red-chroma, thick-stroke, and thin-stroke rewards.The captions identify the three reward settings and datasets but do not state comparative outcomes.Source: Figure 5 · p. 30, Figure 6 · p. 30, Figure 7 · p. 30

F.2 CELEBA

  • Complete qualitative CelebA comparisons cover Natural Toothy Smile, Heavy Makeup, and Elderly Woman rewards, with additional methods including DPS at γ = 0.5, 1, and 2 and CondSNIS with K = 4 proposals per fixed source.Source: §F.2 · p. 31, Figure 8 · p. 31, Figure 9 · p. 32

F ADDITIONAL QUALITATIVE RESULTS

  • The toy figure overlays source and original or tilted targets, then compares terminal samples and trajectories from DSBM and TSBM.It also depicts total variation, sliced W1, and translation cost relative to a Sinkhorn reference.Source: Figure 11 · p. 34

G.1 TOY ILLUSTRATIVE EXPERIMENT

The toy experiment evaluates tilted-target fit and transport cost; further experiments examine CelebA inference costs and Colored MNIST reward steering across stages and corrector variants.

G.1 TOY ILLUSTRATIVE EXPERIMENT

  • TSBM alone closely matches the toy tilted target and approximates the reference bridge’s transport cost.It has the best total-variation and sliced W1 fit and the cost closest to the Sinkhorn reference among evaluated methods.Source: §G.1 · p. 34

G.2 CELEBA-128 INFERENCE TIME AND WORKING MEMORY

  • On CelebA 128 × 128, TSBM has essentially the pretrained bridge’s inference cost; in the single-output protocol, DPS is 4.88× slower and uses 2.229 GiB of working memory versus TSBM’s 0.180 GiB.Source: §G.2 · p. 35, Table 5 · p. 35

G.3 COLORED MNIST

  • Across red-chroma stages, TSBM improves reward and hue error faster than the static-corrector ablation.The qualitative comparison also shows regular TSBM following the reward noticeably better.Source: Figure 12 · p. 36, §G.3 · p. 36, Figure 13 · p. 36

G.4 COLORED MNIST THICK-TARGET. TSBM STAGES ABLATION.

  • On Colored MNIST with the “thick stroke” reward, TSBM improves reward across stages, with most gains in the first several; thickness nearly reaches target early but improves later, while hue error is almost always within 13°–16°.Figure 14 reports mean reward, stroke thickness, and foreground hue error over 1,024 fixed test source images.Source: §G.4 · p. 36, Figure 14 · p. 37

G.5 CELEBA ADDITIONAL QUANTITATIVE ANALYSIS

Additional CelebA evaluations compare prompt alignment and source-to-output change across methods, then test TSBM on aesthetics and thick-framed eyeglasses rewards.

G.5 CELEBA ADDITIONAL QUANTITATIVE ANALYSIS

  • TSBM and DPS γ = 4 are the two best methods by HPS and CLIPScore; TSBM has the largest source-to-output MSE but lower LPIPS than every DPS variant.CondSNIS and CondSMC have smaller input-to-output distances but weaker prompt scores, so distance alone does not indicate better edits.Source: §G.5 · pp. 37–38, Table 6 · p. 38

G.6 ADDITIONAL CELEBA REWARDS

  • TSBM has the largest ImageReward and PickScore means for both additional CelebA rewards, and the largest CLIP-IQA naturalness mean for aesthetics.The eyeglasses task is harder: its CLIP-IQA falls from 0.5446 for pretrained DSBM to 0.4754 for TSBM, while CondSMC reaches 0.5547.Source: Table 7 · p. 39, Table 8 · p. 39, §G.6 · p. 39
  • For eyeglasses with thick frames, TSBM improves reward alignment over the baselines but does not fully solve the task.The passage characterizes adding eyeglasses as particularly challenging under both ImageReward and PickScore.Source: §G.6 · p. 39
Loading 2609.34642v1…