gradient descent solving mu, sigma for generative gaussian

gd_gaussian.py

# -*- coding: utf-8 -*-
"""
generative gaussian model, minimizing <-log p>_data  wrt. mu, sigma
with gradient descent

-log p = log sigma + log sqrt(2pi) + (x - mu)^2 / (2 sigma^2)
So

grad_mu (-log p) = - (x - mu) / sigma^2
grad_sigma (-log p) = 1 / sigma - (x - mu)^2 / sigma^3


"""
import numpy as np

N = 10000
data = (np.random.randn(N, 1) + 0.9) * 1.5
print(f'empirical mean: {np.mean(data)} empirical std: {np.std(data)}')
mu, sigma = 0.1, 0.5
loss = np.mean(np.log(sigma) + np.log(np.sqrt(2*np.pi)) + (data - mu)**2 / (2 * sigma**2))
lr = 1e-4

for i in range(3*10**4):
    grad_mu = np.mean(- (data - mu) / sigma**2)
    grad_sigma = np.mean(1. / sigma - (data - mu)**2 / sigma**3)
    mu -= lr * grad_mu
    sigma -= lr * grad_sigma
    if i % 10**3 == 0:
        print(f'mu: {mu} sigma: {sigma}')