[Python] How to code a neural network that sample from a mixture of gaussian using isotropic...

I am trying to code a neural network that can sample from µ= p·N((1−p)a; In) + (1−p)·N(−pa; In) using isotropic diffusion (Y_t = t*x + W_t, with (W_t)_t Brownian motion. We use a = 1, p = 0.7, n = 128.

Each row is obtained using a different model for the posterior expectation m(y; t), and reports the histogram of ⟨Xt,a⟩/∥a∥2 2 obtained by 1000 independent runs of the generation process. Here Xt = m(Yt; t) is the sample generated at time t. These empirical results are compared with the correct distribution ⟨X,a⟩/∥a∥2 2 under X∼µ.

This is my algorithm :

Function Forward(x∈Rn,α∈[0,π/2]):
φ←(cos(α·i),sin(α·i); i≤20)
s←Lin0(φ)
x1 ←ReLU ◦Lin1(x))
x2 ←Flatten(s⊗x1)
xout ←cos(α)x+ Lin2(x2) ∈Rn
return xout


We encode t in terms of the angle variable α= arctan(1/√t).

We trained on N samples from µ. Parameters where chosen as follows: N = 5,000, 500 epochs, L = 3, m = 256;

We used stochastic gradient descent with batch size 50 over a number of epochs and samples that changes depending on the row of Figures 1, 2. A fixed number of samples is generated, while the noise G and signal-to-noise ratio t are resampled independently at each SGD sample. At training time we sample t= tan(α)−2 with α∼Unif([0,π/2]). At generation time we use a Euler discretization with K equi-spaced values of α, K ∈{200,400}and check that results are insensitive to the value of K.

The projection should at the end look like the image above.

My code is the following:

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal


def density_function(x):
pdf1 = multivariate_normal.pdf(x, mu1, cov)
pdf2 = multivariate_normal.pdf(x, mu2, cov)
return p * pdf1 + (1 - p) * pdf2

class NN1(nn.Module):
def __init__(self, n, m, L):
super(NN1, self).__init__()
self.L0 = nn.Linear(40, L)
self.L1 = nn.Sequential(nn.Linear(n, m), nn.ReLU())
self.L2 = nn.Linear(m * L, n)

def forward(self, x, alpha):
phi = torch.zeros(40, device=x.device)
for i in range(40):
if i % 2 == 0:
phi = torch.cos(alpha * i)
else:
phi = torch.sin(alpha * i)
s = self.L0(phi)
x_1 = self.L1(x)
kron_list = []
for i in range(x_1.size(0)):
kron_prod = torch.kron(s, x_1).flatten()
kron_list.append(kron_prod)
x_2 = torch.stack(kron_list) # shape: (batch_size, L*m)
x_out = torch.cos(alpha) * x + self.L2(x_2) # shape: (batch_size, n)
return x_out

def sample_from_mu(N, p, mu1, mu2, cov, n):
samples = torch.empty((N, n))
for i in range(N):
x = np.random.rand()
if x < p:
sample = np.random.multivariate_normal(mu1, cov)
else:
sample = np.random.multivariate_normal(mu2, cov)
samples = torch.from_numpy(sample).float()
return samples

# Numbers for Figure 1
n = 128
a = np.ones(n)
p = 0.7
mu1 = (1 - p) * a
mu2 = -p * a
cov = np.eye(n)
N = 5000
m = 256
L = 3
epochs = 500
batch_size = 50

x_train = sample_from_mu(N, p, mu1, mu2, cov, n)
train_dataset = TensorDataset(x_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

loss_fn = nn.MSELoss()
model = NN1(n, m, L)
optimizer = optim.SGD(model.parameters(), lr=0.1)

losses = []
for epoch in range(epochs):
for batch in train_loader:
x_batch = batch[0]
alpha = torch.empty(1).uniform_(0, torch.pi/2)
t = 1 / (torch.tan(alpha) ** 2)
W_t = torch.randn(x_batch.size(0), n) * torch.sqrt(t)
Y = t * x_batch + W_t
y_pred = model(Y, alpha)
loss = loss_fn(y_pred, x_batch)
losses.append(loss.item())
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

# Now use the trained NN to predict Xt
t1 = torch.tensor(0.21)
alpha_1 = torch.atan(1 / torch.sqrt(t1))

x_samples = sample_from_mu(1000, p, mu1, mu2, cov, n)
a_tensor = torch.ones(n, dtype=torch.float32)

X_hat = []
for i in range(1000):
W_1 = torch.randn(n) * torch.sqrt(t1)
Y_sample = t1 * x_samples + W_1
x_pred = model(Y_sample.unsqueeze(0), alpha_1)
projection = (x_pred.squeeze(0) @ a_tensor) / n
X_hat.append(projection.item())


I also tried to change my training with :

for epoch in range(epochs):
for batch in train_loader:
x_batch = batch[0] # shape: (batch_size, n)
# Sample a separate alpha for each sample in the batch.
alpha = torch.empty(x_batch.size(0)).uniform_(0, torch.pi/2)
t = 1 / (torch.tan(alpha) ** 2)
W_t = torch.randn(x_batch.size(0), n) * torch.sqrt(t).unsqueeze(1)
Y = t.unsqueeze(1) * x_batch + W_t
y_pred = model(Y, alpha)
loss = loss_fn(y_pred, x_batch)
losses.append(loss.item())
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()


to get more different time, but it still doesn't work.

You can notice that I did not use Euler discretization yet as I first wanted to tried on sample from my distribution and my prediction <X_t, 1>/n should converge to my true distribution.

My problem is that it doesn't. I printed my loss function and it does not converges. I also tried many different learning rate, but it just won't work.

Does anyone have any idea ?

Thank you !

Continue reading...