UNet1DModel does not converge

[MrInformatic]

Describe the bug

I tried to train a UNet1DModel, DDPMScheduler Diffusion Pipeline using AdamW optimizer and mse_loss. No matter what I tried, I never got the model to produce a loss below 0.5. As a sanity check, I also tried to replace the UNet1DModel with a UNet2DModel, which performed significantly better. Both Pipelines should produce silence or a blank image respectively. It seems like something is wrong with the UNet1DModel since this is the only part which was changed. #3203 Also mentions problems with UNet1DModel, but I tried to train my model with different learning rates using HPO allready.

Reproduction

import torch
from diffusers import UNet1DModel, DDPMScheduler, UNet2DModel
from diffusers.utils.torch_utils import randn_tensor
from torch.nn.functional import mse_loss
import matplotlib.pyplot as plt

def test_diffusers(dimensions: int):
    device = torch.device("cpu")
    sample_size = 32
    generator = torch.Generator(device=device)

    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

    if dimensions == 1:
        model = UNet1DModel(
            sample_size=sample_size,
            in_channels=1,
            out_channels=1,
            block_out_channels=(
                64,
            ),
            down_block_types=(
                "DownBlock1D",
            ),
            up_block_types=(
                "UpBlock1D",
            ),
        ).to(device)

        shape = (16, 1, sample_size)
    elif dimensions == 2:
        model = UNet2DModel(
            sample_size=sample_size,
            in_channels=1,
            out_channels=1,
            block_out_channels=(
                64,
            ),
            down_block_types=(
                "DownBlock2D",
            ),
            up_block_types=(
                "UpBlock2D",
            ),
        ).to(device)

        shape = (16, 1, sample_size, sample_size)
    else:
        raise Exception("only 1D and 2D are supported")

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    losses = []
    for i in list(range(100)):
        audio = torch.zeros(shape, device=device)
        noise = randn_tensor(shape, generator=generator, device=device, dtype=audio.dtype)

        batch_size = audio.shape[0]
        time_steps = torch.randint(0, noise_scheduler.num_train_timesteps, (batch_size,), device=device).long()

        noisy_voice = noise_scheduler.add_noise(audio, noise, time_steps)

        target = noise

        pred = model(noisy_voice, time_steps, return_dict=False)[0]

        loss = mse_loss(pred, target)
        loss.backward(loss)
        losses.append(loss.item())

        optimizer.step()
        optimizer.zero_grad()

    return losses

if __name__ == "__main__":
    losses_1d = test_diffusers(1)
    losses_2d = test_diffusers(2)

    plt.plot(losses_1d)
    plt.plot(losses_2d)
    plt.legend(["1D", "2D"], loc="upper right")

    plt.savefig(f"plot.png")

Logs

System Info

diffusers: 0.32.2
torch: 2.6.0
Python: 3.10.14
OS: Manjaro Linux
CPU: AMD Ryzen 5 1600X
GPU: Nvidia RTX 3090 24GB
RAM: 32 GB

Who can help?

No response

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.