PaperlayerView exact revision on arXiv

Source-linked AI summary

Deep operator learning for efficient sampling from invariant measures of stochastic differential equations

Ling Guo, Lei Li, Jingtong Zhang

arXiv:2609.11376v1math.NAcs.LG

TL;DR

Sampling invariant measures across many SDE instances is costly because conventional methods require repeated simulations, especially when mixing is slow. The paper learns an operator-conditioned continuous normalizing flow using coefficient encodings and Lagrangian sensors, yielding competitive regime-dependent accuracy, speedups in slow-mixing settings, and a 64D proof of concept. The framework also provides expressivity and resolution-invariance results, but higher-dimensional experiments remain preliminary.

  • Problem

    Sampling invariant measures for families of SDEs is difficult because traditional sequential methods repeat potentially long mixing simulations for each varying coefficient instance.

  • Method

    The method combines neural operator encoders, flow matching, and continuous normalizing flows to transport a Gaussian reference distribution to invariant measures using random Lagrangian trajectory sensors and attention.

  • Results

    Across experiments, MCMC is stronger for easy fast-mixing settings, while the amortized sampler is useful when slow mixing dominates; it outperforms DeepONet in 2D Constant Noise and demonstrates 64D sampling.

  • Takeaways & Limitations

    The approach supports reusable sampling across related SDEs with high-dimensional, mesh-free inputs and inference costs decoupled from mixing time.

  • Takeaways & Limitations

    The supervised approach requires training samples from target invariant measures, and higher-dimensional experiments are only proof of concept with limited optimization.

Abstract

from arXiv · show

We introduce an amortized neural sampler that combines operator learning with flow methods for sampling. It maps SDE coefficient functions to pushforwards from a reference measure to the invariant measures, enabling efficient sampling across families of stochastic differential equations. Our framework shifts traditional sampling cost to an initial training phase, after which new SDE instances require only one encoder pass and a few ODE solver steps, independent of mixing time. To handle problems in high dimensions, we use Lagrangian trajectory sensors for the coefficient functions and cross attention in the architecture. We also theoretically establish the expressivity and resolution invariance of our framework. Experiments on 1D and 2D SDE families show competitive accuracy with substantial speedups over MCMC in regimes with slow mixing, transfer across sensor counts, and demonstration results on a 64D interacting particle SDE where traditional grid approaches are infeasible.

1 Introduction

The paper frames invariant-measure sampling across SDE families as an operator-learning problem and proposes an amortized flow sampler that reuses training across coefficient instances. Lagrangian trajectory sensors and attention support high-dimensional, mesh-free inputs, while theory addresses expressivity and resolution invariance.

  • 1 Introduction: Traditional MCMC requires sequential, correlated simulations with equilibration and can incur prohibitive mixing costs for multimodal or metastable systems.Repeated long simulations become a bottleneck when drift or diffusion coefficients vary across many instances.
  • 1 Introduction: The proposed sampler combines neural operators with flow matching to map unseen SDE drift and diffusion functions to samples from their invariant measures.It learns a conditional velocity field and uses an encoder-decoder design to reuse coefficient representations during ODE integration.
  • 1 Introduction: Amortization shifts the high construction cost to offline training, leaving unseen SDE instances with inference costs decoupled from mixing time.Once trained, the sampler generates approximate invariant-measure samples without a new long simulation for each coefficient pair.
  • 1 Introduction: Random Lagrangian trajectory sensors with attention accept variable sensor counts, avoid meshes, and support flexible domains in high-dimensional problems.The Perceiver architecture aggregates sensor embeddings into a fixed-dimensional context vector.
  • 1 Introduction: The framework theoretically establishes sampler expressivity for compact admissible coefficient sets and resolution invariance for random Lagrangian attention sensors.The analysis uses continuity of the coefficient-to-flow map and uniform laws of large numbers for compact attention classes.

2 Problem Setting

The paper formulates invariant-measure sampling for function-parameterized SDE families as learning a coefficient-to-distribution map. Its amortized formulation uses a shared sampler from a reference distribution and evaluates quality with Wasserstein population error.

  • 2 Problem Setting: The problem is to sample invariant measures for a parametric SDE family whose parameters are coefficient functions such as drift and diffusion fields.Under stated regularity and dissipativity conditions, each coefficient pair has a unique invariant probability measure.
  • 2 Problem Setting: Quadratic Wasserstein distance measures both the theoretical approximation guarantees and the experimental quality of the learned sampler.The metric is used for population error over the coefficient-family distribution.
  • 2 Problem Setting: The amortized formulation learns one map that pushes a fixed reference distribution forward to approximate the invariant measure for every coefficient pair in the family.This replaces separate per-instance sampling procedures with a shared operator-conditioned sampler.
  • 2 Problem Setting: Once constructed, the sampler produces samples for previously unseen coefficient pairs using forward evaluations without per-instance simulation of the underlying SDE.The formulation exchanges an upfront construction cost for cheap evaluation.

3 The general framework for neural operator sampler

The framework learns a coefficient-conditioned velocity field that transports a reference measure to invariant measures using flow matching. An encoder–decoder amortizes coefficient processing, while alternative transport constructions clarify trainability and data requirements.

  • Neural operator sampler: The neural operator sampler maps SDE coefficients and latent reference samples to an approximate target distribution through a learned conditional velocity field.For each coefficient pair a, Sθ(a, ·) pushes forward γ to an approximate invariant measure.
  • Benamou–Brenier optimal transport flow: Benamou–Brenier transport provides an exact geometrically natural flow, but requires an optimal map or coupling unavailable when training data are unpaired.With an absolutely continuous reference measure, the optimal map is the gradient of a convex potential.
  • Marginal flow matching construction: Flow matching constructs a valid transport flow from any source–target coupling, including the independent coupling that requires only unpaired samples.The conditional velocity is obtained through regression and generates the marginal transport path.
  • Marginal flow matching construction: The default flow-matching path independently pairs X0 ∼ γ and X1 ∼ µa, then uses Xt = (1 −t)X0 + tX1 with velocity target X1 −X0.Alternative empirical or minibatch optimal-transport couplings can produce different valid training targets.
  • Encoder–decoder architecture: An encoder–decoder computes the coefficient representation once and reuses it while the decoder is evaluated during ODE integration.Training uses conditional flow matching; inference encodes each new coefficient pair once before K solver steps.

4 A Lagrangian attention operator sampler

The Lagrangian attention sampler replaces infeasible high-dimensional grids with trajectory sensors, fixed-width embeddings, and a permutation-invariant Perceiver-style encoder. The resulting context conditions a decoder that outputs the flow velocity from state and time.

  • Lagrangian trajectory sensors: Random Lagrangian trajectory sensors address the exponential cost of high-dimensional grids while allowing variable sensor counts and flexible, mesh-free domains.Probe trajectories can follow a cheap exploratory process or the original SDE, and need only support pointwise evaluation of b and σ.
  • Lagrangian trajectory sensors: Each sensor records visited states with local drift and diffusion values, and a 1D CNN converts its trajectory into a fixed-size embedding.Same-padding convolutions capture local temporal patterns, while pooling makes the embedding dimension independent of trajectory length and sensor count.
  • Attention Perceiver encoder: The encoder aggregates sensors through their empirical measure, making the context invariant to sensor permutations and compatible with changing sensor counts.The architecture's measure-based representation underlies its resolution-invariance analysis as the number of sensors grows.
  • Attention Perceiver encoder: Cross attention transfers sensor information into a small latent array, self attention mixes the latent tokens, and an MLP refines them through residual blocks.The encoder uses fixed latent tokens and shared projections for the sensor embeddings.
  • Decoder: The decoder combines the encoded context with state x and flow time t to output the velocity vθ[b, σ](x, t).Time is represented with a fixed sinusoidal embedding before conditioning the decoder blocks.

5 Expressivity of the Lagrangian attention sampler

This section establishes that the sampler can approximate coefficient-dependent transport velocities and thereby invariant-measure sampling uniformly over compact SDE coefficient families, under stated assumptions. It also shows that the Lagrangian attention architecture satisfies the required approximation conditions with high probability.

  • 5.3 Expressivity of the neural operator approximation: The reduction uses the target velocity's Lipschitz continuity, while linear growth ensures global learned-ODE well-posedness and controls exits from the approximation domain.The approximating network's Lipschitz constant is not required for the reduction theorem.
  • 5.3 Expressivity of the neural operator approximation: The sampling-error bound decomposes into neural operator, escape-tail, truncated-velocity, mollification, and projection errors.These terms respectively capture compact-domain approximation, trajectory escape, ODE truncation, Gaussian smoothing, and invariant-measure projection.
  • 5.3 Expressivity of the neural operator approximation: For every η > 0, the sampler can approximate invariant-measure sampling uniformly over compact coefficient families under Assumption 5.1.Theorem 5.1 derives sampler expressivity from compact-domain neural operator approximation, with suitable truncation and smoothing parameters.
  • 5.4 Universal approximation property of the Lagrangian attention neural operator: The Lagrangian attention architecture uniformly approximates continuous permutation-invariant maps and satisfies the neural-operator assumptions when activations are globally Lipschitz.This establishes the finite sensor encoder-decoder property needed by the expressivity theorem.
  • 5.4 Universal approximation property of the Lagrangian attention neural operator: With m > m0(δ, r, ε), random probes satisfy the finite-sensor resolution condition with probability 1 − δ, yielding expressivity with at least the same probability.The probe initial points must be sampled from a distribution with a continuous density that is positive everywhere.

6 Resolution Invariance

This section formalizes random sensor resolution invariance by replacing empirical probe measures with their population law in the fixed attention architecture. Under compactness and decoder regularity assumptions, the finite-sensor terminal law converges to this population-probe law in expectation.

  • 6 Resolution Invariance: Resolution invariance concerns stability under changing sensor resolution for a fixed trained attention architecture, in expectation over random probe trajectories.The population-probe law is formed by replacing the empirical probe measure with the underlying probe law.
  • 6 Resolution Invariance: The encoder treats sensor order as irrelevant by operating on the empirical measure of embedded sensors through cross attention.The same measure formulation defines the population embedded sensor law and the finite-sensor recursion.
  • 6 Resolution Invariance: Compact sensor embeddings and regular decoder behavior provide bounded attention weights, positive denominators, compact reachable latents, and Lipschitz decoder bounds.These consequences support the uniform law of large numbers for the attention class.
  • 6 Resolution Invariance: The attention-class empirical error admits an O(m^-1/2) bound with a constant independent of the coefficient a.The rate follows from finite metric entropy induced by Lipschitz dependence on queries over a compact query set.
  • 6 Resolution Invariance: The finite-sensor terminal law converges to the population-probe terminal law with expected W2 error equal to zero.Theorem 6.1 establishes this for every coefficient a in the compact family K under Assumption 6.1.
  • 6 Resolution Invariance: The population-probe limit is not the target invariant measure but the law produced by the same finite-depth Perceiver with fixed weights.Thus resolution invariance does not by itself imply convergence to the true invariant measure.

7 Numerical Experiments

The experiments evaluate amortized operator sampling across 1D and 2D SDE families, including benign, slow-mixing, sensor-resolution, and high-dimensional settings. The method is less accurate than MCMC in fast mixing but gains substantial efficiency in rare-event regimes and shows promise for structured 64D sampling.

  • 7.1 1D Experiments: The 1D studies train on 65,536 coefficient-function instances and evaluate on 1,024 test functions using Sinkhorn divergence and W2 against ground truth.These experiments identify slow mixing as the regime where amortized operator sampling gains its main advantage.
  • 7.1.1 1D Variable Noise: In the fast-mixing variable-noise regime, MCMC is more accurate while the operator sampler remains comparable in speed and reproduces target laws across 1,024 test functions.The model systematically underestimates invariant-measure variance, and Kolmogorov–Smirnov tests reject for most test functions.
  • 7.1.2 1D Rare Event: The rare-event model achieves lower mean Sinkhorn divergence than converged MCMC with 7× less elapsed time and dominates the MCMC Pareto frontier.MCMC can become trapped in a single mode when diffusion is small, producing highly skewed error distributions and large failures.
  • 2D Constant Noise: The trajectory + Perceiver + AdaLN architecture provides competitive 2D accuracy while supporting fixed or mixed sensor counts during training and different counts at inference without retraining.The comparison includes grid DeepONet, point probes, endpoint probes, trajectory sensors, and alternative aggregation mechanisms.
  • 7.3 64D Interacting Particle: In the 64D interacting-particle proof of concept, correct interaction-law sensors outperform shuffled sensors and Gaussian baselines, while generated samples reproduce pair-distance and energy patterns.The experiment tests whether observations of an unknown pair interaction law yield physically structured samples beyond the confining potential’s one-body scale.

8 Conclusion

The framework combines operator learning with continuous normalizing flows for amortized invariant-measure sampling, with theoretical guarantees and empirical benefits across dimensions and mixing regimes. Its main scope boundaries are training-data cost, limited higher-dimensional optimization, and focus on invariant rather than time-dependent laws.

  • 8 Conclusion: The framework combines operator learning with continuous normalizing flows, using attention-based Lagrangian sensors to avoid grid discretization and proving expressivity and resolution invariance.The sampler encodes SDE coefficients into a latent representation and conditions a flow that transports a Gaussian reference distribution to a target distribution.
  • 8 Conclusion: In 1D and 2D experiments, the sampler is useful when slow mixing dominates computation, while MCMC remains stronger for easy, fast-mixing settings.In 2D Constant Noise, random Lagrangian sensors materially outperform DeepONet.
  • 8 Conclusion: The 64D Interacting Particle experiment demonstrates the sampler’s capability in high dimensions.
  • 8 Conclusion: Training requires target invariant-measure samples and substantial offline simulation effort, limiting usefulness when few related SDE instances or fast mixing make amortization unnecessary.The approach is most useful when many related SDE instances must be solved and mixing is slow.
  • 8 Conclusion: Higher-dimensional experiments remain proof-of-concept, and scaling to larger particle counts and higher-dimensional configuration spaces remains challenging.The network structure and training procedure have not been extensively optimized.

A Proofs of Main Results

The proofs establish stability properties for invariant measures under coefficient convergence, using uniform moment bounds, tightness, weak convergence, and uniqueness. These ingredients support convergence of the associated stationary measures.

  • A Proofs of Main Results: Uniform p-moment bounds yield tightness of the invariant measures, allowing subsequences to converge weakly to a probability measure.The proof uses the Lyapunov function Vp(x) = 1 + |x|p and uniform integrability of second moments.
  • A Proofs of Main Results: Coefficient convergence passes through the defining conditions, so the admissible coefficient family is closed under locally uniform convergence.
  • A Proofs of Main Results: The weak limit satisfies the invariant-measure relation through convergence on test functions and density arguments.The argument extends the relation to C2 test functions.
  • A Proofs of Main Results: Uniqueness identifies every subsequential weak limit with the target stationary measure, implying convergence of the full sequence.
  • A Proofs of Main Results: Uniform integrability combines with weak convergence to justify convergence of second-moment quantities.

A.2 Proofs for the canonical target velocity

The canonical-target proofs establish continuity, growth, and approximation control for the target velocity fields and their learned ODE flows. They use smoothing, truncation, coupling, Grönwall bounds, and sensor-based operator approximation.

  • A.2 Proofs for the canonical target velocity: The canonical velocity construction satisfies a continuity equation and admits uniform lower-density and regularity bounds on compact space-time cylinders.Gaussian-kernel representations support continuity and uniform convergence of velocities and first spatial derivatives.
  • A.2 Proofs for the canonical target velocity: A smooth cutoff gives the truncated velocity field a global Lipschitz bound while preserving continuity in the SDE coefficients.
  • A.2 Proofs for the canonical target velocity: The target and learned ODE paths have uniform moment bounds, which control exit probabilities and support removal of truncation.Linear-growth estimates and Gaussian initial moments are used for both flows.
  • A.2 Proofs for the canonical target velocity: Coupling the exact and learned flows with a common initial point converts velocity approximation into terminal-law error through stopped estimates and Grönwall’s inequality.The proof controls path deviations before exit and then bounds the resulting terminal distributions.
  • A.2 Proofs for the canonical target velocity: Finite-sensor Perceiver–AdaLN networks approximate continuous permutation-invariant sensor maps, yielding uniform coefficient-to-velocity approximation on compact sets.Compactness and partitions of unity reduce the target operator to finitely many local approximations.
  • A.2 Proofs for the canonical target velocity: Random sensor trajectories form increasingly fine nets with high probability, supporting approximation of coefficient functions as sensor counts grow.

A.4 Resolution invariance proofs

The resolution-invariance proof compares finite- and infinite-sensor encoders and their induced flows while keeping the architecture fixed. Uniform approximation, Lipschitz control, and moment bounds show convergence as the number of probes increases.

  • A.4 Resolution invariance proofs: The finite- and infinite-sensor flows start from the same Gaussian reference distribution, enabling pathwise comparison of their trajectories.
  • A.4 Resolution invariance proofs: Layer normalization, self-attention, and MLP blocks provide bounded latent states and Lipschitz maps needed for uniform encoder control.
  • A.4 Resolution invariance proofs: Resolution invariance keeps the parameters and architecture fixed while increasing the number of probes, so only the empirical sensor measure changes.
  • A.4 Resolution invariance proofs: Encoder approximation errors are controlled through uniform laws of large numbers and fixed-depth error propagation.
  • A.4 Resolution invariance proofs: Local Lipschitz estimates and Grönwall’s inequality bound flow discrepancies before exit, while linear-growth moment bounds control exit probabilities.

B Experimental Details

Runtime measurements include the full computational overhead and use comparable-cost GPU and CPU hardware.

  • Elapsed times include data loading, compilation, and I/O overhead.
  • Experiments run on a single NVIDIA RTX 5090 GPU and AMD EPYC 9575F CPU.
  • The listed GPU and CPU systems have comparable prices of $2,000 and $2,200, respectively.

B.1 1D experiments

The 1D experiments evaluate amortized operator sampling on variable-noise and rare-event SDE families, including controlled comparisons with MCMC and solver, sensor, and architecture studies.

  • 1D experiments: The shared 1D architecture uses latent dimension q = 128, a width-32 branch CNN, and a depth-4, width-256 trunk MLP.Training uses conditional flow matching with minibatch optimal transport coupling.
  • 1D Rare Event: The Rare Event experiment uses direct samples from a known Gaussian mixture, making it a controlled test of representation and generalization.Complete SDE sample generation is reserved for the other experiments.
  • 1D experiments: RK4-4 achieves Sinkhorn = 0.131 in 1.4 s, matching dopri5 at Sinkhorn = 0.130 in 7.4 s.The four-step RK4 configuration is used as the default solver in the 1D experiments.
  • 2D Constant Noise: Reference stationary samples in the 2D study use Euler–Maruyama with dt = 0.01, 2,000 equilibration steps, and 4,096 retained samples per function.Reported metrics compare 4,096 generated samples with 4,096 reference samples on 100 evaluation drift fields.
  • Sensor and architecture studies: Trajectory-probe variants retain 256 evaluation sensors while changing sensor budgets and decoder conditioning designs.The variants use latent width 128 and residual flow decoders trained with conditional flow matching.

B.3 64D Interacting Particle

The 64D interacting-particle study uses permutation-aware particle sensing and an equivariant flow architecture across nonsingular interaction kernels, while benchmarking finite-sample diagnostic limits.

  • 64D Interacting Particle: The study covers four nonsingular interaction-kernel families, with Morse and WCA energies capped at UMAX = 10 for MCMC equilibration stability.At β = 2, the associated Boltzmann weight is approximately 2 × 10^-9.
  • 64D Interacting Particle: Each kernel uses 512 isolated pair trajectories, each running 20 steps and recording relative displacement and interaction force.The resulting sensor tensor has shape (512, 20, 4) per kernel.
  • 64D Interacting Particle: The particle sampler combines a 1D CNN sensor encoder, four-token Perceiver aggregator, and equivariant flow network conditioned by AdaLN.The model has 1.78M parameters and is trained with conditional flow matching.
  • 64D Interacting Particle: The correctly conditioned sampler scores 1.7860 mean and 1.0159 median, above the target-split noise floor of 0.2038 mean and 0.1071 median.The passage attributes the remaining gap primarily to finite-sample behavior of sliced 2-Wasserstein in 64 dimensions.
  • 64D Interacting Particle: Permutation-invariant observable checks are evaluated on 100 test kernels with 4,096 samples per kernel.
Loading 2609.11376v1…