Annealed GMM Analytical Log Probabilities

[ahsanMah]

This may be a bit pedantic but, I wonder if there is a bug in the way the log probabilities are calculated for the annealed GMM distribution. My understanding is that annealed sampling uses scores from the ‘noise distribution’ which should reflect a Gaussian perturbed input. I’m assuming you guys achieve this by adding the noise distribution to the original GMM, hence allowing you to calculate the true log probabilities from the same GMM now scaled with the noise sigmas.

However, adding two Gaussian random variables should add the variances from both distributions as shown (taken from Wikipedia)

In this case sigma_x would be the original GMM sigma (which is set to 1), and sigma_y would correspond to the noise level. So shouldn't the calculation in the lines below be using something like sigma = np.sqrt(sigma**2 + self.sigma**2) instead of just the sigma passed in to the function? In my own tests, sampling with these updated sigmas gives more faithful results to the original distribution.

ncsn/models/gmm.py

Lines 45 to 46 in adb98fb

	logps.append((-((samples - self.means[i]) ** 2).sum(dim=-1) / (2 * sigma ** 2) - 0.5 * np.log(
	2 * np.pi * sigma ** 2)) + self.mix_probs[i].log())