Pith. sign in

REVIEW 3 major objections 6 minor 1 cited by

laplax -- Laplace Approximations with JAX

T0 review · 3 major / 6 minor · reviewed 2026-08-06 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read JAX package laplax spans the full Laplace-approximation pipeline for Bayesian neural networks.

desk verdict laplax is a useful JAX Laplace library, but the appendix's LML formula has a sign error that undermines the calibration story. read the letter →

arxiv 2507.17013 v1 pith:LVK7AJEW submitted 2025-07-22 cs.LG cs.AI

Tobias Weber , Bálint Mucsányi , Lenard Rommel , Thomas Christie , Lars Kasüschke , Marvin Pförtner , Philipp Hennig This is my paper
classification cs.LGcs.AI
keywords LaplaceapproximationuncertaintyquantificationBayesianneuralnetworksJAXcurvatureestimationmatrix-freelinearoperatorsmodelselectionopen-sourcesoftware
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

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

The reading

This paper introduces laplax, an open-source Python package built on JAX that implements the entire Laplace approximation pipeline for turning trained neural networks into Bayesian neural networks. The package covers curvature estimation (full, diagonal, and low-rank via Lanczos or LOBPCG), construction of the posterior precision, pushing weight-space uncertainty into output space, hyperparameter calibration, and evaluation. The authors argue that a modular, purely functional design with minimal dependencies makes this the flexible counterpart in JAX to the existing PyTorch-focused implementation, and that it lowers the barrier to research on uncertainty quantification and improved Laplace methods.

What carries the argument

The central abstraction is the matrix-free curvature-vector product: every curvature estimate (Hessian or GGN) is represented as a function $v \mapsto \sum_n \nabla^2_{\theta\theta}\ell(f_\theta(x_n), y_n)v$ (or the corresponding GGN product), never as a stored matrix. On top of this, the posterior is exposed as a function of hyperparameters $(\tau, C)$, mapping a vector $v$ to $(\mathrm{Curv}(C) + \tau I)^{-1}v$, with low-rank approximations replacing the full curvature by $U S U^\top + \tau I$ and using the matrix determinant lemma for the log marginal likelihood. This design makes each component independently testable and swappable, and it is the mechanism that lets a lightweight functional API and low-level building blocks coexist.

What would settle it

Train a small model, compute the full curvature by materializing it, and compare laplax's FULL, DIAGONAL, and LANCZOS posterior means and covariances, and the resulting negative log-likelihoods, against a direct dense implementation; if the low-rank or diagonal results disagree with the dense solve on a well-conditioned problem, the numerical core of the package is not reliable.

Watch

Extended reading notes

Core claim

The central claim is that one coherent, functional library can cover every step of the Laplace approximation without forcing a particular model framework, and that doing so in JAX makes Bayesian deep-learning workflows practical and composable. On the paper's own terms, laplax's contribution is not a new approximation algorithm but a complete, replaceable pipeline: a single `laplace(...)` call turns a model, optimised weights, and data into a posterior function that maps hyperparameters to posterior covariance, with each stage exposed as an independent building block. The demonstration applies this pipeline to a regression task and a classification task, reporting negative log-likelihood and expected calibration error across curvature and calibration choices.

Load-bearing premise

The load-bearing premise is that laplax's numerical routines—matrix-free Hessian and GGN products, Lanczos and LOBPCG eigensolvers, low-rank precision updates, and pushforwards—are correct and stable on realistic models, since the paper does not supply independent verification such as a test suite or comparison against a dense computation.

Editorial extensions

If this is right

  • Any JAX-based model can be made Bayesian by calling `laplace(...)`, so the pipeline applies across JAX deep-learning frameworks without code changes.
  • Because every building block is a standalone function, researchers can replace one stage—say, the eigensolver or the pushforward—without touching the rest of the pipeline.
  • Hyperparameters such as prior precision and observation noise can be calibrated either by marginal-likelihood maximisation or by gradient-based optimisation of a downstream metric, since the whole pipeline is differentiable.
  • The same code path serves both regression and classification, with several logit-to-probability approximations (Monte Carlo bridge, Laplace bridge, mean-field variants) available.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • Editorial inference: because laplax exposes curvature as matrix-vector products, the same machinery could serve as a neutral testbed for comparing curvature approximations on identical models, a benchmark the paper does not itself run.
  • Editorial inference: the low-rank posterior routines could be plugged into continual-learning and sequential-inference loops, where updating a posterior precision without materialising a full matrix is a known bottleneck.
  • Editorial inference: the absent KFAC and exact Fisher products are natural next candidates; if added, laplax would cover nearly every curvature surrogate currently used in Laplace-approximation research.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 6 minor

Summary. The paper presents laplax, an open-source JAX library for Laplace approximations in Bayesian deep learning. It describes the library's design goals (purely functional, modular, minimal dependencies) and its main components: curvature-vector products (Hessian and GGN), curvature approximations (FULL, DIAGONAL, LANCZOS, LOBPCG), posterior precision and scale computation, linear and nonlinear pushforwards, several classification predictive approximations, hyperparameter calibration via log marginal likelihood or downstream metrics, and an evaluation interface. Two small experiments—a 1D regression task and a CIFAR-10 classification task—illustrate functionality, and an appendix gives computational details, including an implementation of FSP-Laplace. The central claim is that laplax is a functional, flexible, researcher-friendly package covering the full Laplace approximation pipeline in JAX.

Significance. If the implementation is correct, laplax fills a real gap by providing a modular JAX-based alternative to Laplace Redux, with matrix-free curvature representations, multiple eigen-solver options, and support for recent extensions such as FSP-Laplace. The paper's strengths include a publicly available codebase, a compact functional API demonstrated in Code Snippet 1, and an explicit enumeration of predictive approximations. However, the empirical evaluation is minimal (no error bars, seeds, or independent benchmarks), and the appendix contains a concrete error in the log marginal likelihood formulas used for calibration. The significance is therefore conditional on correcting that error and providing stronger evidence that the core numerical routines are correct.

major comments (3)
  1. [Appendix A.2/A.2.1/A.2.2] The log marginal likelihood evidence correction in the appendix has the opposite sign to Eq. (2) in the main text. Equation (2) gives the correction as -(1/2)log|H| + (P/2)log(2π), but Appendix A.2 writes -(1/2)(P log(2π) - log|H|), which equals -(P/2)log(2π) + (1/2)log|H|. Additionally, the 'joint' term is written only as the unnormalized training objective, omitting the likelihood normalization -N/2 log(2πσ²) and the prior normalization P/2 log(τ/(2π)), both of which depend on the calibration hyperparameters σ² and τ. Because Section 3 advertises LML-based calibration and Tables 1 and 2 report LML-calibrated results, this is load-bearing: if the code follows the appendix, the calibrated hyperparameters and all LML-based numbers are biased; if the code is correct, the appendix text must be corrected to match the implementation.
  2. [Section 4, Tables 1 and 2] The empirical demonstrations do not report seeds, repeated runs, error bars, or dataset splits, so the numerical differences between curvature and calibration settings (e.g., NLL values of 0.4482 vs. 0.5104 in Table 1) cannot be distinguished from noise. The paper explicitly disclaims comprehensiveness, which is acceptable for a software paper, but the central claim that laplax 'works' would be materially strengthened by a small test suite, a commit hash, and at least one comparison against an independent reference implementation (e.g., Laplace Redux on the same model) to validate the Lanczos/LOBPCG routines and the low-rank posterior precision update of Eq. (5).
  3. [Section 3 / Appendix A.3] The predictive formulas in Appendix A.3, specifically Eqs. (7)-(10), are presented without derivation or a precise source for the λ0 rescaling and the pairwise mean-field forms. Since these predictives are advertised as core features, the paper should either provide concise derivations or point to exact references, so that users can understand the approximations and their validity regimes.
minor comments (6)
  1. [Table 1] The column headers such as 'LML-GS-L' are defined only in the table caption; define the abbreviations in the text or in a table note so the table is self-contained.
  2. [Section 4 / Figure 2] The caption of Figure 2 references 'NLL-GD-L' and 'FSP-Laplace (Matern/Periodic)' without describing the corresponding experimental setup; clarify what was run and which settings are shown.
  3. [Reproducibility] No commit hash, Python version, JAX version, or dependency versions are provided; including these would make the software artifact reproducible and easier to audit.
  4. [Appendix C.1, Algorithm 1] The algorithm signature does not include the number of context points n_C, although line 5 uses n_C; also, the connection between R_FSP^(1), R_FSP^(2) and Eqs. (11)-(13) should be made explicit.
  5. [Section 5] The phrase 'loosely-typed' is unusual; consider using 'dynamically typed' or 'not statically typed' for clarity.
  6. [Section 3 / Appendix A.2.2] The low-rank scale update 'v ↦ τ^(-1/2) v + U diag(S)(U^T v)' with S = (S+τ)^(-1/2) - τ^(-1/2) is stated without derivation; citing the Woodbury identity used would help users implement or verify it.

Circularity Check

0 steps flagged · score 2.0 of 10

No circularity: the central claim is the software artifact itself; self-citations are contextual and non-load-bearing.

full rationale

laplax is a software/implementation paper. The derivation chain for the Laplace approximation (Eq. 1) and log marginal likelihood (Eq. 2) is standard textbook material, and the curvature approximations (Eqs. 3-5) are standard definitions of full, diagonal, and low-rank posterior precisions. The empirical NLL/ECE tables are self-reported outputs of the package itself, not predictions fitted to a subset and then re-predicted, so no fitted-input-called-prediction pattern is present. The package's central claim is backed by the artifact: code snippets, documented APIs, and a reproduction of an existing FSP-Laplace experiment from prior work. Many references are indeed to the authors' own prior work, but none is invoked as a uniqueness theorem, and none carries the load of a claim that is not also directly implemented or demonstrated here. The appendix's log-marginal-likelihood formulas (A.2, A.2.1, A.2.2) contain a sign discrepancy relative to main-text Eq. (2) and drop hyperparameter-dependent normalization constants; this is a substantive correctness risk for LML-based calibration, but it is not circularity, because the evaluation does not reduce to an input by construction. The limitations section honestly states that KFAC and Fisher curvature-vector products are not yet implemented. Overall, no circular step was identified, and the score reflects only the presence of minor, non-load-bearing self-citations.

Assumptions & free parameters 0 free parameters · 4 assumptions · 0 invented entities

The central claim of the paper rests on standard mathematics of the Laplace approximation (Taylor expansion, Gaussian linear algebra) and on the correctness of the library's implementation. No ad hoc fitted constants are introduced; hyperparameters like prior precision τ and observation noise σ² are calibrated within the method itself. No invented scientific entities are postulated. The load-bearing assumptions are the differentiability of the loss, the local optimality of the MAP estimate, and the numerical correctness of the implementation.

assumptions (4)
  • standard math The MAP estimate θ* is a local optimum of the regularized empirical risk L, so the gradient term vanishes in the Taylor expansion.
    Invoked in Section 2 to derive the Gaussian posterior approximation N(θ*, H^{-1}).
  • domain assumption The loss function ℓ and regularizer Ω are twice differentiable with respect to parameters and network outputs, so Hessian and GGN-vector products exist and can be computed via JAX autodiff.
    Required for all curvature approximations in Section 3 and Appendix A.1.
  • domain assumption The prior precision τ is positive, making the posterior precision H = Curv(C) + τI invertible in all curvature approximation pipelines.
    Used in Equations (3)-(5) and the log marginal likelihood formulas in Appendix A.
  • domain assumption For FSP-Laplace, the RKHS norm of the prediction error can be represented via context point evaluations, and the low-rank truncation bound on posterior marginal variance is valid.
    Needed for Algorithm 2 in Appendix C.2, which the paper reproduces from Cinquin et al. (2024).

how reviews work

0 comments
Cite this review

Pith. "Pith review of laplax -- Laplace Approximations with JAX." pith.science (2026). https://pith.science/paper/LVK7AJEW

@misc{pith2026250717013,
  author       = {Pith},
  title        = {Pith review of: laplax -- Laplace Approximations with JAX},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/LVK7AJEW}},
  note         = {Machine review of arXiv:2507.17013}
}
read the original abstract

The Laplace approximation provides a scalable and efficient means of quantifying weight-space uncertainty in deep neural networks, enabling the application of Bayesian tools such as predictive uncertainty and model selection via Occam's razor. In this work, we introduce laplax, a new open-source Python package for performing Laplace approximations with jax. Designed with a modular and purely functional architecture and minimal external dependencies, laplax offers a flexible and researcher-friendly framework for rapid prototyping and experimentation. Its goal is to facilitate research on Bayesian neural networks, uncertainty quantification for deep learning, and the development of improved Laplace approximation techniques.

Figures

Figures reproduced from arXiv: 2507.17013 by the authors.

Figure 1
Figure 1. Linearised Laplace approximation on a two [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 2
Figure 2. Comparison of linearised Laplace with Lanczos-approximated curvature (NLL-GD-L) (left) and FSP-Laplace [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗

Discussion (0). Sign in to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. VOLTA: The Surprising Ineffectiveness of Auxiliary Losses for Calibrated Deep Learning

    cs.LG 2026-04 unverdicted novelty 5.0 of 10

    VOLTA, consisting of a deep encoder with learnable prototypes plus cross-entropy and post-hoc temperature scaling, matches or exceeds ten UQ baselines in accuracy, achieves lower expected calibration error, and perfor...

Reference graph

Works this paper leans on

28 extracted references · 16 canonical work pages · cited by 1 Pith paper

  1. [1]

    1: function FSPL APLACE TRAIN (f, θ(0), GP (µ, Σ), PC, D, b) 2: i ← 0 3: for all minibatch B = (XB, YB) ∼ D of size b do 4: R(1) FSP(θ(i)) ← −n b Pb j=1 log p(y(j) B | f (x(j) B , θ(i))) 5: Sample context points C(i) = {C(i) j }nC j=1 i.i.d. ∼ PC 6: R(2) FSP(θ(i)) ← 1 2 (f (C(i), θ(i)) − µ(C(i)))⊤Σ(C(i), C(i))−1(f (C(i), θ(i)) − µ(C(i))) 7: θ(i+1) ← optim...

    work page 2024

  2. [2]

    Dangel, F., Eschenhagen, R., Ormaniec, W., Fernandez, A., Tatzel, L., and Kristiadi, A

    URL https://arxiv.org/abs/ 2407.13711. Dangel, F., Eschenhagen, R., Ormaniec, W., Fernandez, A., Tatzel, L., and Kristiadi, A. Position: Curvature Matrices Should Be Democratized via Linear Operators,

    arXiv

  3. [4]

    Shaving Weights with Occam's Razor: Bayesian Sparsification for Neural Networks Using the Marginal Likelihood

    URL https: //arxiv.org/abs/2402.15978. 5 laplax – Laplace Approximations with JAX Eschenhagen, R., Daxberger, E., Hennig, P., and Kristiadi, A. Mixtures of Laplace Approximations for Improved Post-Hoc Uncertainty in Deep Learning. In Bayesian Deep Learning Workshop, NeurIPS,

    work page Pith review arXiv

  4. [7]

    Fast Predictive Uncertainty for Classification with Bayesian Deep Networks

    URL https://arxiv.org/abs/2003.01227. Immer, A., Bauer, M., Fortuin, V ., R¨atsch, G., and Khan, M. E. Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning. In Proceedings of the 38th International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, 2021a. URL https://arxiv.org/abs/2104.04975. Immer, A....

    work page Pith review arXiv 2003

  5. [9]

    Learnable Uncertainty under Laplace Approximations

    URL https: //arxiv.org/abs/2010.02720. Kristiadi, A., Immer, A., Eschenhagen, R., and Fortuin, V . Promises and Pitfalls of the Linearized Laplace in Bayesian Optimization. In Fifth Symposium on Advances in Approximate Bayesian Inference,

    work page Pith review arXiv 2010

  6. [11]

    MacKay, D

    URL https://arxiv.org/ abs/2006.07584. MacKay, D. J. A Practical Bayesian Framework for Back- propagation Networks. Neural computation, 4(3):448– 472,

    arXiv 2006

  7. [12]

    Mucs´anyi, B., Kirchhof, M., and Oh, S

    URL https://arxiv.org/abs/ 2406.05072. Mucs´anyi, B., Kirchhof, M., and Oh, S. J. Benchmarking Uncertainty Disentanglement: Specialized Uncertainties for Specialized Tasks. In Proceedings of the 38th Inter- national Conference on Neural Information Processing Systems (NeurIPS),

    arXiv

  8. [13]

    Pinder, T

    URL https://arxiv.org/abs/2402.00809. Pinder, T. and Dodd, D. GPJax: A Gaussian process framework in JAX. Journal of Open Source Software , 7(75):4455,

    arXiv

Show all 28 references
  1. [14]

    URL https://doi.org/10.21105/joss.04455

    doi: 10.21105/joss.04455. URL https://doi.org/10.21105/joss.04455. Rasmussen, C. and Ghahramani, Z. Occam’s Razor. In Advances in Neural Information Processing Systems (NeurIPS), volume 13,

    doi

  2. [16]

    Sliwa, J., Schneider, F., Bosch, N., Kristiadi, A., and Hennig, P

    URL https://arxiv.org/abs/ 2406.03334. Sliwa, J., Schneider, F., Bosch, N., Kristiadi, A., and Hennig, P. Efficient Weight-Space Laplace-Gaussian Filtering and Smoothing for Sequential Deep Learning,

    arXiv

  3. [17]

    6 laplax – Laplace Approximations with JAX Sun, S., Zhang, G., Shi, J., and Grosse, R

    URL http://arxiv.org/abs/2410.06800. 6 laplax – Laplace Approximations with JAX Sun, S., Zhang, G., Shi, J., and Grosse, R. Functional Variational Bayesian Neural Networks,

  4. [19]

    van der Ouderaa, T

    URL http://arxiv.org/abs/ 2310.06131. van der Ouderaa, T. F. A., Nagel, M., van Baalen, M., Asano, Y . M., and Blankevoort, T. The llm surgeon. In International Conference of Learning Representations (ICLR),

    arXiv

  5. [20]

    URL https://arxiv.org/abs/ 2312.17244. 7 laplax – Laplace Approximations with JAX A Computational details In the following section, we provide additional details of how different computations are performed and list more available building blocks of laplax. A.1 Curvatures The p...

    2024 arXiv

  6. [21]

    (8) • MEAN FIELD 0 PREDICTIVE

    approximates the true predictive as follows: ˆp := 1 ˜σ2 1 − 2 C + e ˜µ C2 PC c=1 e−˜µc PC c=1 1 ˜σ2c 1 − 2 C + e ˜µ C2 PC c′=1 e−˜µc′ (7) where ˜µ2 := s p C/2 PC c=1 σ2c µ, ˜σ2 := p C/2 PC c=1 σ2c σ2. (8) • MEAN FIELD 0 PREDICTIVE . A zeroth-order mean-field (probit-style) ap...

    2021

  7. [22]

    B Applications and extensions of the Laplace approximation Section 4 discusses Laplace approximation with the goal of calibrated predictive uncertainty

    or O(C) work; and the Laplace bridge often gives the best calibrated probabilities when variances are large. B Applications and extensions of the Laplace approximation Section 4 discusses Laplace approximation with the goal of calibrated predictive uncertainty. Here, a variety...

    2021

  8. [23]

    (2025) lift the method to the setting of operator learning

    and Magnani et al. (2025) lift the method to the setting of operator learning. Notable work shows that adding some weight space uncertainty fixes overconfidence in classification networks (Kristiadi et al., 2020). In comparison with other uncertainty quantification methods Lap...

    2025

  9. [24]

    for learning layerwise equivariance (van der Ouderaa et al., 2023)

    or model selection, e.g. for learning layerwise equivariance (van der Ouderaa et al., 2023). The analytic uncertainty structure provided by Laplace has also been used to apply filtering techniques to neural network learning with the goal of online/continual learning (Ritter et...

    2023

  10. [25]

    Various other applications exist and this non-extensive list aimed only at provided some pointers for potential use cases

    or to explore Bayesian optimization (Kristiadi et al., 2023). Various other applications exist and this non-extensive list aimed only at provided some pointers for potential use cases. 10 laplax – Laplace Approximations with JAX C FSP-Laplace The BNN literature offers only a m...

    2023

  11. [26]

    This yields a more refined MAP estimate and well-calibrated epistemic uncertainties when prior knowledge is available

    extends the linearised Laplace approximation by placing interpretable Gaussian Process (GP) priors directly in function space, thereby overcoming the non-interpretability of the weight-space prior. This yields a more refined MAP estimate and well-calibrated epistemic uncertain...

    2024

  12. [1992]

    This regulariser (eq

    and standard linearised Laplace approximation (Immer et al., 2021b) in that it requires training the model with a Reproducing Kernel Hilbert Space (RKHS) HΣ regulariser. This regulariser (eq. (12)) is added to the negative log-likelihood (eq. (11)) to form the FSP objective fu...

    2024

  13. [2018]

    Roy, H., Miani, M., Ek, C

    URL http://arxiv.org/abs/1805.07810. Roy, H., Miani, M., Ek, C. H., Hennig, P., Pf ¨ortner, M., Tatzel, L., and Hauberg, S. Reparameterization Invari- ance in Approximate Bayesian Inference. In Advances in Neural Information Processing Systems (NeurIPS) , volume 37,

    arXiv

  14. [2019]

    Tran, B.-H., Rossi, S., Milios, D., and Filippone, M

    URL https://arxiv.org/abs/1903.05779. Tran, B.-H., Rossi, S., Milios, D., and Filippone, M. All You Need is a Good Functional Prior for Bayesian Deep Learning. Journal of Machine Learning Research , 23 (74):1–56,

    1903 arXiv

  15. [2020]

    Kristiadi, A., Hein, M., and Hennig, P

    URL https: //arxiv.org/abs/2002.10118. Kristiadi, A., Hein, M., and Hennig, P. Learnable Uncer- tainty under Laplace Approximations. In Proceedings of the Thirty-Seventh Conference on Uncertainty in Ar- tificial Intelligence (UAI), volume 161 of Proceedings of Machine Learning...

    2002 arXiv

  16. [2021]

    Fortuin, V

    URL https: //arxiv.org/abs/2111.03577. Fortuin, V . Priors in Bayesian Deep Learning: A Review. International Statistical Review , 90(3):563–591,

    arXiv

  17. [2022]

    Hobbhahn, M., Kristiadi, A., and Hennig, P

    URL https://arxiv.org/abs/2105.06868. Hobbhahn, M., Kristiadi, A., and Hennig, P. Fast Predictive Uncertainty for Classification with Bayesian Deep Net- works. In Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence (UAI), volume 180 of Proceed...

    arXiv

  18. [2023]

    Lu, Z., Ie, E., and Sha, F

    URL https: //arxiv.org/abs/2304.08309. Lu, Z., Ie, E., and Sha, F. Mean-Field Approximation to Gaussian-Softmax Integral with Application to Uncer- tainty Estimation,

    arXiv

  19. [2024]

    URL https://docs

    doi: 10.1145/3620665.3640366. URL https://docs. pytorch.org/assets/pytorch2-2.pdf. Botev, A. and Martens, J. KFAC-JAX,

  20. [2025]

    Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P

    URL https://arxiv.org/abs/2501.19183. Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., and Hennig, P. Laplace Redux – Effortless Bayesian Deep Learning. In Advances in Neural Infor- mation Processing Systems (NeurIPS), volume 34, 2021a. URL https://arxiv.or...

    arXiv

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.