Single-seed generation of Brownian paths and integrals for adaptive and high order SDE solvers

2024-05-10

arXiv.org

AI Summary

Adaptive stochastic differential equation (SDE) solvers lack methods for exact sampling of higher-order Brownian time integrals—such as ∫Wᵣ dr—essential for achieving high-order convergence. Method: We propose the extended Virtual Brownian Tree (VBT), the first framework enabling joint exact sampling of both Brownian paths and their time integrals. Implemented in JAX, it employs a single-seed pseudorandom number generator to realize a memory-constant, error-controlled, reproducible binary-search VBT, with theoretical guarantees of strict distributional consistency at ε-separated query points. The implementation is integrated into the Diffrax library. Results: Experiments demonstrate that our method more than doubles the observed convergence order of adaptive high-order SDE solvers on a highly volatile CIR model. In MCMC tasks, a third-order kinetic Langevin solver using our method outperforms NUTS, reducing function evaluations by 90%.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Intelligent Robots: State Estimation

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Virtualization and resource management in Web systems and infrastructures
• Search and Retrieval-Augmented AI: Vertical and domain-specific search

Abstract

Despite the success of adaptive time-stepping in ODE simulation, it has so far seen few applications for Stochastic Differential Equations (SDEs). To simulate SDEs adaptively, methods such as the Virtual Brownian Tree (VBT) have been developed, which can generate Brownian motion (BM) non-chronologically. However, in most applications, knowing only the values of Brownian motion is not enough to achieve a high order of convergence; for that, we must compute time-integrals of BM such as $int_s^t W_r , dr$. With the aim of using high order SDE solvers adaptively, we extend the VBT to generate these integrals of BM in addition to the Brownian increments. A JAX-based implementation of our construction is included in the popular Diffrax library (https://github.com/patrick-kidger/diffrax). Since the entire Brownian path produced by VBT is uniquely determined by a single PRNG seed, previously generated samples need not be stored, which results in a constant memory footprint and enables experiment repeatability and strong error estimation. Based on binary search, the VBT's time complexity is logarithmic in the tolerance parameter $varepsilon$. Unlike the original VBT algorithm, which was only precise at some dyadic times, we prove that our construction exactly matches the joint distribution of the Brownian motion and its time integrals at any query times, provided they are at least $varepsilon$ apart. We present two applications of adaptive high order solvers enabled by our new VBT. Using adaptive solvers to simulate a high-volatility CIR model, we achieve more than twice the convergence order of constant stepping. We apply an adaptive third order underdamped or kinetic Langevin solver to an MCMC problem, where our approach outperforms the No U-Turn Sampler, while using only a tenth of its function evaluations.

Problem

• Adaptive time-stepping for high-order SDE solvers
• Generating Brownian motion and time integrals non-chronologically
• Enabling constant memory usage and experiment repeatability

Innovation

• Extends Virtual Brownian Tree to generate time-integrals
• Uses single PRNG seed for constant memory and repeatability
• Enables adaptive high-order SDE solvers with logarithmic complexity