lnccbrown / lnccbrown/HSSM

Determine the root cause of hierarchical TruncatedNormal sampling failures

Open
#1,282 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
124
Forks
24
Avg merge
19h 32m
Merged PRs (30d)
60

Description

## Purpose

Determine why the rejected centered hierarchical `TruncatedNormal` candidate samples badly, and distinguish a PyMC correctness defect from difficult-but-correct posterior geometry, backend/adaptation behavior, prior calibration, initialization, and HSSM likelihood interactions.

This is now a **research tracker**, not a production release gate:

- PR #1277 is closed unmerged and its exact rejected candidate is frozen on `codex/1269-truncated-normal-experimental-base` at `9e7002f0`.
- PR #1288 is an independent research-only sink based on that frozen branch. It preserves the v2 harness and immutable smoke evidence and must not merge into `main` merely to retain the record.
- The general unmatched-group prior policy remains independent in #1272.
- The only future production intersection is a small evidence-selected default-policy PR based on merged #1272, with no experimental machinery.

## What the canonical smoke established

Canonical run [33410419447](https://github.com/lnccbrown/HSSM/actions/runs/33410419447) on exact commit `8fdd272b` rejected the automatic native centered-TN candidate:

- direct Gaussian PyMC lower-only and two-sided cells failed the frozen divergence screen;
- current Bambi and the HSSM softmax cell also failed;
- the LBA chain was pathological despite zero divergences: R-hat up to 2.97, bulk ESS 2.33, BFMI 0.029, and all retained draws at maximum tree depth;
- other TN smoke cells passed, so this is not universal TN failure;
- completed finite-difference and PyTensor/JAX checks passed at their evaluated starts; the LBA gradient diagnostic did not complete.

Therefore:

- the problem is not HSSM- or Bambi-only, because direct PyMC reproduces it;
- the evidence does **not** yet establish a PyMC density or gradient bug;
- the leading hypothesis is a centered hierarchical funnel sharpened by TN normalization/boundary coupling, with additional calibration/start/likelihood interactions;
- the linked-Normal controls are practical alternatives but change the statistical model, so they cannot establish the root cause.

Do not run the full frozen 720-cell v2 qualification for a candidate already rejected by smoke, and do not weaken its thresholds or reroll its seeds.

## Causal investigation

Preserve all v2 files/results byte-for-byte. Add a small direct-Gaussian v3 study with four representations of the **same natural-scale model**:

1. `native_centered`: PyMC native centered `TruncatedNormal` location and group effects.
2. `manual_centered`: independent TN log density in explicit transformed coordinates using `Flat` variables, a `Potential`, and all Jacobians.
3. `group_icdf_noncentered`: native TN location plus exact conditional inverse-CDF group latents.
4. `full_icdf_noncentered`: exact inverse-CDF non-centering for both the location and group effects.

Every representation retains:

```text
group_location ~ TruncatedNormal(base_mu, 0.25, lower, upper)
group_scale ~ Weibull(1.5, 0.3)
group_effect[g] ~ TruncatedNormal(group_location, group_scale, lower, upper)
y[g, i] ~ Normal(group_effect[g], 0.5)
```

### B1 — exact models and independent correctness oracle

Child branch: `codex/1282-tn-causal-oracle`, stacked on #1288.

- implement a NumPy/SciPy second-order jet oracle with no PyMC/PyTensor/JAX/model-builder imports;
- include the full transformed posterior: priors, Gaussian likelihood, TN normalization, support/positive transforms, and all log-Jacobians;
- compare value, gradient, and Hessian against PyTensor and JAX for lower-only and finite bounds, boundary-near/interior values, and scale regimes from 0.01 to 1.0;
- compile the manual and inverse-CDF graphs under both backend paths before freezing a sampling manifest.

### B2 — prospective causal sampling study

Freeze two exact known-failing Gaussian regimes:

- lower-only, outside anchor, weak information, float64;
- two-sided, near-boundary truth, float32.

Use identical natural data and natural starts across all representations/backends, then map those starts into each coordinate system. Use PyMC and NumPyro:

- smoke: 2 regimes × 4 representations × 2 backends × 1 dataset = 16 fits, 2 chains, 250 tune + 250 draws;
- confirmation: 2 × 4 × 2 × 5 datasets = 80 fits, 4 chains, 1,000 tune + 1,000 draws.

Run each four-representation comparison on one worker with fresh per-cell caches and counterbalanced order. Expected centered failures must not prevent the non-centered cells from running.

### B3 — evidence report

Commit only reviewed aggregates and a causal decision report. Raw chains remain workflow artifacts. The report must include:

- `root_cause_class`;
- `pymc_correctness_bug_supported`;
- `centered_failure_reproduced`;
- `same_model_nc_candidate`;
- `backend_specific`;
- `default_policy_implication`;
- `next_required_experiment`.

## Predeclared causal classifications

- `native-pymc-correctness-defect`: native value/gradient/Hessian disagrees with the independent oracle while manual centered agrees. Stop and produce a minimal upstream PyMC reproducer.
- `native-graph-lowering-or-adaptation`: both centered densities agree with the oracle, but only the native graph samples badly.
- `conditional-centering-root-cause`: both centered forms fail similarly; group-level inverse-CDF NC is healthy and posterior-equivalent.
- `location-and-conditional-centering-root-cause`: only full inverse-CDF NC is healthy.
- `backend-specific`: one representation consistently succeeds on only one backend.
- `tn-or-residual-geometry`: every exact TN representation fails despite oracle agreement.
- `mixed-inconclusive`: no predeclared pattern is met, including a case where shared natural starts make the centered models healthy. The next experiment then isolates initialization/calibration prospectively.

Healthy representations must meet the existing v2 convergence/geometry thresholds and agree on the natural-scale posterior (combined rank R-hat below 1.01 and posterior-mean differences within four combined MCSEs).

A successful inverse-CDF representation advances to HSSM likelihood isolation. It does **not** authorize a new production default without a separate HSSM qualification and product-policy PR.

## Historical context

- HSSM PR #246 had unexplained sampling failures under the old custom truncation wrapper; that is a warning, not this model.
- PyMC #4417 documented a historical wrong-gradient TN defect, but the completed current local probes have not reproduced it.
- Bambi #617 shows that a centered native hierarchy can work in some regimes.
- Bambi #1003 concerns dropped locations under Bambi's Normal non-centering and is a separate compatibility problem.

Contributor guide

Open the contributing guide

Research direction

Start from the frozen v2 harness on codex/1282-tn-causal-oracle, stacked on #1288, and implement the four exact-model representations and independent NumPy/SciPy oracle described in B1. Run the two smoke regimes across both backends before the confirmation study. Done means reviewed aggregates and a causal report containing the seven required decision fields, with healthy representations meeting the stated convergence and posterior-equivalence thresholds.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.