Skip to content

Add RankUpdateEuclideanMetric #443

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add RankUpdateEuclideanMetric #443

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bf633df
Add RankUpdateEuclideanMetric
ErikQQY May 9, 2025
d4f6e9a
Ensure correct resie
ErikQQY May 9, 2025
cc67475
Merge branch 'main' into qqy/RankUpdateEuclideanMetric
ErikQQY May 9, 2025
78ff471
Add some serious tests
ErikQQY May 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions src/metric.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,25 +99,36 @@ function Base.show(io::IO, dem::DenseEuclideanMetric)
end

"""
RankUpdateEuclideanMetric{T,M} <: AbstractMetric
RankUpdateEuclideanMetric{T,AM,AB,AD,F} <: AbstractMetric
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
RankUpdateEuclideanMetric{T,AM,AB,AD,F} <: AbstractMetric
RankUpdateEuclideanMetric{T,AM<:AbstractVecOrMat{T},AB,AD,F} <: AbstractMetric


A Gaussian Euclidean metric whose inverse is constructed by rank-updates.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
A Gaussian Euclidean metric whose inverse is constructed by rank-updates.
A Gaussian Euclidean metric whose inverse is constructed by low-rank updates to a diagonal matrix.


# Fields

$(TYPEDFIELDS)

# Constructors

RankUpdateEuclideanMetric(n::Int)
RankUpdateEuclideanMetric(M⁻¹, B, D)

Construct a Gaussian Euclidean metric of size `(n, n)` with inverse of `M⁻¹`.
- Construct a Gaussian Euclidean metric of size `(n, n)` with `M⁻¹` being diagonal matrix.
- Construct a Gaussian Euclidean metric of `M⁻¹`, where `M⁻¹` should be a full rank positive definite matrix,
and `B` `D` must be chose so that the Woodbury matrix `W = M⁻¹ + B D B^\\mathrm{T}` is positive definite.
Comment on lines +115 to +117
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the other metrics M⁻¹ is defined as the inverse of the metric. Wouldn't it be more consistent to here define M⁻¹ = A + B D B' as the inverse metric and require A be diagonal and positive-definite? (also, "full-rank" for diagonal PD-mat is redundant).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For Pathfinder to safely convert its WoodburyPDMat to this type, we'll need the contents/form of the factorization field to be documented.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
and `B` `D` must be chose so that the Woodbury matrix `W = M⁻¹ + B D B^\\mathrm{T}` is positive definite.
and `B` `D` must be chosen so that the Woodbury matrix `W = M⁻¹ + B D B^\\mathrm{T}` is positive definite.


# Example

```julia
julia> RankUpdateEuclideanMetric(3)
RankUpdateEuclideanMetric(diag=[1.0, 1.0, 1.0])
```

# References

- Ben Bales, Arya Pourzanjani, Aki Vehtari, Linda Petzold, Selecting the Metric in Hamiltonian Monte Carlo, 2019
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The factorized form used here is due to the Pathfinder paper, so I recommend also citing http://jmlr.org/papers/v23/21-0889.html

"""
struct RankUpdateEuclideanMetric{T,AM<:AbstractVecOrMat{T},AB,AD,F} <: AbstractMetric
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

More a question than a request: Is there a reason why AM can be a Vector? Also, is it intentional that AB and AD don't have to have the same element type?

# Diagnal of the inverse of the mass matrix
"Diagnal of the inverse of the mass matrix"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
"Diagnal of the inverse of the mass matrix"
"Diagonal of the inverse of the mass matrix"

M⁻¹::AM
B::AB
D::AD
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
D::AD
D::AD
"Woodbury factorisation of M⁻¹ + B D transpose(B)"

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there some nice explanation of what the purpose of B and D is, why we want this Woodbury factorisation, or how this is important for RankUpdateEuclideanMetric? If not, can just refer to somewhere else for an explanation, but would be good to say something about them.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't seen one place that explains it all. D is not needed, since it can be absorbed into B with any square-root factorization; it's just sometimes more convenient to have it (e.g. Pathfinder directly constructs it). On the other hand, when fitting a factor model, D is usually constrained to be I.

One interpretation comes from the factor model written in #443 (comment). It's useful when most of the structure of the covariance lies on a low-dimensional subspace. One such example is when the a small number of parameters are highly correlated with each other and not with any other parameters, and all remaining parameters are uncorrelated. Typically efficiency of HMC would be hindered by the high correlations, and one would need a dense covariance matrix, but then you end up with quadratic cost for each leapfrog stop. But if you can roughly upper bound the number of correlated parameters as << the number of sampled parameters, you can use a rank-update metric and get linear cost.

It's a mouthful, so I'm not certain the explanation belongs here. I think an explanation could wait until adaptation methods (e.g. Bales') are included here.

Expand Down Expand Up @@ -146,18 +157,23 @@ function RankUpdateEuclideanMetric(::Type{T}, n::Int) where {T}
factorization = woodbury_factorize(M⁻¹, B, D)
return RankUpdateEuclideanMetric(M⁻¹, B, D, factorization)
end

function RankUpdateEuclideanMetric(M⁻¹, B, D)
factorization = woodbury_factorize(M⁻¹, B, D)
return RankUpdateEuclideanMetric(M⁻¹, B, D, factorization)
end

function RankUpdateEuclideanMetric(::Type{T}, sz::Tuple{Int}) where {T}
return RankUpdateEuclideanMetric(T, first(sz))
end
RankUpdateEuclideanMetric(sz::Tuple{Int}) = RankUpdateEuclideanMetric(Float64, sz)

AdvancedHMC.renew(::RankUpdateEuclideanMetric, M⁻¹) = RankUpdateEuclideanMetric(M⁻¹)
renew(::RankUpdateEuclideanMetric, (M⁻¹, B, D)) = RankUpdateEuclideanMetric(M⁻¹, B, D)

Base.size(metric::RankUpdateEuclideanMetric, dim...) = size(metric.M⁻¹.diag, dim...)

function Base.show(io::IO, metric::RankUpdateEuclideanMetric)
print(io, "RankUpdateEuclideanMetric(diag=$(diag(metric.M⁻¹)))")
return nothing
function Base.show(io::IO, ::MIME"text/plain", metric::RankUpdateEuclideanMetric)
return print(io, "RankUpdateEuclideanMetric(diag=$(diag(metric.M⁻¹)))")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's another example where calling the full-rank part of the inverse-metric M⁻¹ leads to a different diagonal being shown than the diagonal of the full inverse-metric. The diagonal entries needed can be efficiently computed, see https://github.com/mlcolab/Pathfinder.jl/blob/f4ca90dc3d91f077f479d13904a2b6bf99e8ee25/src/woodbury.jl#L326-L329

end

# `rand` functions for `metric` types.
Expand Down
6 changes: 6 additions & 0 deletions test/sampler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ end
:UnitEuclideanMetric => UnitEuclideanMetric(D),
:DiagEuclideanMetric => DiagEuclideanMetric(D),
:DenseEuclideanMetric => DenseEuclideanMetric(D),
:RankUpdateEuclideanMetric => RankUpdateEuclideanMetric(D),
)
h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
@testset "$lfsym" for (lfsym, lf) in Dict(
Expand Down Expand Up @@ -104,6 +105,11 @@ end
@test mean(samples) ≈ zeros(D) atol = RNDATOL
end

if metricsym == :RankUpdateEuclideanMetric
# Skip tests with `RankUpdateEuclideanMetric` for `MassMatrixAdaptor`
continue
end

@testset "$adaptorsym" for (adaptorsym, adaptor) in Dict(
:MassMatrixAdaptorOnly => MassMatrixAdaptor(metric),
:StepSizeAdaptorOnly => StepSizeAdaptor(0.8, τ.integrator),
Expand Down
Loading