Self-tuning sampler: AdaptiveMALA (dual-averaging warmup + diagonal metric)

[davidkhjo]

AdaptiveMALA(MALA) — a MALA that removes the step_size sweep.

What

• Dual-averaging warmup (Nesterov, Hoffman & Gelman 2014): drives the smooth acceptance probability ᾱ = mean min(1, e^{logα}) to the MALA-optimal 0.574, then freezes the averaged ε and samples unbiasedly. The passed step_size is only a starting guess.
• Optional diagonal metric (precondition=True): estimates a per-coordinate metric M (geometric mean 1) from a first warmup window, then restarts dual averaging and re-tunes ε under it. The step and its noise rescale together (x ← x − εM∇E + √(2εM)ξ), so the chain still targets exactly p ∝ exp(-E) — ill-conditioned targets mix at a much larger usable step.

Validation (closed-form / distributional)

• Correlated Gaussian: acceptance → 0.574 from a 9×-wrong initial step; covariance recovered.
• diag(25, 0.25) (condition number 100): learned metric ratio ≈ 100; per-axis std recovered; acceptance still on target.

Notes

• NUTS is intentionally deferred to a follow-up. A correct batched NUTS needs per-chain data-dependent termination that's easy to get subtly (silently) biased; shipping that would betray the library's correctness-first stance. AdaptiveMALA already delivers the "self-tuning samplers" goal. The verified NUTS reference spec is preserved for a later PR.
• torch-only; public API preserved; accumulates under the unreleased 0.15.0 (no version bump).

🤖 Generated with Claude Code