Deep Ensembles: Leveraging Ensemble Methods for Uncertainty Estimation in AI (with codes)

Ensemble methods are powerful tools for estimating uncertainty in machine learning models, especially for tasks where predicting confidence levels along with predictions is essential, such as medical diagnosis, autonomous driving, and financial forecasting. By combining predictions from multiple models, ensemble methods can capture a broader range of potential outcomes and help quantify uncertainty in a model’s predictions.

Deep ensembles involve training multiple instances of the same model architecture independently. Each model in the ensemble is initialized with different random seeds and potentially trained on different data subsets, leading to diverse models. This allows the ensemble to capture a wider range of patterns and nuances present within the data, ultimately enhancing the model’s generalization capabilities. By combining predictions from these varied models, deep ensembles mitigate the risk of overfitting to a specific dataset, resulting in more robust and reliable performance across different scenarios and tasks. Additionally, the integration of multiple models can enhance the predictive power significantly, as the ensemble can aggregate the strengths of individual models while compensating for their weaknesses, thus providing a more comprehensive and accurate output.

To use Deep Ensemble for uncertainty estimation: For each prediction, the ensemble outputs predictions across the different models (one prediction from each model in the ensemble). The variance in predictions reflects model uncertainty, while the mean gives the final prediction. While this is simple and intuitive, it requires training multiple models, which can be computationally expensive and memory-intensive. Therefore, for large scale datasets, I would recommend using Monte Carlo Dropout for estimating uncertainty instead.

Import Libraries

First, we import the necessary libraries such as PyTorch, scikit-learn, and numpy.

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np

Load and Preprocess the Dataset

Next, we fetch the California Housing dataset, split it into training and testing sets, and standardize the features.

# Load and preprocess the dataset
data = fetch_california_housing()
X, y = data.data, data.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Standardize the features
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

Convert Data to PyTorch Tensors

We convert the data to PyTorch tensors for compatibility with PyTorch models.

# Convert data to PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

Define a Simple Feedforward Neural Network Model

We define a neural network with three fully connected layers. The forward method defines the forward pass using ReLU activation functions for the hidden layers.

# Define a simple feedforward neural network model
class SimpleNN(nn.Module):
    def __init__(self, input_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

Train a Single Model

We define a function to train the model using Mean Squared Error loss and the Adam optimizer.

# Train a single model
def train_model(model, X_train, y_train, epochs=100, lr=0.001):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs.squeeze(), y_train)
        loss.backward()
        optimizer.step()
    return model

Create an Ensemble of Models

We create an ensemble of 10 models, each trained independently on the same data.

# Define the number of models in the ensemble
ensemble_size = 10
input_dim = X_train.shape[1]
# Create an ensemble of models
ensemble_models = []
for i in range(ensemble_size):
    model = SimpleNN(input_dim)
    trained_model = train_model(model, X_train, y_train, epochs=100, lr=0.001)
    ensemble_models.append(trained_model)

Make Predictions Using Each Model in the Ensemble

We define a function to make predictions using the ensemble, calculating the mean and variance of the predictions.

# Make predictions using each model in the ensemble and calculate the mean and variance
def ensemble_predict(models, X):
    predictions = []
    for model in models:
        model.eval()
        with torch.no_grad():
            predictions.append(model(X).squeeze().numpy())
    predictions = np.array(predictions)
    mean_prediction = np.mean(predictions, axis=0)
    variance_prediction = np.var(predictions, axis=0)  # Uncertainty estimate
    return mean_prediction, variance_prediction

Get Predictions on the Test Set

We use the trained ensemble to get predictions on the test set.

# Get predictions on the test set
mean_pred, var_pred = ensemble_predict(ensemble_models, X_test)

Print the Results

We print the mean predictions and the uncertainty (variance) of the predictions.

# Print the results
print("Mean Predictions:", mean_pred)
print("Uncertainty (Variance of Predictions):", var_pred)

combined codes:

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np
# Load and preprocess the dataset
data = fetch_california_housing()
X, y = data.data, data.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Standardize the features
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
# Convert data to PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)
# Define a simple feedforward neural network model
class SimpleNN(nn.Module):
    def __init__(self, input_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# Train a single model
def train_model(model, X_train, y_train, epochs=100, lr=0.001):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs.squeeze(), y_train)
        loss.backward()
        optimizer.step()
        
    return model
# Define the number of models in the ensemble
ensemble_size = 10
input_dim = X_train.shape[1]
# Create an ensemble of models
ensemble_models = []
for i in range(ensemble_size):
    model = SimpleNN(input_dim)
    trained_model = train_model(model, X_train, y_train, epochs=100, lr=0.001)
    ensemble_models.append(trained_model)
# Make predictions using each model in the ensemble and calculate the mean and variance
def ensemble_predict(models, X):
    predictions = []
    for model in models:
        model.eval()
        with torch.no_grad():
            predictions.append(model(X).squeeze().numpy())
    
    predictions = np.array(predictions)
    mean_prediction = np.mean(predictions, axis=0)
    variance_prediction = np.var(predictions, axis=0)  # Uncertainty estimate
    
    return mean_prediction, variance_prediction
# Get predictions on the test set
mean_pred, var_pred = ensemble_predict(ensemble_models, X_test)
# Print the results
print("Mean Predictions:", mean_pred)
print("Uncertainty (Variance of Predictions):", var_pred)