bug while fitting a multivariate normal

[aldopareja]

I'm trying to fit a multivariate normal and it's not working. This should be trivial since I'm initializing at the optimum.

import jax
from flax import linen as nn
import jax.numpy as jnp
import optax

samples = jax.random.multivariate_normal(jax.random.PRNGKey(12346), jnp.array([5]), jnp.array([[10.0]]), shape=(1000,))

class MultiVariateNormal(nn.Module):
    
    @nn.compact
    def __call__(self, samples, mu, cov_terms):
        mu = self.param('mu',
                    nn.initializers.constant(5.0), # Initialization function
                    (1,)) 
        cov = self.param('cov',
                    nn.initializers.constant(10.0), # Initialization function
                    (1,1)) 
        return jax.vmap(jax.scipy.stats.multivariate_normal.logpdf, in_axes=(0,None,None))(x,mu, cov)

def update_step(apply_fn, samples, opt_state, params, mu, cov_terms):
    def loss(params):
        l = apply_fn(params, samples, mu, cov_terms)
        return -l.sum()
    
    l, grads = jax.value_and_grad(loss, has_aux=False)(params)
    updates, opt_state = tx.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return opt_state, params, l

m = MultiVariateNormal()
params = m.init(jax.random.PRNGKey(567), jnp.ones((10,2)), jnp.array([0.5,0.5]), jnp.array([1.0,0.0,1.0]))
tx = optax.adam(learning_rate=0.01)
opt_state = tx.init(params)


for i in range(1000):
    # opt_state, params, l = jax.jit(update_step, static_argnums=(0,))(m.apply, samples, opt_state, params, jnp.array([0.5,0.5]), jnp.array([1.0,0.0,1.0]))
    opt_state, params, l = update_step(m.apply, samples, opt_state, params, jnp.array([0.5,0.5]), jnp.array([1.0,0.0,1.0]))
    if i % 100 == 0:
        print(l)
        print(params)

This is the output I'm getting:

3086.906
FrozenDict({
    params: {
        cov: DeviceArray([[10.01]], dtype=float32),
        mu: DeviceArray([4.9900002], dtype=float32),
    },
})
2697.4744
FrozenDict({
    params: {
        cov: DeviceArray([[10.756521]], dtype=float32),
        mu: DeviceArray([4.044425], dtype=float32),
    },
})
2456.2666
FrozenDict({
    params: {
        cov: DeviceArray([[10.664117]], dtype=float32),
        mu: DeviceArray([3.224071], dtype=float32),
    },
})
2265.7512
FrozenDict({
    params: {
        cov: DeviceArray([[9.686479]], dtype=float32),
        mu: DeviceArray([2.4944096], dtype=float32),
    },
})
2087.3037
FrozenDict({
    params: {
        cov: DeviceArray([[8.19261]], dtype=float32),
        mu: DeviceArray([1.8418326], dtype=float32),
    },
})
1908.8782
FrozenDict({
    params: {
        cov: DeviceArray([[6.4797406]], dtype=float32),
        mu: DeviceArray([1.2831304], dtype=float32),
    },
})
1709.6995
FrozenDict({
    params: {
        cov: DeviceArray([[4.6283555]], dtype=float32),
        mu: DeviceArray([0.8519289], dtype=float32),
    },
})
1408.9744
FrozenDict({
    params: {
        cov: DeviceArray([[2.5492816]], dtype=float32),
        mu: DeviceArray([0.58346385], dtype=float32),
    },
})
222.00249
FrozenDict({
    params: {
        cov: DeviceArray([[0.12405935]], dtype=float32),
        mu: DeviceArray([0.50057507], dtype=float32),
    },
})
177.4865
FrozenDict({
    params: {
        cov: DeviceArray([[0.0836898]], dtype=float32),
        mu: DeviceArray([0.50000083], dtype=float32),
    },
})

weird enough, when printing the loss output of the found parameters and the optimal parameters it clearly shows that the optimal are better, but somehow the printed loss is different. What could be the problem?

mu = params['params']['mu']
cov = params['params']['cov']
print(jax.vmap(jax.scipy.stats.multivariate_normal.logpdf, in_axes=(0,None,None))(samples, mu, cov).sum())
print(jax.vmap(jax.scipy.stats.multivariate_normal.logpdf, in_axes=(0,None,None))(samples, jnp.array([0.5]), jnp.array([[10.0]])).sum())

-179831.36
-3574.5342

[aldopareja]

I have a bug in the code, using x instead of samples in the forward function. Sorry.