NeuralBayesianNetworks (NBN)

A PyTorch-native Bayesian network library where each mechanism is learnt by a neural network. Every node carries a learnable, batched, GPU-resident conditional distribution; every query is a batched tensor operation; inference and parameter learning both run end-to-end on cuda.

NBN is 9-22× faster than pgmpy on continuous Linear Gaussian inference and 2.3× more accurate on discrete parameter learning at scale, and is the only library in our benchmark suite that handles hybrid (mixed continuous-discrete) networks at scale.

Install

pip install nbn                 # the library
pip install "nbn[neural]"       # + zuko-backed flows / MDNs (NBN's headline mechanisms)
pip install "nbn[bench]"        # + the benchmark suite (`nbn-bench`, `nbn.bench`)

The import name is nbn. The [neural] extra adds zuko-backed flows/MDNs; bench pulls in the benchmark runner's own dependencies (pandas, pyarrow, scipy, yaml, tqdm), the external-baseline libraries (pgmpy, pomegranate) and plotting (matplotlib, seaborn); gp/mcmc add the gpytorch and pyro baselines. All of bench, gp, and mcmc are required for paper-grade benchmark runs — without them the runner silently skips those baselines (cells emit not_supported rather than erroring).

Working from a checkout (development, or reproducing the paper runs, whose YAML configs are addressed by repo-relative path):

git clone https://github.com/Giovannibriglia/NeuralBayesianNetworks.git
cd NeuralBayesianNetworks
pip install -e ".[dev,bench,neural,gp,mcmc]"

dev is the test and docs toolchain. Releases are cut by pushing a v* tag; .github/workflows/publish.yml builds from the tag and uploads to PyPI.

Run the benchmark suite

From the repo root, launch the six paper-scale benchmarks. At most three run in parallel; the next starts as soon as one finishes:

bash scripts/run_all_benchmarks.sh

Override the pool size or device via MAX_PARALLEL=2 bash … or DEVICE=cpu bash … (three GPU benchmarks in flight can exceed an 8 GB card). Each benchmark leaves a run directory under results/; turn it into figures and tables with nbn-bench plot (see Plot the results).

Why NBN

NBN is to Bayesian Networks what GPyTorch is to Gaussian Processes: a torch-native, batchable, autograd-friendly framework where every conditional distribution is a swappable, learnable module, and every query is a batched tensor operation.

Headline results

These numbers come from the canonical paper-data run at tag v0.6c-d (see Reproducibility below).

Inference: 9-22× faster on continuous Linear Gaussian networks

NBN-lg-lw vs pgmpy-lg-predict on continuous Linear Gaussian networks: 22× faster at n=10 (1.9 ms vs 42 ms), 12× at n=1000 (0.72 s vs 8.5 s). Accuracy matches pgmpy within 0.02 W₁ at every n_nodes — speed gain comes without quality regression.

On discrete networks at n=10, NBN-cat-ve runs at 1.4 ms vs pgmpy-mle-ve at 108 ms (75× faster).

Parameter learning: 2.3× more accurate on discrete networks at scale

On discrete Bayesian networks, NBN-cat reaches TV ≈ 0.14 across all n_nodes ≥ 50; pgmpy-mle saturates at TV ≈ 0.34. The quality gap opens at n=50 (0.10 vs 0.25) and persists through n=1000 (0.146 vs 0.340). NBN's gradient-based fitting scales past pgmpy's sample-complexity wall.

On continuous Linear Gaussian networks, NBN matches pgmpy quality (W₁ ≈ 0.083 across all n) at 2× the speed.

Hybrid networks

NBN-hybrid handles mixed continuous-discrete networks across all n_nodes in our benchmark (n ∈ {10, 50, 100, 500, 1000}). Among the external libraries, only pyro covers hybrid inference (Importance sampler); pgmpy, gpytorch, and pomegranate have no applicable hybrid baselines.

Quick start

import torch
from nbn import NeuralBayesianNetwork, TensorVariableElimination
from nbn.mechanisms import CategoricalTableMechanism

# A → B → C, all categorical with cardinality 4
edges = [("A", "B"), ("B", "C")]
model = NeuralBayesianNetwork(
    edges,
    variables={"A": ("discrete", 4), "B": ("discrete", 4), "C": ("discrete", 4)},
)

# Fit each node's mechanism from data
data = {"A": torch.randint(0, 4, (10_000,)),
        "B": torch.randint(0, 4, (10_000,)),
        "C": torch.randint(0, 4, (10_000,))}
for node in model.dag.topological_order():
    parents = model.dag.parents(node)
    pa = torch.stack([data[p] for p in parents], dim=-1).float() if parents else None
    mech = CategoricalTableMechanism()
    mech.fit_local(data[node], pa, parent_cards=[4] * len(parents))
    model.set_mechanism(node, mech)

# Batched query: P(C | A=a) for 4 evidence rows at once
engine = TensorVariableElimination()
posterior = engine.query_batch(model, ["C"], {"A": torch.tensor([0, 1, 2, 3])})
# posterior shape: (B=4, 1, 4) — a distribution over C for each evidence row

For continuous, hybrid, and neural-mechanism examples, see the test suite under tests/integration/.

Gradients

Parent values and do= values are gradient-transparent — a parent computed by your own nn.Module carries autograd back into it:

The last two are deliberate. When you need gradients through an intervention, use model.sample(n, do=...), which applies it against the live parameters; intervene() is for building a mutilated model to query. The contract is pinned by tests/unit/test_parent_gradient_contract.py.

Snapshotting parameters

To take an optimisation step and be able to reject it (backtracking an M-step that decreased the objective, say), use torch's state_dict / load_state_dict — but copy the snapshot:

snap = copy.deepcopy(mech.state_dict())   # NOT mech.state_dict()
...                                       # optimiser step
mech.load_state_dict(snap)                # reverts exactly

state_dict() returns tensors sharing storage with the live parameters, and optimisers update in place, so an uncopied snapshot is mutated by the step it is meant to undo — silently, with nothing raising. Pinned by tests/unit/test_parameter_snapshot_contract.py.

Repository layout

nbn/                Library code (mechanisms, inference, sampling, core).
nbn/bench/          Benchmark suite: runner, baseline adapters, configs, data.
notebooks/          Colab-ready notebooks for the six paper-scale benchmarks.
scripts/            run_all_benchmarks.sh and other operational helpers.
results/            Where nbn-bench writes runs (gitignored).
tests/              Unit + integration tests.
RESEARCH.md         Paper outline and contribution claims.

Reproducibility

The headline numbers above are anchored at tag v0.6c-d (commit 2e0dd32):

git checkout v0.6c-d          # the suite lived at benchmarking/ at this tag
nbn-bench inference \
  --config benchmarking/configs/inference_paper_laptop.yaml
nbn-bench param-learning \
  --config benchmarking/configs/parameter_learning_paper_laptop.yaml

• Hardware: NVIDIA GeForce RTX 4070 Laptop (8 GB VRAM)
• PyTorch: 2.11.0+cu130
• Wall time: 11.4 h inference + 6.5 h parameter-learning
• Paper data anchor: see docs/v0.6c-d/run_summary.md for full headline tables and docs/v0.6c-d/dnf_cells.md for the DNF table

Numerical values vary within MC noise across hardware; STATUS counts and qualitative findings (cluster, speedup, quality gap) are stable. The committed parquets, tables, and figures under results/{raw,tables,figures}/ are the canonical paper artefacts.

Crash tests

NBN ships two crash tests on synthetic Bayesian networks with known ground truth, sweeping network size on the x-axis:

1. Parameter-learning crash test — measures accuracy of fitted CPDs against the true generative process. Speed is not measured.
2. Inference crash test — measures both accuracy and total time for Q conditional queries. NBN uses query_batch(B=Q) (one batched call); other libraries loop over the same Q queries in Python.

Each crash test has a smoke config (CI, < 60s) and a paper config (local reproduction, ~17.9 h on RTX 4070 Laptop 8 GB; CPU not supported for paper-config).

Reproduce

# Smoke (runs in CI):
nbn-bench param-learning --config nbn/bench/configs/synthetic/smoke_tests/parameter_learning_smoke.yaml
nbn-bench inference      --config nbn/bench/configs/synthetic/smoke_tests/inference_smoke.yaml

# Paper (8 GB VRAM, the laptop variant used for v0.6c-d paper data):
nbn-bench param-learning --config nbn/bench/configs/synthetic/complete/parameter_learning_complete_laptop.yaml
nbn-bench inference      --config nbn/bench/configs/synthetic/complete/inference_complete_laptop.yaml

# Paper (≥16 GB VRAM, canonical config without batch reductions):
nbn-bench param-learning --config nbn/bench/configs/synthetic/complete/parameter_learning_complete.yaml
nbn-bench inference      --config nbn/bench/configs/synthetic/complete/inference_complete.yaml

Each invocation writes one run directory under results/; the parquet is the single canonical artefact of a run (figures and tables are never generated automatically):

results/benchmark_<benchmark>_<config_name>_<YYYYMMDD_HHMMSS>/
    <config_name>_metrics.parquet     one row per (cell, metric)
    metrics.jsonl                     the same rows, streamed while running
    run.log                           per-cell log (fit/query phases, errors)

Plot the results

nbn-bench plot turns one (or more) run directories into paper figures and LaTeX tables. The same command serves every benchmark; the figures it emits are decided by what the parquet contains (which metrics have ok rows, whether n_train or batch_size is swept), so you never pick a plotter:

nbn-bench plot results/benchmark_synthetic_learning_curves_20260908_092714 \
  --output-dir results/figures/learning_curves
# options: --aggregation iqm_iqr|mean_std (default iqm_iqr)
#          --benchmark synthetic|bnlearn  (default: every benchmark in the parquet)

Output tree: <output-dir>/<benchmark>/<family>/all/{plots,tables}/, plus common/ or subset<k>/ siblings of all/ when the baselines do not all cover the same problems (_subsets_overview.txt explains the split). A figure that had unfinished cells gets a *_dnf.txt sidecar naming them. The run-directory name uses the config's config_name (complete, scalability_complete, batch_speed, param_learning_complete, learning_curves, bnlearn_complete), not the YAML file name. Per benchmark:

Two things that bite:

• learning_curves.yaml and parameter_learning_complete.yaml declare metrics: log_likelihood and their baselines carry no inference_method, so they must run under param-learning. Under inference the loader refuses them and prints the command to use.
• A figure is only written when at least one ok row exists for its metric in that family. If a plot you expect is missing, nbn-bench plot -v logs skip empty (...) with the reason, and run.log in the run directory has the per-cell error.

Design notes and the full figure/table spec live in docs/v0.13-paper-figures.md.

Configuration

Each config is a YAML file with these fields:

mode:                 'parameter_learning' | 'inference'
families:             list of families ∈ {discrete, continuous_lg,
                      continuous_nongauss, hybrid}
n_nodes:              list of network sizes
n_seeds:              number of seeds per cell (mean ± std reported)
n_queries_per_cell:   number of queries per cell
nbn_batch_size:       B for NBN's query_batch (inference mode only)
baselines:            list of baseline spec dicts, each with required
                      fields {library, mechanism, param_method} plus
                      optional inference_method and device (cpu|cuda|auto)
per_cell_timeout_s:   wall-clock cap per (family, n_nodes, seed, baseline)

See nbn/bench/configs/**/*.yaml for all shipped configs.

Status

Current release: v0.6c-d (paper-data anchor). The library is in publishable empirical state.

Active backlog (v0.7, none paper-blocking):

• Plotter polish: W₁ band lower-clip (#42), parameter-learning accuracy panels for non-discrete families (#44)
• Adapter audits: pgmpy-mle vs pgmpy-bayes / nbn-cat vs nbn-neuralcat fit-path distinctness (#43)
• HybridRouter cuda assert at hybrid n ≥ 10 (#30)
• NeuralCategorical-VE engine refactor (#26)
• pyro inference speedup (v0.8 candidate) — current Importance sampler is Python-bound and CPU-only; GPU is 11× slower at benchmark scale, so speedup requires pyro.plate vectorisation or alternative inference modes (SVI, NUTS). See docs/audits/v0.12-pyro-gpu-investigation.md.

See the open issues for the full v0.7 backlog.

License

Apache License 2.0 — see LICENSE.