Use consistent forward/reverse probabilities in MCMC grow/prune moves

[dwachsmuth]

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]

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?

[dwachsmuth]

Hi Andrew, sure thing!

Here is a six-observation example using the R API. With one binary predictor and min_samples_leaf = 2, there are only two possible tree shapes: an unsplit root, or a split into two groups of three. Fixing both variances at 1 lets us calculate the exact posterior split probability.

library(stochtree)

# Only two possible tree shapes: a root, or a split into two groups of 3.
x <- matrix(rep(c(1, 2), each = 3), ncol = 1)
y <- rep(0, 6)
fit <- bart(
  x, y, num_gfr = 0, num_burnin = 1000, num_mcmc = 100000,
  general_params = list(
    standardize = FALSE, random_seed = 123,
    sample_sigma2_global = FALSE, sigma2_global_init = 1
  ),
  mean_forest_params = list(
    num_trees = 1, alpha = 0.9, beta = 2, min_samples_leaf = 2,
    sample_sigma2_leaf = FALSE, sigma2_leaf_init = 1
  )
)

# Forest and tree indices are zero-based.
split <- vapply(0:99999, function(i) {
  fit$mean_forests$num_leaves(i, 0) == 2
}, logical(1))

# Posterior split/root odds = depth-prior odds * integrated likelihood ratio.
# With zero outcomes and both variances = 1, a leaf contributes 1/sqrt(1+n).
odds <- (0.9 * (1 - 0.9 / 4)^2 / (1 - 0.9)) * (sqrt(7) / 4)
round(c(exact_split_probability = odds / (1 + odds),
        sampled_split_frequency = mean(split)), 4)

Running this same script against the PR's base and its corrected commit gives:

Version	Exact split probability	Sampled split frequency
Base (5c09ac77)	0.7814	0.6418
This PR (5c5f2e8a)	0.7814	0.7813

The exact calculation uses the sampler's existing depth prior and integrated Gaussian likelihood: the prior split/root odds are 0.9 * (1 - 0.9/4)^2 / 0.1, and the likelihood ratio is sqrt(7)/4.

This isolates the reverse-grow probability after pruning to the root. Grow is selected with probability 1 at the root, and prune with probability 1 in the split tree (neither three-observation leaf meets the size threshold for grow). The old prune calculation nevertheless uses a reverse-grow probability of 0.5 because it checks the tree's node count before pruning. Together with the old proposal-ratio direction, this doubles the prune acceptance probability here, giving a stationary split probability of odds / (2 + odds) = 0.6413 instead of 0.7814.

The sampled frequencies above are from R 4.6.1 on macOS ARM64.

[andrewherren]

Thanks for this! I agree with the assessment of our MH ratio bug and the proposed fix is a nice refactoring! I have a bunch of integration tests that check runtime and other empirics for big changes that I'll want to run against this PR before I fully merge it, but I like the fix and appreciate the contribution!

In the meantime, would you mind looking into the failure of the python unit tests? The CI runners all fail this test: test/python/test_bart.py::TestBARTFloat32::test_bart_float32_rfx_matches_float64

[andrewherren]

Just to follow up on this -- it's possible that test_bart_float32_rfx_matches_float64 is brittle to this particular PR which changes (fixes) the MH ratio, so I'm not too worried about that for now, at least until I've had to time to dig into this further. Thanks again!

[andrewherren]

Sorry for the delay here! I've update the python unit tests which were failing on some floating point comparison tests (tl;dr I added an absolute tolerance of 1e-6 in addition to the relative tolerance which is necessary when one of the numbers in a comparison is really small).

I've had a chance to test it locally and it seems to work well, so I anticipate merging this PR today. Thanks again for your contribution and for catching this issue!

[andrewherren]

Looks good! Merging