pith. machine review for the scientific record. sign in

arxiv: 2605.10914 · v1 · submitted 2026-05-11 · 📊 stat.CO

Recognition: 2 theorem links

· Lean Theorem

gemlib.mcmc: composable kernels for Metropolis-within-Gibbs sampling schemes

Alin Morariu, Chris Jewell, Jess Bridgen

Pith reviewed 2026-05-12 03:46 UTC · model grok-4.3

classification 📊 stat.CO
keywords MCMCMetropolis-within-Gibbswriter monadsdata augmentationBayesian inferencecomposable kernelsstate-transition modelsepidemic models
0
0 comments X

The pith

Writer monads enable automatic composition of kernels for Metropolis-within-Gibbs sampling.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper presents a method to simplify building Markov Chain Monte Carlo algorithms for Bayesian inference on state-transition models that have high-dimensional unobserved states. It treats sampling kernels as units that can be combined without the user tracking internal state details at each step. This matters because current tools force a choice between easy automation that limits customization and full flexibility that requires heavy custom coding for each new model. The approach targets fields like epidemiology where joint estimation of parameters and event times is common but implementation is burdensome. If the method works, researchers could reuse and extend advanced sampling schemes across applications with far less effort while retaining correct statistical behavior.

Core claim

The framework formalizes kernel composition using writer monads from category theory, allowing parameter-estimation kernels and data-augmentation kernels to integrate directly into joint sampling schemes for models with intractable likelihoods and missing data, with the composition logic, state management, and performance handling all automated so that developers define only two methods per new kernel.

What carries the argument

The writer monad, a category-theoretic structure that tracks and propagates state updates across successive kernel applications without explicit user-managed bookkeeping.

If this is right

  • Complex Metropolis-within-Gibbs procedures for partially observed models can be expressed as concise chains of kernel operations.
  • New kernels require implementation of only two methods, after which composition and reuse become automatic.
  • Hardware acceleration of the resulting samplers occurs without additional user code.
  • Kernels developed for one epidemic or ecological application can be reused in others without re-deriving the composition logic.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same abstraction pattern could reduce coding effort for other families of iterative statistical algorithms beyond Metropolis-within-Gibbs.
  • Applied modelers might shift time from sampler engineering toward domain-specific questions about latent processes.
  • Empirical tests on diverse real datasets would reveal whether the reduced code volume also lowers the rate of subtle sampling errors.

Load-bearing premise

The monad-based composition produces a mathematically valid overall transition kernel that preserves the required statistical properties of the individual kernels.

What would settle it

A concrete counterexample in which two individually valid kernels, when composed through the framework, yield a chain whose stationary distribution differs from the target posterior would falsify the claim.

Figures

Figures reproduced from arXiv: 2605.10914 by Alin Morariu, Chris Jewell, Jess Bridgen.

Figure 1
Figure 1. Figure 1: Mechanics of performing a MWG transition. The state [PITH_FULL_IMAGE:figures/full_fig_p006_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Comparison of steps between standard MCMC and MWG MCMC algorithms. [PITH_FULL_IMAGE:figures/full_fig_p008_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Computational logistics of MCMC this basic requirement. As such, the mechanics of ensuring the output of an MWG algorithm remains Markovian is a non-trivial task which involves the composition of multiple MCMC kernels. Writer monads are a favourable programming pattern for implementing MCMC. They are functions that take an old state and returns a new state along with a computed value. We adopt a Haskell-li… view at source ↗
Figure 4
Figure 4. Figure 4: Diagrammatic overview of a MWG algorithm step. Each step described maps [PITH_FULL_IMAGE:figures/full_fig_p015_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Visual representation of the evolution of the MCMC kernel [PITH_FULL_IMAGE:figures/full_fig_p016_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Diagram outlining the partial kernel application for a compound algorithm [PITH_FULL_IMAGE:figures/full_fig_p018_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: gemlib constructs the MCMC kernel K2 by grabbing all of the elements it needs from the global scope of the program. The library first computes the conditional target probability distribution π(θ2|θ (0) −2 )) and returns a function (this is analogous to F π in [PITH_FULL_IMAGE:figures/full_fig_p019_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Diagnostics for MWG algorithm on 2-dimensional Gaussian with unknown mean [PITH_FULL_IMAGE:figures/full_fig_p023_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Diagnostics for Metropolis algorithm on 2-dimensional Gaussian with unknown [PITH_FULL_IMAGE:figures/full_fig_p028_9.png] view at source ↗
Figure 11
Figure 11. Figure 11: Simulated epidemic curves for the discrete time, meta-population SIR model with [PITH_FULL_IMAGE:figures/full_fig_p029_11.png] view at source ↗
Figure 12
Figure 12. Figure 12: Output of MWG scheme constructed to fit partially observed SIR model [PITH_FULL_IMAGE:figures/full_fig_p030_12.png] view at source ↗
read the original abstract

State-transition models are essential across epidemiology and ecology, but statistical inference remains challenging owing to high-dimensional latent state spaces, temporal dependence, and intractable likelihood functions. Bayesian inference via Markov Chain Monte Carlo (MCMC) enables joint estimation of model parameters and missing event times through data augmentation, but Metropolis-within-Gibbs (MWG) schemes that combine multiple specialised kernels are notoriously difficult to implement. Current probabilistic programming frameworks face a trade-off: automation sacrifices extensibility, whilst flexibility demands substantial implementation overhead. This divide has created a software landscape characterised by tightly coupled, model-specific implementations that resist reuse and extension. We introduce gemlib.mcmc, an MCMC module designed to bridge methodological and applied communities through principled, composable kernel abstractions. The framework employs writer monads from category theory to formalise kernel composition, enabling seamless integration of parameter-estimation and data-augmentation kernels without manual state management. Built on JAX and TensorFlow Probability for high-performance computation, gemlib.mcmc provides an ergonomic interface -- leveraging Python's right-shift operator for intuitive kernel chaining -- whilst maintaining statistical rigour and transparency. Developers can extend the library by implementing only two methods; composition and hardware acceleration are automated. We demonstrate the framework through parameter inference on partially observed epidemic models, showing how complex inference algorithms can be expressed concisely and reused across applications. By reducing implementation burden we provide access to sophisticated MCMC methods and enable applied researchers to employ state-of-the-art algorithms without reimplementation overhead.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 2 minor

Summary. The paper introduces gemlib.mcmc, a Python library for Metropolis-within-Gibbs MCMC that uses writer monads to formalize composable kernels. It claims to enable seamless integration of parameter-estimation and data-augmentation kernels via an ergonomic interface (Python right-shift operator), with users needing to implement only two methods for extension; composition, state management, and JAX/TFP hardware acceleration are automated while preserving statistical rigour. The framework is demonstrated through parameter inference on partially observed epidemic models, with the goal of reducing implementation burden for complex MWG schemes in epidemiology and ecology.

Significance. If the monad-based composition is shown to preserve detailed balance and stationary distributions, the library could meaningfully lower barriers for applied researchers implementing data-augmentation MCMC in high-dimensional state-transition models, promoting reuse across applications and reducing manual state-management errors.

major comments (1)
  1. Abstract: the central claim that the two-method extension interface 'maintains statistical rigour and transparency' while automating composition is unsupported; the manuscript provides no proofs of correctness, numerical validation against standard MWG implementations, benchmarks, or error analysis to confirm that kernel chaining preserves the target distribution.
minor comments (2)
  1. The manuscript would benefit from explicit code examples showing the right-shift operator for kernel chaining and the two required extension methods.
  2. Clarify the precise interface contract (method signatures and expected return types) for the two extension methods to aid reproducibility.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for their constructive review and for identifying a key area where the manuscript's claims require stronger support. We address the major comment below and will incorporate the necessary revisions to strengthen the presentation of the framework's statistical properties.

read point-by-point responses

  1. Referee: [—] Abstract: the central claim that the two-method extension interface 'maintains statistical rigour and transparency' while automating composition is unsupported; the manuscript provides no proofs of correctness, numerical validation against standard MWG implementations, benchmarks, or error analysis to confirm that kernel chaining preserves the target distribution.

    Authors: We agree that the manuscript as submitted does not contain explicit proofs of correctness for the monad-based composition or numerical validation against reference MWG implementations. The abstract's phrasing regarding maintenance of statistical rigour is therefore not fully substantiated in the current text. In the revised manuscript we will add a new section (likely Section 3 or an appendix) that formally shows preservation of detailed balance: if each constituent kernel satisfies detailed balance with respect to its target conditional distribution, then the writer-monad composition (via the right-shift operator) yields a kernel that satisfies detailed balance with respect to the joint target. The proof will rely on the standard properties of the writer monad and the fact that state updates are performed only through the monadic bind. We will also include a short numerical validation study on a low-dimensional Gaussian target and on a simple SIR epidemic model, comparing posterior samples and acceptance rates obtained from gemlib.mcmc kernels against hand-written Metropolis-within-Gibbs code. Benchmarks measuring wall-clock time and effective sample size per second, together with a brief error analysis (e.g., total-variation distance to a long-run reference chain), will be added to the results section. revision: yes

Circularity Check

0 steps flagged

No significant circularity

full rationale

The paper introduces a software library (gemlib.mcmc) that uses writer monads to compose MCMC kernels, with users implementing only two methods for extension. No equations, fitted parameters, predictions, or first-principles derivations are present that could reduce to inputs by construction. The presentation is an implementation and interface description rather than a mathematical derivation chain; category theory concepts are adopted as external tools without self-referential reduction. No self-citation load-bearing steps, ansatz smuggling, or renaming of known results occur in the provided text.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claim rests on the domain assumption that writer monads can be applied to MCMC kernels to eliminate manual state management while preserving statistical properties; no free parameters or new physical entities are introduced.

axioms (1)
  • domain assumption Writer monads from category theory can formalise kernel composition in MCMC without manual state management
    Invoked in the description of the framework's core abstraction.

pith-pipeline@v0.9.0 · 5569 in / 1242 out tokens · 48556 ms · 2026-05-12T03:46:44.513349+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

Reference graph

Works this paper leans on

31 extracted references · 31 canonical work pages

  1. [1]

    Sampling-Based Approaches to Calculating Marginal Densities

    Alan E. Gelfand and Adrian F. M. Smith. “Sampling-Based Approaches to Calculating Marginal Densities”. In:Journal of the American Statistical Association85.410 (1990), pp. 398–409.issn: 01621459, 1537274X.url: http : / / www . jstor . org / stable / 2289776(visited on 02/27/2026)

    work page 1990

  2. [2]

    Christian P Robert and George Casella.Monte Carlo statistical methods. Vol. 2. Springer, 2004

    work page 2004

  3. [3]

    Brooks et al.Handbook of Markov Chain Monte Carlo

    S. Brooks et al.Handbook of Markov Chain Monte Carlo. Chapman & Hall/CRC Handbooks of Modern Statistical Methods. CRC Press, 2011.isbn: 9781420079425. url:https://books.google.co.uk/books?id=qfRsAIKZ4rIC

    work page 2011

  4. [4]

    The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo

    Matthew D Hoffman, Andrew Gelman, et al. “The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo.” In:J. Mach. Learn. Res.15.1 (2014), pp. 1593–1623

    work page 2014

  5. [5]

    H. I. Freedman.Deterministic Mathematical Models in Population Ecology. Vol. 57. Pure and Applied Mathematics: A Series of Monographs and Textbooks. New York: Marcel Dekker, 1980, p. 254.isbn: 0-8247-6653-9

    work page 1980

  6. [6]

    A contribution to the mathematical theory of epidemics

    WO Kermack and AG McKendrick. “A contribution to the mathematical theory of epidemics”. In:Proc. R. Soc. Lond. A115 (1927), pp. 700–721

    work page 1927

  7. [7]

    Novel coronavirus 2019-nCoV (COVID-19): early estimation of epi- demicological parameters and epidemic size estimates

    JM Read et al. “Novel coronavirus 2019-nCoV (COVID-19): early estimation of epi- demicological parameters and epidemic size estimates”. In:Philos Trans R Soc Lond B Biol Sci376 (1829 2021), p. 20200265

    work page 2019

  8. [8]

    A tutorial introduction to Bayesian inference for stochastic epidemic models using Markov chain Monte Carlo methods

    Philip D O’Neill. “A tutorial introduction to Bayesian inference for stochastic epidemic models using Markov chain Monte Carlo methods”. en. In:Math. Biosci.180.1-2 (Nov. 2002), pp. 103–114

    work page 2002

  9. [9]

    PLOS Com- putational Biology18(9), 1010492 (2022) https://doi.org/10.1371/journal.pcbi

    Sam Moore et al. “Modelling optimal vaccination strategy for SARS-CoV-2 in the UK”. In:PLOS Computational Biology17.5 (2021), e1008849.doi: 10.1371/journal.pcbi. 1008849.url:https://doi.org/10.1371/journal.pcbi.1008849

    work page doi:10.1371/journal.pcbi 2021

  10. [10]

    Cambridge University Press, 2024.url:https://arxiv.org/pdf/2407.12751

    Paul Fearnhead et al.Scalable Monte Carlo for Bayesian Learning. Cambridge University Press, 2024.url:https://arxiv.org/pdf/2407.12751

    work page arXiv 2024

  11. [11]

    F. W. Lawvere.The category of probabilistic mappings - With applications to stocahstic process, statistics, and pattern recognition. 1962.url: https://ncatlab.org/nlab/ files/lawvereprobability1962.pdf

    work page 1962

  12. [12]

    Monad transformers and modular inter- preters

    Sheng Liang, Paul Hudak, and Mark Jones. “Monad transformers and modular inter- preters”. In:Proceedings of the 22nd ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages. POPL ’95. San Francisco, California, USA: Association for Computing Machinery, 1995, 333–343.isbn: 0897916921.doi: 10.1145/199448.199528. url:https://doi.org/10.1145/199448.199528

    work page doi:10.1145/199448.199528 1995

  13. [13]

    Junpeng Lao et al.tfp.mcmc: Modern Markov Chain Monte Carlo Tools Built for Modern Hardware. 2020. arXiv: 2002.01184 [stat.CO] .url: https://arxiv.org/ abs/2002.01184. 36

    work page arXiv 2020

  14. [14]

    Alberto Cabezas et al.BlackJAX: Composable Bayesian inference in JAX. 2024. arXiv: 2402.10797 [cs.MS]

    work page arXiv 2024

  15. [15]

    Apache Software Foundation.OpenXLA Project. Feb. 19, 2010.url: https://github. com/openxla/xla

    work page 2010

  16. [16]

    School of Haskell, FP Complete

    Bartosz Milewski.Basics of Haskell. School of Haskell, FP Complete. Archived from the original on 2016-10-27. Retrieved 2018-07-13. 2013

    work page 2016

  17. [17]

    Functional Programming with Overloading and Higher-Order Polymor- phism

    Mark P. Jones. “Functional Programming with Overloading and Higher-Order Polymor- phism”. In:Advanced Functional Programming. Ed. by Johan Jeuring and Erik Meijer. Vol. 925. Lecture Notes in Computer Science. Berlin, Heidelberg: Springer-Verlag, May 1995

    work page 1995

  18. [18]

    gemlib - probabilistic programming for epidemic models

    Alin Morariu, Jess Bridgen, and C.P. Jewell. “gemlib - probabilistic programming for epidemic models”. In:arXiv preprint arXiv(2025).url: https://arxiv.org/abs/ 2511.08124

    work page arXiv 2025

  19. [19]

    Examples of Adaptive MCMC

    Gareth O. Roberts and Jeffrey S. Rosenthal. “Examples of Adaptive MCMC”. In: Journal of Computational and Graphical Statistics18.2 (2009), pp. 349–367.doi: 10.1198/jcgs.2009.06134 . eprint: https://doi.org/10.1198/jcgs.2009.06134 . url:https://doi.org/10.1198/jcgs.2009.06134

    work page doi:10.1198/jcgs.2009.06134 2009

  20. [20]

    Bayesian Analysis for Emerging Infectious Diseases

    Chris Jewell et al. “Bayesian Analysis for Emerging Infectious Diseases”. In:Bayesian Analysis4 (Sept. 2009), pp. 465–496.doi:10.1214/09-BA417

    work page doi:10.1214/09-ba417 2009

  21. [21]

    A case study in non-centering for data augmentation: Stochastic epidemics

    Peter Neal and Gareth Roberts. “A case study in non-centering for data augmentation: Stochastic epidemics”. In:Statistics and Computing15 (2005), pp. 315–327

    work page 2005

  22. [22]

    PyMC: A Modern and Comprehensive Probabilistic Programming Framework in Python

    Oriol Abril-Pla et al. “PyMC: A Modern and Comprehensive Probabilistic Programming Framework in Python”. In:PeerJ Computer Science9.e1516 (2023).doi: 10.7717/ peerj-cs.1516

    work page 2023

  23. [23]

    Programming With Models: Writing Statistical Algorithms for General Model Structures With NIMBLE

    Perry de Valpine et al. “Programming With Models: Writing Statistical Algorithms for General Model Structures With NIMBLE”. In:Journal of Computational and Graphical Statistics26.2 (Apr. 2017), 403–413.issn: 1537-2715.doi: 10.1080/10618600.2016. 1172487.url:http://dx.doi.org/10.1080/10618600.2016.1172487

    work page doi:10.1080/10618600.2016 2017

  24. [24]

    Real-time decision-making during emergency disease outbreaks

    William J. M. Probert et al. “Real-time decision-making during emergency disease outbreaks”. In:PLOS Computational Biology14.7 (July 2018), pp. 1–18.doi: 10.1371/ journal.pcbi.1006202.url:https://doi.org/10.1371/journal.pcbi.1006202

    work page doi:10.1371/journal.pcbi.1006202 2018

  25. [25]

    A Bayesian approach to identifying the role of hospital structure and staff interactions in nosocomial transmission of SARS-CoV-2

    Jessica R. E. Bridgen et al. “A Bayesian approach to identifying the role of hospital structure and staff interactions in nosocomial transmission of SARS-CoV-2”. In:Journal of The Royal Society Interface21.212 (Mar. 2024), p. 20230525.issn: 1742-5689.doi: 10.1098/rsif.2023.0525 . eprint: https://royalsocietypublishing.org/rsif/ article-pdf/doi/10.1098/rsi...

    work page doi:10.1098/rsif.2023.0525 2024

  26. [26]

    A modelling assessment for the impact of control measures on highly pathogenic avian influenza transmission in poultry in Great Britain

    Christopher N Davis et al. “A modelling assessment for the impact of control measures on highly pathogenic avian influenza transmission in poultry in Great Britain”. In: bioRxiv(2025).doi: 10.1101/2025.04.24.650264 . eprint: https://www.biorxiv. org / content / early / 2025 / 04 / 25 / 2025 . 04 . 24 . 650264 . full . pdf.url: https : //www.biorxiv.org/co...

    work page doi:10.1101/2025.04.24.650264 2025

  27. [27]

    The Journal of Chemical Physics21(6), 1087–1092 (1953) https://doi.org/10.1063/1.1699114

    Nicholas Metropolis et al. “Equation of State Calculations by Fast Computing Machines”. In:The Journal of Chemical Physics21.6 (June 1953), pp. 1087–1092.issn: 0021- 9606.doi: 10.1063/1.1699114. eprint: https://pubs.aip.org/aip/jcp/article- pdf/21/6/1087/18802390/1087\_1\_online.pdf.url: https://doi.org/10.1063/ 1.1699114

    work page doi:10.1063/1.1699114 1953

  28. [28]

    Monte Carlo Sampling Methods Using Markov Chains and Their Applications

    W. K. Hastings. “Monte Carlo Sampling Methods Using Markov Chains and Their Applications”. In:Biometrika57.1 (1970), pp. 97–109.issn: 00063444, 14643510.url: http://www.jstor.org/stable/2334940(visited on 07/01/2025)

    work page arXiv 1970

  29. [29]

    Coupling and Ergodicity of adaptive Markov chain Monte Carlo algorithms

    Gareth Roberts and Jeffrey Rosenthal. “Coupling and Ergodicity of adaptive Markov chain Monte Carlo algorithms”. In:Journal of Applied Probability44 (2007), pp. 458– 475

    work page 2007

  30. [30]

    An adaptive Metropolis algorithm

    H. Haario, E. Saksman, and J. Tamminen. “An adaptive Metropolis algorithm”. In: Bernoulli7.2 (2001), pp. 223–242

    work page 2001

  31. [31]

    SGMCMCJax: a lightweight JAX library for stochastic gradient Markov chain Monte Carlo algorithms

    Jeremie Coullon and Christopher Nemeth. “SGMCMCJax: a lightweight JAX library for stochastic gradient Markov chain Monte Carlo algorithms”. In:Journal of Open Source Software7.72 (2022), p. 4113. 38

    work page 2022