{"id":"ff066702-cc64-4a1a-a793-c5c9ba597472","arxiv_id":"2505.11190","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"low","formal_verification":"none","parameter_count":6,"one_line_summary":"JaxSGMC provides modular, jit-compilable stochastic gradient MCMC samplers for JAX, including pSGLD, SGHMC, replica exchange, AMAGOLD, and SGGMC.","lead":"JaxSGMC is a new JAX library that packages several stochastic gradient MCMC samplers as modular building blocks so they can be swapped and reassembled. It aims to lower the barrier for using Bayesian uncertainty quantification in deep learning.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Verification of built-in samplers is qualitative; a single wrong noise scale or MH acceptance probability would bias the posterior without being caught by the paper's examples.","rationale":"The reader's weakest assumption identifies the same issue: the samplers' fidelity to the intended posterior is only qualitatively or indirectly validated. My reading of the paper confirms this. The linear-regression comparison (fig. 2) is described as 'agree reasonably well,' with no quantitative metric; the CIFAR-10 accuracy results and certainty-threshold table (table 3) show plausible behavior but cannot distinguish a correct SG-MCMC implementation from a biased one, since both would give reasonable accuracy. The paper's own statements in Section 3.1 ('in line with expectations') and Section 3.2 ('comparable performance') are consistent with the reader's assessment that verification is suggestive rather than rigorous. I do not see an internal inconsistency or an obvious error in the described integrators; the concern is about evidential sufficiency for the central claim that the library implements the cited samplers correctly. A targeted posterior-recovery test on a small model would settle this cleanly. The verdict CONDITIONAL remains appropriate: the software engineering contribution is real and the examples are encouraging, but the posterior-fidelity claim needs the quantitative check before the UQ claims should be taken at face value. I agree with the reader's weakest_assumption and see no need to move the verdict.","tokens_in":12804,"tokens_out":1481,"duration_ms":13142,"concrete_test":"Run a targeted posterior-recovery check on a small tractable model, e.g., a Bayesian logistic regression on synthetic data with a known posterior, for each advertised sampler (pSGLD, SGHMC, reSGLD, AMAGOLD, SGGMC). Compare the resulting marginal distributions against an independent reference (NumPyro NUTS or exact enumeration) using a quantitative discrepancy metric such as Wasserstein distance, with multiple random seeds and a fixed step-size schedule. If any sampler's discrepancy is substantially larger than the Monte Carlo error implied by the chain length, the implementation is not faithful to the cited method and the headline claim weakens.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"The central claim is that JaxSGMC faithfully implements state-of-the-art SG-MCMC samplers (SGLD, pSGLD, SGHMC, reSGLD, AMAGOLD, SGGMC) so users obtain posterior samples. The only direct distributional check is the linear-regression example (Section 3.1), which compares pSGLD samples to NumPyro HMC and reports agreement 'reasonably well' without quantitative diagnostics such as Wasserstein distance, effective sample size, or repeated-seed scatter. The CIFAR-10 example (Section 3.2) reports accuracy and a certainty-threshold table, neither of which is a convergence or posterior-fidelity check. The reSGLD and AMAGOLD implementations, which the paper highlights as novel contributions relative to prior JAX libraries, receive no distributional validation at all; a wrong swap-acceptance probability in the tempered-chain exchange (Deng et al. 2020) or an incorrect noise scaling in the AMAGOLD proposal would leave the computed stationary distribution biased while the paper's examples would still pass. Because the library's scientific value is precisely that these samplers sample the intended posterior, the missing quantitative verification is load-bearing, not cosmetic.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper presents JaxSGMC, a JAX library for stochastic gradient Markov chain Monte Carlo (SG-MCMC) that provides pre-built samplers (pSGLD, SGHMC, reSGLD, AMAGOLD, SGGMC) and a modular API for composing custom samplers from building blocks such as integrators, schedulers, and data loaders. The software is designed to be domain-independent, with end-to-end jit compilation and host-device data transfer. The paper demonstrates the library with two examples: a linear regression model using a custom-built pSGLD sampler, compared qualitatively to NumPyro HMC, and a CIFAR-10 image classification task using a pre-built pSGLD sampler, reporting classification accuracy and certainty-threshold behavior. The central claims are that JaxSGMC faithfully implements state-of-the-art SG-MCMC samplers and that its modular structure lowers the barrier to using and developing SG-MCMC methods.","tokens_in":13156,"tokens_out":3784,"duration_ms":35281,"significance":"If the implementation is faithful, JaxSGMC fills a practical gap by offering a modular, JAX-native SG-MCMC library that supports recently proposed samplers (reSGLD, AMAGOLD) not available in other JAX-based libraries, and it could accelerate adoption of Bayesian UQ in deep learning and physical modeling. The paper ships a public repository, documentation, and code listings, which are strengths for a software paper. However, the empirical verification of the central sampling-fidelity claim is currently weak, especially for the novel samplers, and one of the code listings appears to contain a substantive error. These issues are load-bearing for a library whose purpose is to draw samples from the correct posterior.","major_comments":[{"comment":"The validation of pSGLD against NumPyro HMC is qualitative: the text states that the distributions \"agree reasonably well\" without any numerical diagnostics. To substantiate the claim that the sampler targets the correct posterior, report quantitative measures such as Wasserstein or maximum mean discrepancy, effective sample size, or repeated-seed scatter plots.","section":"Section 3.1"},{"comment":"The CIFAR-10 example reports classification accuracy and certainty thresholds but provides no convergence diagnostics (e.g., trace plots, ESS, Gelman-Rubin) or posterior-fidelity checks. Ensemble classification accuracy alone does not demonstrate that the sampler draws from the intended posterior; add calibration or predictive-coverage checks, or explicitly label the example as an illustrative runtime/feature demonstration.","section":"Section 3.2"},{"comment":"The samplers highlighted as novel contributions relative to prior JAX libraries—replica exchange SG-MCMC (reSGLD) and AMAGOLD—are never empirically validated. A subtle implementation error, such as an incorrect swap-acceptance probability in reSGLD or an incorrect noise scaling in the AMAGOLD proposal, would bias the stationary distribution while remaining invisible in the current examples. Add synthetic-data experiments with known target posteriors for these samplers.","section":"Section 2.3 / Section 3"},{"comment":"The `log_prior` function returns `1 / jnp.exp(sample[\"log_sigma\"])`, which is not the log-density of an exponential prior. For σ = exp(log_sigma) and an exponential(1) prior, the log-density is `-jnp.exp(sample[\"log_sigma\"])` up to an additive constant. As written, the example defines a different potential than described in the text, compromising the illustrative linear regression example.","section":"Listing 2, Section 3.1"}],"minor_comments":[{"comment":"The phrase \"an domain-independent library\" should be \"a domain-independent library\".","section":"Section 4"},{"comment":"The caption says \"contour plots of Gaussians obtained from the Hamiltonian Monte Carlo (HMC) method\"; clarify whether the Gaussians are fitted to HMC samples or derived analytically from the linear regression posterior.","section":"Figure 2 caption"},{"comment":"The statement that \"the cost of the whole pSGLD training of 200 epochs is the same order of magnitude as the cost of generating a single sample with the full-batch HMC\" is vague; specify hardware, the HMC hyperparameters (number of leapfrog steps, trajectory length), and the exact runtime comparison.","section":"Section 3.2"},{"comment":"The text refers to \"Stochastic Variation Inference\"; it should be \"Stochastic Variational Inference\".","section":"Section 4"},{"comment":"Table 2 lists the module as `solvers.py` while the text and listings refer to `solver.py`; make the module name consistent.","section":"Table 2"}],"recommendation":"major_revision","confidential_remarks":"The paper is suitable in scope for a software journal and the repository is publicly available with documentation. The main concern is that the empirical validation is too weak to support the central claim that the samplers faithfully target their intended posteriors, particularly for reSGLD and AMAGOLD. The error in Listing 2's log-prior is also concerning and must be fixed. These issues are addressable with additional experiments and code corrections, hence major revision rather than reject."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Short version: this is a legitimate, useful software contribution—a modular JAX library implementing several SG-MCMC samplers (pSGLD, SGHMC, reSGLD, AMAGOLD, SGGMC) with clean separation of integrators, solvers, schedulers, data loaders, and I/O, plus end-to-end jit support. The paper is honest about what it is: a software description, not a methods paper. The main soft spot is the empirical validation, which is qualitative and doesn't directly verify the implemented samplers sample the intended posterior. That matters, but it's a fixable issue for a software paper, not a fatal flaw.\n\nThe architecture is genuinely new relative to SGMCMCJax and other libraries: it supports building custom samplers from building blocks, and it specifically enables replica exchange and amortized MH, which prior JAX libraries did not. The code is public, Apache-2.0, with docs and a version number, and the examples are runnable. Good credit for that.\n\nThe examples: linear regression compares pSGLD to NumPyro HMC and says they 'agree reasonably well' without any quantitative measure—no Wasserstein distance, no ESS, no repeated-seed scatter. The CIFAR-10 example reports accuracy and certainty thresholds, but that is not a posterior-fidelity check. The reSGLD and AMAGOLD implementations, which the paper highlights as novel, get no distributional validation at all. I don't see evidence of a specific bug, but the stress-test concern is right that a subtle error in noise scaling or swap acceptance would go unnoticed by these examples. The paper's value is precisely that these samplers work, so missing verification is a real gap, though not a disaster. Also the runtimes comparison is hand-wavy ('same order of magnitude').\n\nThe citation pattern looks fine; self-citations are to their own applications, which is normal for a library paper. The paper's claims about UQ and Bayesian deep learning are standard and not overstated—they even note pSGLD doesn't fully explore posterior volume.\n\nBottom line: for a researcher who wants to use SG-MCMC in JAX, this library seems like a genuinely useful tool, and the paper is a reasonable software paper. It deserves a serious referee because the code is the main artifact and it's real. I'd recommend acceptance after the authors add at least one quantitative verification of a sampler against a known posterior (e.g., pSGLD on a Gaussian or a simple Bayesian model with exact posterior) and ideally a quick check of reSGLD/AMAGOLD on a small problem. That would close the main gap.","headline":"A solid, genuinely useful JAX library for SG-MCMC with a weak empirical verification section; worth reviewing, but the authors should add quantitative sampler checks.","tokens_in":13590,"tokens_out":2174,"would_cite":true,"duration_ms":20619,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"JaxSGMC packages stochastic-gradient MCMC samplers as modular JAX building blocks.","keywords":["JAX","stochastic gradient MCMC","Bayesian deep learning","uncertainty quantification","SG-MCMC samplers","modular software","Markov chain Monte Carlo","posterior sampling"],"falsifier":"Run one of the library's samplers on a tractable target with a closed-form posterior, such as a Gaussian linear model with known variance, and compare the MCMC sample variance and the posterior predictive coverage to the nominal level. A mismatch beyond Monte Carlo error, especially in the noise injected by the Langevin integrator or in replica-exchange swap acceptance, would show the implementation does not sample the stated distribution.","tokens_in":12602,"feed_emoji":"🎲","tokens_out":6373,"duration_ms":56625,"temperature":0.7,"pith_summary":"The paper presents JaxSGMC, an application-agnostic JAX library that implements several state-of-the-art stochastic gradient MCMC (SG-MCMC) samplers and exposes standard SG-MCMC building blocks for composing custom samplers. Its aim is to lower the barrier to Bayesian deep learning by letting a user with an existing JAX model switch from stochastic optimization to posterior sampling without changing the model code. If the library works as described, uncertainty quantification becomes accessible for large-scale models, including newer samplers such as replica-exchange SGLD and AMAGOLD that existing JAX libraries did not offer. The two worked examples show a pSGLD sampler assembled from blocks matching a gold-standard HMC in linear regression, and giving uncertainty-aware predictions on CIFAR-10.","feed_headline":"JaxSGMC packages stochastic-gradient MCMC samplers as modular JAX blocks","feed_subtitle":"A modular API lets users switch from stochastic optimization to posterior sampling without rewriting their JAX model.","key_machinery":"The central object is the modular sampler architecture: each sampler is built from independent modules---potential.py (log-likelihood/prior potentials), data.py (jit-compatible DataLoaders), adaption.py (RMSProp and covariance preconditioners), integrator.py (Langevin diffusion, leapfrog with friction, OBABO), solver.py (accept/reject or unconditional sample processing), and scheduler.py (step size, burn-in, thinning). The mathematical object that carries the dynamics is the stochastic potential $U(\\theta) \\approx -\\frac{N}{n}\\sum_{i} \\log p(y_i|x_i,\\theta,\\mathcal{M}) - \\log p(\\theta|\\mathcal{M})$, whose mini-batch gradient drives each integrator; composing modules around this potential lets the same building blocks express different samplers.","core_discovery":"On its own terms, JaxSGMC claims that the full variety of modern SG-MCMC algorithms can be organized into a small set of reusable modules---potential evaluation, data batching, preconditioner adaptation, integrators, solvers, and schedulers---and that composing these modules yields both standard samplers like (preconditioned) SGLD, SGHMC, and SGGMC and more recent schemes like replica-exchange SGLD and AMAGOLD under one API. The library is designed so that sampler components can be jit-compiled end-to-end, including the data loading, which keeps Bayesian sampling runtime-competitive with stochastic optimization. This is offered as a practical path for making SG-MCMC a drop-in alternative to optimization in JAX deep-learning workflows.","pith_inferences":["If the modular decomposition is as clean as presented, SG-MCMC research could shift toward composing and benchmarking blocks rather than reimplementing full samplers, making new samplers easier to compare on identical data-loading and scheduling code.","A natural stress test the paper leaves implicit is quantitative convergence checking: running the library's samplers on targets with known posterior moments would let users verify the noise scaling and acceptance steps that the qualitative agreement plot does not settle.","The CIFAR-10 example suggests a direct extension: the same API could be used to ablate which building block (preconditioner, integrator, or swap schedule) most improves posterior coverage, a comparison the paper does not carry out.","The authors indicate that pSGLD may not fully explore posterior volume; a testable next step is whether AMAGOLD or replica-exchange samplers in this library reduce that gap on the same neural-network potential benchmark."],"forward_implications":["A JAX user with a model already written as a function can call a high-level alias to sample the posterior instead of optimizing, with no model rewrite.","Samplers that require advanced building blocks, such as parallel tempering and amortized Metropolis-Hastings acceptance, become available through a common API rather than as stand-alone code.","Custom samplers can be assembled from the documented blocks, letting practitioners tailor proposals, preconditioners, and schedules to a problem.","Because data loading lives inside the jit-compiled loop, SG-MCMC inference can run at a cost comparable to stochastic optimization; the paper reports pSGLD training of 200 epochs on CIFAR-10 at the same order of magnitude as one full-batch HMC proposal."],"supporting_citations":[{"why":"Defines SGLD, the baseline sampler the library implements.","marker":"[13]"},{"why":"Defines SGHMC, one of the implemented samplers.","marker":"[16]"},{"why":"Defines pSGLD with RMSProp preconditioning, the sampler used in both illustrative examples.","marker":"[17]"},{"why":"Defines replica-exchange SG-MCMC, a headline implemented sampler.","marker":"[27]"},{"why":"Defines AMAGOLD with amortized Metropolis-Hastings acceptance, a headline implemented sampler.","marker":"[28]"},{"why":"Defines SGGMC and the OBABO integrator, implemented building blocks.","marker":"[29]"},{"why":"Provides the prior JAX SG-MCMC library whose lack of these building blocks motivates JaxSGMC.","marker":"[26]"},{"why":"Inspires the second API level for composing custom samplers from modular components.","marker":"[37]"}],"fun_headline_variants":["JaxSGMC: Modular stochastic-gradient MCMC for JAX","Build custom SG-MCMC samplers from JAX building blocks","Stochastic-gradient MCMC, modularized: JaxSGMC in JAX","Switch from optimization to posterior sampling in JAX with JaxSGMC","JaxSGMC: Compose SGLD, SGHMC, and more from reusable JAX modules"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The whole contribution rests on the implemented samplers being faithful to the algorithms they cite, so that the Markov chains they produce actually converge to the intended posterior; the paper's verification is qualitative or indirect, and a subtle implementation bug would silently bias the uncertainty estimates.","fun_headline_variants_meta":{"raw":{"variants":["JaxSGMC: Modular stochastic-gradient MCMC for JAX","Build custom SG-MCMC samplers from JAX building blocks","Stochastic-gradient MCMC, modularized: JaxSGMC in JAX","Switch from optimization to posterior sampling in JAX with JaxSGMC","JaxSGMC: Compose SGLD, SGHMC, and more from reusable JAX modules"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000231,"raw_usage":{"total_tokens":1425,"prompt_tokens":827,"completion_tokens":598,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":443,"completion_tokens_details":{"reasoning_tokens":494}},"tokens_in":443,"tokens_out":598,"duration_ms":5804,"temperature":1.0,"reasoning_tokens":494,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T20:55:32.032022+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run one of the library's samplers on a tractable target with a closed-form posterior, such as a Gaussian linear model with known variance, and compare the MCMC sample variance and the posterior predictive coverage to the nominal level. A mismatch beyond Monte Carlo error, especially in the noise injected by the Langevin integrator or in replica-exchange swap acceptance, would show the implementation does not sample the stated distribution.","supporting_citations":[{"cited_title":"Garriga-Alonso, V","cited_arxiv_id":null,"evidence_quote":"Defines SGGMC and the OBABO integrator, implemented building blocks."},{"cited_title":"Zhang, A","cited_arxiv_id":null,"evidence_quote":"Defines AMAGOLD with amortized Metropolis-Hastings acceptance, a headline implemented sampler."},{"cited_title":"Welling, Y","cited_arxiv_id":null,"evidence_quote":"Defines SGLD, the baseline sampler the library implements."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines SGHMC, one of the implemented samplers."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines pSGLD with RMSProp preconditioning, the sampler used in both illustrative examples."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines replica-exchange SG-MCMC, a headline implemented sampler."},{"cited_title":"Coullon, C","cited_arxiv_id":null,"evidence_quote":"Provides the prior JAX SG-MCMC library whose lack of these building blocks motivates JaxSGMC."},{"cited_title":"Babuschkin, K","cited_arxiv_id":null,"evidence_quote":"Inspires the second API level for composing custom samplers from modular components."}],"review_version":1}