Implementation:Pyro ppl Pyro DualAveraging

Property	Value
Module	pyro.ops.dual_averaging
Source	pyro/ops/dual_averaging.py
Lines	80
Classes	DualAveraging
Dependencies	(none -- pure Python)

Overview

This module implements the Dual Averaging optimization scheme adapted for use in Markov chain Monte Carlo (MCMC) algorithms, particularly for adaptively tuning the step size in Hamiltonian Monte Carlo (HMC) and NUTS samplers.

Dual Averaging belongs to a class of subgradient methods that update parameters in primal space using equally-weighted subgradients from a dual space, solving the counter-intuitive aspect of traditional subgradient methods where new subgradients enter with decreasing weights. The scheme is guaranteed to converge to an optimal value under appropriate conditions.

The implementation follows Hoffman & Gelman (2014), introducing free parameters t0 and kappa to control stability and convergence speed, along with gamma for convergence rate.

Code Reference

Class: DualAveraging

Constructor parameters:

• prox_center (float, default 0): The "prox-center" that pulls the primal sequence towards it.
• t0 (float, default 10): Stabilizes initial steps of the scheme.
• kappa (float, default 0.75): Controls weights of steps. Should be in (0.5, 1]. Smaller values forget earlier states faster.
• gamma (float, default 0.05): Controls convergence speed.

Methods:

• reset(): Resets internal state (_x_avg, _g_avg, step counter).
• step(g): Updates the scheme with a new statistic/subgradient g. Computes the primal update x_t = prox_center - sqrt(t) / gamma * g_avg and updates the weighted running average x_avg.
• get_state(): Returns a tuple (x_t, x_avg) of the latest primal value and the weighted average.

I/O Contract

Method	Input	Output
__init__	prox_center: float, t0: float, kappa: float, gamma: float	DualAveraging instance
step(g)	g: float (statistic/subgradient)	(none, mutates state)
get_state()	(none)	Tuple (x_t: float, x_avg: float)

Usage Examples

from pyro.ops.dual_averaging import DualAveraging
import math

# Adapt step size for HMC
# Target acceptance probability of 0.8
target_accept = 0.8
log_step_size = math.log(0.1)  # initial step size
da = DualAveraging(prox_center=log_step_size, t0=10, kappa=0.75, gamma=0.05)

# Simulate adaptation with mock acceptance probabilities
for i in range(100):
    accept_prob = 0.7 + 0.3 * (i / 100)  # improving acceptance
    # The statistic g is typically (target_accept - accept_prob)
    da.step(target_accept - accept_prob)

x_t, x_avg = da.get_state()
adapted_step_size = math.exp(x_avg)
print(f"Adapted step size: {adapted_step_size:.4f}")

Related Pages

• Pyro_ppl_Pyro_VelocityVerlet -- HMC integrator that uses adapted step sizes
• Pyro_ppl_Pyro_WelfordCovariance -- Mass matrix adaptation used alongside step size adaptation