API

compute_metric(
    metric::KernelizedSteinDiscrepancy,
    pc::ParticleContainer,
    ρ;
    ad_backend::ADTypes.AbstractADType
)

Compute the Kernelized Stein Discrepancy (KSD) between the particles in the ParticleContainer and the target log-density.

Arguments

• metric::KernelizedSteinDiscrepancy: A KernelizedSteinDiscrepancy object that defines the kernel and other parameters used for KSD computation.
• pc::ParticleContainer: The particle container holding the current set of particles.
• ρ: A LogDensityProblem representing the target distribution's log-density function.

Keyword Arguments

• ad_backend: The automatic differentiation backend to use for computing gradients required by the KSD.

Returns

• The computed Kernelized Stein Discrepancy (KSD) value, a scalar representing the discrepancy between the particle distribution and the target distribution.

Details

This function calculates the KSD, a measure of the discrepancy between the distribution represented by the particles in pc and the target distribution defined by ρ. It utilizes the kernel specified in the metric object and the provided automatic differentiation backend ad_backend to compute the necessary gradients.

The function extracts the particle positions from pc.P and calls the kernelized_stein_discrepancy function to perform the KSD computation.

This function serves as a metric tracking tool during particle-based inference, allowing monitoring of the convergence of the particle distribution to the target.

get_problem(ctx::LogDensityProblemContext)

Returns LogDensityProblem stored in LogDensityProblemContext.

Arguments

• ctx::LogDensityProblemContext: The problem context

Returns

• the instance of LogDensityProblem ctx.

infer!(
    pc::ParticleContainer,
    ctx::Context{<:AbstractProblemContext, SVGDInferenceContext};
    iters::Integer=10,
    ad_backend::ADTypes.AbstractADType=ADTypes.AutoForwardDiff(),
    verbose::Bool=false,
    track=Dict{String, Any}()
)

Performs Stein Variational Gradient Descent (SVGD) inference to update a particle container.

Arguments

• pc::ParticleContainer: The container holding the particles to be updated.
• ctx::Context{<:AbstractProblemContext, SVGDInferenceContext}: The context containing the problem definition and the SVGD inference settings.
• iters::Integer=10: The number of SVGD iterations to perform.
• ad_backend::ADTypes.AbstractADType=ADTypes.AutoForwardDiff(): The automatic differentiation backend to use for gradient computations.
• verbose::Bool=false: A flag to enable verbose output during inference (currently not implemented).
• track::Dict{String, Any}(): A dictionary specifying metrics to compute and track during inference. The keys are metric names (strings), and the values are metric types (e.g., functions or structs that can be passed to compute_metric).

Returns

• An SVGDInferenceReport containing the tracked metrics and a success flag.

Notes

• This function retrieves the log-density function from the problem context.
• It initializes and updates the particles using the SVGD dynamics specified in the inference context.
• Optionally, it computes and tracks specified metrics at each iteration.

init(
    ρ,
    dynamics::ParticleDynamics;
    particle_initializer=NormalInitializer(),
    n_particles::Integer=16,
    ad_backend::ADTypes.AbstractADType=ADTypes.AutoForwardDiff()
)

Initializes a particle container and the corresponding inference context.

Arguments

• ρ: The target log-density function.
• dynamics::ParticleDynamics: The particle dynamics to be used for inference.
• particle_initializer: An object that initializes the particle positions (default: NormalInitializer()).
• n_particles::Integer=16: The number of particles to initialize (default: 16).
• ad_backend::ADTypes.AbstractADType=ADTypes.AutoForwardDiff(): The automatic differentiation backend to use for gradient computations if needed (default: ADTypes.AutoForwardDiff()).

Returns

• A tuple containing:
  • pc::ParticleContainer: The initialized particle container.
  • ctx::Context: The initialized inference context.

Notes

• This function determines the dimensionality of the problem from the log-density function ρ.
• It ensures that the log-density function is differentiable by wrapping it with an AD backend.
• It initializes the particle positions using the provided particle_initializer.
• It initializes the inference context based on the provided dynamics.

init_context(
    ρ,
    dynamics::ParticleDynamics
)

Initializes a Context struct containing a problem context and an inference context.

Arguments

• ρ: The target log-density function.
• dynamics::ParticleDynamics: The particle dynamics to be used for inference.

Returns

• An instance of Context containing:
  • A problem context initialized from ρ.
  • An inference context initialized from ρ and dynamics.

init_inference_context(ρ, dynamics::SVGD)

Initializes an SVGDInferenceContext with the provided SVGD dynamics. For now ρ is not used.

Arguments

• ρ: The target LogDensityProblem
• dynamics::SVGD: The SVGD dynamics to be used for inference.

Returns

• An instance of SVGDInferenceContext initialized with the given dynamics.

kernel_and_gradient_fn(K::KernelFunctions.Kernel, ad_backend)

Returns a function that computes the kernel value and its gradient with respect to the first argument.

Arguments

• K::KernelFunctions.Kernel: The kernel function from KernelFunctions.jl
• ad_backend::ADTypes.AbstractADType: The automatic differentiation backend to use (e.g., AbstractDifferentiation.ForwardDiffBackend()).

Returns

• A function k_∇k(x, a) that takes two arguments x and a (of compatible types for the kernel K) and returns a tuple containing:
  • The kernel value K(x, a).
  • The gradient of the kernel with respect to x, evaluated at x.

kernelized_stein_discrepancy(P, q, K::KernelFunctions.Kernel; samplesize::Integer, ad_backend::ADTypes.AbstractADType)

Computes the Kernelized Stein Discrepancy (KSD) between a set of samples P and a distribution q.

The KSD measures the discrepancy between two probability distributions by evaluating the expectation of two Stein operators applied to a kernel function.

Arguments

• P: A matrix of samples from the empirical distribution. Each column represents a sample.
• q: A LogDensityProblems.LogDensityProblem representing the target distribution.
• K: A kernel function from KernelFunctions.Kernel.
• samplesize::Integer: The number of sampled particles to evaluate KSD
• ad_backend::ADTypes.AbstractADType: An automatic differentiation backend from DifferentiationInterface.

Returns

The Kernelized Stein Discrepancy (KSD) as a scalar value.

Details

The function calculates the KSD using the following formula:

\[\text{KSD}(P, q) = \frac{1}{n(n-1)} \sum_{i=1}^n \sum_{j=1}^n u(P_i, P_j)\]

\[u(x, y) = \nabla s(x)^T k(x, y) \nabla s(y) + \nabla s(x)^T \nabla_y k(x, y) + \nabla_x k(x, y)^T \nabla s(y) + \text{tr}(\nabla_{x,y} k(x, y))\]

For more details see :

• A Kernelized Stein Discrepancy for Goodness-of-fit Tests, Qiang Liu, Jason Lee, Michael Jordan

particle_velocity(pc::ParticleContainer, ρ, pi, k_∇k, dynamics::SVGD)

Computes the velocity of a single particle based on the Stein Variational Gradient Descent (SVGD) update rule, potentially using a mini-batch of other particles.

Arguments

• pc::ParticleContainer: The container holding the particles.
• ρ: The log-density function (a LogDensityProblem) which must be differentiable.
• pi::Int: The index of the particle for which to compute the velocity.
• k_∇k: A function that takes two particle positions (as vectors) and returns a tuple containing the kernel value and the gradient of the kernel with respect to the first argument. This is generated by kernel_and_gradient_fn.
• dynamics::SVGD: The SVGD dynamics object containing the kernel, step size, and batch size.

Returns

• velocity::Vector: The computed velocity vector for the particle with index pi.

update_particles!(ρ, pc::ParticleContainer, dynamics::SVGD)

Updates the positions of all particles in the ParticleContainer according to the Stein Variational Gradient Descent (SVGD) update rule.

Arguments

• ρ: The log-density function (a LogDensityProblem) that the particles aim to sample from.
• pc::ParticleContainer: The container holding the current positions of the particles. The particle positions are updated in-place.
• dynamics::SVGD: The SVGD dynamics object specifying the kernel, step size (η), and batch size for the update.

Context{T<:AbstractProblemContext, U<:AbstractInferenceContext}

A mutable struct that holds instances of a problem context and an inference context.

Fields

• problem::T: An instance of a subtype of AbstractProblemContext.
• inference::U: An instance of a subtype of AbstractInferenceContext.

KernelizedSteinDiscrepancy <: Metric

A struct representing the Kernelized Stein Discrepancy (KSD) metric.

Fields

• K::KernelFunctions.Kernel: The kernel function used in the KSD computation.
• samplesize::Integer: the size of sample used to evaluate KSD

Description

This struct encapsulates the kernel function needed to compute the Kernelized Stein Discrepancy. The KSD is a measure of discrepancy between two probability distributions, and it relies on a kernel function to define the feature space in which the discrepancy is measured.

By holding the kernel function, this struct provides a convenient way to pass the necessary information for KSD computation to functions that track or evaluate the discrepancy between a particle distribution and a target distribution.

LangevinInitializer(ϵ, steps) <: ParticleInitializer

Initializes particles using a short Langevin dynamics simulation.

This initializer applies a specified number of Langevin dynamics steps to particles.

Fields

• ϵ::Float64: The step size used in the Langevin dynamics simulation. A smaller ϵ generally leads to more stable but slower exploration of the configuration space.
• steps::Int: The number of Langevin dynamics steps to perform for each particle. More steps allow for greater exploration and potentially better sampling of the target distribution.

mutable struct LogDensityProblemContext <: AbstractProblemContext
    ρ
end

A mutable struct representing the problem context for LogDensityProblem.jl

Fields

• ρ: should implement LogDensityProblem interface

mutable struct ParticleContainer{T} <: AbstractParticleContainer{T}
    P::Matrix{T}
    dim
    size
end

A mutable struct that holds a collection of particles.

Fields

• P::Matrix{T}: A matrix where each column represents a particle.
• dim: The dimensionality of each particle (number of rows in P).
• size: The number of particles (number of columns in P).

SVGD <: ParticleDynamics

Stein Variational Gradient Descent (SVGD) particle dynamics.

Fields

• K::KernelFunctions.Kernel: The kernel function used to define the interaction between particles.
• η: The step size or learning rate for updating particle positions.
• batchsize: The number of particles to use in each update step (for mini-batching). If nothing, all particles are used.

Examples

using KernelFunctions

Define a squared exponential kernel

sqexp_kernel = SqExponentialKernel()

Create an SVGD dynamics object with a fixed step size and full batch

svgd_fullbatch = SVGD(K=sqexp_kernel, η=0.1, batchsize=nothing)

Create an SVGD dynamics object with a smaller step size and a batch size of 100

svgd_minibatch = SVGD(K=sqexp_kernel, η=0.05, batchsize=100)

struct SVGDInferenceContext <: AbstractInferenceContext
    dynamics::SVGD
end

A struct representing the inference context for SVGD.

Fields

• dynamics::SVGD: An instance of the SVGD dynamics

SVGDInferenceReport <: InferenceReport

A mutable struct representing the report generated after running Stein Variational Gradient Descent (SVGD) inference.

Fields

• metrics::Dict{String, Vector{Any}}: A dictionary storing various metrics tracked during the SVGD inference. Keys are metric names (strings), and values are vectors of metric values recorded at each iteration or relevant time step.
• success::Bool: A boolean flag indicating whether the inference process completed successfully.