Improve the robustness of ChEES-HMC. (#803)

* Improve the robustness of ChEES-HMC.

* Add a test for halton_sequence to raise ValueError when max_bits is too large.

* Move constant definitions to top of the file.

Author: pipme

blackjax/adaptation/chees_adaptation.py

@@ -17,6 +17,10 @@
 
 # optimal tuning for HMC, see https://arxiv.org/abs/1001.4460
 OPTIMAL_TARGET_ACCEPTANCE_RATE = 0.651
+# Clip the final log-space update like the original implementation in TFP (~log(2)/2 ≈ 0.35).
+LOG_UPDATE_CLIP = 0.35
+# Small constant to avoid division by zero or log of zero
+EPS_FLOAT = 1e-20
 
 
 class ChEESAdaptationState(NamedTuple):
@@ -52,12 +56,24 @@ class ChEESAdaptationState(NamedTuple):
     step: int
 
 
+def weighted_empirical_mean(x, w):
+    # x: (num_chains, dim), w: (num_chains,)
+    x_safe = jnp.where(jnp.isfinite(x), x, 0.0)
+    w = jnp.where(jnp.isfinite(x).all(axis=-1), w, 0.0)
+
+    w_exp = w.reshape((w.shape[0],) + (1,) * (x.ndim - 1))
+    num = jnp.sum(w_exp * x_safe, axis=0)
+    den = jnp.sum(w_exp, axis=0) + EPS_FLOAT
+    return jax.lax.stop_gradient(num / den)
+
+
 def base(
     jitter_generator: Callable,
     next_random_arg_fn: Callable,
     optim: optax.GradientTransformation,
     target_acceptance_rate: float,
     decay_rate: float,
+    max_leapfrog_steps: int,
 ) -> Tuple[Callable, Callable]:
     """Maximizing the Change in the Estimator of the Expected Square criterion
     (trajectory length) and dual averaging procedure (step size) for the jittered
@@ -144,6 +160,8 @@ def compute_parameters(
         harmonic_mean = 1.0 / jnp.mean(
             1.0 / acceptance_probabilities, where=~is_divergent
         )
+        # Replace inf/nan harmonic mean as zero to avoid issues in dual averaging
+        harmonic_mean = jnp.where(jnp.isfinite(harmonic_mean), harmonic_mean, 0.0)
         da_state_ = da_update(da_state, target_acceptance_rate - harmonic_mean)
         step_size_ = jnp.exp(da_state_.log_x)
         new_step_size, new_da_state, new_log_step_size = jax.lax.cond(
@@ -157,9 +175,14 @@ def compute_parameters(
             1.0 - update_weight
         ) * log_step_size_ma + update_weight * new_log_step_size
 
+        w = jnp.where(~is_divergent, acceptance_probabilities, 0.0)
         proposals_mean = jax.tree_util.tree_map(
-            lambda p: jnp.nanmean(p, axis=0), proposed_positions
+            lambda p: weighted_empirical_mean(p, w), proposed_positions
         )
+        # The above weighted mean is presumably better than the simple mean:
+        # proposals_mean = jax.tree_util.tree_map(
+        #     lambda p: jnp.nanmean(p, axis=0), proposed_positions
+        # )
         initials_mean = jax.tree_util.tree_map(
             lambda p: jnp.nanmean(p, axis=0), initial_positions
         )
@@ -177,19 +200,25 @@ def compute_parameters(
 
         trajectory_gradients = (
             jitter_generator(random_generator_arg)
-            * trajectory_length
+            * trajectory_length  # this effectively make this gradient w.r.t. log_trajectory_length
             * jax.vmap(
                 lambda pm, im, mm: (jnp.dot(pm, pm) - jnp.dot(im, im)) * jnp.dot(pm, mm)
             )(proposals_matrix, initials_matrix, momentums_matrix)
         )
+
         trajectory_gradient = jnp.sum(
-            acceptance_probabilities * trajectory_gradients, where=~is_divergent
-        ) / jnp.sum(acceptance_probabilities, where=~is_divergent)
+            acceptance_probabilities * trajectory_gradients,
+            where=~is_divergent,
+        ) / jnp.sum(acceptance_probabilities + EPS_FLOAT, where=~is_divergent)
 
         log_trajectory_length = jnp.log(trajectory_length)
         updates, optim_state_ = optim.update(
             trajectory_gradient, optim_state, log_trajectory_length
         )
+
+        updates = jax.tree_util.tree_map(
+            lambda u: jnp.clip(u, -LOG_UPDATE_CLIP, LOG_UPDATE_CLIP), updates
+        )
         log_trajectory_length_ = optax.apply_updates(log_trajectory_length, updates)
         new_log_trajectory_length, new_optim_state = jax.lax.cond(
             jnp.isfinite(
@@ -204,6 +233,13 @@ def compute_parameters(
         ) * log_trajectory_length_ma + update_weight * new_log_trajectory_length
         new_trajectory_length = jnp.exp(new_log_trajectory_length_ma)
 
+        # clip new trajectory length to avoid too large trajectories, also the
+        # minimum trajectory length is one integrator step
+        new_trajectory_length = jnp.clip(
+            new_trajectory_length,
+            max=max_leapfrog_steps * new_step_size,
+            min=new_step_size,
+        )
         return ChEESAdaptationState(
             new_step_size,
             new_log_step_size_ma,
@@ -278,6 +314,7 @@ def chees_adaptation(
     jitter_amount: float = 1.0,
     target_acceptance_rate: float = OPTIMAL_TARGET_ACCEPTANCE_RATE,
     decay_rate: float = 0.5,
+    max_leapfrog_steps: int = 1000,
     adaptation_info_fn: Callable = return_all_adapt_info,
 ) -> AdaptationAlgorithm:
     """Adapt the step size and trajectory length (number of integration steps / step size)
@@ -376,13 +413,14 @@ def run(
                 jax.random.fold_in(carry_key, i)
             ) * jitter_amount + (1.0 - jitter_amount)
         else:
+            max_bits = np.ceil(np.log2(num_steps + max_sampling_steps))
             jitter_gn = lambda i: dynamic_hmc.halton_sequence(
-                i, np.ceil(np.log2(num_steps + max_sampling_steps))
+                i, max_bits
             ) * jitter_amount + (1.0 - jitter_amount)
 
-        def integration_steps_fn(random_generator_arg, trajectory_length_adjusted):
+        def integration_steps_fn(random_generator_arg, num_leapfrog_steps):
             return jnp.asarray(
-                jnp.ceil(jitter_gn(random_generator_arg) * trajectory_length_adjusted),
+                jnp.ceil(jitter_gn(random_generator_arg) * num_leapfrog_steps),
                 dtype=int,
             )
 
@@ -392,7 +430,12 @@ def integration_steps_fn(random_generator_arg, trajectory_length_adjusted):
         )
 
         init, update = base(
-            jitter_gn, next_random_arg_fn, optim, target_acceptance_rate, decay_rate
+            jitter_gn,
+            next_random_arg_fn,
+            optim,
+            target_acceptance_rate,
+            decay_rate,
+            max_leapfrog_steps,
         )
 
         def one_step(carry, rng_key):
@@ -404,7 +447,7 @@ def one_step(carry, rng_key):
                 logdensity_fn=logdensity_fn,
                 step_size=adaptation_state.step_size,
                 inverse_mass_matrix=jnp.ones(num_dim),
-                trajectory_length_adjusted=adaptation_state.trajectory_length
+                num_leapfrog_steps=adaptation_state.trajectory_length
                 / adaptation_state.step_size,
             )
             new_states, info = jax.vmap(_step_fn)(keys, states)
@@ -432,7 +475,7 @@ def one_step(carry, rng_key):
             one_step, (init_states, init_adaptation_state), keys_step
         )
 
-        trajectory_length_adjusted = jnp.exp(
+        num_leapfrog_steps = jnp.exp(
             last_adaptation_state.log_trajectory_length_moving_average
             - last_adaptation_state.log_step_size_moving_average
         )
@@ -441,7 +484,7 @@ def one_step(carry, rng_key):
             "inverse_mass_matrix": jnp.ones(num_dim),
             "next_random_arg_fn": next_random_arg_fn,
             "integration_steps_fn": lambda arg: integration_steps_fn(
-                arg, trajectory_length_adjusted
+                arg, num_leapfrog_steps
             ),
         }
 

blackjax/mcmc/dynamic_hmc.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Public API for the Dynamic HMC Kernel"""
+
 from typing import Callable, NamedTuple
 
 import jax
@@ -46,7 +47,11 @@ class DynamicHMCState(NamedTuple):
     random_generator_arg: Array
 
 
-def init(position: ArrayLikeTree, logdensity_fn: Callable, random_generator_arg: Array):
+def init(
+    position: ArrayLikeTree,
+    logdensity_fn: Callable,
+    random_generator_arg: Array,
+):
     logdensity, logdensity_grad = jax.value_and_grad(logdensity_fn)(position)
     return DynamicHMCState(position, logdensity, logdensity_grad, random_generator_arg)
 
@@ -154,7 +159,10 @@ def as_top_level_api(
     A ``SamplingAlgorithm``.
     """
     kernel = build_kernel(
-        integrator, divergence_threshold, next_random_arg_fn, integration_steps_fn
+        integrator,
+        divergence_threshold,
+        next_random_arg_fn,
+        integration_steps_fn,
     )
 
     def init_fn(position: ArrayLikeTree, rng_key: Array):
@@ -176,6 +184,14 @@ def step_fn(rng_key: PRNGKey, state):
 
 
 def halton_sequence(i: Array, max_bits: int = 10) -> float:
+    """Generate the (i+1)-th element of the Halton sequence.
+
+    Warning: max_bits should be less than the bit width of i.dtype to prevent integer overflow (e.g., max_bits <= 63 for int64).
+    """
+    if max_bits >= jnp.iinfo(i.dtype).bits:
+        raise ValueError(
+            f"max_bits ({max_bits}) must be less than bit width of dtype {i.dtype} ({jnp.iinfo(i.dtype).bits})"
+        )
     bit_masks = 2 ** jnp.arange(max_bits, dtype=i.dtype)
     return jnp.einsum("i,i->", jnp.mod((i + 1) // bit_masks, 2), 0.5 / bit_masks)
 

tests/adaptation/test_adaptation.py

@@ -61,21 +61,24 @@ def test_adaptation_schedule(num_steps, expected_schedule):
     ],
 )
 def test_chees_adaptation(adaptation_filters):
+    target_mean = jnp.array([0.0, 0.0])
+    target_std = jnp.array([1.0, 10.0])
     logprob_fn = lambda x: jax.scipy.stats.norm.logpdf(
-        x, loc=0.0, scale=jnp.array([1.0, 10.0])
-    ).sum()
+        x, loc=target_mean, scale=target_std
+    ).sum(axis=-1)
 
     num_burnin_steps = 1000
     num_results = 500
     num_chains = 16
     step_size = 0.1
+    target_acceptance_rate = 0.75
 
     init_key, warmup_key, inference_key = jax.random.split(jax.random.key(346), 3)
 
     warmup = blackjax.chees_adaptation(
         logprob_fn,
         num_chains=num_chains,
-        target_acceptance_rate=0.75,
+        target_acceptance_rate=target_acceptance_rate,
         adaptation_info_fn=adaptation_filters["filter_fn"],
     )
 
@@ -84,13 +87,12 @@ def test_chees_adaptation(adaptation_filters):
         warmup_key,
         initial_positions,
         step_size=step_size,
-        optim=optax.adamw(learning_rate=0.5),
+        optim=optax.adam(learning_rate=0.5, b1=0, b2=0.95),
         num_steps=num_burnin_steps,
     )
     algorithm = blackjax.dynamic_hmc(logprob_fn, **parameters)
-
     chain_keys = jax.random.split(inference_key, num_chains)
-    _, (_, infos) = jax.vmap(
+    final_states, (states, infos) = jax.vmap(
         lambda key, state: run_inference_algorithm(
             rng_key=key,
             initial_state=state,
@@ -99,7 +101,9 @@ def test_chees_adaptation(adaptation_filters):
         )
     )(chain_keys, last_states)
 
-    harmonic_mean = 1.0 / jnp.mean(1.0 / infos.acceptance_rate)
+    harmonic_mean = 1.0 / jnp.mean(1.0 / infos.acceptance_rate, axis=0)
+    assert harmonic_mean.shape == (num_results,)
+    harmonic_mean = jnp.mean(harmonic_mean)
 
     def check_attrs(attribute, keyset):
         for name, param in getattr(warmup_info, attribute)._asdict().items():
@@ -119,6 +123,24 @@ def check_attrs(attribute, keyset):
     for i, attribute in enumerate(["state", "info", "adaptation_state"]):
         check_attrs(attribute, keysets[i])
 
-    np.testing.assert_allclose(harmonic_mean, 0.75, atol=1e-1)
-    np.testing.assert_allclose(parameters["step_size"], 1.5, rtol=2e-1)
-    np.testing.assert_array_less(infos.num_integration_steps.mean(), 15.0)
+    # The harmonic mean of the acceptance rate should be close to the target acceptance rate
+    np.testing.assert_allclose(harmonic_mean, target_acceptance_rate, atol=1e-1)
+
+    # These are empirical values that should be roughly correct for this target distribution
+    np.testing.assert_allclose(parameters["step_size"], 1.5, atol=0.3)
+    np.testing.assert_allclose(infos.num_integration_steps.mean(), 9, atol=3)
+
+    # Check that sample means and stds are close to target values
+    draws = states.position.reshape(-1, states.position.shape[-1])
+    empirical_mean = jnp.mean(draws, axis=0)
+    empirical_std = jnp.std(draws, axis=0)
+    np.testing.assert_allclose(empirical_mean, target_mean, atol=0.5)
+    np.testing.assert_allclose(empirical_std, target_std, rtol=0.1)
+
+
+def test_halton_sequence_raise_value():
+    """Test that halton sequence raises value error when max_bits is too large."""
+    from blackjax.mcmc.dynamic_hmc import halton_sequence
+
+    with pytest.raises(ValueError, match="max_bits"):
+        halton_sequence(jnp.array([0], dtype=jnp.int32), max_bits=32)