TensorFlow Uncertainty Estimation

Introduction

When building machine learning models, we often focus on making the most accurate predictions possible. However, in many real-world applications, knowing how confident a model is in its predictions can be just as important as the predictions themselves. This is where uncertainty estimation comes into play.

Uncertainty estimation allows models to express their confidence (or lack thereof) in their predictions. This is crucial for:

1. Safety-critical applications like healthcare or autonomous driving
2. Active learning scenarios where we need to decide what data to collect next
3. Decision support systems where human operators need to know when to trust or override model predictions

In this tutorial, we'll explore how to implement uncertainty estimation in TensorFlow using TensorFlow Probability (TFP), a library that makes it easy to combine probabilistic approaches with deep learning.

Types of Uncertainty

Before diving into code, it's important to understand two fundamental types of uncertainty:

1. Aleatoric Uncertainty: This represents the inherent noise in the data. It can't be reduced by collecting more data of the same kind.

2. Epistemic Uncertainty: This represents the model's ignorance - uncertainty that arises from limited data or knowledge. It can be reduced with more data.

Let's learn how to capture both types using TensorFlow.

Prerequisites

To follow along with this tutorial, you'll need:

python

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import matplotlib.pyplot as plt

Make sure you have TensorFlow Probability installed:

bash

pip install tensorflow-probability

Basic Uncertainty Estimation: Monte Carlo Dropout

One of the simplest ways to estimate uncertainty is using a technique called Monte Carlo (MC) Dropout. This involves:

1. Adding dropout layers to your neural network
2. Keeping dropout active during inference (not just training)
3. Running multiple forward passes to get a distribution of predictions

Let's implement a simple regression model with MC Dropout:

python

def create_mc_dropout_model(input_shape):
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=input_shape),
        tf.keras.layers.Dropout(0.5),  # Dropout layer with 50% drop rate
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dropout(0.5),  # Another dropout layer
        tf.keras.layers.Dense(1)       # Output layer (no activation for regression)
    ])
    
    model.compile(optimizer='adam', loss='mse')
    return model

Now, let's generate some sample data to test our approach:

python

# Generate synthetic data with noise
np.random.seed(42)
X = np.linspace(-3, 3, 100).reshape(-1, 1)
y_true = X**3 - 2*X**2 + 0.5*X
y = y_true + np.random.normal(0, 1, size=y_true.shape)  # Add noise

# Split into train and test
train_size = 70
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]

Now we train our model:

python

model = create_mc_dropout_model((1,))
model.fit(X_train, y_train, epochs=500, verbose=0)

For inference with uncertainty, we need to run multiple forward passes with dropout enabled:

python

def predict_with_uncertainty(model, X, num_samples=100):
    # Enable dropout during inference
    predictions = []
    for _ in range(num_samples):
        y_pred = model(X, training=True)  # Note training=True to enable dropout
        predictions.append(y_pred)
    
    # Stack the predictions
    predictions = tf.stack(predictions, axis=0)
    
    # Calculate mean and standard deviation
    mean_prediction = tf.reduce_mean(predictions, axis=0)
    std_prediction = tf.math.reduce_std(predictions, axis=0)
    
    return mean_prediction, std_prediction

# Make predictions with uncertainty
mean_pred, std_pred = predict_with_uncertainty(model, X_test)

# Visualize the results
plt.figure(figsize=(10, 6))
plt.scatter(X_train, y_train, label='Training Data', alpha=0.5)
plt.plot(X_test, y_test, 'g-', label='True Function (Test Data)', linewidth=2)
plt.plot(X_test, mean_pred, 'r-', label='Prediction Mean', linewidth=2)

# Plot uncertainty bands (±2 standard deviations covers ~95% confidence interval)
plt.fill_between(X_test.flatten(), 
                 (mean_pred - 2 * std_pred).numpy().flatten(), 
                 (mean_pred + 2 * std_pred).numpy().flatten(), 
                 color='red', alpha=0.3, label='Uncertainty (±2σ)')

plt.legend()
plt.title('Prediction with Uncertainty Estimation using MC Dropout')
plt.xlabel('X')
plt.ylabel('y')
plt.show()

The output will show your model's predictions along with uncertainty bands. Notice that uncertainty tends to be higher in regions with less training data!

Bayesian Neural Networks with TensorFlow Probability

For a more principled approach to uncertainty estimation, we can use Bayesian Neural Networks (BNNs). Unlike traditional neural networks, which have fixed weights, BNNs maintain a distribution over weights.

Let's implement a simple BNN using TensorFlow Probability:

python

def create_bnn(input_shape):
    # Define prior weight distribution
    def prior(kernel_size, bias_size, dtype=None):
        n = kernel_size + bias_size
        return {
            'kernel': tfp.distributions.Normal(loc=0., scale=1. / tf.sqrt(n / 2)),
            'bias': tfp.distributions.Normal(loc=0., scale=1.)
        }
    
    # Define variational posterior weight distribution
    def posterior(kernel_size, bias_size, dtype=None):
        n = kernel_size + bias_size
        return {
            'kernel': tfp.distributions.Normal(
                loc=tf.Variable(tf.random.normal([kernel_size], stddev=1. / tf.sqrt(n / 2))),
                scale=tf.nn.softplus(tf.Variable(tf.fill([kernel_size], -3.)))
            ),
            'bias': tfp.distributions.Normal(
                loc=tf.Variable(tf.random.normal([bias_size], stddev=1.)),
                scale=tf.nn.softplus(tf.Variable(tf.fill([bias_size], -3.)))
            )
        }
    
    # Build the model
    model = tf.keras.Sequential([
        tfp.layers.DenseVariational(
            units=64,
            activation='relu',
            input_shape=input_shape,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1/X_train.shape[0]
        ),
        tfp.layers.DenseVariational(
            units=64,
            activation='relu',
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1/X_train.shape[0]
        ),
        tfp.layers.DenseVariational(
            units=1,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1/X_train.shape[0]
        )
    ])
    
    # Define a negative log likelihood loss function
    def neg_log_likelihood(y_true, y_pred):
        return -tfp.distributions.Normal(loc=y_pred, scale=1).log_prob(y_true)
    
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
                 loss=neg_log_likelihood)
    
    return model

Now let's train our Bayesian Neural Network:

python

bnn_model = create_bnn((1,))
bnn_model.fit(X_train, y_train, epochs=1000, verbose=0)

For inference, we perform multiple forward passes to get a distribution of predictions:

python

def predict_with_bnn(model, X, num_samples=100):
    predictions = []
    for _ in range(num_samples):
        y_pred = model(X)
        predictions.append(y_pred)
    
    predictions = tf.stack(predictions, axis=0)
    mean_prediction = tf.reduce_mean(predictions, axis=0)
    std_prediction = tf.math.reduce_std(predictions, axis=0)
    
    return mean_prediction, std_prediction

# Make predictions with the BNN
bnn_mean, bnn_std = predict_with_bnn(bnn_model, X_test)

# Visualize BNN results
plt.figure(figsize=(10, 6))
plt.scatter(X_train, y_train, label='Training Data', alpha=0.5)
plt.plot(X_test, y_test, 'g-', label='True Function (Test Data)', linewidth=2)
plt.plot(X_test, bnn_mean, 'b-', label='BNN Prediction Mean', linewidth=2)

# Plot uncertainty bands
plt.fill_between(X_test.flatten(), 
                 (bnn_mean - 2 * bnn_std).numpy().flatten(), 
                 (bnn_mean + 2 * bnn_std).numpy().flatten(), 
                 color='blue', alpha=0.3, label='BNN Uncertainty (±2σ)')

plt.legend()
plt.title('Prediction with Uncertainty Estimation using Bayesian Neural Network')
plt.xlabel('X')
plt.ylabel('y')
plt.show()

Direct Prediction of Aleatoric Uncertainty

We can also train models to directly predict their uncertainty along with their outputs. This is particularly useful for capturing aleatoric uncertainty:

python

def create_uncertainty_model(input_shape):
    inputs = tf.keras.Input(shape=input_shape)
    x = tf.keras.layers.Dense(64, activation='relu')(inputs)
    x = tf.keras.layers.Dense(64, activation='relu')(x)
    
    # Mean prediction
    mean = tf.keras.layers.Dense(1)(x)
    
    # Log variance prediction (we predict log(σ²) for numerical stability)
    log_var = tf.keras.layers.Dense(1)(x)
    
    # Define a custom loss that incorporates the predicted uncertainty
    def heteroscedastic_loss(y_true, y_pred):
        mean, log_var = y_pred[:, 0:1], y_pred[:, 1:2]
        precision = tf.exp(-log_var)
        return tf.reduce_mean(
            0.5 * precision * tf.square(y_true - mean) + 0.5 * log_var
        )
    
    model = tf.keras.Model(inputs=inputs, outputs=tf.keras.layers.Concatenate()([mean, log_var]))
    model.compile(optimizer='adam', loss=heteroscedastic_loss)
    
    return model

Let's train this model:

python

hetero_model = create_uncertainty_model((1,))
hetero_model.fit(X_train, y_train, epochs=1000, verbose=0)

And visualize the predictions with uncertainty:

python

# Predict using heteroscedastic model
preds = hetero_model.predict(X_test)
mean_preds = preds[:, 0]
var_preds = np.exp(preds[:, 1])  # Convert log variance to standard deviation
std_preds = np.sqrt(var_preds)

# Visualize
plt.figure(figsize=(10, 6))
plt.scatter(X_train, y_train, label='Training Data', alpha=0.5)
plt.plot(X_test, y_test, 'g-', label='True Function (Test Data)', linewidth=2)
plt.plot(X_test, mean_preds, 'c-', label='Heteroscedastic Prediction Mean', linewidth=2)

# Plot uncertainty bands
plt.fill_between(X_test.flatten(), 
                 mean_preds - 2 * std_preds, 
                 mean_preds + 2 * std_preds, 
                 color='cyan', alpha=0.3, label='Predicted Aleatoric Uncertainty (±2σ)')

plt.legend()
plt.title('Direct Prediction of Aleatoric Uncertainty')
plt.xlabel('X')
plt.ylabel('y')
plt.show()

Real-World Application: Medical Diagnosis

Let's apply uncertainty estimation to a more practical example: medical diagnosis using a simplified dataset.

python

# Generate a synthetic medical dataset
# Features: age, blood pressure, glucose level, etc.
np.random.seed(42)
n_samples = 1000
n_features = 5

# Generate synthetic features
X_medical = np.random.normal(0, 1, size=(n_samples, n_features))

# Generate synthetic diagnosis (binary classification)
# True function with some randomness
w = np.random.randn(n_features)
logits = np.dot(X_medical, w)
p_true = 1 / (1 + np.exp(-logits))
y_medical = np.random.binomial(1, p_true)

# Split into train and test
train_size = int(0.8 * n_samples)
X_med_train, X_med_test = X_medical[:train_size], X_medical[train_size:]
y_med_train, y_med_test = y_medical[:train_size], y_medical[train_size:]

# Create a dropout-based uncertainty model for classification
def create_mc_dropout_classifier(input_shape, dropout_rate=0.3):
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=input_shape),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])
    
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    return model

# Train the model
med_model = create_mc_dropout_classifier((n_features,))
med_model.fit(X_med_train, y_med_train, epochs=100, verbose=0)

# Predict with uncertainty
def predict_classification_with_uncertainty(model, X, num_samples=100):
    predictions = []
    for _ in range(num_samples):
        y_pred = model(X, training=True)  # Keep dropout active
        predictions.append(y_pred)
    
    # Stack predictions
    predictions = tf.stack(predictions, axis=0)
    
    # Calculate mean probability and uncertainty
    mean_prob = tf.reduce_mean(predictions, axis=0)
    uncertainty = tf.math.reduce_std(predictions, axis=0)
    
    return mean_prob, uncertainty

# Get predictions
mean_probs, uncertainties = predict_classification_with_uncertainty(med_model, X_med_test)

# Analyze results - let's focus on high uncertainty predictions
threshold = 0.2  # High uncertainty threshold
high_uncertainty_indices = np.where(uncertainties > threshold)[0]

print(f"Total test samples: {X_med_test.shape[0]}")
print(f"Number of high uncertainty predictions: {len(high_uncertainty_indices)}")

# Calculate accuracy for certain vs uncertain predictions
predictions = (mean_probs > 0.5).numpy().flatten()
correct_predictions = (predictions == y_med_test)

certain_indices = np.where(uncertainties <= threshold)[0]
if len(certain_indices) > 0:
    certain_accuracy = np.mean(correct_predictions[certain_indices])
    print(f"Accuracy for certain predictions: {certain_accuracy:.2f}")

if len(high_uncertainty_indices) > 0:
    uncertain_accuracy = np.mean(correct_predictions[high_uncertainty_indices])
    print(f"Accuracy for uncertain predictions: {uncertain_accuracy:.2f}")

# Create a visualization of uncertainty vs prediction correctness
plt.figure(figsize=(10, 6))
plt.scatter(mean_probs.numpy(), uncertainties.numpy(), 
           c=correct_predictions, cmap='coolwarm', alpha=0.7)
plt.axhline(y=threshold, color='r', linestyle='--', label='Uncertainty Threshold')
plt.colorbar(label='Prediction Correct')
plt.xlabel('Predicted Probability')
plt.ylabel('Uncertainty (Standard Deviation)')
plt.title('Prediction Uncertainty vs Correctness')
plt.legend()
plt.show()

This would produce an output showing how the model's uncertainty correlates with its prediction accuracy. In a medical context, cases with high uncertainty could be flagged for review by a human expert.

Summary

In this tutorial, we've covered several approaches to uncertainty estimation in TensorFlow:

1. Monte Carlo Dropout: A simple technique that uses dropout during inference to estimate uncertainty.
2. Bayesian Neural Networks: A more principled approach that maintains distributions over model weights.
3. Direct Uncertainty Prediction: Training models to predict both outputs and their uncertainty.

We've also seen how uncertainty estimation can be applied to real-world problems like medical diagnosis, where knowing when a model is uncertain can be critical.

Remember, uncertainty estimation is not just about improving accuracy—it's about building safer, more trustworthy AI systems that know when they don't know.

Additional Resources

• TensorFlow Probability Documentation
• "Dropout as a Bayesian Approximation" by Gal and Ghahramani
• "Weight Uncertainty in Neural Networks" by Blundell et al.
• TensorFlow Tutorials on Probabilistic Layers

Exercises

1. Modify the MC Dropout model to use different dropout rates and observe how this affects the uncertainty estimates.
2. Apply uncertainty estimation to a real dataset of your choice (e.g., from Kaggle or UCI ML Repository).
3. Implement a model that captures both aleatoric and epistemic uncertainty simultaneously.
4. Create a classification model that outputs calibrated probabilities and evaluate its calibration using reliability diagrams.
5. Build an active learning system that selects new training examples based on model uncertainty.