blackjax-devs / blackjax-devs/tuningfork
Emit path starts all chains from one un-jittered init, making basin assignment a seed lottery and under-feeding the between-chain evidence channel
- Dominant language
- Python
- Stars
- 3
- Forks
- 0
- Avg merge
- 7h 24m
- Merged PRs (30d)
- 4
Description
## Summary
Every warmup family in the emit path starts **all chains from a single init position**, broadcast with no jitter:
```python
# tuningfork/recipes/_emit/_warmup.py:211, 263, 435, 515
lambda x: jnp.broadcast_to(x[None], (num_chains,) + x.shape), init_position
```
Chains are therefore distinguished only by their warmup PRNG keys. All 271 committed recipes carry `init_strategy: None`, i.e. the default single-point path.
This has two consequences, and the second is the reason this is filed rather than left as a robustness nit.
## 1. Basin assignment becomes a seed lottery on multimodal targets
Diagnosed in detail on `lotka_volterra`, whose posterior has a decoy mode (bad ODE trajectory fit absorbed into inflated observation noise, `sigma_obs` 3.86 vs 0.57). Measured by walking a straight line in unconstrained space between the two modes:
| position | log-density |
|---|---|
| reference mode | −65.8 |
| **midpoint** | **−656.3** |
| decoy mode | −236.5 |
The decoy sits ~171 nats below the true mode (negligible posterior mass, correctly excluded) behind a **~420-nat barrier**. No HMC chain crosses that within any budget we would run, so a chain that descends into it is trapped for the life of the run — and it looks healthy there: stable mean across first/last 250 draws, acceptance 0.99, zero divergences.
Which basin each chain reaches is decided by chaotic accumulation over warmup from the single shared init. That init is a `numpyro` `init_to_uniform` draw, and for this model it is hostile: log-density −9557, gradient norm 1.9e4, and **86% of 512 points sampled from the same `[-2,2]^7` box have non-finite log-density** (the posterior region itself is clean: 0 of 256 non-finite).
Same cell, same stack, seed varied:
| seed | R̂ | min-ESS | verdict |
|---|---|---|---|
| 682737 | 11.10 | 2.02 | FAIL (chain 3 in decoy, \|z\| 26.4) |
| 682738 | 1.013 | 325.6 | PASS |
| 682739 | 1.324 | 5.11 | FAIL |
Not a dependency regression: the same cell fails identically under jax 0.10.0 and 0.11.0 (R̂ 11.41 vs 11.10, same chain, same basin, `u0` agreeing to four decimals), with the model's log-density, gradient, ODE trajectory and RNG stream bit-identical or within a few ULP across those versions.
## 2. It under-feeds the between-chain evidence channel the metric controller depends on
This is the more consequential half. The `metric="auto"` controller selects preconditioner structure from **between-chain** evidence, and the method's own description states that this evidence "is informative to the extent that these starts and the subsequent warmup expose distinct regions".
A harness that starts every chain from one point generates between-chain spread only from warmup randomness rather than from dispersed starts, so it systematically weakens the very signal the controller consumes — and every multi-chain number in this catalog was produced that way.
## The fix exists and is unused
`init_strategy` already supports `uniform_perchain` and `zero_perchain`, and `_ENSEMBLE_FRIENDLY_WARMUPS` already permits them for `window_adaptation_{diag,dense,low_rank}_imm` and `mclmc_tuning`. Nothing needs inventing; the default simply does not use it.
Two components are missing:
1. **A validity guard.** Reject and resample any init with non-finite log-density or gradient, deterministic given the seed, bounded attempts, loud failure on exhaustion. Independent of everything else and unambiguously correct given the 86% figure above.
2. **Per-chain dispersion as the default**, each init independently guarded.
## Why this is deferred rather than fixed now
Changing the default init changes every cell's warmup, hence every committed recipe's adapted step size and mass matrix. Doing it properly means re-emitting the corpus — the third full re-baselining this week, after the ESS-estimator switch and the single-dependency-stack re-baseline (#254). The immediate 19 recert failures are instead being addressed per-cell via the existing `init_strategy` field, which leaves this defect in place for the remaining cells.
## Success criteria
- [ ] Validity-guarded init: non-finite log-density or gradient rejected and resampled, deterministic given seed, bounded attempts with loud failure
- [ ] Per-chain dispersion available as the default rather than only opt-in
- [ ] Measured: between-chain spread under single-broadcast vs dispersed inits, on a multimodal model and a well-behaved control
- [ ] Measured: headline shift on currently-passing cells, to establish whether partial re-emission is defensible or the change must be all-or-nothing
- [ ] Decision recorded on whether the corpus is re-emitted wholesale or the mixed-harness state is accepted and documented
## Note on the target
For a genuinely bimodal posterior the correct outcome is not a clean PASS. With dispersed valid starts some chains may still find the decoy, and the controller's documented response to persistent within-/between-chain disagreement is to retain the within-region matrix and advise a population or tempering method. A harness that hides multimodality from the controller tests it on an easier problem than the one it was built for.
— 🤖 Blackjax-devs AI TL
Contributor guide
Research direction
Start in tuningfork/recipes/_emit/_warmup.py at the four broadcast sites and trace the existing init_strategy handling, including uniform_perchain and zero_perchain. Review the 271 committed recipes and the permitted _ENSEMBLE_FRIENDLY_WARMUPS configurations before deciding how validation and defaults affect re-emission. Done means bounded, deterministic rejection of invalid inits, dispersed default starts, measurements for multimodal and control models, and a recorded corpus re-emission decision.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100