Predict the Noise Too

Train a network that predicts a full Gaussian for each input: a mean and a standard deviation that changes with x.

Plain regression with mean squared error assumes the noise is the same everywhere. Here it is not. The data is y = sin(x) + noise, and the noise is small near x = 0 and large near the edges. A model that also predicts sigma(x) can say where it is unsure.

Train with the Gaussian negative log likelihood:

loss = -Normal(mean(x), std(x)).log_prob(y).mean()

Give the network two outputs. Use the first as the mean. Pass the second through softplus and add a small floor, so the standard deviation is always positive.

Return a function predict(X) that gives (mean, std) for new inputs. It is graded on held-out data. The average negative log likelihood must be close to that of the true noise model, and about 95% of the targets must fall inside mean +- 1.96 * std.

Function signature

def train_gaussian_regressor(x_train, y_train):

Example

Input

predict = train_gaussian_regressor(x_train, y_train)
mean, std = predict(torch.tensor([[0.0], [2.5]]))

Output

mean = tensor([0.00, 0.60])    # about sin(0) and sin(2.5)
std  = tensor([0.06, 0.68])    # about 0.05 and 0.675

Reasoning

The true noise standard deviation is 0.05 + 0.25 * |x|, which is 0.05 at x = 0 and 0.675 at x = 2.5. A good model learns both the curve and this spread.

Constraints and notes

• x_train is float32 of shape (n, 1) with values in [-3, 3]. y_train is float32 of shape (n,). n is about 800.
• predict returns two float tensors of shape (m,), with every std > 0.
• Training must take a few seconds at most on a CPU. Set a seed so it is repeatable.