Unified Riemannian Metric API

[THargreaves]

Overview

Insert relevant XKCD — many apologies guys but I felt like another RHMC PR was needed.

I've had a good look through Qingyu, Niko and Jamie's and found them to be a great start with some useful ideas on tidying up the RHMC interface.

I couldn't however help but feel like they were all doing some small (though useful) tweaks without fundamentally addressing the poor interface design of the current RHMC implementation.

This PR builds on Niko's PR (which builds on Yingqu's PR) whilst pulling in some ideas from Jamie's, whilst fundamentally reshaping the RHMC interface so it is philosophically sound.

The high-level takeaway of these changes is that RHMC now works with both SoftAbs eigendecomposition of G as well as Cholesky decompositions of guaranteed PD G, whilst sharing 90% of their code paths.

The key insight is that the gradient computation can be unified by introducing two dispatch functions (logdet_grad_matrix and kinetic_grad_matrix) that return matrices encoding the appropriate chain rule for each metric type.

We also introduce the notion of a pre-metric — for PD metric this is just G, for SoftAbs this is H. This abstraction allows the SoftAbs hamiltonian derivatives to be computed with optimal caching (since we can pull out the derivative part of G into logdet_grad_matrix leaving just H behind as in Betancourt's paper).

I've added tests for energy conservation and verifying that the samplers run but it would be nice to add some validity tests too. Could this be something @J-Price-3 could look at?

Please feel free to tear this PR apart; it's just a draft and I'm open to changes.

New Types

RiemannianMetric

For user-provided positive-definite metrics (e.g., Fisher information matrix):

calc_G = θ -> fisher_information(θ)
calc_∂G∂θ = θ -> autodiff_jacobian(calc_G, θ)
metric = RiemannianMetric((d,), calc_G, calc_∂G∂θ)

SoftAbsRiemannianMetric

For Hessian-based metrics with SoftAbs regularization:

calc_H = θ -> hessian_of_neg_log_density(θ)
calc_∂H∂θ = θ -> autodiff_jacobian(calc_H, θ)
metric = SoftAbsRiemannianMetric((d,), calc_H, calc_∂H∂θ, α=20.0)

SoftAbsEval

Cached result of evaluating a SoftAbs metric, storing:

• Q: Eigenvectors
• softabsλ: Transformed eigenvalues λ * coth(αλ)
• J: Divided difference matrix for gradient chain rule
• M_logdet: Precomputed matrix for logdet gradient

Key Design Decisions

Unified Gradient Computation

The Hamiltonian gradient ∂H/∂θ has three terms:

1. Potential energy gradient: -∂ℓπ/∂θ
2. Logdet term: 0.5 * tr(M_logdet * ∂P/∂θ)
3. Kinetic term: -0.5 * tr(M_kinetic * ∂P/∂θ)

Where P is the "pre-metric" (G for RiemannianMetric, H for SoftAbsRiemannianMetric).

The dispatch functions logdet_grad_matrix(G) and kinetic_grad_matrix(G, r) return M_logdet and M_kinetic respectively, with the J matrix (divided difference formula) absorbed to handle the chain rule through the SoftAbs transformation.

O(n³) Complexity

By precomputing J-dependent matrices, we maintain O(n³) complexity per gradient evaluation rather than O(n⁴). The trace products are computed efficiently using tr(A*B) = sum(A' .* B) in O(n²).

Caching for Fixed-Point Iterations

The generalized leapfrog integrator requires fixed-point iterations where θ is fixed but r varies. RiemannianGradCache stores all θ-dependent computations (eigendecomposition, logdet gradient terms) to avoid redundant work.

Shared Metric Evaluation

neg_energy and ∂H∂r both need the evaluated metric. An optional G_eval parameter allows sharing the computation, and phasepoint uses this to compute both efficiently.

Deprecations

The following types are deprecated with warnings:

• DenseRiemannianMetric → Use RiemannianMetric or SoftAbsRiemannianMetric
• IdentityMap → Use RiemannianMetric directly
• SoftAbsMap → Use SoftAbsRiemannianMetric

Existing code using the deprecated API will continue to work but will emit deprecation warnings — I'd be quite keen to remove these in the long-run.

References

• Girolami & Calderhead (2011). "Riemann manifold Langevin and Hamiltonian Monte Carlo methods"
• Betancourt (2012). "A general metric for Riemannian manifold Hamiltonian Monte Carlo"

[nsiccha]

I couldn't however help but feel like they were all doing some small (though useful) tweaks without fundamentally addressing the poor interface design of the current RHMC implementation.

Makes sense to me. I've started looking at the old PR seriously shortly before my vacation last year, but relatively quickly realized that I wouldn't want to merge it as is, and that we'd need to change a few to a lot of things. I haven't finished the last PR because I had been thinking slightly further than what's currently there and about how to share as much code as possible between the different integrators as well, but haven't implemented the solution yet.

I'll have a look at this PR next week when I'm back!

[J-Price-3]

This sounds like a big improvement. The current implemenation was quite difficult to improve on because the code is quite disjointed. This should make it a lot easier to ensure our performance is good. I'll review the changes and can work on some tests in the next week or so.

[J-Price-3]

For the SoftAbsEval type, it might make sense to somehow optionally provide the resulting matrix from softabs. The old code was always doing this by default using Q * diagm(softabsλ) * Q', which prevents having to do this multiplication more than once. We shouldn't do this every time though, because some cases don't actually need the full multiplied result.

[THargreaves]

Makes sense to me. I've started looking at the old PR seriously shortly before my vacation last year, but relatively quickly realized that I wouldn't want to merge it as is, and that we'd need to change a few to a lot of things. I haven't finished the last PR because I had been thinking slightly further than what's currently there and about how to share as much code as possible between the different integrators as well, but haven't implemented the solution yet.

I'll have a look at this PR next week when I'm back!

Glad to hear we're on the same page about a rewrite being needed. Looking forward to seeing your thoughts!

[THargreaves]

For the SoftAbsEval type, it might make sense to somehow optionally provide the resulting matrix from softabs. The old code was always doing this by default using Q * diagm(softabsλ) * Q', which prevents having to do this multiplication more than once. We shouldn't do this every time though, because some cases don't actually need the full multiplied result.

That's a good thing to consider but I don't think we ever actually use the reconstructed matrix.

I think this was only returned by the old softabs function so that ∂H∂r could do G \ r with dense G (by laziness since Julia doesn't implemented an linear algebra compatible eigendecomp). Now we use SoftAbsEval directly to solve G \ r much more efficiently using

function Base.:\(G::SoftAbsEval, p::AbstractVector)
    return G.Q * ((G.Q' * p) ./ G.softabsλ)
end

[nsiccha]

@THargreaves and @J-Price-3: I think the general pattern that appears here (but which I've encountered elsewhere as well) is that a) there's a simple dependence structure between the various intermediate and final QOIs that downstream algorithms need, b) there's no "easy" way to "code up the math as in the paper" and get an efficient implementation, c) there's obvious potential (though in this case probably marginal) benefit in reusing working memory across computations or iterations, and d) the efficient "computations paths" may vary for different downstream algorithms (e.g. generalized leapfrog or implicit midpoint integrators).

I think there should be a potentially easy and general solution to this combination of problems!

[THargreaves]

@THargreaves and @J-Price-3: I think the general pattern that appears here (but which I've encountered elsewhere as well) is that a) there's a simple dependence structure between the various intermediate and final QOIs that downstream algorithms need, b) there's no "easy" way to "code up the math as in the paper" and get an efficient implementation, c) there's obvious potential (though in this case probably marginal) benefit in reusing working memory across computations or iterations, and d) the efficient "computations paths" may vary for different downstream algorithms (e.g. generalized leapfrog or implicit midpoint integrators).

I think there should be a potentially easy and general solution to this combination of problems!

I agree with this assessment, especially point (b). Given the downstream use-cases of Turing, I would lean towards department from the "maths in the paper" to focus on an efficient abstraction.

[nsiccha]

I'm putting this in here even though it doesn't relly belong into the PR discussion, but RHMC has been one of the inspirations to think about the below pattern again.

Essentially, I'd like to use Julia macros to be able to specify that dependency graph very concisely, and then have functions/algorithms which compute the needed quantities and only the needed quantities efficiently - essentially reactive programming. I didn't quite find a package which would do this the way I'd like to do it. Maybe there's a good reason for that.

This is for example more or less the functionality I'd like (but probably not the final syntax), first for a Euclidean metric using the standard leapfrog integrators:

@reactive euclidean_phasepoint(pot_f, metric, pos, mom) = begin 
    pot, dpot_dpos = pot_f(pos)

    chol_metric = cholesky(metric)
    # The logdet term could be left out as it doesn't change
    kin = .5 * (@node(logdet(chol_metric)) + dot(mom, dkin_dmom))
    dkin_dmom = chol_metric \ mom

    ham = pot + kin
    dham_dpos = dpot_dpos
    dham_dmom = dkin_dmom
end

@reactive leapfrog!(phasepoint; stepsize) = begin 
    @. phasepoint.mom -= .5 * stepsize * phasepoint.dham_dpos
    @. phasepoint.pos += stepsize * phasepoint.dham_dmom
    @. phasepoint.mom -= .5 * stepsize * phasepoint.dham_dpos
end

The first function definition would specify the parse-/compile-time dependency graph, with @node generating an anonymous node if that makes sense. The second function then uses that dependency graph to appropriately recompute the needed QOIs for every line of code whenever needed.

The same functionality, but for the identity map Riemannian metric together with the generalized leapfrog integrator:

@reactive riemannian_phasepoint(pot_f, metric_f, pos, mom) = begin
    pot, dpot_dpos = pot_f(pos)

    # It would be better to have the interface be to compute (pot, dpot_dpos, metric, metric_grad) in one swoop
    metric, metric_grad = metric_f(pos)
    chol_metric = cholesky(metric)
    inv_metric = Symmetric(inv(chol_metric))
    kin = .5 * (@node(logdet(chol_metric)) + dot(mom, dkin_dmom))
    dkin_dmom = chol_metric \ mom
    dkin_dpos = @node(map(eachindex(pos)) do i
        .5 * tr_prod(inv_metric, metric_grad[:, :, i])
    end) .- Base.broadcasted(eachindex(pos)) do i
        .5 * tr_prod(Base.broadcasted(*, dkin_dmom, dkin_dmom'), metric_grad[:, :, i])
    end

    ham = pot + kin
    dham_dpos = dpot_dpos + dkin_dpos
    dham_dmom = dkin_dmom
end

@reactive generalized_leapfrog!(phasepoint; stepsize, n_fi_steps) = begin 
    pos0, mom0 = map(copy, (phasepoint.pos, phasepoint.mom))
    for _ in 1:n_fi_steps 
        @. phasepoint.mom = mom0 - .5 * stepsize * phasepoint.dham_dpos
    end
    dham_dmom0 = copy(phasepoint.dham_dmom)
    for _ in 1:n_fi_steps 
        @. phasepoint.pos = pos0 + .5 * stepsize * (dham_dmom0 + phasepoint.dham_dmom)
    end
    @. phasepoint.mom -= .5 * stepsize * phasepoint.dham_dpos
end

The implementation for the implicit midpoint integrator would then be

@reactive implicit_midpoint!(phasepoint; stepsize, n_fi_steps) = begin 
    pos0, mom0 = map(copy, (phasepoint.pos, phasepoint.mom))
    for _ in 1:n_fi_steps 
        (;dham_dmom, dham_dpos) = phasepoint
        @. phasepoint.pos = pos0 + .5 * stepsize * dham_dmom
        @. phasepoint.mom = mom0 - .5 * stepsize * dham_dpos
    end
    @. phasepoint.pos = 2 * phasepoint.pos - pos0
    @. phasepoint.mom = 2 * phasepoint.mom - mom0
end

And finally, to make all of the above work with the softabs map:

@reactive riemannian_phasepoint(pot_f, premetric_f, metric_f::typeof(softabs), pos, mom) = begin 
    pot, dpot_dpos = pot_f(pos)

    # It would be better to have the interface be to compute (pot, dpot_dpos, premetric, premetric_grad) in one swoop
    premetric, premetric_grad = premetric_f(pos)
    metric, Q, premetric_eigvals, metric_eigvals = metric_f(premetric)

    # inv(metric) could and should be avoided here - this is just for demo purposes
    inv_metric = inv(metric)
    dkin_dmom = inv_metric * mom
    kin = -.5 * (@node(sum(log, metric_eigvals)) + dot(mom, dkin_dmom))
    J = make_J(premetric_eigvals, alpha)
    RJ = Diagonal(diagonal(J) ./ metric_eigvals)
    D = Diagonal((Q' * mom) ./ metric_eigvals)
    dkin_dpos = @node(map(1:dim) do i
        .5 * tr_prod(@node(Q * RJ * Q'), premetric_grad[:, :, i])
    end) - @node(map(1:dim) do i
        .5 * tr_prod(@node(Q * (D * J * D) * Q'), premetric_grad[:, :, i])
    end)

    ham = pot + kin
    dham_dpos = dpot_dpos + dkin_dpos
    dham_dmom = dkin_dmom
end

Thoughts, @THargreaves or @J-Price-3?

---

Notes: I think some signs are wrong, and some of the algorithms may also be wrong - I didn't take care to double check everything for this syntax-POC.

[github-actions]

AdvancedHMC.jl documentation for PR #485 is available at:
https://TuringLang.github.io/AdvancedHMC.jl/previews/PR485/

[codecov]

Codecov Report

❌ Patch coverage is 94.70199% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.73%. Comparing base (4d4aeab) to head (c4314f3).

Files with missing lines	Patch %	Lines
src/riemannian/hamiltonian.jl	91.66%	4 Missing ⚠️
src/riemannian/metric.jl	97.67%	2 Missing ⚠️
src/sampler.jl	50.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #485       +/-   ##
===========================================
+ Coverage   78.45%   90.73%   +12.28%     
===========================================
  Files          21       22        +1     
  Lines        1323     1339       +16     
===========================================
+ Hits         1038     1215      +177     
+ Misses        285      124      -161     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[J-Price-3]

After updating the validity tests, it seems that there may be a problem with the MacOS version. There is a test that seems to consistently fail even after altering hyperparameters. I have confirmed that this should never fail (on windows at least) by using lots of RNG seeds.

[VarunJain268]

I've had a look through this and it seems sensible to me.

I’ve tested the new code against several targets, including synthetic examples such as Neal’s funnel and PosteriorDB targets such as centered eight schools. The diagnostics I checked - QQ plots, R-hat, ESS, etc. - all seem to line up as expected.

I’m not super familiar with the AdvancedHMC codebase, but a few thoughts came up while playing around with the implementation:

1. It might be worth including diagonal SoftAbs as an option, as in A General Metric for Riemannian Manifold Hamiltonian Monte Carlo. This could be a useful metric variant.
2. For the generalised leapfrog integrator, it might be useful to include some kind of convergence check when solving the implicit equations, perhaps with a user-specified tolerance and maximum number of iterations.
3. I also wonder whether fixed-point iteration is the most efficient way to solve the implicit equations in all cases. It may be worth considering something like Anderson acceleration, or at least leaving room for alternative solvers in the interface?

[VarunJain268]

I've had a look through this and it seems sensible to me.

I’ve tested the new code against several targets, including synthetic examples such as Neal’s funnel and PosteriorDB targets such as centered eight schools. The diagnostics I checked - QQ plots, Geweke, R-hat, ESS, etc. - all seem to line up as expected.

I’m not super familiar with the AdvancedHMC codebase, but a few thoughts came up while playing around with the implementation:

1. It might be worth including diagonal SoftAbs as an option, as in A General Metric for Riemannian Manifold Hamiltonian Monte Carlo. This could be a useful metric variant.
2. For the generalised leapfrog integrator, it might be useful to include some kind of convergence check when solving the implicit equations, perhaps with a user-specified tolerance and maximum number of iterations.
3. I also wonder whether fixed-point iteration is the most efficient way to solve the implicit equations in all cases. It may be worth considering something like Anderson acceleration, or at least leaving room for alternative solvers in the interface?

Some confirmatory experimental results below.

Setup: 4 chains, 1000 iterations of step-size adaptation, 5000 post-adaptation samples. The implicit equations for the generalised leapfrog integrator are iterated until convergence to tolerance $10^{-6}$. If convergence fails, or if a non-finite value is encountered, the proposal is rejected.

(i) 10D Neal's Funnel (full SoftAbs metric, $\alpha = 1000$)

QQ-plots:

Rank normalised R-hat (per dimension): min=1.0012, median=1.0016, max=1.0024
Bulk ESS: min=598.5, median=12447.7, max=13052.5
Tail ESS: min=734.9, median=3583.6, max=3870.4
Bulk ESS/sec: min=1.48, median=30.78, max=32.276
Tail ESS/sec: min=1.817, median=8.861, max=9.57

(ii) Eight schools (centred) (full SoftAbs metric, $\alpha = 10$)

QQ-plots:

Rank normalised R-hat (per dimension): min=1.0022, median=1.0035, max=1.0148
Bulk ESS: min=488.8, median=1102.0, max=1278.0
Tail ESS: min=396.7, median=2026.5, max=2654.5
Bulk ESS/sec: min=0.622, median=1.426, max=1.654
Tail ESS/sec: min=0.513, median=2.622, max=3.435

[VarunJain268]

A few comments regarding the implementation/code itself:

• In src/riemannian/integrator.jl, both implicit equations are solved with a fixed number of iterations. The generalised leapfrog is reversible only when these implicit solves have converged. So, if n is too small, an under-converged solve could break reversibility and bias the chain. Perhaps we could iterate to a tolerance, or we could document this behaviour so users are aware of the importance of n.
• I think the step-size initialisation should be fixed. In src/trajectory.jl, find_good_stepsize uses the helper A(h, z, ϵ), which hardcodes step(Leapfrog(ϵ), h, z). This sends a Hamiltonian with a Riemannian metric through the default Euclidean leapfrog in src/integrator.jl, which calls the two-argument ∂H∂r(h, r). However, the Riemannian implementation needs ∂H∂r(h, θ, r; G_eval=...), since ∂H/∂r = G(θ)^{-1}r depends on the current position θ.
• In src/riemannian/hamiltonian.jl, if a user calls just phasepoint(h, θ, r) without keyword arguments, the default ℓπ = ∂H∂θ(h, θ, r) is evaluated before the function body. For a Hamiltonian with a Riemannian metric, this builds the full Riemannian gradient cache, including metric_eval and metric_sensitivity. The function body then computes ℓκ with G_eval = nothing, so it calls metric_eval(h.metric, θ) again. Thus the metric evaluation is shared between neg_energy and ∂H∂r, but not between ℓπ and ℓκ. For SoftAbsRiemannianMetric, this extra work could be significant.