MCMC is extremely slow and OOM-prone — needs optimization

[eprifti]

Problem

MCMC is unusably slow on typical metagenomics datasets (1980 features, 180 samples). With default parameters (n_iter=10000), it gets OOM-killed in Docker. Even with reduced n_iter=1000, it's orders of magnitude slower than GA, ACO, or SA.

Root causes

1. Not parallelized — runs entirely sequentially on one core
2. SBS (Sequential Backward Selection) iterates from all features down to nmin, running full MCMC at each step — O(features × n_iter) total iterations
3. Memory usage — stores trace data for all iterations × all features
4. No GPU support — unlike GA/Beam

Benchmark comparison (Qin2014, 1980 features, 180 samples)

Method	Time	Memory
ILS	0.05s	~10MB
LASSO	0.1s	~10MB
SA	0.3s	~10MB
GA	0.5s	~50MB
ACO	7.1s	~50MB
MCMC	OOM (>256MB) or minutes	>>256MB

Potential fixes

• Parallelize MCMC iterations (each chain independent)
• Reduce memory: don't store full trace, only summary statistics
• Pre-filter features before MCMC (use feature selection to reduce to top 50-100)
• Add progress logging (currently no output during long runs)
• Lower defaults: n_iter=1000, n_burn=500 (done in v0.9.0)

[eprifti]

Deep Performance Analysis

The Core Problem: O(n_features × n_iter × n_features) complexity

The SBS loop (line 893) iterates from n_features down to nmin:

for n in (nmin..=n_features).rev() {     // 340 → 10 = 330 steps
    compute_mcmc(...)                     // each step: n_iter iterations
}

Inside compute_mcmc() (line 1020), each iteration does:

1. 3 Brent optimizations for betas (a, b, c) — up to 100 sub-iterations each
2. n_features coefficient proposals — one Metropolis step per feature
3. Posterior accumulation after burn-in

Total work: (n_features - nmin) × n_iter × (3 × ~50 + n_features)

With Qin2014 (340 features, 1000 iter): 330 × 1000 × 490 = 161M operations. Compare: SA does 5,000 operations (32,000× less).

Memory: The Real OOM Cause

• rng_trace (line 907): clones full RNG at each SBS step — unnecessary
• BayesPred::new() (line 900): references data, but post_mean accumulates across all 330 steps
• data.X HashMap: 180 × 1980 = 356K entries — stays in memory throughout

Proposed Fixes (ordered by impact)

Fix	Speedup	Memory	Effort
LASSO pre-screen (340→50 features)	8×	-80%	Low
Batch SBS (drop 5 at a time when n>50)	5×	-50%	Low
Remove rng_trace cloning	—	significant	Trivial
Adaptive n_iter per step	3×	—	Low
Combined	~100×	-90%	Medium

Implementation Plan

1. LASSO pre-screen: Before SBS, run coordinate descent to rank features. Keep top 50. Reduces 330 SBS steps to 40.

2. Batch elimination: When n > 50, drop the 5 features with highest neutral probability at once.

3. Remove rng_trace: Line 907 rng_trace.push(rng.clone()) — appears unused in return value. Delete it.

4. Adaptive iterations: n_iter_effective = n_iter × min(1.0, 50.0 / n_current) — use fewer iterations for coarse screening.

5. Progress logging: Upgrade SBS loop logging from debug to info with ETA.

Expected result: 340-feature dataset completes in ~5-10s instead of OOM.

[eprifti]

Parallelization Strategy

What can be parallelized

1. Multiple independent chains (biggest win, trivial to implement)

The standard approach used by Stan, JAGS, PyMC: run N chains with different seeds in parallel, combine posteriors.

SBS step (n features remaining):

  Thread 1: Chain(seed=42)  → 1000 iters → posterior_1
  Thread 2: Chain(seed=43)  → 1000 iters → posterior_2     ← rayon par_iter
  Thread 3: Chain(seed=44)  → 1000 iters → posterior_3
  Thread 4: Chain(seed=45)  → 1000 iters → posterior_4

  Combined posterior = average(posterior_1..4)
  → drop feature with highest P(neutral) in combined posterior

Each chain explores different regions of the posterior. Combining gives better coverage and more reliable feature elimination decisions. 4× throughput with zero algorithmic compromise.

Implementation:

// Current (sequential):
let res = compute_mcmc(&bp, param, rng);

// Proposed (parallel):
let chain_seeds: Vec<u64> = (0..n_chains).map(|i| rng.gen()).collect();
let results: Vec<MCMCAnalysisTrace> = chain_seeds
    .par_iter()
    .map(|&seed| {
        let mut chain_rng = ChaCha8Rng::seed_from_u64(seed);
        compute_mcmc(&bp, param, &mut chain_rng)
    })
    .collect();
let combined = MCMCAnalysisTrace::combine(&results);

2. Parallel Brent optimizations for betas (small win)

The 3 beta parameters (a, b, c) are optimized independently at each iteration (lines 1034-1067). These can run in parallel:

Current:  optimize(a) → optimize(b) → optimize(c)   sequential
Proposed: optimize(a) | optimize(b) | optimize(c)    parallel

Small gain (~2-3× on beta phase) since each Brent optimization is fast (~100 1D evaluations).

3. Parallel feature proposals (approximate, advanced)

The coefficient proposal loop (line 1069) is sequential because each accepted proposal changes the state for the next proposal. However:

• Chromatic Gibbs sampling: partition features into independent groups (features that don't interact), update each group in parallel
• Block proposals: propose all features from the same state, accept the one with the highest posterior improvement
• Trade-off: slight approximation but well-studied in the literature (Gonzalez et al., 2011)

Combined optimization plan

                        Before          After           Speedup
                        ──────          ─────           ───────
LASSO pre-screen        340 features    50 features     8×
Batch SBS (drop 5)      330 steps       40 steps        8×
4 parallel chains       1 chain         4 chains        4×
Adaptive n_iter         1000/step       1000→200        3×
Remove rng_trace        330 clones      0 clones        memory

Multiplicative:                                         ~750×

Expected result: Qin2014 (340 features, 180 samples)

• Current: OOM or several minutes
• After: ~2-3 seconds

Implementation order

1. Remove rng_trace — trivial, immediate memory fix
2. LASSO pre-screen — reuse existing lasso.rs coordinate descent
3. Multiple parallel chains — wrap compute_mcmc in par_iter
4. Batch SBS — drop N features per step when n > threshold
5. Adaptive n_iter — reduce iterations proportionally to feature count
6. Progress logging — info-level with ETA at each SBS step

New parameters needed

mcmc:
  n_chains: 4                    # Number of parallel chains (default: num_cpus)
  prescreen_features: 50         # LASSO pre-screen to top N features (0 = disabled)
  batch_elimination: 5           # Drop N features per SBS step when n > 2×nmin
  adaptive_iterations: true      # Scale n_iter by min(1, 50/n_current)

References

• Gelman, A. et al. (2013). Bayesian Data Analysis, 3rd ed. — Chapter 11: multiple chains
• Robert, C. & Casella, G. (2004). Monte Carlo Statistical Methods — parallel MCMC
• Gonzalez, J. et al. (2011). Parallel Gibbs sampling. AISTATS
• Geyer, C.J. (1992). Practical Markov Chain Monte Carlo. Statistical Science

[eprifti]

Fixed in commits ecf33c9 + a2ee201. MCMC now completes Qin2014 (340 features) in ~11 seconds (was OOM/minutes). Key changes: LASSO pre-screen (340→50), parallel chains, batch SBS, golden section search replacing argmin Brent.

[eprifti]

Feedback from Vadim Puller (original MCMC author)

1. Pre-selection with Bayes factor is essential

"J'ai toujours pre-selectionné les features avec facteur Bayes. C'est inutile de le faire avec tout le catalogue... et 10000 d'iterations ne serait pas suffisants dans ce cas."

Features should always be pre-filtered before MCMC — running on the full catalog is wasteful. Our current LASSO pre-screen (top 50) achieves this, but Vadim's original approach used Bayes factor filtering, which is more principled for the Bayesian framework.

Action: Offer both pre-screen methods — LASSO (fast, default) and Bayes factor (principled, feature_minimal_log_abs_bayes_factor already exists in params).

2. Caution about parallel chains

"Je me méfie de la parallélisation, car on jette toujours les premiers n_burn itérations, qui est par défaut la moitié de n_iter — donc on risquerait de finir avec plusieurs chaînes, aucune desquelles ne converge."

Valid concern: each parallel chain needs its own burn-in period. With n_iter=200 and n_burn=100, each chain only has 100 effective samples. Multiple short chains ≠ one long chain for convergence.

Recommendation: Parallel chains should only be used when n_iter is large enough (e.g., n_iter ≥ 2000, so each of 4 chains gets 250 effective samples). For small n_iter, keep single chain.

3. Alternative: joint optimization over feature space (no SBS)

"Un changement possible est de maximiser aussi dans l'espace de nombre de features — sans backward feature selection — ça demanderait plus d'itérations pour converger, mais on ne devrait le refaire pour chaque feature éliminé."

Instead of SBS (which runs a full MCMC at each elimination step), sample feature inclusion/exclusion AS PART of the MCMC itself. This is essentially Gibbs variable selection or reversible jump MCMC:

• At each iteration, propose: (1) change a coefficient, OR (2) add/remove a feature
• The posterior naturally balances model complexity vs fit
• One MCMC run instead of O(n_features) SBS runs

Vadim started a Python prototype. This would be the most principled approach but needs careful implementation and testing.

This is a separate issue — tracked in #72 or a new dedicated issue for Gibbs variable selection.

[eprifti]

Mostly addressed: LASSO pre-screen, adaptive n_iter, batch SBS, golden section optimizer, single-chain default, Gibbs variable selection (PR #74). Remaining: further convergence tuning.