Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

numpyro/NUTS fit_map path is significantly slower than emcee

未关闭
  • 哪里不对？
#34 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

还没有人认领这个 Issue。

评估

难度
5/5
预计耗时
一周以上
新手友好度
35/100
Issue 类型
重构
描述清晰度
需要澄清
活跃度
冷清
技术栈
python
领域
machine-learning, performance

调研方向

Start with the numpyro path and profile the 9-parameter fit described in docs/tutorials/tutorial_6_numpyro.ipynb, comparing it with the emcee path. Read REFACTORING_PLAN.md §5.1b and inspect _make_model, the JIT cache, and NUTS settings. Done means identifying and validating a concrete speed improvement while preserving the reported posterior behavior.

由索引模型根据 Issue 内容生成。

描述

Summary

After the fit_map emcee speedups (f23cf29, 867bc38, 215a45e, cumulatively ~16×), the mcmc='numpyro' path is now wall-clock-slower than mcmc='emcee' on typical fits, even though it is dramatically more sample-efficient (256× fewer samples to reach the same posterior resolution on the validation fit in REFACTORING_PLAN.md §5.1b).

The two emcee speedups (JIT closure + vmap'd batch ln-prob) compile the per-step log-probability into a single XLA dispatch that vmaps across walkers. numpyro can't share that optimisation: NUTS extends its trajectory until a U-turn, which is a per-chain condition, so vmap'ing across chains doesn't help.

Profile of a current fit

9-parameter HD163296 3D fit (docs/tutorials/tutorial_6_numpyro.ipynb), 500 warmup + 500 samples, single chain:

  • Mean leapfrog steps per NUTS iteration: 127
  • Median: 95
  • ~39 % of iterations hit the max_tree_depth=8 cap (256 leapfrog steps)
  • Each gradient evaluation costs ~2× a likelihood evaluation

So one numpyro sample does the gradient work of ~250 emcee evaluations.

Micro-benchmark (50 warmup + 50 sample, post-JIT-warm)

config wall
1 chain, max_tree_depth=8 (previous tutorial default) 40.1 s
1 chain, max_tree_depth=6 (new tutorial default) 23.1 s
4 chains, chain_method='sequential' 197.3 s
4 chains, chain_method='vectorized' 194.3 s

Findings

  • Capping max_tree_depth at 6 gives a ~42 % wall-time speedup with minimal posterior loss — tutorial 6 now recommends 6.
  • chain_method='vectorized' does not help. NUTS' adaptive tree length is per-chain; vmap'd chains must all run to the longest tree on each iteration, cancelling the vmap win. (vectorized is only a win for HMC with fixed L.)

Untried directions

  • target_accept_prob < 0.8 default — accept noisier steps, smaller trajectories. Trade-off: more divergences.
  • dense_mass=True — adapt a full covariance preconditioner. Could help the strongly-correlated z0/psi/r_taper/q_taper block, but the warmup adaptation is slower.
  • More aggressive image downsampling — _make_model cost scales with pixel count; high-SNR pixels dominate the posterior.
  • Audit the numpyro JIT cache: confirm the model is compiled exactly once per fit and not invalidated between iterations.
  • See whether a bounded-horizon scan over leapfrog steps is feasible (probably blocked by the U-turn termination condition, but worth checking).

Workaround

Stay on mcmc='emcee' (the default) for routine fits. Use mcmc='numpyro' when sample efficiency is the actual bottleneck (very expensive single likelihood, very long autocorrelation under emcee, or GPU-available).

主要语言
Python
星标
23
派生
11
平均合并
28 分钟
30 天内合并 PR
1

环境准备

  • 没有 Dockerfile 或 Docker Compose 文件
  • 没有 Pull Request 模板
  • 阅读贡献指南

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

PlanetFormationLab/eddy 的其他 Issue

查看 PlanetFormationLab/eddy 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。