House prices with a mixture density network¶

Backend: TensorFlow, JAX or PyTorch
Dataset: California housing
Techniques: GaussianHead, MixtureGaussianHead

View in Colab GitHub source

Predict the full distribution of house prices, including the spike at the census price cap.

Some targets do not follow a bell curve. In the California housing census, house values were capped at $500,001: about 5% of the districts sit exactly at the cap. A model predicting a single Gaussian has to smear its probability over that spike.

A mixture density network predicts a mixture of Gaussians for every input: several bumps, each with its own weight, center and width. It can put a narrow bump on the cap and a wide one elsewhere. In keras_proba it is the MixtureGaussianHead, trained with MixtureGaussianNLL.

# 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¶

Eight features per district (median income, house age, rooms, location, ...) and the median house value, in units of $100,000.

import matplotlib.pyplot as plt
import numpy as np

import keras
import keras_proba

keras.utils.set_random_seed(0)
CAP = 5.00001  # $500,001.

(x_train, y_train), (x_test, y_test) = keras.datasets.california_housing.load_data(version="large")
y_train, y_test = (y_train / 1e5).astype("float32")[:, None], (y_test / 1e5).astype("float32")[:, None]
mean, std = x_train.mean(0), x_train.std(0)
x_train, x_test = (x_train - mean) / std, (x_test - mean) / std

plt.figure(figsize=(7, 3))
plt.hist(y_train[:, 0], bins=100)
plt.xlabel("median house value ($100k)")
plt.title("The spike at the right is the $500k cap")
plt.show()

Output plot

Train a single Gaussian and a mixture¶

The same network, with two different output layers: GaussianHead predicts one Gaussian; MixtureGaussianHead(num_components=5) predicts a mixture of five.

def train(head, loss):
    model = keras.Sequential(
        [keras.Input((8,)), keras.layers.Dense(128, activation="relu"), keras.layers.Dense(128, activation="relu"), head]
    )
    model.compile(optimizer="adam", loss=loss)
    model.fit(
        x_train,
        y_train,
        validation_split=0.1,
        epochs=100,
        batch_size=128,
        callbacks=[keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)],
        verbose=0,
    )
    return model


gaussian = train(keras_proba.layers.GaussianHead(1), keras_proba.losses.GaussianNLL())
mixture = train(keras_proba.layers.MixtureGaussianHead(num_components=5), keras_proba.losses.MixtureGaussianNLL())

Which model explains the prices better?¶

The negative log-likelihood (NLL) measures how probable the true prices are under each predicted distribution: lower is better. The difference is largest for the districts at the cap.

gaussian_nll = keras.ops.convert_to_numpy(keras_proba.losses.gaussian_nll(y_test, gaussian.predict(x_test, verbose=0)))
mixture_nll = keras.ops.convert_to_numpy(keras_proba.losses.mixture_gaussian_nll(y_test, mixture.predict(x_test, verbose=0)))
capped = y_test[:, 0] >= CAP - 1e-3
print(f"Test NLL, all districts:     Gaussian {gaussian_nll.mean():.2f}, mixture {mixture_nll.mean():.2f}")
print(f"Test NLL, districts at cap:  Gaussian {gaussian_nll[capped].mean():.2f}, mixture {mixture_nll[capped].mean():.2f}")

Test NLL, all districts:     Gaussian 0.60, mixture 0.36
Test NLL, districts at cap:  Gaussian 1.91, mixture 0.28

Look at one predicted distribution¶

For a district at the cap, the single Gaussian can only stretch its one bell curve, which puts probability on values above the cap that never occur. The mixture adds a separate bump at the cap itself.

def density(params, values, mixture_model):
    if mixture_model:
        logits, means, scales = (keras.ops.convert_to_numpy(t)[0] for t in keras_proba.ops.split_mixture_params(params))
        weights = np.exp(logits) / np.exp(logits).sum()
        means, scales = means[:, 0], scales[:, 0]
    else:
        mean, scale = (keras.ops.convert_to_numpy(t)[0] for t in keras_proba.ops.split_gaussian_params(params))
        weights, means, scales = np.ones(1), mean, scale
    gaussians = np.exp(-0.5 * ((values[:, None] - means) / scales) ** 2) / (scales * np.sqrt(2 * np.pi))
    return gaussians @ weights


# One of the districts at the cap where the mixture puts the most probability near the cap.
mixture_params = mixture.predict(x_test[capped], verbose=0)
near_cap = [density(mixture_params[i : i + 1], np.array([CAP]), True)[0] for i in range(len(mixture_params))]
district = np.flatnonzero(capped)[int(np.argmax(near_cap))]
values = np.linspace(0, 7, 500)
plt.figure(figsize=(7, 3))
plt.plot(values, density(gaussian.predict(x_test[district : district + 1], verbose=0), values, False), label="single Gaussian")
plt.plot(values, density(mixture.predict(x_test[district : district + 1], verbose=0), values, True), label="mixture of 5")
plt.axvline(y_test[district, 0], color="k", ls="--", label="actual value (the cap)")
plt.xlabel("median house value ($100k)")
plt.ylabel("predicted density")
plt.legend()
plt.show()

Output plot

What we learned¶

A mixture density network predicts a full distribution, which can have several modes or spikes that a single Gaussian cannot represent.
MixtureGaussianHead and MixtureGaussianNLL are drop-in replacements for GaussianHead and GaussianNLL. keras_proba.ops.mixture_mean_and_stddev gives the predictive mean and spread.
Mixtures can get stuck covering only some of the modes: ReLU features and a modest learning rate train them reliably.