probability-paths

name: probability-paths description: Work with probability path distributions for diffusion models, including bridge path marginals, memoryless sampling, and efficient batch computation. Use when you need to sample from or evaluate the distribution p(x_t | y_1) at intermediate times.

Probability Paths for Diffusion Models

This skill covers the core probability path framework in linsdex for working with the distributional quantities that arise in diffusion-based generative models. While the /diffusion-conversions skill focuses on converting between different parameterizations (y1, score, flow, drift), this skill focuses on the underlying probabilistic objects and efficient sampling.

When to Use

• Sampling from the bridge path marginal p(x_t | y_1) at intermediate times
• Computing all flow quantities (xt, flow, score, drift, noise) jointly for training
• Efficient batch computation of probability path slices at multiple times
• Memoryless trajectory sampling for Reciprocal Adjoint Matching (RAM) training
• Understanding the mathematical structure of diffusion model distributions

Mathematical Background

A diffusion model defines a stochastic bridge between a prior distribution at t=0 and terminal evidence at t=1. The key distributions are:

• Forward marginal p(x_t | x_0): The distribution of the noised state at time t given initial state x_0
• Backward message β_t(x_t) = p(y_1 | x_t): The likelihood of terminal evidence given current state
• Bridge path marginal p_t(x_t | y_1) = ∫ p(x_t | x_0, y_1) p(x_0) dx_0: The marginal of the stochastic bridge

The ProbabilityPathSlice class encapsulates all these quantities at a specific time t, enabling exact computation of scores, flows, and drifts without approximation.

Key Classes

DiffusionModelComponents

Container for the fundamental components defining a diffusion model:

from linsdex.diffusion_model.probability_path import DiffusionModelComponents
from linsdex import BrownianMotion, StandardGaussian, DiagonalMatrix

dim = 64
sde = BrownianMotion(sigma=1.0, dim=dim)

components = DiffusionModelComponents(
  linear_sde=sde,
  t0=0.0,
  x_t0_prior=StandardGaussian(jnp.zeros(dim), DiagonalMatrix.eye(dim)),
  t1=1.0,
  evidence_cov=DiagonalMatrix.eye(dim) * 1e-4
)

ProbabilityPathSlice

Represents all probabilistic quantities at a specific time t:

from linsdex.diffusion_model.probability_path import ProbabilityPathSlice

t = 0.5
path_slice = ProbabilityPathSlice(components, t)

# Access the bridge path marginal p(x_t | y_1) as a functional
functional_marginal = path_slice.functional_pt_given_y1

# Access the backward message β_t as a functional
functional_beta = path_slice.functional_beta_t

# Access precision matrices
beta_precision = path_slice.beta_precision
marginal_precision = path_slice.marginal_precision

Sampling from the Probability Path

Basic Sampling

import jax
import jax.numpy as jnp

# Sample x_t given y_1 (returns a LinearFunctional that maps y_1 -> x_t)
key = jax.random.PRNGKey(0)
functional_xt = path_slice.sample(key)

# Resolve with a specific y_1 value
y1 = jnp.zeros(dim)
xt = functional_xt(y1)  # Concrete sample

# Or use reparameterization with explicit noise
epsilon = jax.random.normal(key, (dim,))
functional_xt = path_slice._sample(epsilon)
xt = functional_xt(y1)

Sampling All Flow Quantities for Training

The _sample_matching_items method returns all quantities needed for training as LinearFunctional objects:

from linsdex.linear_functional.functional_ops import resolve_functional

epsilon = jax.random.normal(key, (dim,))
functional_items = path_slice._sample_matching_items(epsilon)

# functional_items contains:
#   t: time
#   xt: sampled state (LinearFunctional)
#   flow: probability flow velocity (LinearFunctional)
#   score: score function value (LinearFunctional)
#   drift: SDE drift (LinearFunctional)
#   noise: the epsilon used (LinearFunctional)

# Resolve all quantities with a specific y_1
y1 = jnp.zeros(dim)
resolved_items = resolve_functional(functional_items, y1)

# Now resolved_items contains concrete arrays
xt = resolved_items.xt
flow = resolved_items.flow
score = resolved_items.score
drift = resolved_items.drift

Converting to Transition Distribution

# Convert to a GaussianTransition p(x_t | y_1)
epsilon = jax.random.normal(key, (dim,))
transition = path_slice.to_transition(epsilon)

# The transition maps y_1 to x_t
x_given_y1 = transition.condition_on_x(y1)  # StandardGaussian for x_t

Efficient Batch Computation

When computing probability path slices at multiple times, use get_probability_path which requires only 2 ODE solves regardless of the number of times:

from linsdex.diffusion_model.probability_path import get_probability_path

# Efficient: 2 ODE solves total
times = jnp.linspace(0.0, 1.0, 100)
path_slices = get_probability_path(components, times)

# path_slices is a batched ProbabilityPathSlice with batch_size = 100
print(path_slices.batch_size)  # 100

# Access individual slices by indexing
slice_at_t50 = path_slices[50]

# Or use with vmap for parallel operations
def process_slice(path_slice, y1):
  return path_slice.score(path_slice._sample(jax.random.normal(key, (dim,)))(y1))

scores = jax.vmap(process_slice, in_axes=(0, None))(path_slices, y1)

This is much more efficient than computing slices individually:

# Inefficient: 2n ODE solves for n times
path_slices_slow = jax.vmap(lambda t: ProbabilityPathSlice(components, t))(times)

Memoryless Trajectory Sampling

For efficient sampling of entire trajectories conditioned on terminal state y_1, use the memoryless utilities:

from linsdex.diffusion_model.memoryless import (
  sample_memoryless_trajectory,
  MemorylessForwardSDE,
  get_memoryless_projection_adjoint_path,
)

# Sample trajectory X_t | X_1 efficiently
x1 = jnp.zeros(dim)  # Terminal state
ts = jnp.linspace(0.0, 0.99, 50)  # Times to sample (excluding t1)
key = jax.random.PRNGKey(0)

trajectory = sample_memoryless_trajectory(
  components, x1, ts, key,
  method="discretization"  # or "simulation"
)

# trajectory is a TimeSeries with trajectory.times and trajectory.values

MemorylessForwardSDE

The reverse-time SDE that induces a memoryless path distribution:

memoryless_sde = MemorylessForwardSDE(components)

# Get parameters at reverse time s = t1 - t
s = 0.5
F_s, u_s, L_s = memoryless_sde.get_params(s)

# The SDE is: dX_s = (F_s X_s + u_s) ds + L_s dW_s
# in reverse time s where s = 0 corresponds to t = t1

Full Path for Reciprocal Adjoint Matching

For training with RAM, use the precomputed full path:

from linsdex.diffusion_model.memoryless import get_memoryless_projection_adjoint_path

times = jnp.linspace(0.0, 0.99, 50)
full_path = get_memoryless_projection_adjoint_path(components, times)

# Sample a trajectory given y1
key = jax.random.PRNGKey(0)
y1 = jnp.zeros(dim)
trajectory = full_path.sample(key, y1)

# Access precomputed quantities
p_xt_given_y1 = full_path.p_xt_given_y1  # Batched GaussianTransition
p_y1_given_xt = full_path.p_y1_given_xt  # Batched GaussianTransition
base_drifts = full_path.base_drifts  # Batched LinearFunctional
diffusion_coefficients = full_path.diffusion_coefficients  # Batched matrices

Adjoint Simulation Utilities

For computing gradients through SDE solvers using discrete adjoints:

from linsdex.diffusion_model.adjoints import (
  sde_simulation_with_internals,
  ode_simulation_with_internals,
  adjoint_simulation_from_sim_internals,
)
import diffrax

# Simulate SDE while recording solver internals
solver = diffrax.ShARK()

@diffrax.ODETerm
def drift_fn(t, x, args):
  return -x  # Example drift

def diffusion_fn(t, x, args):
  return jnp.eye(dim) * 0.1

sim_state = sde_simulation_with_internals(
  solver, x0, drift_fn, diffusion_fn,
  t0=0.0, t1=1.0,
  key=key, args=None, n_steps=100
)

# sim_state contains:
#   ts: time grid
#   xts: states at each time
#   states_pre: solver internals for adjoint computation

# Compute discrete adjoint
def terminal_cost(xT):
  return jnp.sum(xT**2)

def running_cost(t, x, args):
  return 0.0

adjoint_state = adjoint_simulation_from_sim_internals(
  sim_state, terminal_cost, running_cost, args=None
)

# adjoint_state contains:
#   ats: adjoint states aligned with ts
#   grad_theta: parameter gradients
#   total_cost: accumulated cost

Helper Classes

Affine Mappings

from linsdex.diffusion_model.probability_path import (
  Y1ToBwdMean,
  Y1ToMarginalMean,
  BwdMeanToMarginalMean,
)

# Linear functional mapping y_1 to backward message mean
y1_to_bwd = Y1ToBwdMean(components, t)
bwd_mean = y1_to_bwd(y1)  # Apply mapping

# Linear functional mapping y_1 to marginal mean
y1_to_marginal = Y1ToMarginalMean(components, t)
marginal_mean = y1_to_marginal(y1)

# Mapping between backward and marginal means
bwd_to_marginal = BwdMeanToMarginalMean(components, t)
marginal_mean = bwd_to_marginal(bwd_mean)

Probability Path Transitions

Compute the transition distribution between two times on the probability path:

from linsdex.diffusion_model.probability_path import probability_path_transition

# Compute p(x_t | x_s, y_1) for s < t
s, t = 0.3, 0.7
transition = probability_path_transition(components, components, t, s)

# transition is a GaussianTransition
x_s = jnp.zeros(dim)
p_xt_given_xs = transition.condition_on_x(x_s)  # StandardGaussian

Noise Schedule Drift Correction

When changing noise schedules while preserving marginals:

from linsdex.diffusion_model.probability_path import noise_schedule_drift_correction

def custom_noise_schedule(t, xt):
  return DiagonalMatrix.eye(dim) * 0.5

corrected_drift = noise_schedule_drift_correction(
  components, t, xt, original_drift,
  noise_schedule=custom_noise_schedule
)

Key Imports

# Core probability path classes
from linsdex.diffusion_model.probability_path import (
  DiffusionModelComponents,
  ProbabilityPathSlice,
  get_probability_path,
  probability_path_transition,
  noise_schedule_drift_correction,
  Y1ToBwdMean,
  Y1ToMarginalMean,
  BwdMeanToMarginalMean,
)

# Memoryless sampling utilities
from linsdex.diffusion_model.memoryless import (
  MemorylessForwardSDE,
  MemorylessFullPath,
  sample_memoryless_trajectory,
  get_memoryless_projection_adjoint_path,
  memoryless_noise_schedule,
)

# Adjoint simulation
from linsdex.diffusion_model.adjoints import (
  SimulationState,
  AdjointSimulationState,
  sde_simulation_with_internals,
  ode_simulation_with_internals,
  adjoint_simulation_from_sim_internals,
)

# Resolving LinearFunctional objects
from linsdex.linear_functional.functional_ops import resolve_functional

Integration with Other Skills

• Use /diffusion-conversions for converting between y1, score, flow, and drift representations
• Use /gaussian-distributions for working with the underlying Gaussian distributions
• Use /sde-conditioning for conditioning SDEs on observations
• Use /matrix-operations for efficient linear algebra with structured matrices

Tips

• Use get_probability_path instead of vmapping over ProbabilityPathSlice for efficiency
• The _sample_matching_items method is ideal for training as it returns all quantities jointly
• LinearFunctional objects defer computation until resolved with a specific y_1 value
• For RAM training, use get_memoryless_projection_adjoint_path to precompute all needed quantities
• The discrete adjoint utilities are useful when you need gradients through custom SDE solvers
• For memoryless sampling, avoid times very close to t0 or t1 (use ranges like 0.1 to 0.9) for numerical stability
• When sample_memoryless_trajectory hits ODE solver limits, try method="simulation" with solver_name="euler"