KLMinRepGradProxDescent

This is a convenience constructor for ParamSpaceSGD with the RepGradELBO objective with a proximal operator of the entropy (see here) of location-scale variational families. It implements the stochastic proximal gradient descent-based algorithm described in: [D2020][KMG2024][DGG2023].

AdvancedVI.KLMinRepGradProxDescent — Function
KLMinRepGradProxDescent(adtype; entropy_zerograd, optimizer, n_samples, averager)

KL divergence minimization by running stochastic proximal gradient descent with the reparameterization gradient in the Euclidean space of variational parameters of a location-scale family.

This algorithm only supports subtypes of MvLocationScale. Also, since the stochastic proximal gradient descent does not use the entropy of the gradient, the entropy estimator to be used must have a zero-mean gradient. Thus, only the entropy estimators with a "ZeroGradient" suffix are allowed.

Arguments

  • adtype: Automatic differentiation backend.

Keyword Arguments

  • entropy_zerograd: Estimator of the entropy with a zero-mean gradient to be used. Must be one of ClosedFormEntropyZeroGrad, StickingTheLandingEntropyZeroGrad. (default: ClosedFormEntropyZeroGrad())
  • optimizer::Optimisers.AbstractRule: Optimization algorithm to be used. Only DoG, DoWG and Optimisers.Descent are supported. (default: DoWG())
  • n_samples::Int: Number of Monte Carlo samples to be used for estimating each gradient.
  • averager::AbstractAverager: Parameter averaging strategy. (default: PolynomialAveraging())

Requirements

  • The variational family is MvLocationScale.
  • The target distribution and the variational approximation have the same support.
  • The target LogDensityProblems.logdensity(prob, x) must be differentiable with respect to x by the selected AD backend.
  • Additonal requirements on q may apply depending on the choice of entropy_zerograd.
source
  • D2020Domke, J. (2020). Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning.
  • KMG2024Kim, K., Ma, Y., & Gardner, J. (2024). Linear Convergence of Black-Box Variational Inference: Should We Stick the Landing?. In International Conference on Artificial Intelligence and Statistics (pp. 235-243). PMLR.
  • DGG2023Domke, J., Gower, R., & Garrigos, G. (2023). Provable convergence guarantees for black-box variational inference. Advances in neural information processing systems, 36, 66289-66327.