REVIEW 2 major objections 5 minor 1 cited by
JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference
T0 review · 2 major / 5 minor · reviewed 2026-08-03 · deepseek-v4-flash
Pith's one-line read JADAI shows that adaptive experimental design and Bayesian posterior inference can be trained together, end to end, around a single objective that rewards reducing posterior error at every step of an experiment.
desk verdict A well-engineered joint policy+inference framework with strong empirical results, whose main theoretical hinge for diffusion models is an unproven proxy — worth serious refereeing, not a desk reject. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the incremental posterior-loss utility u_T(θ, h_{0:T}) = Σ_{t=0}^T (ℓ*_{t-1} - ℓ_t), with ℓ_{-1}=0 and ℓ*_t = detach(ℓ_t), where ℓ_t is the per-step posterior estimation loss conditioned on the history summary. For a normalized density model this sum telescopes exactly and reproduces the lower bound on total expected information gain; for implicit models it is used as a proxy. A second mechanism is the detached-history design pass: the policy receives a summary with gradients stopped, so the history network is trained only through the posterior losses while the policy is trained through the observations its designs produce; this removes nested backpropagation throu
What would settle it
In a small Gaussian adaptive-design problem where the true expected information gain is computable in closed form, train JADAI with a deliberately miscalibrated posterior estimator so that posterior loss decreases while the posterior is inaccurate, then compare the resulting policy's estimated information gain or final inference accuracy against a policy trained on the true objective; a measurable gap would show the surrogate is not tracking information gain.
Extended reading notes
Core claim
The paper's central discovery is an objective identity: for a normalized-density posterior estimator, the variational lower bound on total expected information gain telescopes into a sum over experiment steps of per-step differences in posterior loss, log q(θ|h_t) - log q(θ|h_{t-1}). JADAI builds on this by using the same sum, u_T = Σ_{t=0}^T (ℓ*_{t-1} - ℓ_t), as a generic training utility even when the posterior estimator is an implicit diffusion or flow-matching model whose log-density is unavailable, with the previous loss value detached. Although the variational bound interpretation no longer applies, the objective still drives rollouts along which posterior loss decreases at every step.
Load-bearing premise
The claim rests on the detached per-step posterior-loss difference (ℓ_{t-1} - ℓ_t) being a faithful proxy for expected information gain when the posterior is an implicit diffusion or flow-matching model; if that proxy misaligns with true information gain, the policy can be optimal for the surrogate while poor for the actual experiment.
Editorial extensions
If this is right
- At any point in an experiment, JADAI provides an approximate posterior over parameters, not only after the full horizon, enabling active data-acquisition workflows.
- A policy trained to minimize posterior losses over a longer horizon remains competitive at shorter horizons, so longer training horizons can improve early decisions.
- Diffusion- and flow-matching-based posterior estimators can be inserted into the same design objective, making high-dimensional multimodal posteriors tractable in adaptive design.
- When the design policy is replaced by a random policy, performance falls to baseline levels, indicating that most of the gain comes from the learned policy rather than the posterior network.
- Test-time rollouts take milliseconds on low-dimensional tasks and are roughly an order of magnitude faster than prior methods on the high-dimensional benchmark.
Reading between the lines
- Beyond the paper: if the per-step loss proxy is as reliable as the experiments suggest, the same utility should extend to other implicit posterior families, such as consistency models or adversarial generators, without re-deriving an information-theoretic objective.
- Beyond the paper: because the utility is defined entirely through posterior losses, JADAI could be coupled with gradient-free design optimization or policy-gradient estimators to handle black-box simulators; the paper only tests differentiable simulators.
- Beyond the paper: the paper notes but does not implement sharing a transformer backbone across policy and history networks; doing so could improve sample efficiency on harder design spaces.
- Beyond the paper: a testable consequence is that deliberately miscalibrating the posterior estimator, while keeping the posterior loss artificially low, should degrade the learned designs; monitoring this gap would validate the surrogate.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces JADAI, a framework that jointly amortizes Bayesian adaptive design and posterior inference by training a design policy, a history summary network, and a generative posterior estimator end-to-end. The training objective is the generic utility u_T = Σ_{t=0}^T (ℓ*_{t-1} − ℓ_t), where ℓ_t is a per-step posterior loss and ℓ* denotes a detached copy. For a normalized density estimator (e.g., a normalizing flow), this utility is a telescoping form of the Barber–Agakov lower bound on total expected information gain (Eqs. 12–14). For diffusion and flow-matching estimators, ℓ_t is instead a weighted score-matching or flow-matching regression loss (Eqs. 27 and 34), and the paper explicitly states that the Barber–Agakov interpretation no longer applies and appeals to experimental performance. Experiments on Location Finding, CES, and MNIST Image Discovery report competitive or superior sPCE and SSIM/NRMSE relative to recent baselines, with ablations in Appendix B.
Significance. If the empirical findings hold, JADAI is a practically valuable contribution: it extends amortized adaptive design to high-dimensional, multimodal posteriors without reinforcement-learning value networks, and it provides posterior estimates at every acquisition step. The normalized-density case gives a clean connection to TEIG, and the experiments are substantial, including random-policy controls that support the claim that the performance gap in Image Discovery comes primarily from the learned policy (Section 5.3, Figure 6). The paper is also transparent about limitations, including the heuristic status of the implicit-model objective (Section 3.2) and the need for differentiable simulators (Section 6). The main weakness is theoretical: for the diffusion/flow setting that powers the flagship high-dimensional results, the objective is not shown to track or bound expected information gain, so the paper's formal claim is not established for its primary regime.
major comments (2)
- [Section 3.2, Eqs. (16), (27), (34)] The central claim that JADAI maximizes information gain rests on treating ℓ_{t-1}−ℓ_t as incremental information. For a normalized density estimator, Eq. (14) is an identity and the Barber–Agakov bound is clean. For diffusion models (Eq. 27) and flow matching (Eq. 34), ℓ_t is a regression loss against noise/velocity targets; ℓ_{t-1}−ℓ_t is not a log-density ratio and is not shown to bound, estimate, or track incremental EIG. The manuscript concedes this in Section 3.2 and appeals to experiments. This is load-bearing because the policy is trained through this difference. A policy can reduce ℓ_t by steering to histories that are easy for the current posterior estimator without increasing true information gain unless additional calibration or monotonicity holds. I ask for either (a) a formal statement of conditions under which ℓ_{t-1}−ℓ_t is a valid surrogate for incremental information, or
- [Section 5.3, Table 2, Figure 4] For the high-dimensional Image Discovery benchmark, the central evaluation is SSIM/NRMSE against the ground-truth image, not an information-theoretic measure of posterior quality. The random-policy controls in Figure 6 are informative, but SSIM is a perceptual reconstruction metric and can be high even if the posterior is miscalibrated. To support the claim that the learned policy improves inference rather than only reconstruction, please report a posterior calibration or coverage diagnostic at each measurement step, for example simulation-based calibration or conditional coverage on a validation set.
minor comments (5)
- [Section 3.2, Eq. (15)] The notation suggests a score with respect to θ_τ, but in the diffusion derivation the score is with respect to the latent z_τ; please align notation with Eqs. (23)–(24).
- [Section 5.2 / Appendix C.3] The heading 'Constant Elasticity of Substitutions' should be 'Constant Elasticity of Substitution' for consistency with the CES benchmark name.
- [Table 1] Entries marked '−' are not defined in the caption; please specify whether they indicate 'not reported' or 'not run'.
- [Section 5.1, Figure 2] The corner plot axis labels are difficult to read at the current resolution; please ensure they are legible in the final version.
- [Section 6] The requirement of a differentiable simulator is stated clearly, but it should also be mentioned in the abstract or introduction, since it materially constrains the scope of the method.
Circularity Check
No significant circularity: the JADAI objective is a stated heuristic surrogate, not a hidden re-use of its evaluation targets, and the self-citations are background only.
full rationale
The paper's central derivation replaces the exact Barber–Agakov TEIG lower bound, which telescopes only for normalized density estimators (Eqs. 12–14), with the generic loss-difference utility u_T = sum(ℓ*_{t-1} − ℓ_t) for diffusion and flow-matching posterior estimators. The authors explicitly acknowledge that for these implicit models “the Barber–Agakov interpretation is no longer the same” (Section 3.2) and justify the surrogate empirically rather than theoretically. This is a heuristic modeling choice and a limitation, not circular reasoning: the training objective is not defined in terms of the evaluation metrics (sPCE, SSIM, NRMSE), and those metrics are computed externally from the simulator and ground-truth images. The telescoping observation that u_T forward-equals −ℓ_T is explicitly handled by gradient detachment so the method intentionally optimizes all intermediate losses rather than only the final one; this is an algorithmic design, not a circular reduction. The self-citations (e.g., Arruda et al. 2025, Radev et al. 2020, Schmitt et al. 2024) are used for background on diffusion-based SBI and summary networks, not as load-bearing evidence for the paper's new claims. No fitted parameter is relabeled as a prediction, no uniqueness theorem is imported from the authors' prior work, and no ansatz is smuggled in via citation. The derivation chain is therefore self-contained with respect to circularity, and the appropriate finding is no significant circularity.
Assumptions & free parameters
free parameters (3)
- Exploration schedule ρ_n =
Cosine decay from 1 to 0 over first 30% of ID training; ρ=0 for LF/CES after pretraining
- Rollout length schedule R(n) =
Linear from 2 to 6 for ID; fixed T for LF/CES
- Diffusion noise schedule =
Cosine schedule λ(τ) = -2 log(tan(πτ/2))
assumptions (3)
- domain assumption The simulator is differentiable with respect to design variables and random states can be reparameterized.
- ad hoc to paper The per-step posterior loss ℓ_t is a monotonically valid proxy for information gained about θ.
- domain assumption The summary network η_ω produces approximately sufficient statistics even when the history is detached before the policy.
Cite this review
Pith. "Pith review of JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference." pith.science (2026). https://pith.science/paper/YJHUBZ4A
@misc{pith2026251222999,
author = {Pith},
title = {Pith review of: JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference},
year = {2026},
howpublished = {\url{https://pith.science/paper/YJHUBZ4A}},
note = {Machine review of arXiv:2512.22999}
}
read the original abstract
We consider problems of parameter estimation where design variables can be actively optimized to maximize information gain. To this end, we introduce JADAI, a framework that jointly amortizes Bayesian adaptive design and inference by training a policy, a history network, and an inference network end-to-end. The networks minimize a generic loss that aggregates incremental reductions in posterior error along experimental sequences. Inference networks are instantiated with diffusion-based posterior estimators that can approximate high-dimensional and multimodal posteriors at every experimental step. Across standard adaptive design benchmarks, JADAI achieves superior or competitive performance.
Figures
Figures from the paper (4 more)
Forward citations
Cited by 1 Pith paper
-
Bayesian Experimental Design via Score Matching
SCOREBED isolates EIG double intractability in a policy-independent score-matching stage, then trains design policies with a singly intractable gradient estimator, enabling cheap multi-policy selection.
Reference graph
Works this paper leans on
-
[11]
5 Wang, Z. and Li, Q. Information content weighting for perceptual image quality assessment.IEEE Transactions on Image Processing, 20(5):1185–1198, 2011. doi: 10.1 109/TIP.2010.2092435. 16 Wang, Z., Bovik, A., Sheikh, H., and Simoncelli, E. Image quality assessment: from error visibility to structural similarity.IEEE Transactions on Image Processing, 13 (...
arXiv 2011
-
[48]
13 Kleinegesse, S. and Gutmann, M. U. Bayesian experimental design for implicit models by mutual information neural estimation.arXiv [stat.ML], February 2020. URL http: //arxiv.org/abs/2002.08129. 5 Kuck, H., de Freitas, N., and Doucet, A. Smc samplers for bayesian optimal nonlinear design. In2006 IEEE Nonlinear Statistical Signal Processing Workshop, pp....
arXiv 2020
-
[72]
1 Rainforth, T., Foster, A., Ivanova, D
URL http://dx.doi.org/10.1371/journal.pcbi.1009472. 1 Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, F. Modern bayesian experimental design.Statistical Science, 39(1):100–114, 2024. 2, 3 Rezende, D. and Mohamed, S. Variational inference with normalizing flows. InInternational conference on ma- chine learning, pp. 1530–1538. PMLR, 2015. 2 S...
-
[685]
1, 5, 13, 14 Barber, D. and Agakov, F. The IM algorithm: a variational approach to Information Maximization.Neural Informa- tion Processing Systems, pp. 201–208, December 2003. URL https://papers.nips.cc/paper files/paper/2003/file/ a6ea8471c120fe8cc35a2954c9b9c595-Paper.pdf. 3, 4, 5 Blau, T., Bonilla, E. V ., Chades, I., and Dezfouli, A. Op- timizing seq...
arXiv 2003
-
[924]
Method Details A.1
8 12 JADAI: Joint Adaptive Design and Amortized Inference A. Method Details A.1. Diffusion model background In the following, we provide a brief overview of our diffusion model (Ho et al., 2020; Song & Ermon, 2019; Kingma & Gao,
2020
-
[1093]
22 Huan, X. and Marzouk, Y . M. Simulation-based optimal bayesian experimental design for nonlinear systems.Jour- nal of Computational Physics, 232(1):288–317, 2013. 5 Huang, D., Guo, Y ., Acerbi, L., and Kaski, S. Amortized bayesian experimental design for decision-making.Ad- vances in Neural Information Processing Systems, 37: 109460–109486, 2024. 6, 8 ...
arXiv 2013
-
[2001]
Deep residual learn- ing for image recognition
5 He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. InProceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016. 22 Hedman, M., Ivanova, D. R., Guan, C., and Rainforth, T. Step-dad: Semi-amortized policy-based bayesian exper- imental design.International conference on machine lear...
arXiv 2016
-
[2015]
doi: https://doi.org/10.1016/j.ae ue.2014.09.002
ISSN 1434-8411. doi: https://doi.org/10.1016/j.ae ue.2014.09.002. URL https://www.sciencedirect.com/sc ience/article/pii/S1434841114002519. 16 Hahn, C., Lemos, P., Parker, L., R ´egaldo-Saint Blancard, B., Eickenberg, M., Ho, S., Hou, J., Massara, E., Modi, C., Moradinezhad Dizgah, A., and Spergel, D. Cos- mological constraints from non-gaussian and nonli...
Show all 14 references
-
[2018]
21 Song, Y
URL https://arxiv.org/abs/1708.07120. 21 Song, Y . and Ermon, S. Generative Modeling by Estimating Gradients of the Data Distribution.arXiv [cs.LG], July
-
[2019]
2, 13 Søgaard, J., Krasula, L., Shahid, M., Temel, D., Brunnstr¨om, K., and Razaak, M
URL http://arxiv.org/abs/1907.05600. 2, 13 Søgaard, J., Krasula, L., Shahid, M., Temel, D., Brunnstr¨om, K., and Razaak, M. Applicability of existing objective metrics of perceptual quality for adaptive video streaming. volume 2016, 02 2016. doi: 10.2352/ISSN.2470-1173. 2016.1...
1907 arXiv
-
[2023]
r1GL/talcOUHCzu7TmEknpH/Tl8=
used to approximate the posterior p(θ| ·)≈q ψ(θ| ·). More details on diffusion models in a general setting can be found in (Karras et al., 2022) and, for simulation-based inference, in (Arruda et al., 2025). A diffusion model learns how to gradually denoise a sample from a bas...
2022
-
[2024]
1, 5 Sharrock, L., Simons, J., Liu, S., and Beaumont, M
URL https://openreview.net/forum?id=IHjKpKlj yH. 1, 5 Sharrock, L., Simons, J., Liu, S., and Beaumont, M. Sequen- tial neural score estimation: Likelihood-free inference with conditional score based diffusion models.arXiv [stat.ML], October 2022. URL http://arxiv.org/abs/2210 ...
2022
-
[2025]
doi: 10.1146/annure v-statistics-112723-034123
ISSN 2326-8298,2326-831X. doi: 10.1146/annure v-statistics-112723-034123. URL http://dx.doi.org/10.11 46/annurev-statistics-112723-034123. 1 Zhang, X., Huang, D., Kaski, S., and Martinelli, J. PABBO: Preferential amortized black-box optimization.arXiv [stat.ML], March 2025. UR...
2025 doi
-
[2138]
URL http://dx.doi.org/10.1016/j.cma.2024.117457
doi: 10.1016/j.cma.2024.117457. URL http://dx.doi.org/10.1016/j.cma.2024.117457. 5 Drovandi, C. C., McGree, J. M., and Pettitt, A. N. A se- quential monte carlo algorithm to incorporate model un- certainty in bayesian sequential design.Journal of Com- putational and Graphical ...
2024
Reviewed August 3, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.