Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Partially Stochastic Infinitely Deep Bayesian Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes partially stochastic infinitely deep Bayesian neural networks (PSDE-BNNs), where only a subset of weights (vertical time interval or horizontal parameter split) follows an SDE, and proves they are universal conditional distribution approximators while being faster than SDE-BNN.
Problem
Infinitely deep Bayesian neural networks (Neural SDE-based BNNs) are computationally expensive to train and run because all weights are stochastic and must be simulated at every step; partial stochasticity has not been explored in the infinite-depth setting.
Approach
Make only part of the weight dynamics stochastic in the infinite-depth limit: either restrict the SDE to a time sub-interval (vertical cut, rest is ODE) or split the weight vector into stochastic and deterministic subsets (horizontal cut), with a variational SDE posterior whose diffusion matches the prior, and prove universal conditional distribution approximation.
Key results
- On CIFAR-10, PSDE-BNN ODEFirst reaches 87.84% accuracy versus 86.04% for SDE-BNN, with ECE 4.73 vs 9.13 (x10^-2).
- PSDE-BNN ODEFirst cuts epoch time from 371.9s to 289.4s (about 27.5% faster) and inference time from 43.7s to 31.7s (about 25.9% faster).
- Overall training time is reported as 74.1% faster, since PSDE-BNNs need 100 epochs versus 300 for SDE-BNN.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use PSDE-BNN ODEFirst instead of SDE-BNN | SDE-BNN · CIFAR-10 | accuracy (%): 86.04 \(\pm 0.25\) -> **87.84** \(\pm\) **0.08** | Helped |
| use PSDE-BNN SDEFirst instead of SDE-BNN | SDE-BNN · CIFAR-10 | ECE (× 10^-2): 9.13 \(\pm 0.28\) -> **3.56** $\pm$ **0.15** | Helped |
| replace SDE-BNN with PSDE-BNN ODEFirst | SDE-BNN · CIFAR-10 | epoch time (s): 371.9 -> **289.4** | Helped |
| replace SDE-BNN with PSDE-BNN ODEFirst | SDE-BNN · CIFAR-10 | inference time for 10000 predictions (s): 43.7 -> **31.7** | Helped |
| increase ODEFirst stochasticity ratio from 10% to 70% (ablation) | PSDE-BNN ODEFirst · CIFAR-10 | accuracy (%): 87.84 \(\pm\) 0.08 -> 84.47 \(\pm 0.15\) | Hurt |
| increase SDEFirst stochasticity ratio from 10% to 30% (ablation) | PSDE-BNN SDEFirst · CIFAR-10 | accuracy (%): 85.34 \(\pm 0.21\) -> 87.19 \(\pm 0.24\) | Helped |
| increase SDEFirst stochasticity ratio from 10% to 50% (ablation) | PSDE-BNN SDEFirst · CIFAR-10 | accuracy (%): 85.34 \(\pm 0.21\) -> 82.24 \(\pm 0.32\) | Hurt |
| increase stochasticity ratio (ablation) | PSDE-BNN vertical-cut models · CIFAR-10 | epoch time (s): 289.4 -> 310.3 -> 331.9 -> 353.2; 295.1 -> 314.5 -> 334.9 -> 413.1 | Hurt |
| increase KL coefficient order to 10^-2 (ablation) | PSDE-BNN · image classification | learning rate | Helped |
| increase KL coefficient order to 10^-2 (ablation) | PSDE-BNN · image classification | training stability | Hurt |
| decrease KL coefficient order to 10^-4 (ablation) | PSDE-BNN · image classification | convergence speed | Hurt |
| use PSDE-BNN vertical cuts instead of SDE-BNN | CIFAR-10 OOD detection · CIFAR-10 ID vs OOD | AUC: 0.84 -> 0.88 | Helped |
Surprising findings
- Ablations show no clear correlation between stochasticity ratio and downstream accuracy or uncertainty metrics.
- The SDE-BNN baseline, when trained for the same wall-clock time, achieves much lower accuracy (69.94% on CIFAR-10), showing ECE only looks good at low accuracy.
Limitations
- Horizontal cut lacks an optimised numerical solver, so efficiency was only measured for vertical cut models.
- Performance on CIFAR-10 is below ResNet32 and ODENet point estimates in accuracy, and some baselines (HMC, Laplace) remain better calibrated.
- Hyperparameters chosen by grid search with a KL coefficient sensitive to numerical stability.
Experimental setup
Datasets: MNIST CIFAR-10
Benchmarks: MNIST CIFAR-10 CIFAR10-C
Metrics: accuracy expected calibration error predictive entropy AUC epoch time inference time
Baselines: SDE-BNN ODENet ResNet32 MFVI ODENet MFVI HyperODENet Deep Ensemble HMC ResNet32 + LL Laplace
Builds on: SDE-BNN Augmented Neural ODEs Neural ODE Scalable gradients for SDEs Do Bayesian neural networks need to be fully stochastic?
Closest papers by meaning
- Understanding and Exploring the Network with Stochastic Architectures (NeurIPS 2020 · similarity 0.87)
- On the Expressiveness of Approximate Inference in Bayesian Neural Networks (NeurIPS 2020 · similarity 0.87)
- Why bigger is not always better: on finite and infinite neural networks (ICML 2020 · similarity 0.87)
- Exact marginal prior distributions of finite Bayesian neural networks (NeurIPS 2021 · similarity 0.86)
- Simple and Effective Stochastic Neural Networks (AAAI 2021 · similarity 0.86)
- A Study of Bayesian Neural Network Surrogates for Bayesian Optimization (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.