Skip to content

Implemented algorithms and scientific boundaries

LaFA 0.1.0 implements the six requested families with explicit variants. The single bibliography contains the original papers and source-review provenance. The validation report distinguishes numerical checks, finite-run experiments and outstanding GPU execution. Exa research remains outstanding under the approved primary-source fallback.

All potentials use U = -log π; a batch of states has shape (chains, dimensions) and a batched potential returns (chains,). Float32 and float64 are supported. Specialized inference drivers use their stated ensemble, latent and observation axes. A finite interacting ensemble is not a collection of independent MCMC chains.

Methods

Names below are public imports from the indicated module. All rows have executable implementations and CPU checks; CUDA cases are collected but have not yet run on a real GPU.

Module / API Definition and source Implemented boundary and principal checks
lafa.samplers.ULA x' = x - h M ∇U + √(2hTM) ξ; R01, R02, R20 Constant positive diagonal mobility; finite-step bias. Analytic proposal/noise and Gaussian stationary covariance checks.
lafa.samplers.MALA ULA proposal and full reverse/forward MH ratio; R01, R19, R20 Fixed step, temperature and mobility; whole-chain acceptance. Ratio, rejection, invalid proposals and continuation checks.
lafa.samplers.LipULA, LipMALA Local Lipschitz gradient-difference rule and square-root growth cap; R01 Continuous adaptation is approximate. adapt_steps freezes production; finite maximum step safeguards flat targets. Paper interpretation documents signs, initialization, rejection and normalization decisions.
lafa.specialized.AnchoredLangevin Reference drift and noise scaled by exp(U-U0); R03 Exact supplied reference potential, Euler discretization and random operational clock. No arbitrary attraction or clipping. Monte Carlo Gaussian-smoothed and reflected extensions are not implemented.
lafa.samplers.KineticLangevin Euler–Maruyama for dx=p/m dt, dp=-∇U dt-γp dt+√(2γmT)dW; R05, R06 Explicit Euler, not the exact frozen-force scheme analyzed in R06. Momentum and fluctuation–dissipation tests.
lafa.samplers.BAOAB B–A–O–A–B splitting with exact OU substep; R05 Positive diagonal mass; nonnegative damping. No universal fourth-order claim. Momentum covariance and one-step checks.
lafa.specialized.HOLA Tamed/untamed higher-order overdamped scheme with correlated noise; R07, Eqs. (2)–(4), (8) Distinct tamings of gradient, Hessian, Hessian-gradient contraction and Laplacian gradient. Dense small-dimensional reference; exact derivative-provider path for large dimension. Two-noise covariance and provider agreement tests.
lafa.specialized.AnnealedLangevin, NoiseAnnealedLangevin Known-potential temperature annealing; separately, noise-conditioned score annealing R08 Explicit schedules and finite budgets. Temperature scheduling is a documented variant, not a claim to reproduce R08's learned score model. Endpoint particles are approximate.
lafa.specialized.ReplicaExchange MALA within fixed-temperature replicas and alternating neighboring swaps; R09, R19 Corrected local transitions plus log-domain swap ratio. R09 motivates exchanges; this implementation does not claim its exact discretization or acceleration theorem.
lafa.specialized.MYIPLA, PIPULA, PIPGLA Joint parameter/latent proximal methods; R10, Eqs. (7)–(14) Shared parameter, latent particles, parameter noise √(2h/N). Caller supplies the joint prox; full-potential prox for PIPULA. Exact fixture KKT checks and independent integrated-likelihood comparison. Parameter learning is not ordinary Bayesian parameter-posterior sampling.
lafa.particles.SVGD, WeightSpaceSVGD Kernel attraction plus source-gradient repulsion; R11 Weight-space name is an explicit alias of SVGD. Exact blocked interactions, optional normalized weights. Dense/autograd and permutation oracles.
lafa.particles.ProjectedSVGD Likelihood-gradient informed SVD basis; R12 Whitened isotropic Gaussian prior; frozen-complement profile approximation, not exact marginalized-profile integration. Basis refresh and continuation checks; no dense parameter covariance.
lafa.particles.SteinVariationalNewton Published block-diagonal Newton system, R13 Eq. (17), followed by kernel interpolation Matrix-free HVPs and CG with true residual/curvature checks. Exact or caller-supplied PSD Gauss–Newton curvature, explicit damping. Full coupled SVN is not implemented. R13.
lafa.particles.FunctionSpaceSVGD Function-space Stein field followed by each destination's VJP; R14 Explicit single-state function map and batched function potential. The user defines the function-space target/prior. Nonlinear VJP finite differences and prediction reference checks. Optional chunked vmap outputs/VJPs preserve the loop default.
lafa.diagnostics.ksd Langevin Stein kernel for RBF/IMQ; R15 and R01 Appendix A Exact blocked weighted U/V statistics. U estimates of squared KSD may be negative; no clamp masquerading as an unbiased estimator. No approximate pair estimator. Dense and autograd oracles; missed-mode counterexample.
lafa.particles.RepulsiveEnsemble, FunctionSpaceRepulsiveEnsemble KDE-WGD / function-space KDE-WGD; R16 Own-score attraction and KDE repulsion; function-space version uses a VJP. These are not renamed SVGD. Stable log-density reductions and weighted references.
lafa.specialized.SALEnKF Stochastic approximation with Langevinized ensemble filtering; R17 Linear Gaussian observations, differentiable possibly nonlinear transition, full observation blocks, backward sampling from retained previous-stage pool. Learns point parameters; state posterior is approximate. No nonlinear-observation augmentation or minibatching. Kalman reference and transition-score tests.
lafa.specialized.BirthDeathLangevin KL-driven WFR transport plus sequential KDE birth/death; R18 Fixed population with equal weights, explicit ancestry and recomputed reaction rates. Positive rate kills; negative rate duplicates. Finite step/KDE approximation; reference sequential reactions can be expensive.

The root imports lafa.ULA, MALA, LipULA, LipMALA, BAOAB are torch.optim.Optimizer adapters. Tensor-based drivers live in lafa.samplers. The optimizer MALA decision uses the complete parameter state; dynamic addition of parameter groups is explicitly unsupported. Closures must represent a deterministic full potential, with gradients recomputed on each call. Arbitrary sparse, complex or stateful models are not automatically supported.

Computation and storage

For N states of dimension D, simple transitions retain O(ND) state and perform target/gradient work plus O(ND) arithmetic. Default step results own detached snapshots. Native step(fields=...) selects outputs and advance(steps) suppresses result allocations while retaining validation. Opt-in compiled potential/gradient evaluation has whole-sampler CPU timing evidence; eager validation can still synchronize CUDA. Chunked ModulePotential and optional vectorized function VJPs bound simultaneous transformed evaluations, without changing model or target definitions.

Exact pair methods cost O(N²D) arithmetic. Pair tiles use O(B² + BD) workspace and avoid a full N×N×D tensor. Matrix-free SVN avoids D×D Hessians but performs multiple HVPs and CG solves. pSVGD uses a thin SVD of the N×D likelihood-gradient matrix. Dense HOLA intentionally has a dimension guard; derivative providers must supply exact contractions and a Hessian spectral norm. SA-LEnKF retains a complete post-burn-in pool for a bounded number of stages, so memory depends on particles, dimension and per-stage budget. Birth/death rate recomputation prioritizes the specified sequential algorithm over a claim of fast large-N scaling.

The recording guide describes selected fields, stride, online moments, chunk readers and checkpoints. Tensor staging/queue memory is bounded; serialization overhead and manifest metadata are additional. GPU-to-host copies are currently synchronous, even when disk writing is asynchronous. None of the algorithms accumulates an unrequested trajectory.

Diagnostics, fixtures and remaining extensions

autocorrelation, raw/bulk/tail ESS, mcse_mean, mcse_quantile and rank_normalized_split_rhat accept (chains, draws, dimensions) after warm-up. Raw ESS uses pooled Geyer positive/monotone sums and a conservative total-draw cap [R22]. Bulk ESS uses split rank normalization; tail ESS uses quantile indicators. Mean and quantile MCSE estimate finite-run variability. R-hat uses rank normalization, splitting and folding [R21]. New split ESS/MCSE APIs require at least 10 draws; undefined constant inputs yield NaN. OnlineMoments uses Welford's recurrence [R23]. These tools cannot establish convergence alone.

Fixtures include analytic Gaussians, a multimodal mixture, the supplied paper's normalizable Rosenbrock density, a normalized banana, actual sparse straight-ray 2D traveltime tomography, nonlinear transmission through a positive 2D field, a nonsmooth latent model and a linear Gaussian state-space model. The reduced nonlinear model has a refined two-coefficient quadrature reference. Separate full-grid experiments infer all 9/16 latent cells and compare with finite elliptical-slice references [R24], reporting reference mixing/MCSE and coverage across independent simulated datasets.

The initial implementation does not provide reflected anchored dynamics, finite-Monte-Carlo anchor estimation, full coupled SVN, arbitrary nonlinear-observation SA-LEnKF, approximate KSD pairs, distributed particles, fused CUDA kernels or general mixed-precision inference. These are extensions with separate derivation/performance requirements. Broader calibration, larger neural-network/full-grid problems and stronger reference mixing remain extensions beyond the measured small examples. Real GPU execution remains a release-validation gate, documented in RunPod setup.