菜单
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.
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
基于 SOC 职业分类
A JAX-based library for linear stochastic differential equations, state-space models, and Gaussian inference. Use when working with time series interpolation, diffusion models, Kalman filtering, or probabilistic modeling with linear-Gaussian systems.
Perform inference in chain-structured Gaussian Conditional Random Fields using efficient message passing. Use for discrete-time probabilistic modeling, computing marginals, sampling joint distributions, or Kalman-style filtering and smoothing.
Convert between diffusion model representations including clean data predictions (y1), scores, probability flows, and drifts. Use when building or training diffusion-based generative models.
Work with Gaussian distributions in three parameterizations for numerical stability and efficiency. Use when you need to sample, combine distributions, or convert between mean/covariance and precision/natural forms.
Use specialized matrix types with symbolic tags for efficient linear algebra. Use when working with diagonal, block, or tagged matrices to avoid unnecessary dense computations.
Condition Linear SDEs on observations to interpolate sparse data, perform Bayesian inference on time series, or create bridges between boundary conditions. Use when working with time series interpolation, state estimation, or posterior sampling.
| 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. |
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.
A diffusion model defines a stochastic bridge between a prior distribution at t=0 and terminal evidence at t=1. The key distributions are:
The ProbabilityPathSlice class encapsulates all these quantities at a specific time t, enabling exact computation of scores, flows, and drifts without approximation.
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
)
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
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)
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
# 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
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)
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
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
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
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
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)
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
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
)
# 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
/diffusion-conversions for converting between y1, score, flow, and drift representations/gaussian-distributions for working with the underlying Gaussian distributions/sde-conditioning for conditioning SDEs on observations/matrix-operations for efficient linear algebra with structured matricesget_probability_path instead of vmapping over ProbabilityPathSlice for efficiency_sample_matching_items method is ideal for training as it returns all quantities jointlyLinearFunctional objects defer computation until resolved with a specific y_1 valueget_memoryless_projection_adjoint_path to precompute all needed quantitiessample_memoryless_trajectory hits ODE solver limits, try method="simulation" with solver_name="euler"