PaperlayerView exact revision on arXiv

Source-linked AI summary

Generative Modeling via Nonlinear Schr\"odinger Bridges and Branching Diffusions

Mathias Grau, Pierre Henry-Labordère

arXiv:2610.29476v1cs.LGcs.AIstat.ML

TL;DR

Existing Schrödinger-bridge methods primarily address linear entropic transport, while law-dependent costs yield coupled semilinear dynamics and high-dimensional optimization challenges. The paper formulates SBNL and uses branching diffusions to compute its coupled potentials. Experiments cover bridge recovery, nonlinear interactions, and latent image transport.

  • Problem

    Existing bridge methods primarily target linear entropic transport, and available formulations often lack objectives suited to unbiased stochastic optimization in high dimensions.

    Source: §1 · p. 1

  • Method

    SBNL characterizes semilinear mean-field Schrödinger bridges with coupled potential PDEs and uses branching diffusions for stochastic representations of the potentials.

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

  • Results

    Experiments cover bridge recovery, Gross-Pitaevskii-type couplings, potential interactions, and latent transport.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The formulation extends the Schrödinger-bridge setting beyond linear entropic transport to law-dependent running and terminal costs.

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

    The authors do not claim systematic superiority over mature linear diffusion Schrödinger-bridge pipelines, and nonlinear training requires careful hyperparameter tuning.

    Source: §6 · p. 10

Abstract

from arXiv · show

We propose a new generative modeling framework grounded in forward-backward partial differential equations (PDEs). These equations arise naturally in mean-field game (MFG) theory, and are closely related to nonlinear Schr\"odinger systems. Our main contribution is to reinterpret this connection within the setting of generative modeling and to develop an efficient numerical methodology for its practical implementation. Specifically, we derive and characterize a nonlinear extension of the classical Schr\"odinger bridge problem through a coupled system of PDEs, yielding a model we call SBNL (Schr\"odinger Bridge NonLinear). To compute its solution, we design a simulation algorithm based on branching diffusions and the probabilistic representation of semilinear PDEs. The resulting method is embarrassingly parallel, scalable to high dimensions, and applicable to a broad class of nonlinear interaction functionals. We validate the approach through numerical experiments, including recovery of the classical Schr\"odinger bridge, nonlinear extensions with convex penalties (Gross-Pitaevskii) as well as potential interaction terms, and high-dimensional generative tasks such as image synthesis on the MNIST and FFHQ human face dataset. Furthermore, the same framework embeds score-based generative modeling, diffusion models, and tractable gaussian Schr\"odinger bridges.

1. Introduction

Schrödinger bridges connect generative modeling to controlled stochastic transport, but existing methods largely target linear entropic transport and do not directly address law-dependent costs. This section introduces SBNL and a branching-diffusion approach for coupled semilinear bridge equations.

1. Introduction

  • Existing SB methods mainly target linear entropic transport, while law-dependent running costs induce coupled semilinear dynamics that current controlled-dynamics approaches are not suited to handle.Alternating fitting can be computationally expensive and accumulate errors; Gaussian-mixture approaches can have limited high-dimensional expressivity.Source: §1 · p. 1
  • Recent MFG and generalized bridge formulations clarify forward–backward potential coupling but often lack objectives suited to theoretically grounded, unbiased stochastic optimization in high dimensions..Source: §1 · p. 1
  • SBNL formulates semilinear mean-field Schrödinger bridges with law-dependent running and terminal costs and characterizes optimality through coupled potential PDEs.The PDEs are equivalent to a Lasry–Lions system.Source: §1 · pp. 1–2
  • The proposed marked-branching stochastic algorithm targets the coupled two-potential system and supports joint training rather than alternating refitting.The contributions also report unbiasedness and L2-integrability under explicit conditions, with branching during training and sampling.Source: §1 · p. 2
  • Reported experiments include bridge recovery, soft terminal penalties, Gross-Pitaevskii-type couplings, and high-dimensional latent image transport.The introduction also identifies potential interactions among the experimental settings.Source: §1 · p. 2

1.1. Related work.

  • Related work spans linear bridge fitting, Gaussian-mixture potential solvers, MFG couplings, and branching representations; this work combines law-dependent costs with jointly learned potentials and branching Monte Carlo.The related-work discussion identifies limited high-dimensional expressivity and costly alternating refits as concerns for some existing approaches.Source: §1.1 · p. 2

2. Semilinear Schrödinger bridges

The section extends the classical Schrödinger bridge to law-dependent costs and characterizes the optimizer through coupled forward-backward potentials and nonlinear terminal matching. A soft terminal penalty recovers the classical bridge in the hard-constraint limit.

2.1. Formulation.

  • The variational objective combines quadratic control cost, running marginal-law costs, and terminal-law preferences; soft constraints avoid exact calibration to finite empirical measures.Source: §2.1 · pp. 2–3

2.2. Main result.

  • The optimal process is characterized by coupled potentials whose product gives its law, with terminal conditions set by a nonlinear Sinkhorn system.Source: §2.2 · p. 3
  • As λ→∞, the soft terminal penalty enforces m(T, ·) = µT and the nonlinear Sinkhorn system recovers the classical Schrödinger bridge.Source: §2.2 · p. 4
  • Special cases include Gaussian semilinear bridges and exact score-based and DDPM models when F ≡ 0 and φ̄ ≡ 1.Source: §2.2 · p. 4

3. Branching diffusions

Marked branching diffusions turn nonlinear terms in the semilinear PDEs into particle-reproduction rules, providing stochastic representations for high-dimensional simulations. Under explicit conditions, the resulting tree estimator is unbiased and square-integrable.

3. Branching diffusions

  • The branching representation offers a stochastic alternative to deterministic Sinkhorn-type PDE iteration, which suffers from the curse of dimensionality.The marked branching approach is described as embarrassingly parallel and effective for high-dimensional problems.Source: §3 · pp. 4–5
  • Particles follow Brownian motion and reproduce according to the nonlinearity, yielding unbiased Monte Carlo estimators that parallelize for high-dimensional simulations.For local analytic interactions, truncating the power series at order K gives a polynomial nonlinearity amenable to branching.Source: §3 · p. 5

3.1. Branching representation.

  • Under explicit bounds on terminal potentials and branching coefficients, the tree is almost surely finite and its payoff expectation equals the mild PDE solution.The proposition also establishes L2 integrability of the payoff.Source: §3.1 · p. 6

4. Stochastic-optimization form of the nonlinear Sinkhorn system

The nonlinear Sinkhorn conditions are enforced through an MMD objective estimated with branching trees and optimized stochastically. Once the potentials are learned, sampling simulates the controlled process using a branching-based score estimate.

4. Stochastic-optimization form of the nonlinear Sinkhorn system

  • For a characteristic kernel, MMD is zero exactly when the two probability laws are equal.MMD is expressed using within-law kernel expectations and a cross-law expectation.Source: §4 · p. 6
  • The nonnegative MMD objective equals zero exactly when the terminal potentials satisfy the nonlinear Sinkhorn conditions almost everywhere.The objective is used to match the measures associated with those conditions.Source: §4 · p. 6

4.1. Meta-algorithm for training semilinear Schrödinger bridges.

  • The training estimators use white- or black-rooted branching trees initialized at samples from an easy-to-sample reference density.The corresponding one-tree estimators are defined for the terminal potentials.Source: §4.1 · p. 6
  • Independent branching-tree estimates provide an unbiased objective estimate, which is averaged over distinct minibatch samples and minimized by stochastic gradient descent.The construction uses independent trees attached to the reference-law samples.Source: §4.1 · p. 7

4.2. Meta-algorithm for sampling with semilinear Schrödinger bridges.

  • After learning the potentials, sampling simulates the optimally controlled process using a score estimate from branching-tree estimates of φθ.The estimated score is ∇bφ/bφ, consistent as the number of trees M tends to infinity; the dynamics can be discretized by Euler–Maruyama.Source: §4.2 · p. 7

5. Numerical experiments

Experiments test bridge recovery, nonlinear costs, and generative tasks across low-dimensional and latent spaces. The reported comparisons show competitive mixture transport results and qualitative image-translation outcomes, alongside implementation settings for the learned potentials and sampling.

5.1. Implementation details.

  • The potentials are positive feedforward networks, and joint stochastic-gradient training uses Adam with learning rate 10^-4 and batch size 64.Networks use 3–8 fully connected layers; training runs for 300 epochs, and sampling uses 50–200 Euler steps with 10^3–10^4 branching trees per score evaluation.Source: §5.1 · p. 7

5.2. Results.

  • With F=0 and a hard terminal constraint, the method recovers the classical Schrödinger bridge, reducing the problem to entropic optimal transport.The temporal-evolution figure traces Law(Xt) from Moons source to Swiss target.Source: §5.2 · p. 7, Figure 2 · p. 8
  • SBNL captures all modes and achieves the best mixture W2, while BranchSBM has the lowest empirical W2; SBNL remains competitive across evaluation metrics.The comparison uses a Gaussian-to-eight-Gaussian-mixture benchmark; Table 3b reports terminal distances.Source: §5.2 · p. 8
  • On the d=50 Gaussian-mixture benchmark, terminal mean, standard deviation, and skewness converge to target values within 50 training epochs, with no observed variance collapse.The authors attribute the absence of variance collapse to the consistent branching estimator of the score.Source: §5.2 · p. 8
  • For soft terminal conditions, terminal KL divergence decreases monotonically as λ increases, consistent with stronger penalties enforcing closer target-law alignment.The reported comparison varies the finite terminal penalty parameter λ.Source: §5.2 · p. 8
  • The Gross-Pitaevskii penalty increases branching with larger κ and c, while a potential running cost routes bridge mass around a high-potential obstacle.The potential-cost results include evolving samples and estimated-density snapshots.Source: §5.2 · p. 9
  • For MNIST digits 1 and 3 in a d=32 VAE latent space, decoded intermediate reconstructions remain visually smooth and digit-like.Figure 6 presents the temporal evolution of latent variables and corresponding decoded representations.Source: §5.2 · p. 9, Figure 6 · p. 9
  • In d=512 ALAE latent space, SBNL preserves identity and attributes in FFHQ young-to-old transfer, while several baselines collapse to a near-canonical face.The comparison uses 3×10^4 samples per age group.Source: §5.2 · p. 9, Figure 7 · p. 10

6. Discussion

The paper formulates semilinear mean-field Schrödinger bridges through coupled forward–backward PDEs and reports a branching-diffusion approach. It presents the work as a new variational and optimality system, while noting that its numerics do not establish systematic superiority in the linear setting.

6. Discussion

  • The framework characterizes law-dependent semilinear mean-field Schrödinger bridges through coupled PDEs, equivalently a Lasry–Lions system.The formulation extends beyond the linear entropic setting of Sinkhorn and iterative Markovian fitting.Source: §6 · p. 10
  • The authors do not claim systematic superiority over mature diffusion Schrödinger-bridge pipelines in the linear setting.Nonlinear-regime training requires careful tuning, and branching parallelization and Monte Carlo variance leave room for improvement.Source: §6 · p. 10

Appendix A. Proof of Theorem 2.1

This section develops the framework’s theoretical foundations and extensions, establishing its optimality system and connecting the forward-backward dynamics to imaginary-time semilinear Schrödinger dynamics. It also covers assumptions, sample-based training forms, links to mean-field control and soft terminal matching, and structured bridge diffusions arising from confining potentials and Coxeter-domain constraints.

Appendix A. Proof of Theorem 2.1

  • The proof verifies that the feedback control is optimal and that the primal and dual values coincide.Convexity yields the verification argument; the proof also establishes uniqueness of the optimal control and transforms the system into the nonlinear Sinkhorn equations.Source: Appendix A · pp. 14–15

Appendix B. Stochastic-optimization forms of the nonlinear Sinkhorn system

  • The nonlinear Sinkhorn boundary equations are recast as sample-based objectives for neural terminal potentials, estimated with branching trees and optimized by stochastic gradient descent.The appendix describes MMD matching of the reweighted measures and notes that a mixture of Gaussian kernels is a stable default.Source: Appendix B · p. 15

Appendix C. RKHS and MMD

  • For the finite measures used in the matching objective, MMD is nonnegative and equals zero exactly when the measures agree.The appendix uses this property to show that the matching loss vanishes exactly when both nonlinear Sinkhorn boundary equations hold.Source: Appendix C · p. 17

Appendix D. Standing assumptions for well-posedness

  • The verification argument assumes structural conditions and a regular forward-backward solution; existence is not proved, and local couplings may require separate arguments.The cited existence results cover regularizing nonlocal couplings under regularity and growth conditions; local examples κm and κ log m are outside those conditions.Source: Appendix D · p. 17

Appendix E. State dependence and relation to mean-field games

  • Explicit state dependence can be absorbed into dependence on the law, yielding a mean-field control problem whose optimality system couples HJB and Fokker–Planck equations while retaining entropy minimization.The appendix relates this functional to Schrödinger’s entropic optimal transport problem.Source: Appendix E · p. 18

Appendix F. Soft terminal constraint via Wasserstein cost

  • A soft terminal constraint can use an optimal-transport discrepancy, with its conjugate expressed through a cost transform; entropic relaxation offers a more computationally tractable alternative.The transport cost considered is based on |x − y|^p.Source: Appendix F · p. 18

Appendix G. Link to the (complex) semilinear Schrödinger equation and Gross-Pitaevskii

  • After Wick rotation, identifying the complex wavefunction pair with the bridge potentials recovers the paper’s forward-backward system.For the cubic Gross-Pitaevskii case, c > 0 gives a monotone mean-field interaction, whereas c < 0 gives negative monotonicity and may permit concentration or collapse.Source: Appendix G · p. 19

H.1. Confining potentials.

  • With a potential-only interaction, the system is linear and its optimal drift has a Feynman-Kac representation; a confining potential restricts diffusion to D as Λ → ∞.The limiting diffusion uses the heat kernel on D with absorbing (Dirichlet) boundary conditions.Source: §H.1 · p. 19

H.1.1. Coxeter domains and Dyson-type diffusions.

  • Coxeter chambers impose ordering or positivity constraints, and the reflection principle expresses their heat kernel as a signed sum of free Gaussian kernels.The chamber heat kernel generalizes the Karlin–McGregor determinant formula; types A and C provide ordered and positive-ordered domains.Source: §H.1.1 · p. 20
  • In type A, the long-time bridge drift converges to the Dyson Brownian motion interaction, describing eigenvalue dynamics of Hermitian Brownian ensembles (GUE).This result is for the hard constraint λ = ∞ and coincident point-mass endpoint distributions.Source: §H.1.1 · p. 21
  • In type C, the limiting diffusion is the type-C analogue of Dyson Brownian motion: three-dimensional Bessel processes conditioned never to collide.Its drift is 1/x, and the process is supported on the positive ordered chamber.Source: §H.1.1 · p. 21

H.2. Score-based generative modeling as an exact linear-SB solution (F ≡0, noised-data prior).

The paper identifies score-based diffusion sampling as an exact classical Schrödinger-bridge solution for a noised-data prior, and describes Gaussian and lifted-bridge specializations. It also gives an ansatz-based proof for the Gaussian semilinear bridge.

H.2. Score-based generative modeling as an exact linear-SB solution (F ≡0, noised-data prior).

  • With F ≡ 0 and a noised-data prior, the Schrödinger-bridge optimal drift is the score of Gaussian-smoothed data, yielding the reverse-time score-based sampler.The passage identifies this with DDPM / score-based diffusion when the second potential is 1.Source: §H.2 · p. 21
  • For Gaussian endpoint laws under the stated log couplings, the optimal diffusion has Gaussian marginals Law(X_t) = N(ˆµ_t, ˆΣ_t).The mean and covariance follow the coupled ODE system, assuming it has a solution on [0, T].Source: §H.2 · p. 22

Appendix I. Proof of Proposition H.1

  • The Gaussian semilinear-bridge proof uses Gaussian and quadratic ansätze, matches coefficients in the forward Fokker–Planck and backward HJB equations, then verifies optimality.The affine feedback is the unique optimal control under the stated assumptions.Source: §H.2 · p. 22, Appendix I · p. 23

H.4. Lifted (semilinear) Schrödinger bridge and Gyöngy’s Markov projection.

  • Lifting endpoint laws to an augmented space, solving the bridge there, and projecting back yields a Markov diffusion with the original endpoint marginals.The projected diffusion is generally not itself a Schrödinger bridge; under bounded ∇x log φ, every solution has the stated marginals.Source: §H.2 · p. 22, §H.4 · p. 23

Appendix J. Proof of Proposition H.2

The proof constructs a Markov diffusion by projecting the lifted bridge onto the original coordinates. It establishes the projected process’s marginals and notes that the time-reflected product in the bridge system is isolated by a scalar branching benchmark.

Appendix J. Proof of Proposition H.2

  • The projected drift is the conditional expectation of the lifted bridge drift given the original coordinate, and Gyöngy’s theorem supplies a Markov diffusion with matching one-time marginals.The construction uses the finite-energy assumption; bounded ∇x log φ ensures every solution has those marginals.Source: Appendix J · p. 23
  • Integrating out the lifted coordinate gives a Fokker–Planck equation whose initial and terminal marginals are µ0 and µT.The projection therefore preserves the endpoint marginals of the lifted bridge.Source: Appendix J · p. 24

Appendix K. Numerical sanity check: branching with time-reflected coupling

  • The scalar branching model isolates u(t, x)u(T − t, x) as an analytic benchmark for the bridge system’s time-reflected coupling.The bridge system pairs φ(t, ·) with ¯φ(T − t, ·).Source: Appendix K · p. 24

K.1. Sanity check.

A one-dimensional semilinear PDE provides a numerical sanity check for the branching implementation. Its Monte Carlo estimate closely matches the known analytic solution, with the heat kernel included to show the nonlinear term’s impact.

K.1. Sanity check.

  • The section introduces a one-dimensional semilinear PDE with nonlinearity u(t, x)u(T − t, x) and an explicit solution.The closed-form solution makes u(0, ·) available for comparison with the numerical estimate.Source: §K.1 · p. 24
  • The example uses a single particle family and rewrites u(t, x)u(T − t, x) in the canonical form for branching.At a branching time, a particle either restarts alone or is replaced by two particles born at s and T − s.Source: §K.1 · p. 24
  • The branching Monte Carlo estimate closely matches the analytic solution u(0, x) for T = 1.The heat-kernel solution with the same terminal condition is also shown to assess the nonlinear term’s impact.Source: §K.1 · p. 24, Figure 8 · p. 24

Appendix L. Proof of Proposition 3.1

The proof represents coupled semilinear PDE potentials as expectations over two-type branching trees and establishes conditions for well-defined, square-integrable estimates. The accompanying diagrams connect perturbative branching orders to the classical linear Schrödinger bridge.

Appendix L. Proof of Proposition 3.1

  • If cT K̄ < 1, the expected tree size is bounded and the tree is finite almost surely; if cT K̄ ≥ 1, expected size is unbounded.The bound is N(t) = 1 + c(T − t)(1 + 2K̄)/(1 − cT K̄).Source: Appendix L · p. 25
  • A root of the majorant fixed-point equation yields bounded tree expectations; without a root, at least one expectation is unbounded.When a root exists, the smallest root is obtained by iterating Ψ from zero.Source: Appendix L · p. 26
  • Square integrability requires the majorant condition for qk = b_k^2/p_k; failure can leave trees finite while allowing infinite estimator variance.For Gross-Pitaevskii experiments, condition (i) holds when cT p1 < 1, while condition (ii) depends on trained-potential sup norms and must be checked afterward.Source: Appendix L · pp. 26–27
  • The branching-tree expectations form a mild solution of the coupled PDE system, with the stated bounds justifying the averaging used in the proof.The proof also establishes uniqueness among mild solutions dominated by the specified majorants.Source: Appendix L · pp. 26–27

Appendix M. Feynman-Galton-Watson diagrams

  • For the cubic nonlinearity, each branching vertex splits one incoming line into three, and n branchings contribute order c^n.The order-0, no-branching case reduces to the classical Schrödinger bridge and linear forward-backward heat equations.Source: Appendix M · pp. 27–28, Figure 9 · p. 28

Appendix N. Implementation details

The implementation normalizes data to stabilize branching-based training and score estimation. It uses covariance whitening for Gaussian-like data and pretrained latent representations for high-dimensional images.

Appendix N. Implementation details

  • Input normalization addresses large variance in branching-estimator terminal-factor products and ill-conditioned score differentiation for differently scaled or heavy-tailed coordinates.The implementation selects the normalization mechanism according to dimensionality and data modality.Source: Appendix N · p. 28

N.1. Covariance whitening in data space.

  • For Gaussian endpoints, affine whitening maps the initial distribution to N(0, I) and the terminal distribution to N(0, σ²I), after which generated samples are mapped back.The bridge is trained on the transformed endpoint laws; T = 2 is used in practice to retain nontrivial interpolation.Source: §N.1 · p. 28

N.2. Latent representations for imaging.

  • For high-dimensional images, singular or poorly conditioned pixel covariances make pixel-space whitening inappropriate, so images are encoded into pretrained low-dimensional latent codes.MNIST uses a VAE and FFHQ uses an ALAE; the bridge is trained in latent space and decoded for visualization.Source: §N.2 · pp. 28–29

O.1. 2D synthetic experiments.

These experiments examine classical and nonlinear bridges across synthetic distributions, potential landscapes, and image latents. They also track distributional evolution and compare methods on a two-dimensional Gaussian-mixture benchmark.

O.1. 2D synthetic experiments.

  • The recovered classical bridges show the source progressively deforming into the target in two synthetic 2D cases.The cases are Gaussian to a Gaussian mixture and Moon curves to Swiss curves.Source: §O.1 · p. 29

O.2. Comparison on the 2D benchmark.

  • The Gaussian-to-eight-Gaussian comparison reports terminal samples, mixture W2 distance, and per-mode coverage across SBNL and four baselines.The methods are LightSB, LightSBB-M, DSBM, and BranchSBM; the supplied passages describe the comparison metrics but do not report its results.Source: §O.2 · p. 29

O.3. Potential running cost.

  • A state-dependent potential V shapes the transport path, with trajectories deflecting around high-potential regions while matching both marginals.The potential is evaluated during training and sampling; the appendix also shows gated-maze, S-curve, valley, spiral, and three-Gaussian examples.Source: §O.3 · p. 32, Figure 16 · p. 33, Figure 17 · p. 33

O.4. 50D synthetic experiments.

  • The 50D experiment visualizes generated and target mixture marginals and projected trajectories using the first two PCA directions.The PCA basis is fit on pooled validation draws from the source and target distributions.Source: §O.4 · p. 35, Figure 21 · p. 35, Figure 22 · p. 35

O.5. MNIST in latent space.

  • The MNIST latent bridge uses d=32 codes and visualizes both their PCA projection and decoded generated and target digits.The projected terminal latents are overlaid on the target, with point colors showing local structure.Source: §O.5 · p. 36, Figure 24 · p. 36

O.7. Comparison on FFHQ.

  • The FFHQ appendix panel compares SBNL and four baselines on young-to-old transfer in ALAE latent space.Each row includes a source child, method-generated terminal samples, and a randomly drawn adult reference.Source: §O.7 · p. 38, Figure 26 · p. 38

Appendix P. Additional figures

  • Additional figures track terminal-law convergence and the evolution of both Schrödinger potentials for Moon-to-Swiss curves.The terminal law is shown as approaching the target law as training epochs increase.Source: §O.1 · p. 29, Figure 27 · p. 39, Figure 28 · p. 39

P.1. Evolution of the potentials during training.

  • The flow-evolution figure compares λ values of 105, 104, 5×10^3, 10^3, and 10^2.Source: Figure 30 · p. 40

Related papers

Loading 2610.29476v1…