Two Pressures on One Batch

A VAE's encoder does not hand over a code. For each input it hands over a centre and a spread for every code entry, and the code is drawn from that little cloud. Training then scores the result with two terms that pull in opposite directions, and one number $\beta$ that sets the balance between them.

In this problem the random draws have already been made, so everything is deterministic. For one input $x$ with $P$ positions and a code with $d$ entries:

step	rule
draw the code	$z_j = c_j + s_j\,\varepsilon_j$
decode	$\hat{x}_i = \sum_j W_{ij}\, z_j + b_i$
reconstruction term	$\dfrac{1}{P}\sum_{i=1}^{P} (\hat{x}_i - x_i)^2$
KL term	$\dfrac{1}{2}\sum_{j=1}^{d}\left(c_j^2 + s_j^2 - 1 - \ln s_j^2\right)$

Here $c_j$ is the centre, $s_j$ the spread (always positive), and $\varepsilon_j$ the draw from the standard shape at the origin (a standard normal), all for code entry $j$ ; $\ln$ is the natural logarithm. The decoder is a single linear layer: weights $W$ with one row per position and one column per code entry, and bias $b$ .

Task: write vae_loss(inputs, centres, spreads, noise, weights, bias, beta).

Row $n$ of centres, spreads and noise belongs to row $n$ of inputs. The batch may hold a single input.
Average the reconstruction term over the batch, and separately average the KL term over the batch.
The loss is the average reconstruction term $+\ \beta\ \times$ the average KL term, with $\beta$ given by beta.
Return [reconstruction, kl, loss], each rounded to 4 decimal places. Round only at the very end: the loss is built from the unrounded averages.

The two terms want different things. Reconstruction would like every input's cloud small and far from the others, so a draw can never be mistaken for a neighbour's. The KL term charges every cloud for straying from one standard shape at the origin. $\beta$ is the exchange rate, and the whole behaviour of a VAE lives in that trade.

Two Pressures on One Batch

Examples

Hints