Keras Probability

Uncertainty awareness for deep learning in Keras 3.

keras_proba gives Keras models calibrated, distance-aware uncertainty estimates. Every layer, loss and metric runs on the TensorFlow, JAX and PyTorch backends, plugs into model.compile / model.fit, and saves with model.save.

The distance-aware and Bayesian layers are ports of TensorFlow Model Garden and TensorFlow Probability, tested for numerical parity with them.

---

What's inside

Need	Use
A classifier that knows when inputs are far from the training data	SNGP: SpectralNormalization + RandomFeatureGaussianProcess, or the all-in-one GaussianProcessClassificationHead
Uncertainty over the weights	Bayesian layers: DenseFlipout, DenseReparameterization, Conv{1,2,3}DFlipout, ...
Regression with a predicted noise level	GaussianHead with GaussianNLL
Multi-modal or skewed regression targets	MixtureGaussianHead with MixtureGaussianNLL
Measuring calibration	ExpectedCalibrationError, BrierScore
Fixing an overconfident trained classifier	TemperatureScaling
Intervals or sets with a coverage guarantee	SplitConformalRegressor, SplitConformalClassifier

---

Installation

pip install keras-proba

Keras 3 needs one backend; install it with an extra, e.g. pip install "keras-proba[jax]" ([torch], [jax] or [tensorflow]).

---

A first example

A classifier whose predictions become uncertain away from the training data:

import numpy as np
import keras
import keras_proba

x = np.random.normal(size=(512, 2)).astype("float32")
y = (x[:, 0] * x[:, 1] > 0).astype("int32")

inputs = keras.Input((2,))
hidden = keras.layers.Dense(64, trainable=False)(inputs)  # A frozen random projection.
hidden = keras_proba.layers.SpectralNormalization(keras.layers.Dense(64, activation="relu"))(hidden)
logits = keras_proba.layers.GaussianProcessClassificationHead(
    inner_dim=0,
    num_classes=2,
    temperature=np.pi / 8,
    normalize_input=False,
    scale_random_features=True,
    custom_random_features_initializer=keras.initializers.RandomNormal(stddev=1.0),
)(hidden)
model = keras.Model(inputs, logits)

model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy", keras_proba.metrics.ExpectedCalibrationError(from_logits=True)],
)
model.fit(x, y, epochs=10, verbose=0, callbacks=[keras_proba.callbacks.ResetCovarianceCallback()])

probs = keras.ops.softmax(model.predict(np.array([[0.5, 0.5], [8.0, -8.0]], "float32"), verbose=0))
# Confident near the data, close to uniform far from it.

Continue with Getting Started, jump to a guide, or run an example notebook on Google Colab.

---

Verified against the references

Each component is tested on all three backends for:

• Parity with the original implementation (Model Garden, TensorFlow Probability, PyTorch, SciPy), feeding both the same weights and random draws.
• Training: the method does its job on real fits, e.g. out-of-distribution uncertainty, improved calibration, or the theoretical conformal coverage.
• Serialization: model.save / keras.models.load_model round trips.

See Testing for details, including the bugs in the references that these tests surfaced.