General

ParamSpaceSGD SGD is a general algorithm for leveraging automatic differentiation and SGD. Furthermore, it operates in the space of variational parameters. Consider the case where each member $q_{\lambda} \in \mathcal{Q}$ of the variational family $\mathcal{Q}$ is uniquely represented through a collection of parameters $\lambda \in \Lambda \subseteq \mathbb{R}^p$. That is,

\[\mathcal{Q} = \{q_{\lambda} \mid \lambda \in \Lambda \},\]

Then, as implied by the name, ParamSpaceSGD runs SGD on $\Lambda$, the (Euclidean) space of parameters.

Any algorithm that operates by iterating the following steps can easily be implemented via ParamSpaceSGD:

Obtain an unbiased estimate of the target objective.
Obtain an estimate of the gradient of the objective by differentiating the objective estimate with respect to the parameters.
Perform gradient descent with the stochastic gradient estimate.

After some simplifications, each step of ParamSpaceSGD can be described as follows:

function step(rng, alg::ParamSpaceSGD, state, callback, objargs...; kwargs...)
    (; adtype, problem, objective, operator, averager) = alg
    (; q, iteration, grad_buf, opt_st, obj_st, avg_st) = state
    iteration += 1

    # Extract variational parameters of `q`
    params, re = Optimisers.destructure(q)

    # Estimate gradient and update the `DiffResults` buffer `grad_buf`.
    grad_buf, obj_st, info = estimate_gradient!(...)

    # Gradient descent step.
    grad = DiffResults.gradient(grad_buf)
    opt_st, params = Optimisers.update!(opt_st, params, grad)

    # Apply operator
    params = apply(operator, typeof(q), opt_st, params, re)

    # Apply parameter averaging
    avg_st = apply(averager, avg_st, params)

    # Updated state
    state = ParamSpaceSGDState(re(params), iteration, grad_buf, opt_st, obj_st, avg_st)
    state, false, info
end

The output of ParamSpaceSGD is the final state of averager. Furthermore, operator can be anything from an identity mapping, a projection operator, a proximal operator, and so on.

`ParamSpaceSGD`

The constructor for ParamSpaceSGD is as follows:

AdvancedVI.ParamSpaceSGD — Type

ParamSpaceSGD(
    objective::AbstractVariationalObjective,
    adtype::ADTypes.AbstractADType,
    optimizer::Optimisers.AbstractRule,
    averager::AbstractAverager,
    operator::AbstractOperator,
)

This algorithm applies stochastic gradient descent (SGD) to the variational objective over the (Euclidean) space of variational parameters.

The trainable parameters in the variational approximation are expected to be extractable through Optimisers.destructure. This requires the variational approximation to be marked as a functor through Functors.@functor.

Note

Different objective may impose different requirements on adtype, variational family, optimizer, and operator. It is therefore important to check the documentation corresponding to each specific objective. Essentially, each objective should be thought as forming its own unique algorithm.

Arguments

objective: Variational Objective.
adtype: Automatic differentiation backend.
optimizer: Optimizer used for inference.
averager : Parameter averaging strategy.
operator : Operator applied to the parameters after each optimization step.

Output

q_averaged: The variational approximation formed from the averaged SGD iterates.

Callback

The callback function callback has a signature of

callback(; rng, iteration, restructure, params, averaged_params, restructure, gradient)

The 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.
params: Current variational parameters.
averaged_params: Variational parameters averaged according to the averaging strategy.
gradient: The estimated (possibly stochastic) gradient.

source

Objective Interface

To define an instance of a ParamSpaceSGD algorithm, it suffices to implement the AbstractVariationalObjective interface. First, we need to define a subtype of AbstractVariationalObjective:

AdvancedVI.AbstractVariationalObjective — Type

AbstractVariationalObjective

Abstract type for the VI algorithms supported by AdvancedVI.

Implementations

To be supported by AdvancedVI, a VI algorithm must implement AbstractVariationalObjective and estimate_objective. Also, it should provide gradients by implementing the function estimate_gradient. If the estimator is stateful, it can implement init to initialize the state.

source

In addition, we need to implement some methods associated with the objective. First, each objective may maintain a state such as buffers, online estimates of control variates, batch iterators for subsampling, and so on. Such things should be initialized by implementing the following:

AdvancedVI.init — Method

init(rng, obj, adtype, prob, params, restructure)

Initialize a state of the variational objective obj given the initial variational parameters λ. This function needs to be implemented only if obj is stateful.

Arguments

rng::Random.AbstractRNG: Random number generator.
obj::AbstractVariationalObjective: Variational objective.

adtype::ADTypes.AbstractADType`: Automatic differentiation backend.

params: Initial variational parameters.
restructure: Function that reconstructs the variational approximation from λ.

source

If this method is not implemented, the state will be automatically be nothing.

Next, the key functionality of estimating stochastic gradients should be implemented through the following:

AdvancedVI.estimate_gradient! — Function

estimate_gradient!(rng, obj, adtype, out, params, restructure, obj_state)

Estimate (possibly stochastic) gradients of the variational objective obj targeting prob with respect to the variational parameters λ

Arguments

rng::Random.AbstractRNG: Random number generator.
obj::AbstractVariationalObjective: Variational objective.
adtype::ADTypes.AbstractADType: Automatic differentiation backend.
out::DiffResults.MutableDiffResult: Buffer containing the objective value and gradient estimates.
params: Variational parameters to evaluate the gradient on.
restructure: Function that reconstructs the variational approximation from params.
obj_state: Previous state of the objective.

Returns

out::MutableDiffResult: Buffer containing the objective value and gradient estimates.
obj_state: The updated state of the objective.
stat::NamedTuple: Statistics and logs generated during estimation.

source

AdvancedVI only interacts with each variational objective by querying gradient estimates. In a lot of cases, however, it is convinient to be able to estimate the current value of the objective. For example, for monitoring convergence. This should be done through the following:

AdvancedVI.estimate_objective — Function

estimate_objective([rng,] obj, q, prob; kwargs...)

Estimate the variational objective obj targeting prob with respect to the variational approximation q.

Arguments

rng::Random.AbstractRNG: Random number generator.
obj::AbstractVariationalObjective: Variational objective.
prob: The target log-joint likelihood implementing the LogDensityProblem interface.
q: Variational approximation.

Keyword Arguments

Depending on the objective, additional keyword arguments may apply. Please refer to the respective documentation of each variational objective for more info.

Returns

obj_est: Estimate of the objective value.

source