Merge branch 'stoch_grad_hmc' of sivapvarma/AdvancedHMC.jl into stoch_grad_hmc

Author: sivapvarma

src/trajectory.jl

@@ -747,6 +747,93 @@ end
 
 
 
+###
+### Stochastic Gradient Langevin Dynamics sampler.
+###
+"""
+Stochastic Gradient Langevin Dynamics with fixed number of steps.
+"""
+mutable struct SGLD{
+    I<:AbstractIntegrator,
+    F<:AbstractFloat
+} <: AbstractTrajectory{I}
+    integrator      :: I
+    n_steps         :: Int  # number of samples
+    ϵ               :: F    # constant scale factor of the learning rate
+    i               :: Int  # iteration counter
+    γ               :: F    # scaling constant
+end
+
+function transition(
+    rng::AbstractRNG,
+    τ::SGLD,
+    h::Hamiltonian,
+    z::PhasePoint
+) where {T<:Real}
+    # z′ = step(rng, τ.integrator, h, z, τ.n_steps)
+    DEBUG && @debug "compute current step size..."
+    # γ = .35
+    τ.i += 1
+    ϵ_t = τ.ϵ / τ.i ^ τ.γ # NOTE: Choose γ=.55 in paper
+
+    DEBUG && @debug "recording old variables..."
+    θ = z.θ
+    grad = -z.ℓπ.gradient
+
+    DEBUG && @debug "update latent variables..."
+    θ .+= ϵ_t .* grad ./ 2 .+ rand.(Normal.(zeros(length(θ)), sqrt(ϵ_t)))
+
+    # no M-H step
+    z = PhasePoint(h, θ, -z.r)
+    stat = (
+        step_size=τ.integrator.ϵ,
+        n_steps=τ.n_steps,
+        log_density=z.ℓπ.value,
+        hamiltonian_energy=energy(z),
+        )
+    return Transition(z, stat)
+end
+
+##
+## Stochastic Gradient Hamilton Samplers
+##
+
+###
+### Stochastic Gradient Hamiltonian Monte Carlo sampler.
+###
+"""
+Stochastic Gradient HMC with fixed number of steps.
+"""
+struct SGHMC{
+    I<:AbstractIntegrator,
+    F<:AbstractFloat
+} <: AbstractTrajectory{I}
+    integrator      :: I
+    n_steps         :: Int  # number of samples
+    η               :: F    # learning rate
+    α               :: F    # momentum decay
+end
+
+function transition(
+    rng::AbstractRNG,
+    τ::SGHMC,
+    h::Hamiltonian,
+    z::PhasePoint
+) where {T<:Real}
+    z′ = step(rng, τ.integrator, h, z, τ.n_steps)
+    # no M-H step
+    z = PhasePoint(z′.θ, z′.r, z′.ℓπ, z′.ℓκ)
+    stat = (
+        step_size=τ.integrator.ϵ,
+        n_steps=τ.n_steps,
+        log_density=z.ℓπ.value,
+        hamiltonian_energy=energy(z),
+        )
+    return Transition(z, stat)
+end
+
+
+
 ###
 ### Stochastic Gradient Langevin Dynamics sampler.
 ###