MediumCalculusca-193

ELBO Gradient Step

Problem

Take one gradient ascent step on the ELBO for a one-sample Gaussian variational posterior. With z = mu + sigma * eps and dropped constants, the objective is

L = -(z - x)^2 / 2 - z^2 / 2 + log(sigma) + eps^2 / 2

whose gradients are dL/dmu = x - 2z and dL/dsigma = eps * (x - 2z) + 1/sigma. Return [mu + lr*dL/dmu, sigma + lr*dL/dsigma].

Hint

The sampled value z depends on mu and sigma, so both partials flow through it.

Test cases

case 1

in: [0,1,1,0.5,0.1]

expected: [0,1.1]

case 2

in: [0,2,3,-1,0.05]

expected: [0.35000000000000003,1.675]

case 3

in: [1,1,1,0,0.1]

expected: [0.9,1.1]

case 4

in: [0,1,0,2,0.25]

expected: [-1,-0.75]

Starter codePython

def elbo_gradient_step(mu, sigma, x, eps, lr):
    # Return [mu_new, sigma_new]
    # Your code here
    pass

4 test cases · run your code against them in the editor