suq 0.1.0__tar.gz

suq-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 AaltoML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

suq-0.1.0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: suq
+Version: 0.1.0
+Summary: Streamlined Uncertainty Quantification (SUQ)
+Home-page: https://github.com/AaltoML/SUQ
+Author: Rui Li, Marcus Klasson, Arno Solin, Martin Trapp
+License: MIT
+Classifier: Programming Language :: Python :: 3
+License-File: LICENSE
+Requires-Dist: torch>=1.10
+Requires-Dist: numpy>=1.21
+Requires-Dist: tqdm>=4.60
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary

suq-0.1.0/README.md

@@ -0,0 +1,137 @@
+# SUQ: Streamlined Uncertainty Quantification
+
+![image](suq.png)
+
+This repository contains an open-source library implementation of Streamlined Uncertainty Quantification (SUQ) used in the paper *Streamlining Prediction in Bayesian Deep Learning* accepted at ICLR 2025.
+
+<table>
+<tr>
+	<td>
+   		<strong> Streamlining Prediction in Bayesian Deep Learning</strong><br>
+            Rui Li, Marcus Klasson, Arno Solin, Martin Trapp<br>
+		<strong>International Conference on Learning Representations (ICLR 2025)</strong><br>
+		<a href="https://arxiv.org/abs/2411.18425"><img alt="Paper" src="https://img.shields.io/badge/-Paper-gray"></a>
+		<a href="https://github.com/AaltoML/suq"><img alt="Code" src="https://img.shields.io/badge/-Code-gray" ></a>
+		</td>
+    </tr>
+</table>
+
+## SUQ Library
+### 📦 Installation
+Install the stable version with `pip`:
+```bash
+pip install suq
+```
+
+Or install the latest development version from source:
+```bash
+git clone https://github.com/AaltoML/SUQ.git
+cd SUQ
+pip install -e .
+```
+
+### 🚀 Simple Usage
+#### Streamline Whole Network
+```python
+from suq import streamline_mlp, streamline_vit
+
+# Load your model and estimated posterior
+model = ...
+posterior = ...
+
+# Wrap an MLP model with SUQ
+suq_model = streamline_mlp(
+    model=model,
+    posterior=posterior,
+    covariance_structure='diag',       # currently only 'diag' is supported
+    likelihood='classification'        # or 'regression'
+)
+
+# Wrap a Vision Transformer with SUQ
+suq_model = streamline_vit(
+    model=model,
+    posterior=posterior,
+    covariance_structure='diag',      # currently only 'diag' is supported
+    likelihood='classification',      
+    MLP_deterministic=True,
+    Attn_deterministic=False,
+    attention_diag_cov=False,
+    num_det_blocks=10
+)
+
+# Fit scale factor
+suq_model.fit(train_loader, scale_fit_epoch, scale_fit_lr)
+
+# Make a prediction
+pred = suq_model(X)
+```
+
+📄 See [`examples/mlp_la_example.py`](examples/mlp_la_example.py), [`examples/vit_la_example.py`](examples/vit_la_example.py), [`examples/mlp_vi_example.py`](examples/mlp_vi_example.py), and [`examples/vit_vi_example.py`](examples/vit_vi_example.py) for full, self-contained examples that cover:
+- Training the MAP model
+- Estimating the posterior with Laplace or IVON (mean field VI)
+- Wrapping the model into a streamlined SUQ version
+
+
+> ❗ **Note on Vision Transformer Support**  
+Currently, SUQ only supports Vision Transformers implemented in the same style as [`examples/vit_model.py`](examples/vit_model.py). If you're using a different ViT implementation, compatibility is not guaranteed.
+
+#### Streamline Individual Layers
+
+In addition to wrapping full models like MLPs or ViTs, SUQ allows you to manually wrap individual layers in your own networks.
+
+You can directly import supported modules from `suq.streamline_layer`.
+
+Supported Layers:
+
+| Layer Type         | SUQ Wrapper                   |
+|--------------------|-------------------------------|
+| `nn.Linear`        | `SUQ_Linear_Diag`             |
+| `nn.ReLU`, etc.    | `SUQ_Activation_Diag`         |
+| `nn.BatchNorm1d`   | `SUQ_BatchNorm_Diag`          |
+| `nn.LayerNorm`     | `SUQ_LayerNorm_Diag`          |
+| `MLP (Transformer block)`    | `SUQ_TransformerMLP_Diag`     |
+| `Attention`    | `SUQ_Attention_Diag`          |
+| `Transformer block`  | `SUQ_Transformer_Block_Diag`  |
+| `Final classifier`   | `SUQ_Classifier_Diag`         |
+
+Example:
+
+```python
+from suq.streamline_layer import SUQ_Linear_Diag
+
+# Define a standard linear layer
+linear_layer = nn.Linear(100, 50)
+# Provide posterior variances for weights and biases
+w_var = torch.rand(50, 100)
+b_var = torch.rand(50)
+
+# Wrap the layer with SUQ's linear module
+streamlined_layer = SUQ_Linear_Diag(linear_layer, w_var, b_var)
+
+# Provide input mean and variance (e.g., from a previous layer)
+input_mean = torch.randn(32, 100)
+input_var = torch.rand(32, 100)
+
+# Forward pass through the streamlined layer
+pred_mean, pred_var = streamlined_layer(input_mean, input_var)
+```
+
+### 🛠️ TODO
+- Extend support to other Transformer implementations
+- Add Kronecker covariance
+- Add full covariance
+
+
+## Citation
+
+```bibtex
+@inproceedings{li2025streamlining,
+  title = {Streamlining Prediction in Bayesian Deep Learning},
+  author = {Rui Li, Marcus Klasson, Arno Solin and Martin Trapp},
+  booktitle = {International Conference on Learning Representations ({ICLR})},
+  year = {2025}
+}
+```
+
+## License
+This software is provided under the MIT license.

suq-0.1.0/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"

suq-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

suq-0.1.0/setup.py

@@ -0,0 +1,19 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='suq',
+    version='0.1.0',
+    description='Streamlined Uncertainty Quantification (SUQ)',
+    author='Rui Li, Marcus Klasson, Arno Solin, Martin Trapp',
+    url='https://github.com/AaltoML/SUQ',
+    packages=find_packages(exclude=["examples*", "tests*"]),
+    install_requires=[
+        'torch>=1.10',
+        'numpy>=1.21',
+        'tqdm>=4.60'
+    ],
+    license='MIT',
+    classifiers=[
+        'Programming Language :: Python :: 3',
+    ],
+)

suq-0.1.0/suq/SUQ_MLP.py

@@ -0,0 +1,10 @@
+from .diag_suq_mlp import SUQ_MLP_Diag
+
+def streamline_mlp(model, posterior, covariance_structure, likelihood, scale_init = 1.0):
+    if covariance_structure == 'diag':
+        return SUQ_MLP_Diag(org_model = model, 
+                            posterior_variance = posterior, 
+                            likelihood = likelihood,
+                            scale_init = scale_init)
+    else:
+        raise NotImplementedError(f"Covariance structure '{covariance_structure}' is not implemented.")

suq-0.1.0/suq/SUQ_ViT.py

@@ -0,0 +1,14 @@
+from .diag_suq_transformer import SUQ_ViT_Diag
+
+def streamline_vit(model, posterior, covariance_structure, likelihood, MLP_deterministic, Attn_deterministic, attention_diag_cov, num_det_blocks, scale_init = 1.0):
+    if covariance_structure == 'diag':
+        return SUQ_ViT_Diag(ViT = model, 
+                            posterior_variance = posterior, 
+                            MLP_determinstic = MLP_deterministic, 
+                            Attn_determinstic = Attn_deterministic, 
+                            likelihood = likelihood,
+                            attention_diag_cov = attention_diag_cov, 
+                            num_det_blocks = num_det_blocks,
+                            scale_init = scale_init)
+    else:
+        raise NotImplementedError(f"Covariance structure '{covariance_structure}' is not implemented.")

suq-0.1.0/suq/__init__.py

@@ -0,0 +1,4 @@
+from .SUQ_MLP import streamline_mlp
+from .SUQ_ViT import streamline_vit
+
+__all__ = ["streamline_mlp", "streamline_vit"]

suq-0.1.0/suq/base_suq.py

@@ -0,0 +1,181 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from tqdm import tqdm
+from torch.distributions import Categorical
+from torch.distributions.normal import Normal
+from torch.utils.data import DataLoader
+
+from suq.utils.utils import torch_dataset
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+class SUQ_Base(nn.Module):
+    """
+    Base class for SUQ models.
+
+    Provides core functionality for:
+    - Managing likelihood type (regression or classification)
+    - Probit-based approximation for classification
+    - NLPD-based fitting of the scale factor
+
+    Inputs:
+        likelihood (str): Either 'classification' or 'regression'
+        scale_init (float): Initial value for the scale factor parameter
+    """
+    
+    def __init__(self, likelihood, scale_init):
+        super().__init__()
+        
+        if likelihood not in ['classification', 'regression']:
+            raise ValueError(f"Invalid likelihood type {likelihood}")
+        
+        self.likelihood = likelihood
+        self.scale_factor = nn.Parameter(torch.Tensor([scale_init]).to(device))
+    
+    def probit_approximation(self, out_mean, out_var):
+        """
+        Applies a probit approximation to compute class probabilities from the latent Gaussian distribution.
+        
+        Inputs:
+            out_mean (Tensor): Latent function mean, shape [B, C]
+            out_var (Tensor): Latent function variance, shape [B, C] or [B, C, C]
+
+        Outputs:
+            posterior_predict_mean (Tensor): Predicted class probabilities, shape [B, C]
+        """
+
+        if out_var.dim() == 3:
+            kappa = 1 / torch.sqrt(1. + np.pi / 8 * out_var.diagonal(dim1=1, dim2=2))
+        else:
+            kappa = 1 / torch.sqrt(1. + np.pi / 8 * out_var)
+            
+        posterior_predict_mean = torch.softmax(kappa * out_mean, dim=-1)
+        return posterior_predict_mean
+
+    def fit_scale_factor(self, data_loader, n_epoches, lr, speedup = True, verbose = False):
+        """
+        Fits the scale factor for predictive variance using negative log predictive density (NLPD).
+
+        Inputs:
+            data_loader (DataLoader): Dataloader containing (input, target) pairs
+            n_epoches (int): Number of epochs for optimization
+            lr (float): Learning rate for scale optimizer
+            speedup (bool): If True (classification only), caches forward pass outputs to accelerate fitting
+            verbose (bool): If True, prints NLPD at each epoch
+
+        Outputs:
+            total_train_nlpd (List[float]): Average NLPD per epoch over training data
+        """
+        print("fit scale factor")
+        optimizer = torch.optim.Adam([self.scale_factor], lr)
+        total_train_nlpd = []
+        
+        # store intermediate result and pack it into a data loader, so we only need to do one forward pass
+        if speedup:
+            
+            if self.likelihood == 'regression':
+                raise ValueError(f"Speed up not supported for regression atm")
+            
+            if self.likelihood == 'classification':
+
+                f_mean = []
+                f_var = []
+                labels = []
+
+                for (X, y) in tqdm(data_loader, desc= "packing f_mean f_var into a dataloader"):
+                    out_mean, out_var = self.forward_latent(X.to(device))
+                    f_mean.append(out_mean.detach().cpu().numpy())
+                    f_var.append(out_var.detach().cpu().numpy())
+                    if y.dim() == 2:
+                        labels.append(y.numpy().argmax(1).reshape(-1, 1))
+                    if y.dim() == 1:
+                        labels.append(y.numpy().reshape(-1, 1))
+
+                f_mean = np.vstack(f_mean)
+                f_var = np.vstack(f_var)
+                labels = np.vstack(labels)
+
+                scale_fit_dataset = torch_dataset(f_mean, f_var, labels)
+                scale_fit_dataloader = DataLoader(scale_fit_dataset, batch_size=16, shuffle=True)
+
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for data_pair in scale_fit_dataloader:
+                        x_mean, x_var_label = data_pair
+                        num_class = x_mean.shape[1]
+                        x_mean = x_mean.to(device)
+                        x_var, label = x_var_label.split(num_class, dim=1)
+                        x_var = x_var.to(device)
+                        label = label.to(device)
+
+                        optimizer.zero_grad()
+                        # make prediction
+                        x_var = x_var / self.scale_factor.to(device)
+                        posterior_predict_mean = self.probit_approximation(x_mean, x_var)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Categorical(posterior_predict_mean)
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                    total_train_nlpd.append(running_nlpd / len(scale_fit_dataloader))
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+                        
+                del scale_fit_dataloader
+                del scale_fit_dataset
+        
+        else:
+            
+            if self.likelihood == 'classification':
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for (data, label) in data_loader:
+                        
+                        data = data.to(device)
+                        label = label.to(device)
+
+                        optimizer.zero_grad()
+                        # make prediction
+                        posterior_predict_mean = self.forward(data)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Categorical(posterior_predict_mean)
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                    total_train_nlpd.append(running_nlpd / len(data_loader))
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+
+
+            if self.likelihood == 'regression':
+                for epoch in tqdm(range(n_epoches), desc="fitting scaling factor"):
+                    running_nlpd = 0
+                    for (data, label) in data_loader:
+                        data = data.to(device)
+                        label = label.to(device)
+                        
+                        optimizer.zero_grad()
+                        # make prediction
+                        posterior_predict_mean, posterior_predict_var = self.forward(data)
+                        # construct log posterior predictive distribution
+                        posterior_predictive_dist = Normal(posterior_predict_mean, posterior_predict_var.sqrt())
+                        # calculate nlpd and update
+                        nlpd = -posterior_predictive_dist.log_prob(label).mean()
+                        nlpd.backward()
+                        optimizer.step()
+                        # log nlpd
+                        running_nlpd += nlpd.item()
+                
+                    total_train_nlpd.append(running_nlpd / len(data_loader))
+                    
+                    if verbose:
+                        print(f"epoch {epoch + 1}, nlpd {total_train_nlpd[-1]}")
+                    
+        return total_train_nlpd