Key Probabilistic Layers in probly

1. BayesLinear & BayesConv2d

What they do: These layers replace standard Linear and Conv2d layers to build a Bayesian Neural Network (BNN).

How they work: Instead of using fixed weights, each weight is modeled as a probability distribution (e.g., Gaussian). On every forward pass, weights are sampled from these distributions.

Result: The model can explicitly represent uncertainty in its own parameters.

---

2. DropConnectLinear

What it does: Implements the DropConnect technique as a neural network layer.

How it works: During training (and during inference when performing uncertainty quantification), individual weights are randomly set to zero. This is a more general form of Dropout.

Result: Running multiple forward passes produces a distribution of outputs that captures model uncertainty.

---

3. NormalInverseGammaLinear (for Evidential Regression)

What it does: A specialized output layer designed for Evidential Regression.

How it works: Instead of predicting a single value, the layer outputs the four parameters of a Normal-Inverse-Gamma (NIG) distribution:

• gamma

• nu

• alpha

• beta

Result: This distribution directly models:

• The predictive mean

• The predictive variance

• The model’s confidence in its own predictions

# Example 1: BayesLinear and BayesConv2d (via bayesian transformation)
import torch
from torch import nn

from probly.transformation import bayesian

# Create a standard model with Linear and Conv2d layers
model = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3), nn.ReLU(), nn.Flatten(), nn.Linear(8 * 26 * 26, 32), nn.ReLU(), nn.Linear(32, 10)
)

print("Original model:")
print(model)

# Transform to Bayesian - replaces Linear -> BayesLinear, Conv2d -> BayesConv2d
bnn_model = bayesian(model)

print("\nBayesian model:")
print(bnn_model)

# Test with dummy input
dummy_input = torch.randn(2, 1, 28, 28)
output = bnn_model(dummy_input)

print(f"\nOutput shape: {output.shape}")

# Example 2: DropConnectLinear (via dropconnect transformation)
import torch
from torch import nn

from probly.transformation import dropconnect

# Create a standard model
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))

print("Original model:")
print(model)

# Transform to DropConnect with p=0.25 (25% of weights dropped)
dc_model = dropconnect(model, p=0.25)

print("\nDropConnect model:")
print(dc_model)

# Test stochastic behavior
dummy_input = torch.randn(4, 10)

dc_model.train()  # Enable stochastic weight dropping
output1 = dc_model(dummy_input)
output2 = dc_model(dummy_input)

print("\nSame input, different outputs due to random weight dropping:")
print(f"Output 1: {output1[0, :3].detach().numpy()}")
print(f"Output 2: {output2[0, :3].detach().numpy()}")
print(f"Difference: {(output1[0] - output2[0]).abs().mean().item():.4f}")

# Example 3: NormalInverseGammaLinear (via evidential_regression transformation)
import torch
from torch import nn

from probly.transformation import evidential_regression

# Create a standard regression model
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))

print("Original model (outputs single value):")
print(model)

# Transform to evidential regression
# Replaces final Linear with NormalInverseGammaLinear
evid_model = evidential_regression(model)

print("\nEvidential regression model (outputs 4 NIG parameters):")
print(evid_model)

# Test output
dummy_input = torch.randn(2, 10)
output = evid_model(dummy_input)

print("\nOutput is a dictionary with 4 parameters:")
print(f"Keys: {output.keys()}")
print(
    f"Shapes: gamma={output['gamma'].shape}, nu={output['nu'].shape}, "
    f"alpha={output['alpha'].shape}, beta={output['beta'].shape}"
)