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).

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 its output under results/:

results/figures/{prefix}_total_time_vs_size.{pdf,svg,png}
results/figures/{prefix}_accuracy_vs_size.{pdf,svg,png}
results/raw/{prefix}_metrics.parquet
results/raw/{prefix}_{timestamp}.log         (gitignored)
results/raw/{prefix}_{timestamp}.run.json    (gitignored)
results/tables/{prefix}_summary.{csv,md,parquet,tex}

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.