层归一化应用于扩散模型，通过归一化神经网络层中的激活，尤其是在迭代噪声预测过程中，来稳定训练。 扩散模型的工作原理是逐渐向数据添加噪声，然后学习通过一系列去噪步骤来逆转这一过程。 核心架构，通常是 U-Net，在每个步骤中处理带噪输入，层归一化有助于在这些步骤中保持一致的信号尺度。 与依赖批统计的批量归一化不同，层归一化使用每个样本的统计信息，使其适用于扩散训练中常见的具有不同噪声水平和小批量大小的场景。
在实践中，层归一化被集成到 U-Net 架构的残差块中。 例如，在每个下采样或上采样块中，归一化应用于卷积层或注意力机制之前。 一个典型的序列可能包括：（1）卷积层，（2）层归一化，（3）非线性激活（例如，SiLU），以及（4）跳跃连接。 在基于 transformer 的扩散变体中，层归一化应用于多头自注意力或前馈层之前，以标准化输入。 时间步长嵌入，将模型置于当前去噪步骤中，通常在归一化之后注入，以避免破坏归一化的特征统计信息。 这种放置确保了在数百或数千个去噪步骤中，反向传播期间的稳定梯度流。
层归一化在扩散模型中的主要好处是它能够处理时间步长中不同的噪声幅度。 例如，早期的去噪步骤处理高度噪声的数据，而后面的步骤处理几乎干净的数据。 层归一化适应这些变化，而不依赖于批次范围的统计信息，这可能不一致。 与批量归一化等替代方案相比，这有助于更快的收敛和更稳定的训练。 虽然一些扩散实现（例如，DDPM）使用组归一化来实现计算效率，但层归一化仍然是一个灵活的选择，尤其是在结合了卷积层和注意力层的架构中，其中组件之间的特征尺度可能差异很大。