distprop - Uncertainty Quantification via Stable Distribution Propagation through Neural Networks

This repository includes the official implementation of our ICLR 2024 Paper "Uncertainty Quantification via Stable Distribution Propagation" (Paper @ ArXiv).

distprop is a Python 3.6+ and PyTorch 1.13.1+ based package for propagating distributions through neural networks and can be installed with:

pip install distprop

💻 Usage

from distprop import SDP
import torch

model = torch.nn.Sequential(
        torch.nn.Linear(4, 100),
        torch.nn.ReLU(),
        torch.nn.Linear(100, 100),
        torch.nn.ReLU(),
        torch.nn.Linear(100, 3),
    )
data = torch.randn(10, 4)

sdp_model = SDP(model, std=0.1, num_outputs=3)
output_mean, output_cov = sdp_model(data)
print(output_mean.shape, output_cov.shape)

📖 Citing

@inproceedings{petersen2024uncertainty,
  title={{Uncertainty Quantification via Stable Distribution Propagation}},
  author={Petersen, Felix and Mishra, Aashwin and Kuehne, Hilde and 
          Borgelt, Christian and Deussen, Oliver and Yurochkin, Mikhail},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2024}
}

📜 License

distprop is released under the MIT license. See LICENSE for additional details about it.