要在 PyTorch 中实现一个基本的扩散模型，首先定义核心组件：前向过程（向数据添加噪声）、反向过程（去除噪声）以及预测噪声的神经网络。前向过程使用预定义的噪声时间表在多个时间步长上逐步破坏输入数据（例如图像）。例如，你可以定义一个 beta 值（噪声水平）的线性时间表，该时间表在 1000 个步骤中从 1e-4 增加到 0.2。每一步使用 x_t = sqrt(alpha_bar[t]) * x_0 + sqrt(1 - alpha_bar[t]) * epsilon 计算带噪声的样本，其中 alpha_bar 是 1 - beta 的累积乘积。
接下来，构建 U-Net 模型来预测反向过程中的噪声。一个最小的 U-Net 可能包括带有残差连接的卷积块、注意力层以及下采样/上采样步骤。例如，使用 PyTorch 的 nn.Module 创建一个包含编码器和解码器块的网络。编码器在增加通道的同时减小空间维度，而解码器则执行相反的操作。包含时间嵌入层以将时间步长信息注入到网络中，这通常通过正弦位置编码或学习到的嵌入来完成。模型将带噪声的图像和时间步长作为输入，输出预测的噪声。
对于训练，采样一批数据，生成随机时间步长，向数据添加噪声，然后训练模型以预测添加的噪声。使用 Adam 优化器和预测噪声与实际噪声之间的均方误差 (MSE) 损失。在推理（采样）期间，从纯噪声开始，并通过在每个时间步长预测和减去噪声来迭代去噪。使用 ddpm_sampler 函数应用反向扩散步骤，根据预测的噪声和噪声时间表调整样本。保持代码模块化，分离模型、训练循环和采样逻辑以提高清晰度。