Skip to content

Use consistent forward/reverse probabilities in MCMC grow/prune moves - #426

Open
dwachsmuth wants to merge 1 commit into
StochasticTree:mainfrom
dwachsmuth:codex/fix-mcmc-acceptance
Open

Use consistent forward/reverse probabilities in MCMC grow/prune moves#426
dwachsmuth wants to merge 1 commit into
StochasticTree:mainfrom
dwachsmuth:codex/fix-mcmc-acceptance

Conversation

@dwachsmuth

Copy link
Copy Markdown

Hi,

Here's the second of two issues uncovered with my current modelling run. Again, fix and reprex are both included.

Overview

The grow/prune acceptance formulas use proposal probabilities that differ from those used to select moves. This can violate detailed balance even with finite numeric covariates. This PR computes move probabilities from shared whole-tree state and uses reverse-over-forward probabilities in both acceptance ratios.

This branch is based directly on upstream main at 5c09ac77e4125d078a9ae0ad5b29e71b33615a5c. It contains only the acceptance fix and its tests. The numeric-extrema correction is a separate PR; the exact-posterior validation described below includes that correction to isolate acceptance logic.

Problems

  • The prune formula uses log(prob_prune_old) - log(prob_grow_new); the proposal ratio requires the opposite signs.
  • The reverse grow probability after pruning uses the pre-prune NumNodes(). This gives 0.5 when pruning to a root should give 1, and allocated-node counts can include deleted nodes.
  • Growing a leaf always assumes the number of reverse-prune choices increases by one. If its sibling is a leaf, the old parent ceases to be a leaf parent, so the count stays unchanged.
  • Acceptance calculations inspect the affected children and feature constancy, whereas move dispatch uses the sizes of all leaves. Unaffected leaves can therefore make the calculated and actual move probabilities disagree.
  • Dispatch's strict n > 2 * min_leaf cutoff excludes valid equal-size boundary splits. An undersized root currently causes a fatal error despite having a valid no-op transition.

Reproductions and intended ratio

Let L(T) be the number of leaves, B(T) the number of prune choices, and g(T) / p(T) the actual probabilities of choosing grow / prune. For a grow T -> T+, the move-type/node-selection ratio is

[p(T+) / B(T+)] / [g(T) / L(T)]

The corresponding prune must use its reciprocal. The target-density ratio still includes the existing integrated likelihood and depth-prior terms. As in the existing sampler, cancellation of the conditional split-rule density assumes that its prior and proposal match; this PR does not change that prior.

Concrete examples:

  1. Minimum leaf size 5, a root split into a constant left child of size 5 and a varying right child of size 15: the old acceptance code predicts reverse prune probability 1, while actual dispatch uses 0.5.
  2. Start with two leaves and grow one: there remains one leaf parent, while the old grow ratio uses two reverse-prune choices.
  3. Minimum leaf size 2, a four-observation root split into two children of size 2: the split is legal. Grow at the root and prune in the resulting tree both have probability 1. The original dispatch cannot propose that root split.

Changes

  • Add a proposal-state summary of leaves, leaf parents and size-eligible leaves.
  • Use its shared grow/prune probabilities for dispatch and acceptance.
  • Compute hypothetical post-move state without changing the live tree, accounting for unaffected leaves and changing leaf-parent membership.
  • Use reverse-over-forward log proposal ratios in both directions.
  • Permit the equality boundary and retain an undersized root without a tree move.

Grow still selects uniformly among all leaves; invalid choices are rejected. Size eligibility is not a promise that a grow succeeds: maximum-depth and constant-feature proposals can still be rejected. Feature selection, cutpoint sampling, likelihoods and depth-prior terms are otherwise unchanged.

Validation

On this independent branch, all 33 C++ tests pass on macOS ARM64 with Apple Clang, C++17 and OpenMP disabled (28 existing plus five new tests). The earlier combined extrema-plus-acceptance build passed 36 tests.

Five deterministic C++ tests cover proposed versus actual whole-tree state, forward/reverse ratios, unaffected leaves, leaf-parent replacement, equality at the size cutoff, root probabilities after node deletion, and undersized-root handling. They use the real tree and tracker classes.

In supplementary one-feature checks with the extrema correction present in both versions, exact transition matrices were enumerated for 15, 6 and 23 valid trees (4/min1, 6/min2 and 8/min2 observations). The original acceptance formulas violate
detailed balance; corrected matrices preserve the target to approximately 1e-16. These calculations integrate unit-width intervals between integer covariates under the conditional uniform cutpoint prior, with alpha=0.6, beta=2, zero outcomes, and unit Gaussian leaf-prior and observation variances.

A direct C++ sampler check on the four-observation example, with 10,000 burn-in and 500,000 retained iterations (seed 20260907), gives:

Leaves Exact posterior Old acceptance logic Corrected logic
1 50.183% 42.630% 50.386%
2 42.176% 54.351% 42.036%
3 7.196% 3.020% 7.137%
4 0.446% 0.000% 0.441%

@andrewherren

Copy link
Copy Markdown
Collaborator

Hi @dwachsmuth, thank you for the detailed report and PR! I won’t be able to give this a proper review until Sunday at the earliest, but if there’s a bug in the MH ratios as you noted, I am definitely eager to fix that.

In the meantime, you mentioned a reprex in your PR, are you able to share that as well?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants