PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Bayesian Networks for Modular and Scalable Decision-Making Under Uncertainty

Giovanni Briglia, Marco Lippi, Stefano Mariani, Franco Zambonelli

arXiv:2610.16902v1cs.AIcs.LG

TL;DR

Classical BNs constrain conditional representations and evaluate sequentially, limiting expressiveness and scale. NBNs preserve the BN graph and factorization while using tensor-valued mechanisms; the benchmark confirms potential in expressiveness and scalability.

  • Problem

    Classical BN conditionals constrain expressiveness, while sequential evaluation limits scalability across samples and queries.

    Source: §2.1 · p. 3

  • Method

    NBNs preserve the BN graph and factorization while using differentiable, tensor-valued conditional mechanisms that run in batches.

    Source: §1 · p. 1, §2 · p. 2

  • Results

    The benchmark confirms NBNs' potential for scalability and expressiveness; on non-Gaussian networks, MDNs have between one third and 60% of a linear-Gaussian model's W1 and better held-out likelihood.

    Source: §1 · p. 2, §4 · p. 5

  • Takeaways & Limitations

    The released benchmark fixes problems, queries, methods, and budgets to support reproducible comparisons of Bayesian-network methods.

    Source: §3 · p. 4

    The work assumes a given graph and complete i.i.d. data, and leaves structure learning and planning-trajectory sampling to future work.

    Source: §2.1 · p. 3

Abstract

from arXiv · show

Bayesian Networks (BNs) are established tools for letting agents reason under uncertainty in sequential decision-making. They split a joint distribution into local conditional mechanisms arranged on a directed acyclic graph, which makes them modular, interpretable, and ready to perform probabilistic inference. Their classical formulation, however, has limitations in both expressiveness and scale: each conditional is usually stored as a table or a linear-Gaussian form, and inference evaluates the graph one node and one query at a time. In this paper we introduce Neural Bayesian Networks (NBNs), in which each conditional is a differentiable, tensor-valued mechanism, while the graph and its factorization are unchanged. NBNs retain the modularity and local mechanisms of BNs, can represent any learnable conditional distribution, and run natively on GPUs. We also release NBN, a PyTorch-based library for learning and inference in NBNs, together with a public benchmark. Our evaluation shows gains in both expressiveness and scale: NBNs fit distributions that classical BN libraries struggle to represent, and on large networks and large batches of queries they run two orders of magnitude faster thanks to GPU parallelization. NBN is available at https://github.com/Giovannibriglia/NeuralBayesianNetworks.

1 Introduction

Neural Bayesian Networks keep the Bayesian network graph and factorization while replacing scalar conditionals with differentiable, tensor-valued mechanisms. This addresses limits in classical networks’ expressiveness and sequential evaluation, and supports batched GPU computation.

1 Introduction

  • Factored BN transition models support agents reasoning about uncertain outcomes, especially when high-dimensional states make flat models impractical.The paper identifies planning, reinforcement learning, and diagnosis as settings that can use learned transition models queried under evidence.Source: §1 · p. 1
  • Classical BNs face limits in expressiveness and scale: conditional tables grow combinatorially with parent count and are undefined for continuous parents, while inference evaluates one node and sample at a time.These representational and sequential-evaluation constraints also affect mainstream BN software packages.Source: §1 · pp. 1–2
  • NBNs preserve the BN graph and factorization while each differentiable, tensor-valued mechanism returns conditional distributions for batches of parent configurations.Mechanisms may be tabular, parametric, neural, or non-parametric, and can be evaluated a graph layer at a time on a GPU.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • The authors release NBN, a PyTorch library, and a benchmark comparing BN frameworks across discrete, Gaussian, and non-Gaussian networks; the benchmark indicates scalability and expressiveness potential, though NBN does not win every comparison.The library includes tabular, linear-Gaussian, neural, and non-parametric mechanisms.Source: §1 · pp. 1–2

2 Neural Bayesian Networks

Neural Bayesian Networks retain a Bayesian network’s graph and factorization while replacing scalar conditionals with learnable modules that support batch learning and inference. This preserves BN properties while broadening conditional representations and reducing the dependence of inference time on batch size until device saturation.

2 Neural Bayesian Networks

  • An NBN retains the BN graph and factorization but replaces each conditional with a learnable module returning a full distribution for each parent configuration, allowing mixed tabular, parametric, non-parametric, and neural mechanisms.The NBN library is built from a DAG and per-node mechanism assignments, then used for fitting and queries.Source: §2 · p. 2, §2.2 · p. 3

2.3 Learning and Inference Methods

  • Differentiable conditionals support batch gradient fitting, with closed-form fits where available; VE, LW, and AVI provide exact, sampling-based, and amortized query engines.Mechanism options include tables, linear-Gaussian forms, MDNs, KNN, and FlexCode.Source: §2 · p. 2, §2.2 · p. 3, §2.3 · p. 4

2.1 Motivation

  • The setup assumes a given graph, complete i.i.d. data, and discrete finite or real-valued scalar variables; classical scalar conditionals constrain learning and many-query inference.Tables grow combinatorially with parents and do not support continuous parents, while linear-Gaussian forms cannot represent multimodal or heteroscedastic conditionals.Source: §2.1 · p. 3

2.2 Formal definition and Properties

  • Inference uses depth-wise tensor operations over active nodes, removing dependence on sample and evidence-batch counts until device saturation while preserving modularity and arbitrary-query use.Each mechanism has its own parameters and can be inspected, refit, or replaced without changing the others.Source: §2.2 · pp. 3–4

3 Benchmark Design

The benchmark addresses the lack of a fully reproducible BN comparison by fixing problems, queries, methods, and budgets and releasing its configuration and harness. It spans synthetic and published networks, multiple query scenarios, and reference-based accuracy and timing measures.

3 Benchmark Design

  • The benchmark fixes problems, queries, methods, and budgets, and records a row per problem, seed, method, query, and metric, including failures.It is released with NBN to make results more comparable.Source: §3 · p. 4
  • Discrete accuracy uses total variation and continuous accuracy uses 1-Wasserstein distance, but finite-sample references impose a shared accuracy floor.Exact references are left for future work because they are unavailable for every family and size.Source: §3 · pp. 4–5
  • The protocol sets separate 1200-second fit and query budgets, memory caps, and status reporting, and compares NBN engines with pgmpy, pomegranate, and Pyro baselines.Figures report the interquartile mean over solved seeds and mark failures; NBN and pomegranate run on GPU while pgmpy runs on CPU.Source: §3 · p. 5
  • Scenarios vary network family and size, query load, and batching, while real-network queries target hubs, cut nodes, and terminals in prediction and diagnosis directions.Synthetic problems cover discrete, linear-Gaussian, and Gaussian-mixture conditionals; the accuracy scenario uses 128 random queries and sizes from 10 to 1000 nodes.Source: §3 · pp. 4–5

4 Experimental Results

Across the experiments, NBN preserves strong results on classical model families, represents non-Gaussian conditionals more effectively, and gains substantial throughput on large query batches. The results also show where accuracy, fitting cost, or memory constrain its performance.

4 Experimental Results

  • Across tasks, NBN matches classical accuracy where their model families apply, improves non-Gaussian modeling, and delivers two-order-of-magnitude speedups on large query batches.Sampling engines achieve these batch speedups, while NBN still has accuracy and resource limits in some settings.Source: §4 · p. 5

4.1 Inference accuracy (RQ1, RQ2)

  • One-third to 60%: on non-Gaussian networks, MDN has this fraction of the W1 error of a linear-Gaussian model; on discrete networks, cat+VE matches pgmpy at every size.Figure 3 plots inference accuracy against network size; lower TV and W1 are better.Source: §4 · p. 5, §4.1 · p. 5, Figure 3 · p. 6

4.2 Scalability

  • At 10,000 nodes, cat+LW and LG+LW finish every seed, while cat+VE exceeds the GPU’s 12 GB on two of five seeds.The heaviest-query comparison extends to 10,000 nodes.Source: Table 1 · p. 6, §4.2 · p. 6

4.3 Batched queries (RQ3)

  • At B=1024, cat+AVI gains 456× and LG+AVI 227×, reaching 0.09 and 0.05 ms per query, respectively; cat+VE gains only 4×.The batch-time comparison uses 50-node networks with full evidence.Source: §4.3 · p. 6, Figure 4 · p. 7

4.4 Real networks

  • On CLG networks, NBN’s best mechanism trails Pyro on all three; MDN wins on sangiovese and mehra, while kNN wins on healthcare.On discrete networks, cat+VE has the lowest TV on five networks, including the two largest.Source: Table 2 · p. 7, §4.4 · p. 7

4.5 Parameter learning and sample efficiency

  • On 50-node non-Gaussian networks, linear-Gaussian NLL stays at 1.26–1.27 from 4096 to 65 536 rows, while MDN reaches 1.20 with 4096 rows.When the model is correct, NBN’s table and linear-Gaussian fits match the corresponding estimators; at 1,000 nodes, their fit times are 2.1 s and 2.6 s, versus pgmpy’s 9.8 s and 4.0 s. MDN fitting exceeds budget at 65 536 rows and from 100 nodes.Source: Table 3 · p. 7, §4.5 · p. 8

5 Related Works

NBNs extend Bayesian networks along expressiveness and scalability while preserving their graph and factorization. Unlike approaches tied to particular conditional families or inference procedures, NBNs use tensor-valued mechanisms with shared batched inference.

Expressiveness

  • NBNs preserve the BN graph while allowing tabular, linear-Gaussian, neural, and non-parametric mechanisms to be mixed and trained together.The same batched engine can query any mechanism that samples and evaluates a density.Source: §5 · p. 8
  • Earlier expressive approaches either fix conditional families and inference procedures or hide dependencies in a single network without directly available posteriors for arbitrary evidence.NBNs instead retain inspectable, replaceable local conditionals.Source: §5 · p. 8

Efficiency and scalability

  • Prior scalable exact inference uses GPU parallelism or tensor contractions for discrete networks but remains exponential in treewidth, while probabilistic circuits trade per-variable conditionals for latent sum and product units.Likelihood weighting can waste samples when evidence is unlikely; amortized inference trains proposals that agree with evidence and reuses them.Source: §5 · p. 8
  • NBNs do not introduce a new estimator; they evaluate same-depth nodes in one tensor operation across particles and evidence assignments.This execution strategy makes query cost stop growing with sample and query counts, as stated in Proposition 2.3.Source: §5 · p. 8

Libraries

  • Existing libraries split capabilities across classical CPU BN packages, a GPU package limited to discrete loopy belief propagation, and probabilistic programming systems without an explicit BN object.The benchmark compares representative systems from these classes; the stated gap is an explicit DAG combined with arbitrary differentiable conditionals and GPU-batched learning and inference.Source: §5 · p. 8, Table 4 · p. 8

6 Conclusions and Future Works

NBNs combine explicit graphical structure with learnable tensor-valued conditionals for learning and GPU-batched inference. The paper presents the library and benchmark as tools for evaluating trade-offs and identifies several directions for future work.

Conclusions

  • NBNs preserve Bayesian-network factorization and modularity while bringing heterogeneous conditional models into a common framework for learning and GPU-batched inference.Their conditionals are learnable tensor-valued mechanisms within an explicit graph.Source: §6 · p. 8
  • The PyTorch library and reproducible benchmark provide a basis for evaluating trade-offs among expressiveness, accuracy, and computational cost.Separately replaceable mechanisms support adapting individual model components, while batched inference suits evaluating many candidate actions or hypotheses.Source: §6 · p. 9

Future works

  • Future work includes evaluating NBNs in agent decision-making tasks, removing the assumption that the graph is given, assessing sampling, and considering CPU execution alongside GPU execution.Source: §6 · p. 9

A Benchmark Design – Extended Version

The extended benchmark design specifies synthetic and real-network data, query selection, reference distributions, learning metrics, and execution settings. It also documents conditions that affect score interpretation and batching measurements.

Synthetic generator

  • Synthetic networks use seeded DAGs with at most four parents per node, and the benchmark includes discrete, linear-Gaussian, and non-Gaussian conditional families.The non-Gaussian family uses three-component Gaussian mixtures with skewed component offsets; the parent cap bounds CPT size as network size grows.Source: A Benchmark Design – Extended Version · p. 11

Data

  • The benchmark varies training-set sizes across scenarios and real networks, uses 4096 test rows in most settings, and evaluates five seeds.Real-network training sets range from 10,240 to 81,920 rows, with sizes selected by network parameter count.Source: A Benchmark Design – Extended Version · p. 11

Query selectors

  • Queries are selected as random, topological, or heaviest cases, with targets and evidence assigned by specified structural roles and patterns.The batching scenario asks each selected query for evidence values from 1024 rows.Source: A Benchmark Design – Extended Version · p. 11

Reference

  • Discrete references use finite samples, while continuous references use likelihood-weighted draws from the true conditional posterior; the resulting error floor is not a bound.The floor measures exact answers against sampled references, so approximate methods can score slightly better than exact engines.Source: A Benchmark Design – Extended Version · p. 11

Learning metrics

  • Learning is scored by parameter-recovery TV for discrete networks and held-out per-node NLL for continuous families, which do not share a common parameter space.PIT-KS calibration alone can favor a model that ignores parent variables, so the benchmark prefers NLL for continuous-family comparisons.Source: A Benchmark Design – Extended Version · p. 11

Execution settings

  • Benchmark settings include 2048 sampling particles in most scenarios, scenario-specific fit and query budgets, and a memory cap tied to available machine memory.The batching scenario uses 16 particles and a 600-second query budget; empty-evidence batches repeat the same marginal query and are not counted as a batching gain.Source: A Benchmark Design – Extended Version · p. 11

B Complete results

The appendix reports all methods across all scenarios, including additional mechanisms, engines, and timing results. It specifies how figures and tables encode performance and failures.

Complete results

  • The appendix reports every method for every scenario, including mechanisms and engines omitted from the main text and fit and query times where the main text reports accuracy only.It includes normalizing-flow, kernel-density, and neural categorical conditionals, plus the AIS engine.Source: B Complete results · p. 13

Figures

  • Figures show interquartile means over solved seeds, with error bars spanning contributing values; axes are logarithmic and failed seeds or absent results are indicated.For NBN, figure color identifies the mechanism and hatch identifies the engine.Source: B Complete results · p. 13

Tables

  • Tables report IQM ± half the interquartile range and failure counts, with the best method in each column selected first by fewest failures and then by IQM.Source: B Complete results · p. 13

B.1 Inference accuracy

The inference-accuracy evaluation reports accuracy and query time across network sizes, alongside fit time. It covers discrete, linear-Gaussian, and non-Gaussian networks.

B.1 Inference accuracy

  • Inference accuracy is evaluated against network size using TV for discrete targets and W1 for continuous targets.The figures summarize all methods; tables provide discrete and linear-Gaussian results.Source: Figure 5 · p. 14, Table 5 · p. 17, Table 8 · p. 17
  • Total query time is compared across network sizes for discrete and linear-Gaussian networks, with additional non-Gaussian query-time results.The figure reports total query time across all methods.Source: Figure 6 · p. 15, Table 6 · p. 17, Table 9 · p. 18
  • Fit time is reported against network size for discrete and linear-Gaussian networks, with non-Gaussian fit-time results also included.The fit-time figure compares all methods.Source: Figure 7 · p. 16, Table 7 · p. 17, Table 10 · p. 18

B.2 Scalability

The scalability evaluation compares accuracy and query time against network size across discrete, linear-Gaussian, and non-Gaussian networks. It includes both evidence modes and full-evidence results.

B.2 Scalability

  • For discrete networks, scalability is assessed using TV and time per query against network size, for both evidence modes and full evidence.The figures compare accuracy and per-query time across all methods.Source: Figure 8 · p. 20, Figure 10 · p. 22, Table 14 · p. 24
  • For linear-Gaussian networks, scalability is assessed using W1 and time per query against network size, for both evidence modes and full evidence.The corresponding figures compare accuracy and per-query time across all methods.Source: Figure 8 · p. 20, Figure 10 · p. 22, Table 18 · p. 24
  • For non-Gaussian networks, scalability is assessed using W1 and time per query against network size, for both evidence modes and full evidence.The corresponding figures compare accuracy and per-query time across all methods.Source: Figure 8 · p. 20, Figure 10 · p. 22, Table 22 · p. 26

B.3 Batched queries

The evaluation covers batched-query timing, query accuracy and time on real networks, and parameter learning across three network types. The reported measures vary by task and network type.

B.3 Batched queries

  • Batched-query evaluation compares time per query against batch size for discrete, linear-Gaussian, and non-Gaussian networks, with full or empty evidence.Figures summarize all methods; tables report the network-specific timings.Source: Figure 12 · p. 28, Figure 13 · p. 29, Table 26 · p. 30

B.4 Real networks

  • On real networks, the evaluation reports discrete-network TV and Gaussian- and CLG-network W1, alongside total query time.The figures show accuracy and total query time against network for all methods.Source: Figure 14 · p. 32, Figure 15 · p. 33, Table 32 · p. 34

B.5 Parameter learning

  • For discrete-network parameter learning, the evaluation reports parameter-recovery TV, held-out negative log-likelihood per node, and fit time against network size.These measures are presented in the corresponding tables and figures.Source: Figure 16 · p. 36, Figure 18 · p. 38, Table 38 · p. 40
  • For linear-Gaussian and non-Gaussian parameter learning, the evaluation reports PIT-KS, held-out negative log-likelihood per node, and fit time against network size.The figures cover all methods; the tables separate the two network types.Source: Figure 17 · p. 37, Table 41 · p. 40, Table 42 · p. 40

B.6 Sample efficiency

The sample-efficiency evaluation examines parameter recovery, calibration, held-out likelihood, and fit time as training-set size varies. Its tables and figures cover discrete, linear-Gaussian, and non-Gaussian networks.

B.6 Sample efficiency

Related papers

Loading 2610.16902v1…