Tilted Schrödinger Bridge Matching

Sergei Kholkin; Evgeny Burnaev; Alexander Korotin

cs.LGSep 28, 2026

Tilted Schrödinger Bridge Matching

Authors: Sergei Kholkin, Evgeny Burnaev, Alexander Korotin

Organizations: Applied AI Institute

Abstract

Schrödinger 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ödinger bridges. We introduce Tilted Schrödinger 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.

Figures & tables

Figure 1: Reward tilting a pretrained Schrödinger bridge. The pretrained SB $P$ transports $p_{0}$ to $p_{1}$ with drift $b_{t}$ ( brown ). TSBM learns an incremental control $u_{t}$ , resulting in drift $b_{t}+\sigma_{t}u_{t}$ ( blue ), to target the reward $r(x)$ tilted marginal, i.e., $p_{1}^{r}(x)\propto p_{1}(x)r(x)$ .

Figure 2: Qualitative results on Colored MNIST. Rows show the source input, pretrained DSBM, TSBM, DPS with

\gamma=1

, and CondSMC with

N=8

particles for the thin-stroke, thick-stroke, and red-chroma targets. Comparisons with additional baselines are in Appendix F .

Figure 3: Quantitative analysis on Colored MNIST. Reward $\uparrow$ is the mean unscaled terminal reward $r(x)$ . The two stroke panels use logarithmic spacing by reward magnitude. Faithfulness $\downarrow$ is foreground hue error relative to the input for thin and thick strokes, and mean thickness deviation from the input for red chroma. The red dotted lines connect the non-TSBM Pareto points, while pale red shading marks values below them.

Natural Toothy Smile				Heavy Makeup				Elderly Woman
TSBM (ours)	DPS $\gamma=1$	DPS $\gamma=4$	CondSMC	TSBM (ours)	DPS $\gamma=1$	DPS $\gamma=4$	CondSMC	TSBM (ours)	DPS $\gamma=1$	DPS $\gamma=4$	CondSMC

Figure 4: Qualitative results for CelebA experiment. Each row uses the same source across the pretrained DSBM and three CelebA prompt groups. Within each group, we compare TSBM, DPS with

\gamma\in\{1,4\}

, and CondSMC. Red borders highlight selected outputs with visible artifacts.

	Natural Toothy Smile			Heavy Makeup			Elderly Woman
Method	ImageReward (train) $\uparrow$	PickScore (eval) $\uparrow$	CLIP-IQA $\uparrow$	ImageReward (train) $\uparrow$	PickScore (eval) $\uparrow$	CLIP-IQA $\uparrow$	ImageReward (train) $\uparrow$	PickScore (eval) $\uparrow$	CLIP-IQA $\uparrow$
Pretrained	-1.2348	17.6109	0.5950	-1.4054	17.3777	0.5950	-1.4960	18.1592	0.5950
TSBM (ours)	0.1398	18.2345	0.6595	0.9146	18.0390	0.3508	-0.0204	18.8200	0.6340
DPS ∗ $\gamma=1$	0.0622	17.7347	0.6192	0.8576	17.5803	0.3789	-0.7328	18.2559	0.5623
DPS ∗ $\gamma=2$	0.5993	17.6831	0.6312	1.2507	17.6124	0.3401	-0.2527	18.3329	0.5218
DPS ∗ $\gamma=4$	0.9783	17.4688	0.6164	1.5151	17.6236	0.3133	0.2239	18.3346	0.4618

Table 1: Quantitative results for CelebA experiment . Means are over 1024 fixed inputs per prompt, with TSBM evaluated after stage 5. ImageReward and PickScore use translation, cropping, and horizontal-flip augmentation; CLIP-IQA measures naturalness. All metrics are higher-is-better; bold marks the largest mean. ( ∗ ) indicates the appearance of visual artifacts.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Table 2 — Table 2: Colored MNIST terminal energies and desired outputs. Each pair shows a source $x_{0}$ (left) and the terminal TSBM generation from the same column of the corresponding comparison grid (right).

Hyperparameter	Toy 2D	Colored MNIST	CelebA
Controller gradient steps / stage	1,000	1,000	1,000
Corrector gradient steps / stage	782	1,000	1,000
Training batch size	128	32	32
Controller replay size	None (fresh batch for every step)	512 (red); 1,024 (strokes)	1,024
Controller replay refresh	Every step	Every 50 steps	Every 100 steps
Corrector data-set size / stage	100,000 endpoints	4,096 endpoints	8,192 endpoints

Appendix

Table 3: TSBM hyperparamters.

Hyperparameter	Toy 2D	Colored MNIST	CelebA
Forward / backward gradient steps, initial IMF iter	10,000	100,000	200,000
Forward / backward gradient steps	10,000	5,000	20,000
Training batch size	128	128	64
Coupling cache size	100,000 pairs	10,000 pairs	1,000 pairs
Coupling cache refresh	Once / direction / cycle	Every 1,000 steps	Every 1,000 steps
EMA rate	None	0.999	0.999

Appendix

Table 4: DSBM pretraining hyperparamters.

Input

TSBM (ours)

DPS

\gamma=0.5

DPS

\gamma=1

DPS

\gamma=2

CondSNIS

CondSMC

Appendix

Figure 5: Full red-chroma qualitative comparison on Colored MNIST.

Input

TSBM (ours)

DPS

\gamma=0.5

DPS

\gamma=1

DPS

\gamma=2

CondSNIS

CondSMC

Appendix

Figure 6: Full Thick-stroke qualitative comparison on Colored MNIST.

Input

TSBM (ours)

DPS

\gamma=0.5

DPS

\gamma=1

DPS

\gamma=2

CondSNIS

CondSMC

Appendix

Figure 7: Full Thin-stroke qualitative comparison on Colored MNIST.


0
1
2
3
4
5

Appendix

Figure 8: Full qualitative comparison for ”Natural Toothy Smile” reward on CelebA. Gray outputs and visual artifacts are retained from the saved samples.


0
1
2
3
4
5

Appendix

Figure 9: Full qualitative comparison for ”Heavy Makeup” reward on CelebA. Gray outputs and visual artifacts are retained from the saved samples.


0
1
2
3
4
5

Appendix

Figure 10: Full qualitative comparison for ”Elderly Woman” reward on CelebA. Gray outputs and visual artifacts are retained from the saved samples.

Panel of Figure 11 — Figure 11: Two-dimensional toy experiment. (a) The first two panels overlay the red source with the blue original and tilted targets; the remaining panels show terminal samples and generation trajectories from DSBM and TSBM.Red and blue dots with black outlines mark trajectory starts and endpoints, respectively; black dotted lines trace the trajectories. (b) Total variation and sliced $W_{1}$ are measured against the ground-truth tilted distribution. Cost denotes $\mathbb{E}\lVert\hat{x}_{1}-x_{0}\rVert_{2}$ . Bold marks the best distributional metrics and the cost closest to the Sinkhorn reference.

Method	$K$	Single-output time (s)	Batch-16 s/output	Relative to TSBM	Working memory (GiB)
Pretrained DSBM	1	$1.468\pm 0.010$	0.355	$1.00\times$	0.180
TSBM (ours)	1	$1.460\pm 0.005$	0.357	$1.00\times$	0.180
DPS ∗ $\gamma=1$	1	$7.116\pm 0.011$	1.836	$4.88\times$	2.229
SNIS	4	$2.089\pm 0.004$	1.452	$1.43\times$	0.379
SNIS	8	$3.094\pm 0.001$	2.905	$2.12\times$	0.758
SNIS	16	$5.794\pm 0.002$	5.811	$3.97\times$	1.512

Appendix

Table 5: Inference cost on CelebA

128\times 128

. Single-output wall-clock time is the mean

\pm

standard deviation over 16 trials. Batch-16 time is the median total time divided by 16 source outputs. Relative time uses the single-output TSBM mean. An asterisk ( ∗ ) indicates visual artifacts.

Input

TSBM

TSBM with static corrector

Appendix

Figure 13: Qualitative comparison of TSBM and TSBM with static corrector samples on Colored MNIST dataset with ”red chroma” reward.

Panel of Figure 14 — Figure 14: Quantitative analysis of different TSBM stages on Colored MNIST dataset with ”thick stroke” reward. The panels show mean reward, stroke thickness, and foreground hue error over $1,024$ fixed test source images. The gray marker is the frozen pretrained DSBM. The dotted line marks the stroke-thickness target.

Natural Toothy Smile
Method	MSE input-output	LPIPS input-output	HPS v2.1 $\uparrow$	CLIPScore $\uparrow$
Pretrained DSBM	0.1677	0.2983	0.1806	19.6923
TSBM (ours)	0.1923	0.3074	0.2012	21.5642
DPS $\gamma=0.5$	0.1815	0.3363	0.1867	20.2365
DPS ∗ $\gamma=1$	0.1827	0.3409	0.1928	20.8415
DPS ∗ $\gamma=2$	0.1843	0.3483	0.1982	21.3689

Appendix

Table 6: Additional CelebA quantitative results. MSE and LPIPS compare each output with its input, smaller values indicate less change. HPS v2.1 and CLIPScore use the same prompts as in Appendix D.3 with the same translation, crop, and horizontal-flip augmentation, larger values indicate better prompt alignment. Bold and underline marks the largest and second largest HPS or CLIPScore mean within each prompt. ( ∗ ) denotes the DPS settings marked for visual artifacts in Table 1 .

Method	ImageReward $\uparrow$	PickScore $\uparrow$	CLIP-IQA $\uparrow$
Pretrained DSBM	-0.2353	18.0223	0.5446
TSBM (ours)	0.6283	18.6727	0.5738
DPS ∗ $\gamma=0.5$	0.1858	18.1216	0.5453
DPS ∗ $\gamma=1$	0.2383	17.9490	0.4901
DPS ∗ $\gamma=2$	0.5524	17.5891	0.3833
DPS ∗ $\gamma=4$	0.5592	17.2191	0.2649

Appendix

Table 7: Aesthetics reward on CelebA-128. ImageReward and PickScore use the positive prompt with translation, crop, and horizontal-flip augmentation; CLIP-IQA measures naturalness. Means are over 128 fixed inputs, with TSBM evaluated after stage 5. All metrics are higher-is-better. Bold marks the largest mean. ( ∗ ) indicates visible artifacts in the saved samples.

Method	ImageReward $\uparrow$	PickScore $\uparrow$	CLIP-IQA $\uparrow$
Pretrained DSBM	-1.8730	17.2151	0.5446
TSBM (ours)	-0.0468	18.1322	0.4754
DPS ∗ $\gamma=0.5$	-1.8729	17.1511	0.5538
DPS ∗ $\gamma=1$	-1.8260	17.0452	0.5499
DPS ∗ $\gamma=2$	-1.7740	16.9616	0.5365
DPS ∗ $\gamma=4$	-1.7151	16.8643	0.4952

Appendix

Table 8: Eyeglasses with thick frames reward on CelebA-128. Metrics and augmentation follow Table 7 . Means are over 128 fixed inputs, with TSBM evaluated after stage 5. All metrics are higher-is-better. Bold marks the largest mean. ( ∗ ) indicates visible artifacts in the saved samples.


0
1
2
3
4
5

Appendix

Figure 15: Full qualitative comparison for the aesthetics reward on CelebA. Rows show the first 16 fixed sources; columns show the input, pretrained DSBM, TSBM (ours) at stage 10, DPS with

\gamma\in\{0.5,1,2,4\}

, CondSNIS, and CondSMC. Gray outputs and visual artifacts are retained from the saved samples.


0
1
2
3
4
5

Appendix

Figure 16: Full qualitative comparison for the eyeglasses with thick frames reward on CelebA. Rows show the first 16 fixed sources; columns show the input, pretrained DSBM, TSBM (ours) at stage 5, DPS with

\gamma\in\{0.5,1,2,4\}

, CondSNIS, and CondSMC. Gray outputs and visual artifacts are retained from the saved samples.

Explore similar work

Twisted Schrödinger Bridge Matching
Jul 18, 2026Maxence Noble, Marie Scheid, Yazid Janati +2Schrödinger Bridges Differentiable Optimal Transport
CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI +1
Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schrödinger Bridge Matching
Feb 17, 2026Jeongwoo Shin, Jinhwan Sul, Joonseok Lee +2Schrödinger Bridges Diffusion Models
Seoul National University · Georgia Institute of Technology · Sungkyunkwan University
Reflected Schrödinger Bridge Matching
Jul 3, 2026Marcus Häggbom, Viktor Nilsson, Pierre Nyquist +1Schrödinger Bridges Generative Models
SEB Group and Department of Mathematics KTH Royal Institute of Technology Stockholm, Sweden · Department of Mathematical Sciences Chalmers University of Technology and University of Gothenburg and Department of Mathematics KTH Royal Institute of Technology

Jul 18, 2026, stat.MLJ/K move · Enter open · S save

Twisted Schrödinger Bridge Matching

Maxence Noble, Marie Scheid, Yazid Janati, Eric Moulines+1

CMAP, Ecole polytechnique · Institute of Foundation Models · MBZUAI+1

Schrödinger Bridges Differentiable Optimal Transport Stochastic Differential Equations

Open paper AskPDF arXiv

Over the past few years, diffusion-based Schrödinger bridge models have been proposed to approximate optimal transport dynamics between two prescribed boundary distributions, with successful applications to generative modeling. More precisely, these methods aim to estimate a path measure whose initial and terminal marginals match the two boundary distributions, while minimizing the Kullback-Leibler divergence with respect to a reference Markov process. In this work, we consider the generalized Schrödinger bridge problem, in which the reference process is a twisted Brownian motion, that is, a Feynman-Kac transform of a Brownian motion induced by a time-dependent differentiable potential. Building on the Iterative Markovian Fitting (IMF) paradigm, and in particular on its special case Diffusion Schrödinger Bridge Matching (DSBM), which corresponds to the zero potential case, we introduce Twisted Schrödinger Bridge Matching (TSBM), a diffusion-based method designed to handle both continuous- and discrete-time potentials. Unlike previous approaches, TSBM provides a rigorous extension of the IMF scheme to the generalized Schrödinger bridge problem. This derivation leads to a new bridge-matching loss that depends explicitly on the gradient of the potential and recovers the DSBM objective when the potential vanishes, yielding improved performance. We further introduce trajectory-based variance-reduction techniques that substantially stabilize optimization and may be useful beyond the present setting. Finally, we empirically demonstrate the benefits of TSBM for trajectory inference across increasingly high-dimensional settings, including crowd navigation and single-cell data. Code available at https://github.com/maxencenoble/twisted-sb-matching.

Abstract

Explore similar work

Twisted Schrödinger Bridge Matching

Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schrödinger Bridge Matching

Reflected Schrödinger Bridge Matching