Bayesian weather forecasting: what the model knows and what it doesn't¶
Backend: TensorFlow, JAX or PyTorch
Dataset: Jena climate
Techniques: Conv1DFlipout, DenseFlipout, GaussianHead
Separate the noise in the weather from the model's own uncertainty, with Bayesian layers.
A forecast can be uncertain for two different reasons:
- Aleatoric uncertainty: the weather is noisy. More data would not remove it.
- Epistemic uncertainty: the model is unsure, e.g. because the input is unlike anything it was trained on. More data would reduce it.
A Bayesian layer learns a distribution
over its weights instead of a single value. Every prediction samples new weights, so repeating a
prediction gives slightly different forecasts: their spread is the epistemic uncertainty. Adding
a GaussianHead on top also predicts the aleatoric part.
# Install keras-proba, then choose a backend: "jax", "tensorflow" or "torch".
!pip install -q keras-proba
import os
os.environ["KERAS_BACKEND"] = "jax"
Load the data¶
As in the probabilistic forecasting
example: predict the temperature 24 hours ahead from the last 48 hours of the
Jena climate data, split in time. Here the inputs keep
their time axis, (48 hours, 5 measurements), for a 1D convolution.
import glob
import matplotlib.pyplot as plt
import numpy as np
import keras
import keras_proba
keras.utils.set_random_seed(0)
LOOKBACK, HORIZON = 48, 24 # Hours.
path = keras.utils.get_file(
origin="https://storage.googleapis.com/tensorflow/tf-keras-datasets/jena_climate_2009_2016.csv.zip",
extract=True,
)
csv = glob.glob(os.path.join(path, "**", "*.csv"), recursive=True)[0]
data = np.genfromtxt(csv, delimiter=",", skip_header=1, usecols=(1, 2, 4, 5, 12), dtype="float32")[5::6]
data[:, 4] = np.maximum(data[:, 4], 0) # -9999 marks missing wind measurements.
temperature = data[:, 1].copy()
train_end, val_end = int(0.7 * len(data)), int(0.85 * len(data))
features = (data - data[:train_end].mean(0)) / data[:train_end].std(0)
def windows(start, stop):
"""Inputs: the last 48 hours of features. Targets: the temperature 24 hours later."""
ends = np.arange(max(start, LOOKBACK), stop - HORIZON + 1)
x = np.stack([features[t - LOOKBACK : t] for t in ends])
return x, temperature[ends + HORIZON - 1][:, None]
x_train, y_train = windows(0, train_end)
x_val, y_val = windows(train_end, val_end)
x_test, y_test = windows(val_end, len(data))
A Bayesian model¶
Conv1DFlipout and DenseFlipout are the Bayesian versions of Conv1D and Dense. Training
balances fitting the data against keeping the weight distributions close to a prior; kl_weight
sets how much the prior counts, and 1 / number of training examples is the standard choice.
Bayesian layers learn more slowly than standard ones, so we train for longer, with a higher
learning rate.
kl_weight = 1.0 / len(x_train)
model = keras.Sequential(
[
keras.Input(x_train.shape[1:]),
keras_proba.layers.Conv1DFlipout(32, 5, padding="same", activation="relu", kl_weight=kl_weight),
keras_proba.layers.Conv1DFlipout(32, 5, strides=2, padding="same", activation="relu", kl_weight=kl_weight),
keras.layers.Flatten(),
keras_proba.layers.DenseFlipout(64, activation="relu", kl_weight=kl_weight),
keras_proba.layers.GaussianHead(1),
]
)
model.compile(optimizer=keras.optimizers.Adam(3e-3), loss=keras_proba.losses.GaussianNLL())
model.fit(
x_train,
y_train,
validation_data=(x_val, y_val),
epochs=40,
batch_size=256,
callbacks=[keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True)],
verbose=0,
)
Repeat the forecast to measure what the model doesn't know¶
Each predict call samples new weights. Over 30 calls:
- the average of the predicted variances is the aleatoric uncertainty;
- the variance of the predicted means is the epistemic uncertainty;
and their sum is the total predictive variance.
def forecast(x, num_samples=30):
means, variances = [], []
for _ in range(num_samples):
mean, std = keras_proba.ops.split_gaussian_params(model.predict(x, batch_size=1024, verbose=0))
means.append(keras.ops.convert_to_numpy(mean)[:, 0])
variances.append(keras.ops.convert_to_numpy(std)[:, 0] ** 2)
means, variances = np.array(means), np.array(variances)
return means.mean(0), np.sqrt(variances.mean(0)), np.sqrt(means.var(0))
mean, aleatoric, epistemic = forecast(x_test)
total = np.sqrt(aleatoric**2 + epistemic**2)
print(f"Average uncertainty (°C): aleatoric {aleatoric.mean():.2f}, epistemic {epistemic.mean():.2f}")
print(f"90% intervals contain the actual temperature {np.mean(np.abs(y_test[:, 0] - mean) <= 1.645 * total):.1%} of the time")
hours = np.arange(24 * 10)
plt.figure(figsize=(9, 3))
plt.fill_between(hours, (mean - 1.645 * total)[hours], (mean + 1.645 * total)[hours], alpha=0.3, label="90% interval (total)")
plt.fill_between(hours, (mean - 1.645 * epistemic)[hours], (mean + 1.645 * epistemic)[hours], alpha=0.5, label="epistemic part")
plt.plot(hours, y_test[hours, 0], "k", lw=1, label="actual")
plt.xlabel("hour")
plt.ylabel("temperature (°C)")
plt.legend()
plt.show()
Average uncertainty (°C): aleatoric 3.04, epistemic 0.77
90% intervals contain the actual temperature 92.2% of the time
Inputs the model has never seen¶
Epistemic uncertainty should grow when the inputs leave the range of the training data, e.g. after a sensor fault. We shift every measurement by 3 standard deviations: a standard network would forecast with unchanged confidence.
_, _, epistemic_shifted = forecast(x_test[:2000] + 3.0)
print(f"Epistemic uncertainty (°C): normal inputs {epistemic[:2000].mean():.2f}, shifted inputs {epistemic_shifted.mean():.2f}")
What we learned¶
- Bayesian layers are drop-in replacements for
DenseandConv1D(alsoConv2D,Conv3D). - Repeating predictions splits the uncertainty into noise in the data (aleatoric) and the model's own uncertainty (epistemic).
- Epistemic uncertainty flags inputs unlike the training data, which a deterministic model,
even with a
GaussianHead, cannot do.