diff --git a/HISTORY.md b/HISTORY.md index 440f765d4..1c4e19251 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,10 @@ +# Release 0.5.1 + +This update adds new variational inference algorithms in light of the flexibility added in the v0.5 update. +Specifically, the following measure-space optimization algorithms have been added: + + - `KLMinWassFwdBwd` + # Release 0.5 ## Default Configuration Changes diff --git a/Project.toml b/Project.toml index 5d31be04b..a7ff8a27a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedVI" uuid = "b5ca4192-6429-45e5-a2d9-87aec30a685c" -version = "0.5.0" +version = "0.5.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/docs/make.jl b/docs/make.jl index 3c1d3a2d7..2cacf5fcb 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -23,9 +23,10 @@ makedocs(; "Normalizing Flows" => "tutorials/flows.md", ], "Algorithms" => [ - "KLMinRepGradDescent" => "klminrepgraddescent.md", - "KLMinRepGradProxDescent" => "klminrepgradproxdescent.md", - "KLMinScoreGradDescent" => "klminscoregraddescent.md", + "`KLMinRepGradDescent`" => "klminrepgraddescent.md", + "`KLMinRepGradProxDescent`" => "klminrepgradproxdescent.md", + "`KLMinScoreGradDescent`" => "klminscoregraddescent.md", + "`KLMinWassFwdBwd`" => "klminwassfwdbwd.md", ], "Variational Families" => "families.md", "Optimization" => "optimization.md", diff --git a/docs/src/index.md b/docs/src/index.md index f82087b00..0abd8c05f 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -17,3 +17,4 @@ For using the algorithms implemented in `AdvancedVI`, refer to the corresponding - [KLMinRepGradDescent](@ref klminrepgraddescent) (alias of `ADVI`) - [KLMinRepGradProxDescent](@ref klminrepgradproxdescent) - [KLMinScoreGradDescent](@ref klminscoregraddescent) (alias of `BBVI`) + - [KLMinWassFwdBwd](@ref klminwassfwdbwd) diff --git a/docs/src/klminwassfwdbwd.md b/docs/src/klminwassfwdbwd.md new file mode 100644 index 000000000..4c1911fff --- /dev/null +++ b/docs/src/klminwassfwdbwd.md @@ -0,0 +1,77 @@ +# [`KLMinWassFwdBwd`](@id klminwassfwdbwd) + +## Description + +This algorithm aims to minimize the exclusive (or reverse) Kullback-Leibler (KL) divergence by running proximal gradient descent (also known as forward-backward splitting) in Wasserstein space[^DBCS2023]. +(This algorithm is also sometimes referred to as "Wasserstein VI".) +Since `KLMinWassFwdBwd` is a measure-space algorithm, its use is restricted to full-rank Gaussian variational families (`FullRankGaussian`) that makes the measure-valued operations tractable. + +```@docs +KLMinWassFwdBwd +``` + +The associated objective value, which is the ELBO, can be estimated through the following: + +```@docs; canonical=false +estimate_objective( + ::Random.AbstractRNG, + ::KLMinWassFwdBwd, + ::MvLocationScale, + ::Any; + ::Int, +) +``` + +[^DBCS2023]: Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In *International Conference on Machine Learning*. PMLR. +## [Methodology](@id klminwassfwdbwd_method) + +This algorithm aims to solve the problem + +```math + \mathrm{minimize}_{q \in \mathcal{Q}}\quad \mathrm{KL}\left(q, \pi\right) +``` + +where $\mathcal{Q}$ is some family of distributions, often called the variational family. +Since we usually only have access to the unnormalized densities of the target distribution $\pi$, we don't have direct access to the KL divergence. +Instead, we focus on minimizing a surrogate objective, the *free energy functional*, which corresponds to the negated evidence lower bound[^JGJS1999], defined as + +```math + \mathcal{F}\left(q\right) \triangleq \mathcal{U}\left(q\right) + \mathcal{H}\left(q\right), +``` + +where + +```math +\begin{aligned} + \mathcal{U}\left(q\right) &= \mathbb{E}_{\theta \sim q} -\log \pi\left(\theta\right) + &&\text{(``potential energy'')} + \\ + \mathcal{H}\left(q\right) &= \mathbb{E}_{\theta \sim q} \log q\left(\theta\right) . + &&\text{(``Boltzmann entropy'')} +\end{aligned} +``` + +For solving this problem, `KLMinWassFwdBwd` relies on proximal stochastic gradient descent (PSGD)---also known as "forward-backward splitting"---that iterates + +```math + q_{t+1} = \mathrm{JKO}_{\gamma_t \mathcal{H}}\big( + q_{t} - \gamma_t \widehat{\nabla_{\mathrm{BW}} \mathcal{U}} (q_{t}) + \big) , +``` + +where $$\widehat{\nabla_{\mathrm{BW}} \mathcal{U}}$$ is a stochastic estimate of the Bures-Wasserstein measure-valued gradient of $$\mathcal{U}$$, the JKO (proximal) operator is defined as + +```math +\mathrm{JKO}_{\gamma_t \mathcal{H}}(\mu) += +\argmin_{\nu \in \mathcal{Q}} \left\{ \mathcal{H}(\nu) + \frac{1}{2 \gamma_t} \mathrm{W}_2 {(\mu, \nu)}^2 \right\} , +``` + +and $$\mathrm{W}_2$$ is the Wasserstein-2 distance. +When $$\mathcal{Q}$$ is set to be the Bures-Wasserstein space of $$\mathbb{R}^d$$, this algorithm is referred to as the Jordan-Kinderlehrer-Otto (JKO) scheme[^JKO1998], which was originally developed to study gradient flows under Wasserstein metrics. +Within this context, `KLMinWassFwdBwd` can be viewed as a numerical realization of the JKO scheme by restricting $$\mathcal{Q}$$ to be a tractable parametric variational family. +Specifically, Diao *et al.*[^DBCS2023] derived the JKO update for multivariate Gaussians, which is implemented by `KLMinWassFwdBwd`. +`KLMinWassFwdBwd` also exactly corresponds to the measure-space analog of [KLMinRepGradProxDescent](@ref klminrepgradproxdescent). + +[^JKO1998]: Jordan, R., Kinderlehrer, D., & Otto, F. (1998). The variational formulation of the Fokker--Planck equation. *SIAM Journal on Mathematical Analysis*, 29(1). +[^JGJS1999]: Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., & Saul, L. K. (1999). An introduction to variational methods for graphical models. Machine learning, 37, 183-233. diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 8fb85d7ca..2b82b7b79 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -352,4 +352,10 @@ include("algorithms/common.jl") export KLMinRepGradDescent, KLMinRepGradProxDescent, KLMinScoreGradDescent, ADVI, BBVI +# Other Algorithms + +include("algorithms/klminwassfwdbwd.jl") + +export KLMinWassFwdBwd + end diff --git a/src/algorithms/constructors.jl b/src/algorithms/constructors.jl index 32c85b228..ed2203816 100644 --- a/src/algorithms/constructors.jl +++ b/src/algorithms/constructors.jl @@ -21,12 +21,12 @@ KL divergence minimization by running stochastic gradient descent with the repar # Output - `q_averaged`: The variational approximation formed by the averaged SGD iterates. -# Callback -The callback function `callback` has a signature of +# Callback Signature +The `callback` function supplied to `optimize` needs to have the following signature: callback(; rng, iteration, restructure, params, averaged_params, restructure, gradient) -The arguments are as follows: +The keyword arguments are as follows: - `rng`: Random number generator internally used by the algorithm. - `iteration`: The index of the current iteration. - `restructure`: Function that restructures the variational approximation from the variational parameters. Calling `restructure(params)` reconstructs the current variational approximation. @@ -100,12 +100,12 @@ Thus, only the entropy estimators with a "ZeroGradient" suffix are allowed. # Output - `q_averaged`: The variational approximation formed by the averaged SGD iterates. -# Callback -The callback function `callback` has a signature of +# Callback Signature +The `callback` function supplied to `optimize` needs to have the following signature: callback(; rng, iteration, restructure, params, averaged_params, restructure, gradient) -The arguments are as follows: +The keyword arguments are as follows: - `rng`: Random number generator internally used by the algorithm. - `iteration`: The index of the current iteration. - `restructure`: Function that restructures the variational approximation from the variational parameters. Calling `restructure(params)` reconstructs the current variational approximation. @@ -178,11 +178,11 @@ KL divergence minimization by running stochastic gradient descent with the score - `q_averaged`: The variational approximation formed by the averaged SGD iterates. # Callback -The callback function `callback` has a signature of +The `callback` function supplied to `optimize` needs to have the following signature: callback(; rng, iteration, restructure, params, averaged_params, restructure, gradient) -The arguments are as follows: +The keyword arguments are as follows: - `rng`: Random number generator internally used by the algorithm. - `iteration`: The index of the current iteration. - `restructure`: Function that restructures the variational approximation from the variational parameters. Calling `restructure(params)` reconstructs the current variational approximation. diff --git a/src/algorithms/klminwassfwdbwd.jl b/src/algorithms/klminwassfwdbwd.jl new file mode 100644 index 000000000..f834b539a --- /dev/null +++ b/src/algorithms/klminwassfwdbwd.jl @@ -0,0 +1,213 @@ + +""" + KLMinWassFwdBwd(n_samples, stepsize, subsampling) + KLMinWassFwdBwd(; n_samples, stepsize, subsampling) + +KL divergence minimization by running stochastic proximal gradient descent (forward-backward splitting) in Wasserstein space[^DBCS2023]. + +Denoting the target log-density as \$\$ \\log \\pi \$\$ and the current variational approximation as \$\$q\$\$, the original algorithm requires estimating the quantity \$\$ \\mathbb{E}_q \\nabla^2 \\log \\pi \$\$. If the target `LogDensityProblem` associated with \$\$ \\log \\pi \$\$ has second-order differentiation [capability](https://www.tamaspapp.eu/LogDensityProblems.jl/dev/#LogDensityProblems.capabilities), we use the sample average of the Hessian. If the target has only first-order capability, we use Stein's identity. + +# (Keyword) Arguments +- `n_samples::Int`: Number of samples used to estimate the Wasserstein gradient. (default: `1`) +- `stepsize::Float64`: Step size of stochastic proximal gradient descent. +- `subsampling::Union{Nothing,<:AbstractSubsampling}`: Optional subsampling strategy. + +!!! note + The `subsampling` strategy is only applied to the target `LogDensityProblem` but not to the variational approximation `q`. That is, `KLMinWassFwdBwd` does not support amortization or structured variational families. + +# Output +- `q`: The last iterate of the algorithm. + +# Callback Signature +The `callback` function supplied to `optimize` needs to have the following signature: + + callback(; rng, iteration, q, info) + +The keyword arguments are as follows: +- `rng`: Random number generator internally used by the algorithm. +- `iteration`: The index of the current iteration. +- `q`: Current variational approximation. +- `info`: `NamedTuple` containing the information generated during the current iteration. + +# Requirements +- The variational family is [`FullRankGaussian`](@ref FullRankGaussian). +- The target distribution has unconstrained support (\$\$\\mathbb{R}^d\$\$). +- The target `LogDensityProblems.logdensity(prob, x)` has at least first-order differentiation capability. +""" +@kwdef struct KLMinWassFwdBwd{Sub<:Union{Nothing,<:AbstractSubsampling}} <: + AbstractVariationalAlgorithm + n_samples::Int = 1 + stepsize::Float64 + subsampling::Sub = nothing +end + +""" + gaussian_expectation_gradient_and_hessian!(rng, q, n_samples, grad_buf, hess_buf, prob) + +Estimate the expectations of the gradient and Hessians of the log-density of `prob` taken over the Gaussian `q`. For estimating the expectation of the Hessian, if `prob` has second-order differentiation capability, this function uses the sample average of the Hessian. Otherwise, it uses Stein's identity. + +# Arguments +- `rng::Random.AbstractRNG`: Random number generator. +- `q::MvLocationScale{<:LowerTriangular,<:Normal,L}`: Gaussian to take expectation over. +- `n_samples::Int`: Number of samples used for estimation. +- `grad_buf::AbstractVector`: Buffer for the gradient estimate. +- `hess_buf::AbstractMatrix`: Buffer for the Hessian estimate. +- `prob`: `LogDensityProblem` associated with the log-density gradient and Hessian subject to expectation. +""" +function gaussian_expectation_gradient_and_hessian!( + rng::Random.AbstractRNG, + q::MvLocationScale{<:LowerTriangular,<:Normal,L}, + n_samples::Int, + grad_buf::AbstractVector{T}, + hess_buf::AbstractMatrix{T}, + prob, +) where {T<:Real,L} + logπ_avg = zero(T) + fill!(grad_buf, zero(T)) + fill!(hess_buf, zero(T)) + + if LogDensityProblems.capabilities(typeof(prob)) ≤ + LogDensityProblems.LogDensityOrder{1}() + # Use Stein's identity + d = LogDensityProblems.dimension(prob) + u = randn(rng, T, d, n_samples) + z = q.scale*u .+ q.location + for b in 1:n_samples + zb, ub = view(z, :, b), view(u, :, b) + logπ, ∇logπ = LogDensityProblems.logdensity_and_gradient(prob, zb) + logπ_avg += logπ/n_samples + grad_buf += ∇logπ/n_samples + hess_buf += ub*(∇logπ/n_samples)' + end + return logπ_avg, grad_buf, hess_buf + else + # Use sample average of the Hessian. + z = rand(rng, q, n_samples) + for b in 1:n_samples + zb = view(z, :, b) + logπ, ∇logπ, ∇2logπ = LogDensityProblems.logdensity_gradient_and_hessian( + prob, zb + ) + logπ_avg += logπ/n_samples + grad_buf += ∇logπ/n_samples + hess_buf += ∇2logπ/n_samples + end + return logπ_avg, grad_buf, hess_buf + end +end + +struct KLMinWassFwdBwdState{Q,P,S,Sigma,GradBuf,HessBuf} + q::Q + prob::P + sigma::Sigma + iteration::Int + sub_st::S + grad_buf::GradBuf + hess_buf::HessBuf +end + +function init( + rng::Random.AbstractRNG, + alg::KLMinWassFwdBwd, + q_init::MvLocationScale{<:LowerTriangular,<:Normal,L}, + prob, +) where {L} + sub = alg.subsampling + n_dims = LogDensityProblems.dimension(prob) + capability = LogDensityProblems.capabilities(typeof(prob)) + if capability < LogDensityProblems.LogDensityOrder{1}() + throw( + ArgumentError( + "`KLMinWassFwdBwd` requires at least first-order differentiation capability. The capability of the supplied `LogDensityProblem` is $(capability).", + ), + ) + end + sub_st = isnothing(sub) ? nothing : init(rng, sub) + grad_buf = Vector{eltype(q_init.location)}(undef, n_dims) + hess_buf = Matrix{eltype(q_init.location)}(undef, n_dims, n_dims) + return KLMinWassFwdBwdState(q_init, prob, cov(q_init), 0, sub_st, grad_buf, hess_buf) +end + +output(::KLMinWassFwdBwd, state) = state.q + +function step( + rng::Random.AbstractRNG, alg::KLMinWassFwdBwd, state, callback, objargs...; kwargs... +) + (; n_samples, stepsize, subsampling) = alg + (; q, prob, sigma, iteration, sub_st, grad_buf, hess_buf) = state + + m = mean(q) + Σ = sigma + η = convert(eltype(m), stepsize) + iteration += 1 + + # Maybe apply subsampling + prob_sub, sub_st′, sub_inf = if isnothing(subsampling) + prob, sub_st, NamedTuple() + else + batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) + prob_sub = subsample(prob, batch) + prob_sub, sub_st′, sub_inf + end + + # Estimate the Wasserstein gradient + logπ_avg, grad_buf, hess_buf = gaussian_expectation_gradient_and_hessian!( + rng, q, n_samples, grad_buf, hess_buf, prob_sub + ) + + m′ = m - η * (-grad_buf) + M = I - η*Hermitian(-hess_buf) + Σ_half = Hermitian(M*Σ*M) + + # Compute the JKO proximal operator + Σ′ = (Σ_half + 2*η*I + sqrt(Hermitian(Σ_half*(Σ_half + 4*η*I))))/2 + q′ = MvLocationScale(m′, cholesky(Σ′).L, q.dist) + + state = KLMinWassFwdBwdState(q′, prob, Σ′, iteration, sub_st′, grad_buf, hess_buf) + elbo = logπ_avg + entropy(q′) + info = merge((elbo=elbo,), sub_inf) + + if !isnothing(callback) + info′ = callback(; rng, iteration, q, info) + info = !isnothing(info′) ? merge(info′, info) : info + end + state, false, info +end + +""" + estimate_objective([rng,] alg, q, prob; n_samples) + +Estimate the ELBO of the variational approximation `q` against the target log-density `prob`. + +# Arguments +- `rng::Random.AbstractRNG`: Random number generator. +- `alg::KLMinWassFwdBwd`: Variational inference algorithm. +- `q::MvLocationScale{<:Any,<:Normal,<:Any}`: Gaussian variational approximation. +- `prob`: The target log-joint likelihood implementing the `LogDensityProblem` interface. + +# Keyword Arguments +- `n_samples::Int`: Number of Monte Carlo samples for estimating the objective. (default: Same as the the number of samples used for estimating the gradient during optimization.) + +# Returns +- `obj_est`: Estimate of the objective value. +""" +function estimate_objective( + rng::Random.AbstractRNG, + alg::KLMinWassFwdBwd, + q::MvLocationScale{S,<:Normal,L}, + prob; + n_samples::Int=alg.n_samples, +) where {S,L} + obj = RepGradELBO(n_samples; entropy=MonteCarloEntropy()) + if isnothing(alg.subsampling) + return estimate_objective(rng, obj, q, prob) + else + sub = alg.subsampling + sub_st = init(rng, sub) + return mapreduce(+, 1:length(sub)) do _ + batch, sub_st, _ = step(rng, sub, sub_st) + prob_sub = subsample(prob, batch) + estimate_objective(rng, obj, q, prob_sub) / length(sub) + end + end +end diff --git a/test/algorithms/klminwassfwdbwd.jl b/test/algorithms/klminwassfwdbwd.jl new file mode 100644 index 000000000..3832017b7 --- /dev/null +++ b/test/algorithms/klminwassfwdbwd.jl @@ -0,0 +1,158 @@ + +@testset "KLMinWassFwdBwd" begin + begin + modelstats = normal_meanfield(Random.default_rng(), Float64; capability=2) + (; model, n_dims, μ_true, L_true) = modelstats + + alg = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3) + L0 = LowerTriangular(Matrix{Float64}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(Float64, n_dims), L0) + + @testset "callback" begin + T = 10 + callback(; iteration, kwargs...) = (iteration_check=iteration,) + _, info, _ = optimize(alg, T, model, q0; callback, show_progress=PROGRESS) + @test [i.iteration_check for i in info] == 1:T + end + + @testset "estimate_objective" begin + q_true = FullRankGaussian(μ_true, LowerTriangular(Matrix(L_true))) + + obj_est = estimate_objective(alg, q_true, model) + @test isfinite(obj_est) + + obj_est = estimate_objective(alg, q_true, model; n_samples=10^5) + @test obj_est ≈ 0 atol=1e-2 + end + + @testset "determinism" begin + seed = (0x38bef07cf9cc549d) + rng = StableRNG(seed) + T = 10 + + q_avg, _, _ = optimize(rng, alg, T, model, q0; show_progress=PROGRESS) + μ = q_avg.location + L = q_avg.scale + + rng_repl = StableRNG(seed) + q_avg, _, _ = optimize(rng_repl, alg, T, model, q0; show_progress=PROGRESS) + μ_repl = q_avg.location + L_repl = q_avg.scale + @test μ == μ_repl + @test L == L_repl + end + end + + @testset "error low capability" begin + modelstats = normal_meanfield(Random.default_rng(), Float64; capability=0) + (; model, n_dims) = modelstats + + alg = KLMinWassFwdBwd(; n_samples=10, stepsize=1.0) + + L0 = LowerTriangular(Matrix{Float64}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(Float64, n_dims), L0) + @test_throws "first-order" optimize(alg, 1, model, q0) + end + + @testset "type stability type=$(realtype), capability=$(capability)" for realtype in [ + Float64, Float32 + ], + capability in [1, 2] + + modelstats = normal_meanfield(Random.default_rng(), realtype; capability) + (; model, μ_true, L_true, n_dims, strong_convexity, is_meanfield) = modelstats + + alg = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3) + T = 10 + + L0 = LowerTriangular(Matrix{realtype}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(realtype, n_dims), L0) + + q, _, _ = optimize(alg, T, model, q0; show_progress=PROGRESS) + + @test eltype(q.location) == eltype(μ_true) + @test eltype(q.scale) == eltype(L_true) + end + + @testset "convergence capability=$(capability)" for capability in [1, 2] + modelstats = normal_meanfield(Random.default_rng(), Float64; capability) + (; model, μ_true, L_true, n_dims, strong_convexity, is_meanfield) = modelstats + + T = 1000 + alg = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3) + + q_avg, _, _ = optimize(alg, T, model, q0; show_progress=PROGRESS) + + Δλ0 = sum(abs2, q0.location - μ_true) + sum(abs2, q0.scale - L_true) + Δλ = sum(abs2, q_avg.location - μ_true) + sum(abs2, q_avg.scale - L_true) + + @test Δλ ≤ 0.1*Δλ0 + end + + @testset "subsampling" begin + n_data = 8 + + @testset "estimate_objective batchsize=$(batchsize)" for batchsize in [1, 3, 4] + modelstats = subsamplednormal(Random.default_rng(), n_data) + (; model, n_dims, μ_true, L_true) = modelstats + + L0 = LowerTriangular(Matrix{Float64}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(Float64, n_dims), L0) + + subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) + alg = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3) + alg_sub = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3, subsampling) + + obj_full = estimate_objective(alg, q0, model; n_samples=10^5) + obj_sub = estimate_objective(alg_sub, q0, model; n_samples=10^5) + @test obj_full ≈ obj_sub rtol=0.1 + end + + @testset "determinism" begin + seed = (0x38bef07cf9cc549d) + rng = StableRNG(seed) + + modelstats = subsamplednormal(Random.default_rng(), n_data) + (; model, n_dims, μ_true, L_true) = modelstats + + L0 = LowerTriangular(Matrix{Float64}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(Float64, n_dims), L0) + + T = 10 + batchsize = 3 + subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) + alg_sub = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-3, subsampling) + + q, _, _ = optimize(rng, alg_sub, T, model, q0; show_progress=PROGRESS) + μ = q.location + L = q.scale + + rng_repl = StableRNG(seed) + q, _, _ = optimize(rng_repl, alg_sub, T, model, q0; show_progress=PROGRESS) + μ_repl = q.location + L_repl = q.scale + @test μ == μ_repl + @test L == L_repl + end + + @testset "convergence capability=$(capability)" for capability in [1, 2] + modelstats = subsamplednormal(Random.default_rng(), n_data; capability) + (; model, n_dims, μ_true, L_true) = modelstats + + L0 = LowerTriangular(Matrix{Float64}(I, n_dims, n_dims)) + q0 = FullRankGaussian(zeros(Float64, n_dims), L0) + + T = 1000 + batchsize = 1 + subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) + alg_sub = KLMinWassFwdBwd(; n_samples=10, stepsize=1e-2, subsampling) + + q, stats, _ = optimize(alg_sub, T, model, q0; show_progress=PROGRESS) + + Δλ0 = sum(abs2, q0.location - μ_true) + sum(abs2, q0.scale - L_true) + Δλ = sum(abs2, q.location - μ_true) + sum(abs2, q.scale - L_true) + + @test Δλ ≤ 0.1*Δλ0 + end + end +end diff --git a/test/models/normal.jl b/test/models/normal.jl index 3db1af08d..a8de0f336 100644 --- a/test/models/normal.jl +++ b/test/models/normal.jl @@ -1,7 +1,8 @@ -struct TestNormal{M,S} +struct TestNormal{M,S,C} μ::M Σ::S + cap::C end function LogDensityProblems.logdensity(model::TestNormal, θ) @@ -16,15 +17,23 @@ function LogDensityProblems.logdensity_and_gradient(model::TestNormal, θ) ) end +function LogDensityProblems.logdensity_gradient_and_hessian(model::TestNormal, θ) + return ( + LogDensityProblems.logdensity(model, θ), + ForwardDiff.gradient(Base.Fix1(LogDensityProblems.logdensity, model), θ), + ForwardDiff.hessian(Base.Fix1(LogDensityProblems.logdensity, model), θ), + ) +end + function LogDensityProblems.dimension(model::TestNormal) return length(model.μ) end -function LogDensityProblems.capabilities(::Type{<:TestNormal}) - return LogDensityProblems.LogDensityOrder{1}() +function LogDensityProblems.capabilities(::Type{TestNormal{M,S,C}}) where {M,S,C} + return C() end -function normal_fullrank(rng::Random.AbstractRNG, realtype::Type) +function normal_fullrank(rng::Random.AbstractRNG, realtype::Type; capability::Int=1) n_dims = 5 σ0 = realtype(0.3) @@ -32,19 +41,33 @@ function normal_fullrank(rng::Random.AbstractRNG, realtype::Type) L = Matrix(σ0 * I, n_dims, n_dims) Σ = Hermitian(L * L') - model = TestNormal(μ, PDMat(Σ, Cholesky(L, 'L', 0))) + cap = if capability == 1 + LogDensityProblems.LogDensityOrder{1}() + elseif capability == 2 + LogDensityProblems.LogDensityOrder{2}() + else + LogDensityProblems.LogDensityOrder{0}() + end + model = TestNormal(μ, PDMat(Σ, Cholesky(L, 'L', 0)), cap) return TestModel(model, μ, LowerTriangular(L), n_dims, 1 / σ0^2, false) end -function normal_meanfield(rng::Random.AbstractRNG, realtype::Type) +function normal_meanfield(rng::Random.AbstractRNG, realtype::Type; capability::Int=1) n_dims = 5 σ0 = realtype(0.3) μ = Fill(realtype(5), n_dims) σ = Fill(σ0, n_dims) - model = TestNormal(μ, Diagonal(σ .^ 2)) + cap = if capability == 1 + LogDensityProblems.LogDensityOrder{1}() + elseif capability == 2 + LogDensityProblems.LogDensityOrder{2}() + else + LogDensityProblems.LogDensityOrder{0}() + end + model = TestNormal(μ, Diagonal(σ .^ 2), cap) L = Diagonal(σ) diff --git a/test/models/subsamplednormals.jl b/test/models/subsamplednormals.jl index 86be0dc1a..a254bbf24 100644 --- a/test/models/subsamplednormals.jl +++ b/test/models/subsamplednormals.jl @@ -1,14 +1,17 @@ -struct SubsampledNormals{D<:Normal,F<:Real} +struct SubsampledNormals{D<:Normal,F<:Real,C} dists::Vector{D} likeadj::F + capability::C end -function SubsampledNormals(rng::Random.AbstractRNG, n_normals::Int) +function SubsampledNormals(rng::Random.AbstractRNG, n_normals::Int, capability) μs = randn(rng, n_normals) σs = ones(n_normals) dists = Normal.(μs, σs) - return SubsampledNormals{eltype(dists),Float64}(dists, 1.0) + return SubsampledNormals{eltype(dists),Float64,typeof(capability)}( + dists, 1.0, capability + ) end function LogDensityProblems.logdensity(m::SubsampledNormals, x) @@ -31,19 +34,30 @@ function LogDensityProblems.logdensity_gradient_and_hessian(m::SubsampledNormals ) end -function LogDensityProblems.capabilities(::Type{<:SubsampledNormals}) - return LogDensityProblems.LogDensityOrder{2}() +function LogDensityProblems.dimension(::SubsampledNormals) + return 1 end -function AdvancedVI.subsample(m::SubsampledNormals, idx) - n_data = length(m.dists) - return SubsampledNormals(m.dists[idx], n_data/length(idx)) +function LogDensityProblems.capabilities(::Type{SubsampledNormals{D,F,C}}) where {D,F,C} + return C() end -function subsamplednormal(rng::Random.AbstractRNG, n_data::Int) - model = SubsampledNormals(rng, n_data) +function AdvancedVI.subsample(m::SubsampledNormals, idx) + n_data = length(m.dists) + return SubsampledNormals(m.dists[idx], n_data/length(idx), m.capability) +end + +function subsamplednormal(rng::Random.AbstractRNG, n_data::Int; capability::Int=1) + cap = if capability == 1 + LogDensityProblems.LogDensityOrder{1}() + elseif capability == 2 + LogDensityProblems.LogDensityOrder{2}() + else + LogDensityProblems.LogDensityOrder{0}() + end + model = SubsampledNormals(rng, n_data, cap) n_dims = 1 μ_true = [mean([mean(dist) for dist in model.dists])] - L_true = Diagonal(ones(1)) + L_true = Diagonal([sqrt(1/n_data)]) return TestModel(model, μ_true, L_true, n_dims, 1, true) end diff --git a/test/runtests.jl b/test/runtests.jl index 47bf0846b..ab67247b3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -64,6 +64,8 @@ if GROUP == "All" || GROUP == "GENERAL" include("families/location_scale.jl") include("families/location_scale_low_rank.jl") + + include("algorithms/klminwassfwdbwd.jl") end if GROUP == "All" || GROUP == "AD"