Technical Summary: Deterministic Decomposition of Stochastic Generative Dynamics

Problem Statement

Modern generative modeling is increasingly viewed through the lens of probability transport, where samples evolve from a simple base distribution $\pi_0$ to a target data distribution $\pi_1$. A fundamental distinction exists between deterministic models (e.g., Normalizing Flows, Flow Matching), which evolve samples via Ordinary Differential Equations (ODEs) governed by a velocity field, and stochastic models (e.g., Diffusion Models), which evolve samples via Stochastic Differential Equations (SDEs) involving both drift and diffusion.

While deterministic models offer tractable velocity fields, stochastic models capture richer density evolution. However, when stochastic dynamics are reformulated into deterministic representations (such as probability flow ODEs), the distinct mechanisms of drift (directional transport) and diffusion (density reshaping) are often compressed into a single effective velocity field. This collapse obscures the specific roles of deterministic evolution versus stochastic fluctuation, making it difficult to analyze, control, or independently recombine these effects during sampling. Existing approaches typically treat the stochasticity as an external noise source or absorb it entirely into a unified drift, failing to provide a structured deterministic representation that separates transport from diffusion-induced effects.

Methodology: Transport–Osmotic Decomposition

The authors propose a theoretical framework that decomposes the deterministic velocity field of a stochastic generative process into two distinct, complementary components: a Transport field ($u_t$) and an Osmotic field ($d_t$).

1. Theoretical Derivation

The decomposition arises from the intrinsic bidirectional structure of stochastic processes. For a given marginal density path $\{\pi_t\}_{t \in [0,1]}$, the forward drift $b^\to_t$ and the backward drift $b^\leftarrow_t$ are coupled through the marginal score $\nabla \log \pi_t$. Under canonical regularity conditions, the relationship is:
$$b^\to_t(x) - b^\leftarrow_t(x) = \beta \nabla \log \pi_t(x)$$
where $\beta$ is the diffusion coefficient.

The authors define the decomposition as:

• Transport Field ($u_t$): The symmetric component, defined as the average of forward and backward drifts:
  $$u_t(x) = \frac{1}{2} (b^\to_t(x) + b^\leftarrow_t(x))$$
  This field governs the marginal probability transport via the continuity equation ($\partial_t \pi_t + \nabla \cdot (\pi_t u_t) = 0$). It represents the deterministic backbone of the dynamics, moving probability mass without altering the local density geometry through diffusion.
• Osmotic Field ($d_t$): The antisymmetric component, defined as half the difference between forward and backward drifts:
  $$d_t(x) = \frac{1}{2} (b^\to_t(x) - b^\leftarrow_t(x)) = \frac{\beta}{2} \nabla \log \pi_t(x)$$
  This field is entirely determined by the marginal score. It captures diffusion-induced, density-dependent effects (analogous to osmotic velocity in classical diffusion theory). It does not represent sample-wise randomness but rather how diffusion reshapes the marginal density.

The original forward drift is reconstructed as $b^\to_t = u_t + d_t$.

2. Bridge Matching (BM) Framework

To learn these decomposed fields, the authors propose Bridge Matching, a flow-based training framework. Instead of learning a single effective velocity, BM learns two separate neural networks: $u_\theta$ (transport) and $d_\phi$ (osmotic).

The training objective is a regression loss that matches the learned fields to target fields:
$$\mathcal{L}_{BM} = \mathbb{E}_{t, x_t} \left[ \| u_\theta(x_t, t) - u^*(x_t, t) \|^2 + \lambda_d \| d_\phi(x_t, t) - d^*(x_t, t) \|^2 \right]$$

The paper introduces two formulations for constructing the targets $u^*$ and $d^*$:

• Marginal Bridge Matching (MBM): Directly targets the marginal decomposition. The transport target $u^*$ satisfies the continuity equation for the prescribed path, and the osmotic target $d^*$ is derived from the marginal score $\nabla \log \pi_t$. This requires estimating the marginal score, often via Kernel Density Estimation (KDE).
• Conditional Bridge Matching (CBM): Uses a tractable conditional distribution $p_t(x|z)$ (e.g., a Gaussian interpolation between endpoints) to avoid explicit marginal score estimation. The targets are defined conditionally: $d^*_t(x, z) = \frac{\beta_t}{2} \nabla_x \log p_t(x|z)$ and $u^*_t = v^*_t - d^*_t$. The expectation over the conditioning variable $z$ recovers the marginal decomposition.

3. Sampling and Recombination

A key feature of the framework is the ability to recombine the learned fields at sampling time. The forward sampling dynamics are defined as:
$$b^{BM}_t(x) = u_\theta(x, t) + \lambda_d d_\phi(x, t)$$
where $\lambda_d$ is a tunable weight. This allows for controllable sampling where the contribution of the diffusion-induced osmotic effect can be adjusted independently of the transport backbone.

Key Contributions

1. Transport–Osmotic Decomposition: The paper derives a formal decomposition of stochastic generative dynamics, demonstrating that deterministic and stochastic models share a common marginal transport backbone ($u_t$) but differ in how diffusion-induced effects ($d_t$) are integrated into sample dynamics.
2. Bridge Matching Framework: The authors propose a training framework that learns these decomposed fields using both marginal and conditional formulations, effectively connecting score-based information with flow-style regression objectives.
3. Controllable and Interpretable Sampling: Empirical results show that recombining the learned components with varying weights ($\lambda_d$) enables interpretable control over the generation process. Adjusting $\lambda_d$ allows for trading off between sharpness and smoothness in generated images and improves endpoint matching in specific 2D transport tasks.

Experimental Results

The authors evaluated Bridge Matching on 2D transport tasks (checkerboard, moons, Gaussian mixtures) and image generation benchmarks (CIFAR-10, ImageNet-32, ImageNet-64).

• 2D Transport: On 2D tasks, BM variants achieved competitive endpoint matching compared to standard Flow Matching (FM) baselines. The optimal osmotic weight $\lambda_d$ was found to be task-dependent; for complex, multimodal targets, increasing the osmotic contribution improved matching, while for simpler structures, it could degrade performance.
• Image Generation: On CIFAR-10 and ImageNet subsets, the CBM model achieved generation quality comparable to or better than the CFM-Diffusion baseline. Crucially, the optimal $\lambda_d$ varied across evaluation metrics (FID, KID, IS), suggesting that the decomposition provides a mechanism to trade off different aspects of generation quality. Notably, the best performance was not always achieved at $\lambda_d = 1$ (the theoretical reconstruction), indicating that sampling-time recombination offers a useful degree of freedom for finite-capacity models.
• Visualization: Visualizations of the vector fields confirmed that $u_t$ drives coherent, large-scale transport, while $d_t$ exhibits structured, smaller-scale patterns responsible for local density reshaping. Particle-level simulations illustrated how $d_t$ shifts trajectories along the probability path without dominating the global transport trend.

Significance and Claims

The paper claims to provide a unified perspective on deterministic and stochastic generative dynamics. By explicitly separating the transport backbone from diffusion-induced density shaping, the work offers an interpretable view of how stochasticity shapes probability transport.

The authors emphasize that this decomposition is not merely a theoretical curiosity but a practical tool for controllability. It allows practitioners to adjust the "osmotic contribution" at sampling time without retraining, enabling fine-grained control over sample characteristics (e.g., sharpness vs. smoothness) and potentially facilitating interactive sampling procedures (e.g., pausing, reversing, or regenerating from intermediate states).

The paper remains modest regarding its scope, noting that the decomposition applies most directly to field-based generative dynamics and that learning separate fields introduces additional modeling and computational overhead. Future work is suggested to focus on reducing this overhead while preserving the interpretability and controllability benefits in downstream applications.