Interactive demo · Bayesian and Probabilistic Machine Learning · Chapter 5

Mixture density networks: when the inverse problem has several answers

Data: x uniform on (0, 1), t = x + 0.3 sin(2πx) + noise. The forward problem (predict t from x) has one answer per x. In the inverse problem the roles are swapped, so in the middle of the range one input has three answers. A network trained with the sum-of-squares error predicts the conditional mean, which lies between the branches. A mixture density network (MDN) outputs the mixing coefficients πk(x), means μk(x) and widths σk(x) of a Gaussian mixture and is trained by maximum likelihood (5.153).

Problem

Model

Mixture components K

Hidden units M

Try these

N = 250 points, seeded. The sum-of-squares network has 5 tanh hidden units; the MDN has M tanh hidden units and 3K outputs: softmax for the mixing coefficients, exp for the standard deviations (5.150)–(5.152). Both are trained with Adam, full batch (step 0.02), from N(0, 1) first-layer and N(0, 1/4) second-layer weights. In the heat map the sum-of-squares network is shown with a Gaussian of the residual standard deviation. The conditional mode is approximated, as in the book, by the mean of the most probable component. MDN training can end in a poor local optimum: try “New initial weights”.

In the lecture notes