Add "low-rank" variational families

docs/src/families.md

@@ -243,9 +243,9 @@ As we can see, `LowRankGaussian` converges faster than `FullRankGaussian`.
 While `FullRankGaussian` can converge to the true solution since it is a more expressive variational family, `LowRankGaussian` gets there faster.
 
 !!! info
+
     `MvLocationScaleLowRank` tend to work better with the `Optimisers.Adam` optimizer due to non-smoothness.
     Other optimisers may experience divergences.
-
 
 ### API
 

src/families/location_scale.jl

@@ -30,8 +30,8 @@ This is necessary to guarantee stable convergence.
 function MvLocationScale(
     location::AbstractVector{T},
     scale::AbstractMatrix{T},
-    dist::ContinuousDistribution;
-    scale_eps::T=sqrt(eps(T)),
+    dist::ContinuousUnivariateDistribution;
+    scale_eps::T=eps(T)^(1//4),
 ) where {T<:Real}
     @assert minimum(diag(scale)) ≥ scale_eps "Initial scale is too small (smallest diagonal value is $(minimum(diag(scale)))). This might result in unstable optimization behavior."
     return MvLocationScale(location, scale, dist, scale_eps)
@@ -143,7 +143,7 @@ Construct a Gaussian variational approximation with a dense covariance matrix.
 - `scale_eps`: Smallest value allowed for the diagonal of the scale. (default: `sqrt(eps(T))`).
 """
 function FullRankGaussian(
-    μ::AbstractVector{T}, L::LinearAlgebra.AbstractTriangular{T}; scale_eps::T=sqrt(eps(T))
+    μ::AbstractVector{T}, L::LinearAlgebra.AbstractTriangular{T}; scale_eps::T=eps(T)^(1//4)
 ) where {T<:Real}
     q_base = Normal{T}(zero(T), one(T))
     return MvLocationScale(μ, L, q_base, scale_eps)
@@ -162,7 +162,7 @@ Construct a Gaussian variational approximation with a diagonal covariance matrix
 - `scale_eps`: Smallest value allowed for the diagonal of the scale. (default: `sqrt(eps(T))`).
 """
 function MeanFieldGaussian(
-    μ::AbstractVector{T}, L::Diagonal{T}; scale_eps::T=sqrt(eps(T))
+    μ::AbstractVector{T}, L::Diagonal{T}; scale_eps::T=eps(T)^(1//4)
 ) where {T<:Real}
     q_base = Normal{T}(zero(T), one(T))
     return MvLocationScale(μ, L, q_base, scale_eps)

src/families/location_scale_low_rank.jl

@@ -35,8 +35,8 @@ function MvLocationScaleLowRank(
     location::AbstractVector{T},
     scale_diag::AbstractVector{T},
     scale_factors::AbstractMatrix{T},
-    dist::ContinuousDistribution;
-    scale_eps::T=sqrt(eps(T)),
+    dist::ContinuousUnivariateDistribution;
+    scale_eps::T=eps(T)^(1//4),
 ) where {T<:Real}
     @assert minimum(scale_diag) ≥ scale_eps "Initial scale is too small (smallest diagonal scale value is $(minimum(scale_diag)). This might result in unstable optimization behavior."
     @assert size(scale_factors, 1) == length(scale_diag)
@@ -156,7 +156,7 @@ function update_variational_params!(
 end
 
 """
-    LowRankGaussian(location, scale_diag, scale_factors; check_args = true)
+    LowRankGaussian(μ, D, U; scale_eps)
 
 Construct a Gaussian variational approximation with a diagonal plus low-rank covariance matrix.
 
@@ -169,7 +169,7 @@ Construct a Gaussian variational approximation with a diagonal plus low-rank cov
 - `scale_eps`: Smallest value allowed for the diagonal of the scale. (default: `sqrt(eps(T))`).
 """
 function LowRankGaussian(
-    μ::AbstractVector{T}, D::Vector{T}, U::Matrix{T}; scale_eps::T=sqrt(eps(T))
+    μ::AbstractVector{T}, D::Vector{T}, U::Matrix{T}; scale_eps::T=eps(T)^(1//4)
 ) where {T<:Real}
     q_base = Normal{T}(zero(T), one(T))
     return MvLocationScaleLowRank(μ, D, U, q_base; scale_eps)

test/families/location_scale.jl

@@ -1,6 +1,7 @@
 
 @testset "interface LocationScale" begin
-    @testset "$(string(covtype)) $(basedist) $(realtype)" for basedist in [:gaussian, :gaussian_nonstd],
+    @testset "$(string(covtype)) $(basedist) $(realtype)" for basedist in
+                                                              [:gaussian, :gaussian_nonstd],
         covtype in [:meanfield, :fullrank],
         realtype in [Float32, Float64]
 
@@ -143,7 +144,8 @@
     end
 
     @testset "scale positive definite projection" begin
-        @testset "$(string(covtype)) $(realtype) $(bijector)" for covtype in [:meanfield, :fullrank],
+        @testset "$(string(covtype)) $(realtype) $(bijector)" for covtype in
+                                                                  [:meanfield, :fullrank],
             realtype in [Float32, Float64],
             bijector in [nothing, :identity]
 

test/families/location_scale_low_rank.jl

@@ -1,6 +1,7 @@
 
 @testset "interface LocationScaleLowRank" begin
-    @testset "$(basedist) rank=$(rank) $(realtype)" for basedist in [:gaussian, :gaussian_nonstd],
+    @testset "$(basedist) rank=$(rank) $(realtype)" for basedist in
+                                                        [:gaussian, :gaussian_nonstd],
         n_rank in [1, 2],
         realtype in [Float32, Float64]
 
@@ -56,7 +57,7 @@
         @testset "statistics" begin
             @testset "mean" begin
                 @test eltype(mean(q)) == realtype
-                @test mean(q) == mean(q_true)
+                @test mean(q) ≈ mean(q_true)
             end
             @testset "var" begin
                 @test eltype(var(q)) == realtype
@@ -81,7 +82,7 @@
                 @test cov(z_samples; dims=2) ≈ cov(q_true) rtol = realtype(1e-2)
 
                 z_sample_ref = rand(StableRNG(1), q)
-                @test z_sample_ref == rand(StableRNG(1), q)
+                @test z_sample_ref ≈ rand(StableRNG(1), q)
             end
 
             @testset "rand batch" begin
@@ -96,7 +97,7 @@
                 @test cov(z_samples; dims=2) ≈ cov(q_true) rtol = realtype(1e-2)
 
                 samples_ref = rand(StableRNG(1), q, n_montecarlo)
-                @test samples_ref == rand(StableRNG(1), q, n_montecarlo)
+                @test samples_ref ≈ rand(StableRNG(1), q, n_montecarlo)
             end
 
             @testset "rand! AbstractVector" begin
@@ -107,7 +108,7 @@
                 end
                 z_samples = mapreduce(first, hcat, res)
                 z_samples_ret = mapreduce(last, hcat, res)
-                @test z_samples == z_samples_ret
+                @test z_samples ≈ z_samples_ret
                 @test dropdims(mean(z_samples; dims=2); dims=2) ≈ mean(q_true) rtol = realtype(
                     1e-2
                 )
@@ -121,13 +122,13 @@
 
                 z_sample = Array{realtype}(undef, n_dims)
                 rand!(StableRNG(1), q, z_sample)
-                @test z_sample_ref == z_sample
+                @test z_sample_ref ≈ z_sample
             end
 
             @testset "rand! AbstractMatrix" begin
                 z_samples = Array{realtype}(undef, n_dims, n_montecarlo)
                 z_samples_ret = rand!(q, z_samples)
-                @test z_samples == z_samples_ret
+                @test z_samples ≈ z_samples_ret
                 @test dropdims(mean(z_samples; dims=2); dims=2) ≈ mean(q_true) rtol = realtype(
                     1e-2
                 )
@@ -141,7 +142,7 @@
 
                 z_samples = Array{realtype}(undef, n_dims, n_montecarlo)
                 rand!(StableRNG(1), q, z_samples)
-                @test z_samples_ref == z_samples
+                @test z_samples_ref ≈ z_samples
             end
         end
     end

test/inference/repgradelbo_distributionsad.jl

@@ -14,7 +14,8 @@ if @isdefined(Enzyme)
 end
 
 @testset "inference RepGradELBO DistributionsAD" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
+                                                                     [Float64, Float32],
         (modelname, modelconstr) in Dict(:Normal => normal_meanfield),
         n_montecarlo in [1, 10],
         (objname, objective) in Dict(

test/inference/repgradelbo_locationscale.jl

@@ -14,7 +14,8 @@ if @isdefined(Enzyme)
 end
 
 @testset "inference RepGradELBO VILocationScale" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
+                                                                     [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
         n_montecarlo in [1, 10],

test/inference/repgradelbo_locationscale_bijectors.jl

@@ -14,7 +14,8 @@ if @isdefined(Enzyme)
 end
 
 @testset "inference RepGradELBO VILocationScale Bijectors" begin
-    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in [Float64, Float32],
+    @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
+                                                                     [Float64, Float32],
         (modelname, modelconstr) in
         Dict(:NormalLogNormalMeanField => normallognormal_meanfield),
         n_montecarlo in [1, 10],