AdaptiveStepSize

phasic.svgd.AdaptiveStepSize(base_step=0.01, kl_target=0.1, adjust_rate=0.1)

Adaptive step size based on particle spread (KL divergence proxy).

Increases step size when particles are too concentrated (low KL), decreases when particles are too dispersed (high KL).

Parameters

base_step : float = 0.01

Base step size

kl_target : float = 0.1

Target KL divergence (in log-space particle spread)

adjust_rate : float = 0.1

Rate of adjustment (0 = no adjustment, 1 = immediate)

Examples

>>> schedule = AdaptiveStepSize(base_step=0.01, kl_target=0.1)
>>> particles = jnp.array([[1.0], [1.1], [0.9]])  # concentrated
>>> schedule(10, particles)  # will increase step size
0.011