I am learning about VI and am implementing a GMM model for clustering using variational inference. However, my implementation is not fitting the data at all, even when initializing the cluster means to their true values.
I have generated the following data using two mixture components, with means N(1, 1) and N(7, 2), 100 points each, plotted with jitter.
I have the following code for dealing with variational parameters
class VParamNormal:
def __init__(self, size, m=None, log_s=None):
if m is None: m = torch.randn(size) * 0.01
if log_s is None: log_s = torch.randn(size) * 0.01
self.vp = torch.stack([m, log_s])
self.vp.requires_grad = True
self.size = size
def dist(self):
return torch.distributions.Normal(self.vp[0], self.vp[1].exp())
def rsample(self, n=torch.Size([])):
return self.dist().rsample(n)
def log_q(self, real):
return self.dist().log_prob(real).sum()
and have done the same for both the Gamma and Dirichlet distribution, using exp(param) for alpha and beta in the Gamma distribution and the exp for the concentration parameters in the Dirichlet.
To calculate the ELBO and likelihood, I use the following. I have skipped priors here because I wanted to keep things as simple as possible, and adding them doesn't make it work either.
def log_priors(real_params, params):
return 0
def log_q(model_params, real_params):
out = 0
for key in model_params:
out += model_params[key].log_q(real_params[key])
return out
def elbo(x, model_params):
params = {}
for key in model_params:
params[key] = model_params[key].rsample()
clust_assignments = (
torch.log(params["w"]) +
torch.distributions.Normal(params["means"], params["vars"]).log_prob(x[:, None].repeat(1, 2))
)
out = clust_assignments.sum()
out += log_priors(params, params)
out -= log_q(model_params, params)
return out
Because VI can't handle discrete variables, I marginalize out the categorical, drawn from the Dirichlet. Found in clust_assignments. Then, if I wanted to get the actual point clusters, I'd have to softmax these values (if my understanding is correct).
Then, optimization is simple, using standard techniques
model_params = {
"means": VParamNormal(size=(2,), m=torch.tensor([1., 7.])),
"vars": VParamGamma(size=(2,)),
"w": VParamDirichlet([1., 1.])
}
optimizer = torch.optim.Adam([model_params[key].vp for key in model_params], lr=0.1)
elbo_hist = []
max_iter = 3000
for t in iters:
loss = -elbo(x, model_params)
elbo_hist.append(-loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
However, when I fit this, the means do find their true values. Or in this case, I initialized them to the true means, and they drift together to something that might be the mean of all the points, as indicated by the red crosses.
I am a bit baffled by this and can't really see where I went wrong. Any help would be appreciated.
