Simple Diffusion Model with Checkerboard Data

I have been trying to build a simple diffusion model with checker board data which is made from dots (x, y). So it is 2d data.

def get_noise_level(step, total_steps, max_noise_level=1.0):
    # Define the mean and standard deviation for the Gaussian distribution
    mean = total_steps / 2
    std_dev = total_steps / 4  # This can be adjusted based on desired spread

    # Calculate the Gaussian noise level
    gaussian_noise_level = np.exp(-((step - mean) ** 2) / (2 * std_dev ** 2))

    # Scale the noise level to be between 0 and max_noise_level
    return gaussian_noise_level * max_noise_level

Above is my func. to get the noise level

# Model definition
class DiffusionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

Above is my simple model

dataset = TensorDataset(data)
data_loader = DataLoader(dataset, batch_size=512, shuffle=True)

model = DiffusionModel()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
total_steps = 200

for epoch in range(num_epochs):
    epoch_loss = 0.0
    for i, (batch,) in enumerate(data_loader):
        current_step = epoch * len(data_loader) + i
        noise_level = get_noise_level(current_step, total_steps)
        noise = torch.randn_like(batch) * noise_level
        noisy_data = batch + noise

        # Train model
        optimizer.zero_grad()
        predicted_noise = model(noisy_data)
        loss = nn.MSELoss()(predicted_noise, noise)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        
    average_epoch_loss = epoch_loss / len(data_loader)
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {average_epoch_loss:.4f}")

Above is the simple training code.

Here, I observed a weird loss decrease:

Epoch 1/10, Loss: 0.2600
Epoch 2/10, Loss: 0.0017
Epoch 3/10, Loss: 0.0000
Epoch 4/10, Loss: 0.0000
Epoch 5/10, Loss: 0.0000
Epoch 6/10, Loss: 0.0000
Epoch 7/10, Loss: 0.0000
Epoch 8/10, Loss: 0.0000
Epoch 9/10, Loss: 0.0000
Epoch 10/10, Loss: 0.0000

Also when I try to generate samples with the model, all the values I generate are NaN. Any idea about the problem here? I got the data from here.

source https://stackoverflow.com/questions/77671619/simple-diffusion-model-with-checkerboard-data