# Partially Stochastic Infinitely Deep Bayesian Neural Networks

arXiv: https://arxiv.org/abs/2402.03495 · HTML: https://mlsanity.com/paper/2402.03495/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 50 references that are in the map.
- First posted on arXiv 5 Feb 2024.

## Builds on (most cited first, 28 of 28)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Error bounds for approximations with deep ReLU networks](https://mlsanity.com/paper/1610.01145/) (2016)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit](https://mlsanity.com/paper/1905.09883/) (2019)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise](https://mlsanity.com/paper/1906.02355/) (2019)
- [Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference](https://mlsanity.com/paper/1703.09194/) (2017)
- [Bayesian Deep Learning via Subnetwork Inference](https://mlsanity.com/paper/2010.14689/) (2020)
- [Do Bayesian Neural Networks Need To Be Fully Stochastic?](https://mlsanity.com/paper/2211.06291/) (2022)
- [Approximation Capabilities of Neural ODEs and Invertible Residual Networks](https://mlsanity.com/paper/1907.12998/) (2019)
- [Bayesian Learning via Neural Schr"odinger-F"ollmer Flows](https://mlsanity.com/paper/2111.10510/) (2021)
- [Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations](https://mlsanity.com/paper/2102.06559/) (2021)
- [Sampling-based inference for large linear models, with application to linearised Laplace](https://mlsanity.com/paper/2210.04994/) (2022)
- [Revealing hidden dynamics from time-series data by ODENet](https://mlsanity.com/paper/2005.04849/) (2020)
- [Efficient Online Bayesian Inference for Neural Bandits](https://mlsanity.com/paper/2112.00195/) (2021)
- [Accounting For Informative Sampling When Learning to Forecast Treatment Outcomes Over Time](https://mlsanity.com/paper/2306.04255/) (2023)
- [Compositional uncertainty in deep Gaussian processes](https://mlsanity.com/paper/1909.07698/) (2019)
- [DeepVol: Volatility Forecasting from High-Frequency Data with Dilated Causal Convolutions](https://mlsanity.com/paper/2210.04797/) (2022)
- [The Missing U for Efficient Diffusion Models](https://mlsanity.com/paper/2310.20092/) (2023)
- [Graph Neural Stochastic Differential Equations](https://mlsanity.com/paper/2308.12316/) (2023)

## Built on by (most cited first, 6 of 6)

- [On Vanishing Gradients, Over-Smoothing, and Over-Squashing in GNNs: Bridging Recurrent and Graph Learning](https://mlsanity.com/paper/2502.10818/) (2025)
- [Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations](https://mlsanity.com/paper/2408.16115/) (2024)
- [A Gaussian Process View on Observation Noise and Initialization in Wide Neural Networks](https://mlsanity.com/paper/2502.01556/) (2025)
- [Low Rank Based Subspace Inference for the Laplace Approximation of Bayesian Neural Networks](https://mlsanity.com/paper/2502.02345/) (2025)
- [Partial Trace-Class Bayesian Neural Networks](https://mlsanity.com/paper/2511.01628/) (2025)
- [Improving Infinitely Deep Bayesian Neural Networks with Nesterov's Accelerated Gradient Method](https://mlsanity.com/paper/2603.25024/) (2026)

## Abstract (arXiv)

In this paper, we present Partially Stochastic Infinitely Deep Bayesian Neural Networks, a novel family of architectures that integrates partial stochasticity into the framework of infinitely deep neural networks. Our new class of architectures is designed to improve the computational efficiency of existing architectures at training and inference time. To do this, we leverage the advantages of partial stochasticity in the infinite-depth limit which include the benefits of full stochasticity e.g. robustness, uncertainty quantification, and memory efficiency, whilst improving their limitations around computational complexity. We present a variety of architectural configurations, offering flexibility in network design including different methods for weight partition. We also provide mathematical guarantees on the expressivity of our models by establishing that our network family qualifies as Universal Conditional Distribution Approximators. Lastly, empirical evaluations across multiple tasks show that our proposed architectures achieve better downstream task performance and uncertainty quantification than their counterparts while being significantly more efficient. The code can be found at urlhttps://github.com/Sergio20f/partstochinfdeep