Learning to Bound: Monotonic, Calibrated Neural Upper Bounds for MPE Inference

Akshay Vyas; Hailiang Dong; Shivvrat Arya; Brij Malhotra; Tahrima Rahman; Vibhav Gogate; Nicholas Ruozzi

cs.LGOct 29, 2026

Learning to Bound: Monotonic, Calibrated Neural Upper Bounds for MPE Inference

Authors: Akshay Vyas, Hailiang Dong, Shivvrat Arya, Brij Malhotra, Tahrima Rahman, Vibhav Gogate, Nicholas Ruozzi

Organizations: The University of Texas at Dallas

Abstract

Probabilistic graphical models (PGMs) compactly encode structured joint distributions, but Most Probable Explanation (MPE), which finds the highest-probability assignment to unobserved variables, remains computationally intractable in large, densely structured models. Anytime systematic search, such as branch-and-bound (B&B), is a popular approach, but its efficiency depends on tight admissible upper bounds that are often prohibitively expensive to compute using classical methods like Mini-Bucket Elimination (MBE). We introduce Learning to Bound (L2B), a data-driven framework that amortizes bound computation cost by training a neural predictor. L2B targets two key properties required for effective search bounds: admissibility, which is achieved via statistical calibration techniques, and progressive tightening as evidence accumulates, which is enforced through structural architectural constraints that ensure monotonicity by construction. Across challenging benchmarks, L2B delivers tighter, monotonic bounds at orders-of-magnitude lower cost on every PGM family, and yields measurable improvements when integrated with branch-and-bound on high-treewidth instances.

Figures & tables

Figure 1

Dataset	MPE ⋆	JGLP	$\gamma=0.95$				$\gamma=0.999$
		JGLP	SMNN (Ours)		UMLP (Ours)		SMNN (Ours)		UMLP (Ours)
		Relative Gap	Cov	Relative Gap	Cov	Relative Gap	Cov	Relative Gap	Cov	Relative Gap
BN_9	-23.81*	0.3649 $\pm$ 0.1745	0.97	0.0683 $\pm$ 0.0351	0.99	0.0696 $\pm$ 0.0357	1.00	0.1285 $\pm$ 0.0396	1.00	0.1331 $\pm$ 0.0414
BN_12	-19.71*	0.3438 $\pm$ 0.1896	0.97	0.1113 $\pm$ 0.0615	0.93	0.0727 $\pm$ 0.0365	0.99	0.1989 $\pm$ 0.0858	1.00	0.1447 $\pm$ 0.0423
BN_14	-26.51*	0.3411 $\pm$ 0.1954	0.92	0.0656 $\pm$ 0.0385	0.95	0.0667 $\pm$ 0.0365	0.99	0.1165 $\pm$ 0.0469	1.00	0.1204 $\pm$ 0.0419
BN_30	-14.66	0.3017 $\pm$ 0.1558	1.00	0.0143 $\pm$ 0.0137	0.95	0.0187 $\pm$ 0.0104	1.00	0.0897 $\pm$ 0.0142	1.00	0.0846 $\pm$ 0.0206

Table 1: Calibrated upper-bound quality on representative PGMs at

i

-bound

{=}4

, with stratified-conformal calibration,

\gamma

, on test data. MPE ⋆ : true MPE log-likelihood value. Relative Gap :

|\text{bound}-\text{true}|/|\text{true}|

. Cov : fraction of queries with

\text{bound}\geq\text{true MPE}

Family	Pairs	UMLP		SMNN (Ours)
Family	Pairs	Raw	Calib.	Raw	Calib.
Bayesian networks	600	203	213	0	0
Grids	700	145	150	0	0
Segmentation	300	65	66	0	0
Random uniform	700	110	111	0	0
WCSP	300	117	114	0	0

Table 2: Path-monotonicity violations, aggregated by family. For each of the 30 PGMs we sample 100 parent

\to

child evidence pairs (child

=

parent plus one assigned variable) and count pairs with

g(\text{child})>g(\text{parent})

, which breaks admissibility of the heuristic along an AOBB path. Raw is the direct network output; Calib. is the deployed bound after monotone-stratified conformal calibration at

\gamma=0.999

. SMNN is monotone by construction and stays at zero end-to-end; the unconstrained UMLP violates on every family.

Instance	$w^{*}$	Method	$\gamma=0.95$			$\gamma=0.999$
Instance	$w^{*}$	Method	1 m	5 m	10 m	1 m	5 m	10 m
rus2_1	105	SMNN	+0.0000	+0.0225	+0.0225	-0.2629	+0.0225	+0.0225
rus2_1	105	UMLP	-0.1416	-0.1541	-0.1399	-0.3020	-0.1490	-0.1399
rus3_1	96	SMNN	+0.0013	+0.0013	+0.0013	+0.0017	+0.0017	+0.0017
rus3_1	96	UMLP	+0.0017	+0.0020	+0.0020	+0.0017	+0.0017	+0.0020
rus3_3	98	SMNN	+0.0032	+0.0032	+0.0032	+0.0032	+0.0032	+0.0032

Table 3: Anytime AOBB MPE improvement over MBE on high-treewidth instances.

\Delta=(\text{NN MPE}-\text{MBE MPE})/|\text{MBE MPE}|

; positive values indicate the NN-augmented search found a strictly better MPE within the time budget. Bold marks

\Delta>10^{-4}

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 1 — Figure 1: (a) SMNN architecture: positive weights, batch normalization without affine parameters, saturating activation, and output negation give monotonicity in evidence by construction . (b) Empirical effect along AOBB search paths: SMNN exhibits zero violations across 30 benchmark PGMs; an unconstrained MLP violates monotonicity on 28.6% of parent–child pairs.

Figure 2: L2B pipeline for calibrated neural upper bounds. Top: A binary-encoded partial assignment $x_{E}$ is processed through MLP $f_{\theta}$ with three [Linear $\to$ BatchNorm $\to$ ReLU $\to$ Dropout] blocks, producing raw prediction $\hat{h}(x_{E})$ . Middle: Stratified conformal calibration computes nonconformity scores $s_{i}=h^{*}(x_{E_{i}})-\hat{h}(x_{E_{i}})$ , stratifies by $|x_{E}|$ into $K$ bins, and extracts per-bin quantiles $\hat{q}^{1}>\cdots>\hat{q}^{K}$ adapted to prediction difficulty. Bottom: At inference, the calibrated bound $g(x_{E})=\hat{h}(x_{E})+\hat{q}^{B(x_{E})}$ satisfies per-stratum coverage $\mathbb{P}[g(x_{E})\geq h^{*}(x_{E})\mid B(x_{E})=b]\geq\gamma$ for all $b$ .

Dataset	$n$	$m$	$d_{\max}$	$w^{*}$	Source
Bayesian Networks
BN_9, BN_12, BN_14	90–155	90–115	2	12–18	UAI
BN_30, BN_32	1156–1444	1156–1444	2	15–22	UAI
BN_70	2315	2315	36	-	UAI
Grid Models
grid10 $\times$ 10	100	300	2	10	Ising

Appendix

Table 4: Benchmark PGM statistics.

Figure 3: Comparison of calibration methods: (a) Gaussian, (b) standard conformal, (c) stratified conformal prediction.

	Gaussian	Conformal	Stratified Conformal
Calibrated bound	$g(x_{E})=\hat{h}(x_{E})+z_{\gamma}\cdot\sigma$	$g(x_{E})=\hat{h}(x_{E})+\hat{q}_{\gamma}$	$g(x_{E})=\hat{h}(x_{E})+\hat{q}_{\gamma}^{B(x_{E})}$
Parameters	$\sigma=\sqrt{\mathbb{E}[(\hat{h}-h^{*})^{2}]}$	$\hat{q}_{\gamma}=\text{Quantile}_{\gamma}(\{s_{i}\})$	$\hat{q}_{\gamma}^{b}=\text{Quantile}_{\gamma}(\{s_{i}\}_{i\in b})$
	$z_{\gamma}$ = Gaussian quantile	$s_{i}=h^{*}(x_{E_{i}})-\hat{h}(x_{E_{i}})$	$B(x_{E})=\text{bin}(\|x_{E}\|)$
Assumption	Gaussian errors	Exchangeability	Exchangeability
Adapts to $\|x_{E}\|$	No	No	Yes
Guarantee	Asymptotic	Finite-sample	Finite-sample (per-stratum)

Appendix

Table 5: Comparison of calibration methods. Stratified conformal adapts margins to evidence size

|x_{E}|

, yielding tighter bounds at the same coverage level.

Dataset	#Variables	#Factors	Max domain
BN_9	105	105	2
BN_12	90	90	2
BN_14	115	115	2
BN_30	1156	1156	2
BN_32	1444	1444	2
BN_70	2315	2315	36

Appendix

Table 6: PGM details.

Figure 4: Per-query bound computation time across seven PGM families, grouped by method (log scale). Bars report mean wall-clock time per test query, with

n

instances per family in parentheses: BN (6), Grid (7), Seg (3), rus (7), WCSP (3), Object Detection (2), Protein Folding (2). MB and JGLP run in seconds to minutes; SMNN (Section 4 ) and UMLP compute bounds in milliseconds—3–5 orders of magnitude faster—on every family. SMNN is marginally slower than UMLP due to its positive-weight reparameterization and batch-norm-without-affine layers, but both remain below

10^{-2}

Instance	$w^{*}$	Method	$\gamma=0.95$			$\gamma=0.999$
Instance	$w^{*}$	Method	1m	5m	10m	1m	5m	10m
rus2_1	105	SMNN	+0.0000	+0.0225	+0.0225	-0.2629	+0.0225	+0.0225
rus2_1	105	UMLP	-0.1416	-0.1541	-0.1399	-0.3020	-0.1490	-0.1399
rus3_1	96	SMNN	+0.0013	+0.0013	+0.0013	+0.0017	+0.0017	+0.0017
rus3_1	96	UMLP	+0.0017	+0.0020	+0.0020	+0.0017	+0.0017	+0.0020
rus3_3	98	SMNN	+0.0032	+0.0032	+0.0032	+0.0032	+0.0032	+0.0032

Appendix

Table 11: Anytime AOBB MAP improvement over MBE, high-treewidth bucket (21 PGMs).

\Delta=(\text{NN MAP}-\text{MBE MAP})/|\text{MBE MAP}|

; positive values mean the NN-augmented search found a strictly better MAP within the wall-clock budget. SMNN and UMLP run on identical evidence queries and AOBB orderings per PGM, isolating the architectural effect. Each cell is the mean over 3 candidate queries at two calibration levels (deep budget 600 s); bold marks

\Delta>10^{-4}

. Negative UMLP entries on rus and BN_70 reflect a configuration constraint discussed in Section 5 .

Instance	$w^{*}$	Method	$\gamma=0.95$			$\gamma=0.999$
Instance	$w^{*}$	Method	1m	5m	10m	1m	5m	10m
BN_9	22	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000
BN_9	22	UMLP	-0.0125	-0.0125	-0.0125	-0.0125	-0.0125	-0.0125
BN_12	28	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000
BN_12	28	UMLP	-0.0153	-0.0153	-0.0153	-0.0153	-0.0153	-0.0153
BN_14	26	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000

Appendix

Table 12: Anytime AOBB MAP improvement over MBE, moderate / low-treewidth bucket (9 PGMs). Setup and notation as in Table 11 .

Instance	$w^{*}$	Method	$\gamma=0.95$			$\gamma=0.999$
Instance	$w^{*}$	Method	20s	40s	1m	20s	40s	1m
BN_9	22	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000
BN_9	22	UF	-0.0125	-0.0125	-0.0125	-0.0125	-0.0125	-0.0125
BN_12	28	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000
BN_12	28	UF	-0.0153	-0.0153	-0.0153	-0.0153	-0.0153	-0.0153
BN_14	26	SMNN	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000	+0.0000

Appendix

Table 13: Anytime AOBB MAP improvement over MBE for SMNN (Mode 7) and UMLP on the moderate/low-treewidth bucket (9 PGMs, deep timeout 60 s). Same protocol as the high-treewidth table: identical evidence query and AOBB configuration per PGM, two calibration levels side-by-side. Reported on the 3 candidate queries per PGM (mean over the 3); bold = strict win over MBE. BN_9/12/14 and grid10x10 saturate at

\Delta=0

because MBE-i4 already finds the optimum at root, leaving no headroom for either NN.

Dataset	$n_{\text{pairs}}$	UMLP violations		SMNN violations ( monotone calib )
Dataset	$n_{\text{pairs}}$	Raw	Calibrated	Raw	Calibrated
wcsp_11	100	61/100	61/100	0/100	0/100
wcsp_12	100	41/100	41/100	0/100	0/100
wcsp_15	100	15/100	12/100	0/100	0/100
rus_2_1	100	9/100	9/100	0/100	0/100
rus_3_1	100	13/100	13/100	0/100	0/100

Appendix

Table 14: Monotonicity-along-AOBB-paths verification. For each PGM we sample 100 random parent

\to

child evidence pairs (child = parent + one additional variable) and count pairs where bound(child)

>

bound(parent), which would violate the admissibility of upper-bounding heuristics in AOBB. Raw = direct neural prediction (the architectural property). Calibrated = stratified-conformal at

\gamma=0.999

(the deployed bound). SMNN raw violations are zero by construction (monotone-by-design with one-hot encoding); the previous stratified-conformal wrapper introduced 63/3000 calibrated violations from per-bin

q

jumps at evidence-size bin boundaries. The new monotone-cap calibrator

q_{\mathrm{mono}}[b]=\min(q_{\mathrm{mono}}[b-1],q_{\mathrm{strat}}[b])

enforces a non-increasing

q

schedule across bins, eliminating these wrapper-induced violations and yielding 0/3000 calibrated SMNN violations across all 30 PGMs . UMLP is unconstrained and frequently violates the property at both raw and calibrated levels.

Probabilistic Graphical Models

Probabilistic graphical models represent joint distributions with graphs to support structured inference, learning, and reasoning under uncertainty.

See all 54 papers

ABSOL: Aggregated Bayesian Subsampling Orchestrated with LLMs
Sep 14, 2026Jackson Hassell, Chen Shen, Estevam HruschkaGraph Structure Learning Bayesian Network Structure Learning
Megagon Labs
VeriBound: PAC-Bayesian Generalization Bounds for Process Reward Models Trained with Formal Verification Tools
Jun 17, 2026Amirul Rahman, Mohammed Sabih AlsharariPAC-Bayesian Generalization Bounds Process Reward Models
University of Malaya
Train for Many, Update with One: Improving Multi-Sample Reasoning via Fenchel-Bregman Posterior Learning
Oct 13, 2026Sungjun Lim, Adrian Robert Minut, Nico Carlsen Daheim +3LLM Reasoning
RIKEN AIP · Yonsei University · Sapienza University of Rome +1

Sep 14, 2026J/K move · Enter open · S save

ABSOL: Aggregated Bayesian Subsampling Orchestrated with LLMs

Jackson Hassell, Chen Shen, Estevam Hruschka

Megagon Labs

Graph Structure Learning Bayesian Network Structure Learning Bayesian Networks Probabilistic Graphical Models LLMs for Graph Learning

Open paper AskPDF arXiv

Large language models are increasingly used as natural-language interfaces to structured data, yet they remain unreliable when answers require consistent evidence conditioning, dependency-aware reasoning, and uncertainty estimation. Bayesian networks provide an explicit probabilistic reasoning layer, but learning useful structures from data remains costly and fragile at scale. We introduce ABSOL, a hybrid LLM-guided Bayesian network structure-learning framework that uses LLMs as bounded semantic guides. Across five discrete BN benchmarks spanning 27 to 1041 nodes, ABSOL is the only evaluated method to produce a viable graph on every benchmark, and achieves the highest Edge F_1 on every benchmark larger than 27 nodes with GPT-5.4. The four LLM augmentations, which contribute complementary semantic evidence to the statistical backbone, improve Edge F_1 over the non-LLM aggregation backbone by +0.23 on average. Complementary post-hoc refinement experiments suggest that these gains depend in part on limiting the LLM's authority over the final structure. Together, these results show that language-derived semantic knowledge can substantially improve scalable probabilistic structure learning when used as bounded guidance within a statistically grounded reasoning pipeline. The code for ABSOL is available at github.com/megagonlabs/absol-bn.

Abstract

Probabilistic Graphical Models

Explore similar work

ABSOL: Aggregated Bayesian Subsampling Orchestrated with LLMs

VeriBound: PAC-Bayesian Generalization Bounds for Process Reward Models Trained with Formal Verification Tools

Train for Many, Update with One: Improving Multi-Sample Reasoning via Fenchel-Bregman Posterior Learning