A JAX toolbox solves optimal transport problems for point clouds and histograms.
problem Optimal transport problems between point clouds and histograms.
method Automatic and custom reverse mode differentiation, vectorization, just-in-time compilation, and accelerators support.
result Solves a wide range of optimal transport problems including regularized OT, barycenters, Gromov-Wasserstein, and low-rank solvers.
We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional c…
jinns is a JAX library for physics-informed neural networks.
problem Physics-informed neural networks for forward and inverse problems.
method Physics-informed neural networks using JAX ecosystem.
result Efficient prototyping and extensions for real problems.
JAX-LOB simulates thousands of LOBs for RL training.
problem Efficient simulation of large LOBs for RL training.
method GPU-accelerated JAX implementation of LOB simulator.
result JAX-LOB processes thousands of LOBs in parallel with reduced processing time.
PGMax automates PGM inference on GPUs, improving quality and speed.
problem Efficient inference in complex discrete PGMs.
method Factor graph specification and loopy belief propagation in JAX.
result Higher-quality inference with up to 3x speedups.
GenSBI offers JAX-native SBI methods for natural sciences.
problem Lack of native SBI libraries in JAX for natural sciences.
method Flow and diffusion generative models in JAX.
result Near-ideal mean C2ST scores on SBIBM tasks.
Fast-vollib offers high-performance option pricing and IV computation.
problem Efficiently pricing and computing implied volatility for financial models.
method Open-source Python library with PyTorch, JAX, and CUDA backends, implementing Halley and LBR algorithms.
result High-performance option pricing and IV computation with vectorized implementations.
NIFTy.re accelerates imaging models and expands Gaussian processes and variational inference.
problem Slow performance and limited inference strategies in NIFTy.
method Rewritten NIFTy with new modeling principles, inference strategies, and JAX integration.
result Dramatic acceleration of models and new inference capabilities.
JaxSGMC simplifies SG-MCMC for Bayesian deep learning.
problem Uncertainty quantification in deep learning models.
method Modular stochastic gradient MCMC in JAX.
result Facilitates trustworthy neural network predictions.
JAXFit speeds up curve fitting on GPUs.
problem Nonlinear least squares curve fitting problems.
method Trust region method on GPU with automatic differentiation.
result Significantly faster than CPU and other GPU libraries.
BlackJAX simplifies Bayesian inference with modular, fast implementations.
problem Efficient and modular Bayesian inference methods for complex models.
method Functional approach to implementing sampling and variational inference algorithms in JAX.
result Ease of use, speed, and modularity in implementing Bayesian inference.
Geometric AD framework simplifies derivative computation in JAX.
problem Efficient and accurate automatic differentiation.
method Jet functors and Weil algebras for geometric analysis.
result Unified view of derivative propagation with algebraic exactness.
NumPyro is a lightweight library that provides an alternate NumPy backend to the Pyro probabilistic programming language with the same modeling interface, language primitives and effect handling abstractions. Effect handlers allow Pyro's modeling API to be extended to NumPyro despite its being built atop a fundamentall…
EagerPy simplifies writing code for multiple deep learning frameworks.
problem Code duplication and framework lock-in when using different deep learning libraries.
method Integrates multiple frameworks into a single Python framework.
result Automatic compatibility across PyTorch, TensorFlow, JAX, and NumPy.
Python package for fast simulation-based inference.
problem Intractable likelihood functions in Bayesian inference.
method Uses neural networks as surrogate models for Bayesian inference.
result Highly efficient and user-friendly for constructing SBI estimators.
New acquisition function for extreme rewards in bandits.
problem Online decision making with extreme payoffs in multi-armed bandits.
method Modeling payoffs as Gaussian processes and using a novel UCB acquisition function.
result Demonstrated benefits across synthetic and real-world benchmarks.
JaxMARL-HFT accelerates MARL for HFT with 240x speedup.
problem Heavy computational cost in MARL for HFT.
method GPU-accelerated JAX framework for multi-agent RL.
result Agents learn to outperform benchmarks in HFT.
Myriad offers a testbed for integrating machine learning and trajectory optimization.
problem Real-world trajectory optimization challenges.
method JAX-based testbed with machine learning and trajectory optimization integration.
result End-to-end learning and planning with neural ODEs.
New method for robust fixed-point smoothing without state augmentation.
problem Estimating initial states in Gaussian smoothing algorithms.
method Cholesky-based formulation without state augmentation.
result Matches runtime and robustness of existing methods.
Machine learning finds Z/2 eigenfunctions on a sphere.
problem Finding Z/2 eigenfunctions on the sphere.
method Created a multivalued neural network and used JAX to implement it. Fixed branch points at tetrahedron and cube vertices, and allowed AI to move them in the third case.
result Found Z/2 eigenfunctions for three cases.
Study develops electronic phenotypes of ICU patient acuity.
problem Limited time for patient acuity assessments and imprecise clinical trajectory prediction.
method Developed electronic phenotypes using automated variable retrieval in electronic health records.
result Identified three phenotypes: persistently stable, persistently unstable, and transitioning from unstable to stable.
Study developed phenotypes for ICU patients' brain dysfunction states.
problem Underdiagnosis of acute brain dysfunction in ICU patients.
method Created algorithms to quantify and cluster brain dysfunction states.
result Developed three phenotypes of ICU patients' brain dysfunction states.
Unified framework for efficient Gaussian process inference.
problem Efficient inference in non-conjugate Gaussian process models.
method Combines expectation propagation with linearization for improved efficiency.
result Unified view of various inference schemes, including classical smoothers and EP.
GraphGP: Scalable Gaussian Processes with Vecchia's Approximation
problem Naive Gaussian Process computation limits practical use
method GPU algorithm for Vecchia's approximation
result Linear time and memory requirements for nearly a billion parameters
New method speeds up Gaussian process inference for large datasets.
problem Numerical instability and inefficiency in approximate inference methods for non-Gaussian likelihoods.
method Conjugate-computation variational inference with Kalman recursions.
result Linear-time inference with fast and stable variational inference for state-space GP models.
New solver avoids memory issues for long differential equations.
problem Memory constraints in adaptive probabilistic ODE solvers.
method Fixed memory demands adaptive probabilistic solver using robust state estimation.
result Eliminates memory issues for long time series simulations.
Efficiently differentiate functions of large matrices using new adjoint systems.
problem Differentiating functions of large matrices in scientific and probabilistic machine learning models.
method Deriving and implementing new adjoint systems for Lanczos and Arnoldi iterations in JAX.
result Efficient differentiation of PDEs, Gaussian process models, and Bayesian neural networks.
This work optimizes MCMC algorithms for modern accelerators without synchronization overheads.
problem Synchronization overheads in vectorized MCMC algorithms.
method Designing single-chain MCMC algorithms as finite state machines (FSMs) to avoid synchronization.
result Achieves up to an order of magnitude speed-up in experiments.
New method speeds up Bayesian inference for complex simulators.
problem Challenges in Bayesian inference for complex stochastic simulators with intractable likelihood functions.
method Optimization Monte Carlo framework reformulated as deterministic optimization problems with gradient-based methods.
result Accurate posterior inference with reduced runtimes compared to existing methods.
Cookbook transforms constrained statistical inference into unconstrained problems.
problem Transforming constrained statistical inference into unconstrained problems.
method Bijective and diffeomorphisms parametrizations.
result Maintains statistical inference properties like identifiability.
A drone catches another agile drone using competitive reinforcement learning.
problem Intercepting an agile drone with another agile drone.
method Formulated as a Competitive Reinforcement Learning problem, trained with PPO, using a high-fidelity simulation environment.
result Trained policies outperform common heuristic baselines in catch rate, time to catch, and crash rate.
CoLA automates efficient numerical linear algebra for complex matrix structures.
problem Efficiently solving large-scale linear algebra problems with complex matrix structures.
method Combining linear operator abstraction with compositional dispatch rules.
result Automatic and efficient numerical algorithms for various linear algebra operations.
Improved exploration in cooperative multi-agent reinforcement learning.
problem Limited expressiveness of Gaussian policies in DecSPG hinders effective exploration.
method Proposes decentralized diffusion policy learning (DDPL) with denoising diffusion probabilistic models.
result Consistently improved performance on various MARL benchmarks.
Efficiently quantifies uncertainty in DeepONets for function spaces.
problem Uncertainty quantification in deep operator networks.
method Randomized prior ensembles for frequentist inference.
result Improved robustness and accuracy, reliable uncertainty estimates, out-of-distribution detection, and model bias quantification.
Regularized Stein thinning improves MCMC output approximations.
problem Pathologies in Stein thinning leading to poor approximations.
method Theoretical analysis and regularization to improve KSD.
result Regularized Stein thinning alleviates pathologies and improves efficiency.
Ray tracing sampler improves neural network sampling efficiency and resilience.
problem Sampling neural network posterior distributions efficiently and robustly.
method Markov Chain Monte Carlo using ray tracing through likelihood space.
result Significantly higher resilience to gradient heating compared to HMC.
Improves numerical solution of ill-conditioned linear systems for machine learning.
problem Wastefulness and instability in solving ill-conditioned linear systems.
method autonugget combines Richardson extrapolation to determine the solution of the ill-conditioned system, improving accuracy over a single nugget.
result Improves accuracy of numerical solution of ill-conditioned linear systems.
GoTube verifies neural networks over time, scaling to large horizons.
problem Verifying the robustness of time-continuous neural networks.
method Solves Go problems to construct a conservative execution set.
result Substantially outperforms existing tools in size, speed, and scalability.
Fast algorithm for rescaling vectors with clipping, improving training efficiency.
problem Efficiently rescale vectors to a desired length while maintaining them within a domain after clipping.
method Analytical solution for optimal rescaling using fast and differentiable algorithm.
result Optimal rescaling can be found analytically, improving training efficiency for neural networks.
Energy-based diffusion models improve molecular sampling and simulation.
problem Inconsistency between diffusion model scores and equilibrium distributions.
method Fokker-Planck regularization to enforce consistency.
result Improved consistency and efficient sampling of biomolecular systems.
Advocates for user-friendly RL problem descriptions to improve usability and generalization.
problem Usability and generalization challenges in RL for non-engineers.
method Development of user-friendly description languages for RL problems.
result Improved ability of RL algorithms to generalize to new problems.
New algorithms compute Volterra signature efficiently for time series analysis.
problem Efficient computation of Volterra signature with matrix-valued kernels.
method Decomposed Chen-type convolution relation, introduced FFT-based and exact recursion algorithms.
result Efficient algorithms for Volterra signature computation with various complexities.
This work introduces new ways to compare adversarial robustness of classifiers globally.
problem The limitations of point-wise measures in comparing adversarial robustness.
method Robustness curves and scale analysis to uncover global properties of robustness.
result Point-wise measures fail to capture important global properties of adversarial robustness.
Unified framework for Bayesian online learning in changing conditions.
problem Probabilistic online learning in non-stationary environments.
method BONE framework with three modelling choices and two algorithmic choices.
result Framework allows for reinterpreting and proposing new methods.
New method extracts cosmological information from dark matter halo catalogues using graph neural networks.
problem Quantifying cosmological information from large-scale structure data.
method Implicit likelihood approach with Information Maximising Neural Networks (IMNNs) on graph representations of dark matter halo catalogues.
result Graph neural network summaries can extract information from noisy catalogues and improve parameter constraints.
A GPU-based workflow for building physics emulators of hypersonic flows
problem Resolving complex physical phenomena in hypersonic flows
method Fully GPU-based workflow integrating accelerated data generation and neural emulators
result Physics emulators remain reliable beyond their training distribution
THRML uses energy-based models for index tracking, reducing portfolio tracking error and improving returns.
problem NP-hard combinatorial optimization in portfolio optimization under cardinality constraints.
method THRML reformulates index tracking as probabilistic inference on an Ising Hamiltonian, using GPU-accelerated block Gibbs sampling.
result THRML achieves 4.31 percent annualized tracking error compared to 5.66-6.30 percent for baselines, with 128.63 percent total return.
Generative model predicts financial market order flow with high accuracy.
problem Creating realistic order flow models for financial markets.
method Token-level autoregressive generative model using deep state space layers.
result Model generates high-quality order flow data with low perplexity.