Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget
如何最大限度地利用有限的预算:在微型预算下从头开始进行扩散训练

We denote the data distribution by $\mathcal{D}$, such that $({\boldsymbol{x}},{\boldsymbol{c}})\sim\mathcal{D}$, where ${\boldsymbol{x}}$ and ${\boldsymbol{c}}$ represent the image and the corresponding natural language caption, respectively. We assume that $p({\boldsymbol{x}};\sigma)$ is the image distribution obtained after adding $\sigma^{2}$-variance Gaussian noise.
我们用 $𝒟\backslash mathcal\{D\}$ 表示数据分布,使得 $(𝒙,𝒄)\sim 𝒟(\{\backslash boldsymbol\{x\}\},\{\backslash boldsymbol\{c\}\})\backslash sim\backslash mathcal\{D\}$ ,其中 $𝒙\{\backslash boldsymbol\{x\}\}$ 和 $𝒄\{\backslash boldsymbol\{c\}\}$ 分别表示图像和相应的自然语言字幕。我们假设 $p(𝒙;\sigma )p(\{\backslash boldsymbol\{x\}\};\backslash sigma)$ 是在添加 ${\sigma }^2\backslash sigma^\{2\}$ 方差高斯噪声后获得的图像分布。

Diffusion-based probabilistic models use a forward-reverse process-based approach, where the forward process corrupts the input data and the reverse process tries to reconstruct the original signal by learning to reconstruct the data degradation at each step. Though multiple variations of diffusion models exists (Ho et al., 2020; Song et al., 2020b; Bansal et al., 2024), it is common to use time-dependent Gaussian noise ($\sigma(t)$) for model data corruption. Under deterministic sampling, both the forward and reverse processes in diffusion models can be modeled using an ordinary differential equation (ODE).
基于扩散的概率模型使用基于正向-逆向过程的方法,其中正向过程会破坏输入数据,而逆向过程则试图通过学习在每一步重建数据退化来重建原始信号。尽管存在多种扩散模型的变体(何等, 2020; 宋等, 2020b; 班萨尔等, 2024),但使用时间依赖高斯噪声( $\sigma (t)\backslash sigma(t)$ )来对模型数据进行破坏是常见的。在确定性采样下,扩散模型中的正向和逆向过程都可以使用常微分方程(ODE)进行建模。

where $\dot{\sigma}(t)$ represents the time derivative and $\nabla_{\boldsymbol{x}}\log p({\boldsymbol{x}};\sigma(t))$ is the score of the underlying density function. Thus the move towards (or away from) higher density regions in the diffusion process is proportional to both the scores of the density function and changes in the noise distribution ($\sigma(t)$) with time. Going beyond the deterministic sampling based on the ODE formulation in Equation 1, the sampling process in diffusion models can also be formulated as a stochastic differential equation,
其中 $\dot{\sigma }(t)\backslash dot\{\backslash sigma\}(t)$ 表示时间导数, ${\nabla }_{𝒙}\log p(𝒙;\sigma (t))\backslash nabla\_\{\backslash boldsymbol\{x\}\}\backslash log\; p(\{\backslash boldsymbol\{x\}\};\backslash sigma(t))$ 为基础密度函数的分数。因此,在扩散过程中向更高密度区域移动(或远离)与密度函数的分数和噪声分布 ( $\sigma (t)\backslash sigma(t)$ ) 随时间的变化成正比。超越方程 1 中基于 ODE 的确定性采样,扩散模型中的采样过程也可以表述为一个随机微分方程。

where $\text{d}w_{t}$ is standard the Wiener process and the parameter $\beta(t)$ controls the rate of injection of additional noise. The generation process in diffusion models starts by sampling an instance from $p({\boldsymbol{x}},\sigma_{\text{max}})$ and iteratively denoising it using the ODE or SDE-based formulation.
其中 $\text{d}{w}_t\backslash text\{d\}w\_\{t\}$ 是标准的维纳过程,参数 $\beta (t)\backslash beta(t)$ 控制额外噪声的注入速率。扩散模型的生成过程从采样 $p(𝒙,{\sigma }_{\text{max}})p(\{\backslash boldsymbol\{x\}\},\backslash sigma\_\{\backslash text\{max\}\})$ 的一个实例开始,然后使用基于 ODE 或 SDE 的公式进行迭代去噪。

Training. Diffusion model training is based on parameterizing the score of the density function, which doesn’t depend on the generally intractable normalization factor, using a denoiser, i.e., $\nabla_{\boldsymbol{x}}\log p({\boldsymbol{x}};\sigma(t))\approx\left(F_{\theta}({\boldsymbol{x}},\sigma(t))-{\boldsymbol{x}}\right)/\sigma(t)^{2}$. The denoising function $F_{\theta}({\boldsymbol{x}},\sigma(t))$ is generally modeled using a deep neural network. For text-to-image generation, the denoiser is conditioned on both noise distribution and natural language captions and trained using the following loss function,
训练。扩散模型训练是基于参数化密度函数的分数,该分数不依赖于通常难以处理的归一化因子,使用去噪器,即 ${\nabla }_{𝒙}\log p(𝒙;\sigma (t))\approx \left(F_{\theta }(𝒙,\sigma (t))-𝒙\right)/\sigma {(t)}^2\backslash nabla\_\{\backslash boldsymbol\{x\}\}\backslash log\; p(\{\backslash boldsymbol\{x\}\};\backslash sigma(t))\backslash approx\backslash left(F\_\{\backslash theta\}(\{\backslash boldsymbol\{x\}\},\backslash sigma(t))-\{\backslash boldsymbol\{x\}\}\backslash right)/\backslash sigma(t)^\{2\}$ 。去噪函数 $F_{\theta }(𝒙,\sigma (t))F\_\{\backslash theta\}(\{\backslash boldsymbol\{x\}\},\backslash sigma(t))$ 通常使用深度神经网络建模。对于文本到图像的生成,去噪器同时依赖于噪声分布和自然语言字幕,并使用以下损失函数进行训练。

The noise distribution during training is lognormal ($\text{ln}(\sigma)\sim\mathcal{N}(P_{\text{mean}},P_{\text{std}})$), where the choice of the mean and standard deviation of the noise distribution strongly influences the quality of generated samples. For example, the noise distribution is shifted to the right to sample larger noise values on higher resolution images, as the signal-to-noise ratio increases with resolution (Teng et al., 2023).
训练期间的噪声分布是对数正态分布( $\text{ln}(\sigma )\sim 𝒩(P_{\text{mean}},P_{\text{std}})\backslash text\{ln\}(\backslash sigma)\backslash sim\backslash mathcal\{N\}(P\_\{\backslash text\{mean\}\},P\_\{\backslash text\{std\}\})$ ),其中噪声分布的均值和标准差的选择会严重影响生成样本的质量。例如,对于分辨率更高的图像,信噪比会随着分辨率的增加而提高,因此我们将噪声分布向右偏移,以采样更大的噪声值(Teng 等人, 2023)。

Classifier-free guidance. To increase alignment between input captions and generated images, classifier-free guidance (Ho & Salimans, 2022) modifies the image denoiser to output a linear combination of denoised samples in the presence and absence of input captions.
无分类器引导。 为了增加输入标题和生成图像之间的对齐度,无分类器引导(Ho & Salimans, 2022)修改了图像去噪器,以输出在存在和不存在输入标题的情况下的去噪样本的线性组合。

where $w\geq 1$ controls the strength of guidance. During training, a fraction of image captions (commonly set to 10%) are randomly dropped to learn unconditional generation.
控制指引的强度。在训练期间,图像标题的一部分(通常设置为 10%)被随机删除以学习无条件生成。>

Latent diffusion models. In contrast to modeling higher dimensional pixel space ($\mathbb{R}^{h\times w\times 3}$), latent diffusion models (Rombach et al., 2022) are trained in a much lower dimensional compressed latent space ($\mathbb{R}^{\frac{h}{n}\times\frac{w}{n}\times c}$), where $n$ is the compression factor and $c$ is the number of channels in the latent space. The image-to-latent encoding and its reverse mapping are performed by a variational autoencoder. Latent diffusion models achieve faster convergence than training in pixel space and have been heavily adopted in large-scale diffusion models (Podell et al., 2024; Betker et al., 2023; Chen et al., 2023; Balaji et al., 2022).
潜在扩散模型。 与建模更高维的像素空间( ${ℝ}^{h\times w\times 3}\backslash mathbb\{R\}^\{h\backslash times\; w\backslash times\; 3\}$ )相比,潜在扩散模型 (Rombach et al., 2022) 在低维压缩的潜在空间( ${ℝ}^{\frac{h}{n}\times \frac{w}{n}\times c}\backslash mathbb\{R\}^\{\backslash frac\{h\}\{n\}\backslash times\backslash frac\{w\}\{n\}\backslash times\; c\}$ )中进行训练,其中 $nn$ 是压缩系数, $cc$ 是潜在空间的通道数。图像到潜在编码及其反向映射由变分自编码器完成。潜在扩散模型的收敛速度比在像素空间中训练快,并已被大规模扩散模型大量采用(Podell et al., 2024; Betker et al., 2023; Chen et al., 2023; Balaji et al., 2022)。

Diffusion transformers. We consider that the transformer model ($F_{\theta}$) consists of $k$ hierarchically stacked transformer blocks. Each block consists of a multi-head self-attention layer, a multi-head cross-attention layer, and a feed-forward layer. For simplicity, we assume that all images are converted to a sequence of patches to be compatible with the transformer architecture. In the transformer architecture, we refer to the width of all linear layers as $d$.
扩散型变换器。 我们认为 transformer 模型( $F_{\theta }F\_\{\backslash theta\}$ ) 由 $kk$ 层次堆叠的 transformer 块组成。每个块由一个多头自注意力层、一个多头交叉注意力层和一个前馈层组成。为简单起见,我们假设所有图像都转换为与 transformer 架构兼容的一系列块。在 transformer 架构中,我们将所有线性层的宽度称为 $dd$ 。

Assuming transformer architectures where linear layers account for the majority of the computational cost, the total training cost is proportional to $M\times N\times S$, where $M$ is the number of samples processed, $N$ is the number of model parameters, and $S$ is the number of patches derived from one image, i.e., sequence size for transformer input (Figure 2a). As our goal is to train an open-world model, we believe that a large number of parameters are required to support diverse concepts during generation; thus, we opt to not decrease the parameter count significantly. Since diffusion models converge slowly and are often trained for a large number of steps, even for small-scale datasets, we keep this choice intact (Rombach et al., 2022; Balaji et al., 2022). Instead, we aim to exploit any redundancy in the input sequence to reduce computational cost while simultaneously not drastically degrade performance.
$M\times N\times SM\backslash times\; N\backslash times\; S$ ，其中 $MM$ 是处理的样本数量， $NN$ 是模型参数数量， $SS$ 是从一个图像衍生的 patch 数量，即 transformer 输入的序列大小（图2a）。由于我们的目标是训练一个开放世界模型，我们认为需要大量的参数来支持生成过程中的多样化概念;因此，我们选择不显著减少参数数量。由于扩散模型收敛缓慢，即使对于小规模数据集，也需要进行大量步骤的训练,我们保留了这一选择（Rombach 等人，2022；Balaji 等人，2022）。相反，我们旨在利用输入序列中的冗余来降低计算成本,同时不大幅降低性能。

A. Using larger patch size (Figure 2b). In a vision transformer, the input image is first transformed into a sequence of non-overlapping patches with resolution $p\times p$, where $p$ is the patch size (Dosovitskiy et al., 2020). Though using a higher patch size quadratically reduces the number of patches per image, it can significantly degrade performance due to the aggressive compression of larger regions of the image into a single patch.
使用更大的块大小（图 2b）。在视觉 Transformer 中,输入图像首先被转换为一系列不重叠的分块,分辨率为 $p\times pp\backslash times\; p$ ,其中 $pp$ 是块大小(Dosovitskiy 等人,2020)。尽管使用更高的块大小可以将每个图像的分块数量减少到四分之一,但由于将图像的较大区域压缩为单个分块,这可能会严重降低性能。

B. Using patch masking (Figure 2c). A contemporary approach is to keep the patch size intact but drop a large fraction of patches at the input layer of the transformer (He et al., 2022). Naive token masking is similar to training on random crops in convolutional UNets, where often upsampling diffusion models are trained on smaller random crops rather than the full image (Nichol et al., 2021; Saharia et al., 2022). However, in contrast to random crops of the images, patch masking allows training on non-continuous regions of the image²²2We observe consistent degradation in the quality of image generation with block masking, i.e., retaining continuous regions of the image (Appendix C). Thus, we argue that random patch masking in transformers offers a more powerful masking paradigm compared to random image crops in a convolutional network under identical computational cost.
我们观察到使用块遮罩导致图像生成质量的持续恶化,即保留图像的连续区域(附录C)。因此,我们认为与卷积网络中的随机图像裁剪相比,变压器中的随机分块遮蔽提供了更强大的遮蔽范式,计算成本相同。
使用补丁掩码（图 2c）。一种当代方法是保持补丁大小不变,但在变换器的输入层丢弃大部分补丁(He 等人,2022 年)。简单的标记掩码类似于在卷积 UNets 上训练随机裁剪,其中通常将扩散模型训练在较小的随机裁剪上而不是完整的图像(Nichol 等,2021 年;Saharia 等,2022 年)。然而,与图像的随机裁剪相比,补丁掩码允许在图像的非连续区域上进行训练。. Due to its effectiveness, masking-based training of transformers has been adopted across both vision and language domains (Devlin et al., 2018; He et al., 2022).
由于其有效性,基于掩蔽的变换器训练已广泛应用于视觉和语言领域(Devlin et al., 2018; He et al., 2022)。

C. Masking patches with autoencoding (MaskDiT - Figure 2d). In order to also encourage representation learning from masked patches, Zheng et al. (2024) adds an auxiliary autoencoding loss that encourages the reconstruction of masked patches. This formulation was motivated by the success of the autoencoding loss in learning self-supervised representations (He et al., 2022). MaskDiT first appends a lightweight decoder, a small-scale transformer network, to the larger backbone transformer. Before the decoder, it expands the backbone transformer output by inserting spurious learnable patches in place of masked patches.
使用自编码掩蔽补丁(MaskDiT - 图 2d)。为了鼓励从遮蔽的补丁中学习表示,郑等人(2024)增加了一个辅助自编码损失,鼓励重建被遮蔽的补丁。这一表述是受到自编码损失在学习自监督表示方面成功的启发(何等人,2022)。MaskDiT 首先在更大的主干变换器网络上附加一个轻量级解码器,即一个小规模的变换器网络。在解码器之前,它通过在被遮蔽的补丁位置插入虚构的可学习补丁来扩展主干变换器的输出。

Assuming that the binary mask $m$ represents the masked patches, the final training loss is the following,
假设二进制掩码 $mm$ 表示掩蔽补丁,最终的训练损失如下:

where $\bar{F}_{\theta}$ represents the sequential backbone transformer and decoder module, and $\gamma$ is a hyperparameter to balance the influence of the masked autoencoding loss with respect to the diffusion training loss.
${\overline{F}}_{\theta }\backslash bar\{F\}\_\{\backslash theta\}$ 表示顺序骨架变压器和解码器模块, $\gamma \backslash gamma$ 是一个超参数,用于平衡掩码自动编码损失与扩散训练损失的影响。

To drastically reduce the computational cost, patch masking requires dropping a large fraction of input patches before they are fed into the backbone transformer, thus making the information from masked patches unavailable to the transformer. High masking ratios, e.g., 75% masking, significantly degrade the overall performance of the transformer. Even with MaskDiT, we only observe a marginal improvement over naive masking, as this approach also drops the majority of image patches at the input layer itself.
为了大幅降低计算成本,补丁遮罩需要在将输入补丁传递到主干变换器之前丢弃大部分补丁,从而使被遮罩补丁的信息对变换器不可用。高遮罩比率(例如 75%遮罩)会严重降低变换器的整体性能。即使采用 MaskDiT,我们也只观察到对朴素遮罩的微小改善,因为该方法也在输入层自身就丢弃了大部分图像补丁。

Deferred masking to retain semantic information of all patches. As high masking ratios remove the majority of valuable learning signals from images, we are motivated to ask whether it is necessary to mask at the input layer. As long as the computational cost is unchanged, it is merely a design choice and not a fundamental constraint. In fact, we uncover a significantly better masking strategy, with nearly identical cost to the existing MaskDiT approach. As the patches are derived from non-overlapping image regions in diffusion transformers (Peebles & Xie, 2023; Dosovitskiy et al., 2020), each patch embedding does not embed any information about other patches in the image. Thus, we aim to preprocess the patch embeddings before masking, such that non-masked patches can embed information about the whole image. We refer to the preprocessing module as patch-mixer.
保留所有补丁的语义信息的延迟遮蔽。 由于高遮蔽率会移除图像中大部分有价值的学习信号,我们被激励去问是否有必要在输入层进行遮蔽。只要计算成本不变,这仅仅是一个设计选择,而不是一个基本约束。事实上,我们发现了一种显著更好的遮蔽策略,其成本与现有的 MaskDiT 方法几乎相同。由于补丁来自扩散变压器中的非重叠图像区域(Peebles & Xie, 2023; Dosovitskiy et al., 2020),每个补丁嵌入不会嵌入关于图像中其他补丁的任何信息。因此,我们的目标是在遮蔽之前预处理补丁嵌入,使得未被遮蔽的补丁可以嵌入有关整个图像的信息。我们将这个预处理模块称为补丁混合器。

Training diffusion transformers with a patch-mixer. We consider a patch mixer to be any neural architecture that can fuse the embeddings of individual patches. In transformer models, this objective is naturally achieved with a combination of attention and feedforward layers. So we use a lightweight transformer comprising only a few layers as our patch mixer. We mask the input sequence tokens after they have been processed by the patch mixer (Figure 2e). Assuming a binary mask $m$ for masking, we train our model using the following loss function,
使用 patch-mixer 训练扩散变换器。我们认为 patch mixer 是任何能够融合单个 patch 的嵌入的神经架构。在 transformer 模型中,这一目标通过注意力和前馈层的组合自然实现。因此,我们使用由少量层组成的轻量级 transformer 作为我们的 patch mixer。我们在 patch mixer（图 2e）处理后屏蔽输入序列令牌。假设有一个用于掩码的二进制掩码 $mm$ ,我们使用以下损失函数训练我们的模型。

where $M_{\phi}$ is the patch-mixer model and $F_{\theta}$ is the backbone transformer. Note that compared to MaskDiT, our approach also simplifies the overall design by not requiring an additional loss function and corresponding hyperparameter tuning between two losses during training. We don’t mask any patches during inference.
其中 $M_{\varphi }M\_\{\backslash phi\}$ 是 patch-mixer 模型, $F_{\theta }F\_\{\backslash theta\}$ 是 backbone transformer。值得注意的是,与 MaskDiT 相比,我们的方法还通过在训练过程中不需要额外的损失函数及相应的超参数调整来简化了整体设计。我们在推理过程中不会遮蔽任何 patches。

Unmasked finetuning. As a very high masking ratio can significantly reduce the ability of diffusion models to learn global structure in the images and introduce a train-test distribution shift in sequence size, we consider a small degree of unmasked finetuning after the masked pretraining. The finetuning can also mitigate any undesirable generation artifacts due to the use of patch masking. Thus, it has been crucial in previous work (Zheng et al., 2024) to recover the sharp performance drop from masking, especially when classifier-free guidance is used in sampling. However, we don’t find it completely necessary, as even in masked pretraining, our approach achieves comparable performance to the baseline unmasked pretraining. We only employ it in large-scale training to mitigate any unknown-unknown generation artifacts from a high degree of patch masking.
无掩蔽微调。 由于非常高的掩蔽率可能会显著降低扩散模型学习图像全局结构的能力,并在序列大小上引入训练-测试分布偏移,我们在掩蔽预训练之后考虑进行少量的无掩蔽微调。微调还可以缓解由于使用分块掩蔽而产生的任何不受欢迎的生成伪影。因此,在之前的工作(Zheng 等人, 2024)中,它在恢复由于掩蔽而导致的性能下降方面至关重要,尤其是在采样中使用无分类引导时。然而,我们发现它并非完全必要,因为即使在掩蔽预训练中,我们的方法也达到了与基线无掩蔽预训练相当的性能。我们只在大规模训练中使用它来缓解由于高度分块掩蔽而产生的任何未知-未知的生成伪影。

Improving backbone transformer architecture with mixture-of-experts (MoE) and layer-wise scaling. We also leverage innovations in transformer architecture design to boost the performance of our model under computational constraints. We use mixture-of-experts layers as they increase the parameters and expressive power of our model without significantly increasing the training cost (Shazeer et al., 2017; Fedus et al., 2022). We use a simplified MoE layer based on expert-choice routing (Zhou et al., 2022), where each expert determines the tokens routed to it, as it doesn’t require any additional auxiliary loss functions to balance the load across experts (Shazeer et al., 2017; Zoph et al., 2022). We also consider a layer-wise scaling approach, which has recently been shown to outperform canonical transformers in large-language models (Mehta et al., 2024). It linearly increases the width, i.e., hidden layer dimension in attention and feedforward layers, of transformer blocks. Thus, deeper layers in the network are assigned more parameters than earlier layers. We argue that since deeper layers in visual models tend to learn more complex features (Zeiler & Fergus, 2014), using higher parameters in deeper layers would lead to better performance. We provide background details on both techniques in Appendix B. We describe the overall architecture of our diffusion transformer in Figure 3.
使用专家组合(MoE)和逐层缩放改进 transformer 架构。 我们还利用 transformer 架构设计中的创新来提高模型在计算约束下的性能。我们使用专家组合层,因为它们可以增加模型的参数和表达能力,而不显著增加训练成本(Shazeer et al., 2017; Fedus et al., 2022)。我们使用基于专家选择路由的简化 MoE 层(Zhou et al., 2022),其中每个专家确定路由到其中的令牌,因为它不需要任何额外的辅助损失函数来平衡专家之间的负载(Shazeer et al., 2017; Zoph et al., 2022)。我们还考虑采用逐层缩放方法,该方法最近被证明在大型语言模型中优于经典的 transformers(Mehta et al., 2024)。它线性增加注意力和前馈层中的宽度,即隐藏层维度。因此,网络中更深层的层被分配了更多的参数。我们认为,由于视觉模型中更深层的层会学习更复杂的特征(Zeiler & Fergus, 2014),在更深层使用更高的参数会带来更好的性能。我们在附录B中提供了这两种技术的背景细节。我们在图3中描述了我们的扩散 transformer 的整体架构。

A key limitation of patch masking strategies is poor performance when a large fraction of patches are masked. For example, Zheng et al. (2024) observed large degradation in even MaskDiT performance beyond a 50% masking ratio. Before optimizing the performance of our approach with a rigorous ablation study, we evaluate its out-of-the-box performance with common training parameters for up to $87.5\%$ masking ratios. As a baseline, we train a network with no patch-mixer, i.e., naive masking (Figure 2c) for each masking ratio. For deferred masking, we use a four-transformer block patch-mixer that comprises less than $10$% of the parameters of the backbone transformer.
郑等人(2024)观察到即使 MaskDiT 在 50% 以上遮挡比率下性能也会大幅下降,这是遮挡掩码策略的一个关键局限性。在对我们的方法进行严格的消融研究来优化性能之前,我们使用常见的训练参数评估了其达到最高 $87.5\%87.5\backslash \%$ 遮挡比率时的开箱即用性能。作为基准,我们训练了一个没有补丁混合器的网络,即简单遮挡掩码(图 2c)。对于延迟遮挡,我们使用一个包含不到主干变压器参数 $1010$ % 的四个变压器块补丁混合器。

We use commonly used default hyperparameters to simulate the out-of-the-box performance for both our approach and the baseline. We train both models with the AdamW optimizer with an identical learning rate of $1.6\times 10^{-4}$, $0.01$ weight decay, and a cosine learning rate schedule. We set $(P_{\text{mean}},P_{\text{std}})$ to $(-1.2,1.2)$ following the original work (Karras et al., 2022). We provide our results in Figure 4. Our deferred masking approach achieves much better performance across all three performance metrics. Furthermore, the performance gaps widen with an increase in masking ratio. For example, with a 75% masking ratio, naive masking degrades the FID score to $16.5$ while our approach achieves $5.03$, much closer to the FID score of $3.79$ with no masking.
我们使用常用的默认超参数来模拟我们的方法和基线方法的开箱即用性能。我们使用相同的学习率 $1.6\times {10}^{-4}1.6\backslash times\; 10^\{-4\}$ 、 $0.010.01$ 权重衰减和余弦学习率调度训练了两个模型。我们将 $(P_{\text{mean}},P_{\text{std}})(P\_\{\backslash text\{mean\}\},P\_\{\backslash text\{std\}\})$ 设置为 $(-1.2,1.2)(-1.2,1.2)$ ,遵循原始工作(Karras et al., 2022)。我们在图4中给出了结果。我们的延迟掩蔽方法在所有三个性能指标上都取得了更好的表现。此外,随着掩蔽率的增加,性能差距也越来越大。例如,当掩蔽率为 75%时,简单掩蔽会将 FID 得分降低到 $16.516.5$ ,而我们的方法达到了 $5.035.03$ ,接近无掩蔽时的 FID 得分 $3.793.79$ 。

We ablate design choices across masking, patch mixer, and training hyperparameters. We summarize the techniques that improved out-of-the-box performance in Table 5(a). We provide supplementary results and discussion of the ablation study in Appendix C.
我们对遮罩、补丁混合器和训练超参数的设计选择进行消融。我们总结了在表 5(a) 中提高开箱即用性能的技术。我们在附录 C 中提供了消融研究的补充结果和讨论。

Comparison with training hyperparameters of LLMs. Since the diffusion architectures are very similar to transformers used in large language models (LLMs), we compare the hyperparameter choices across the two tasks. Similar to common LLM training setups (Touvron et al., 2023; Jiang et al., 2023; Chowdhery et al., 2022), we find that SwiGLU activation (Shazeer, 2020) in the feedforward layer outperforms the GELU (Hendrycks & Gimpel, 2016) activation function. Similarly, higher weight decay leads to better image generation performance. However, we observe better performance when using a higher running average coefficient for the AdamW second moment ($\beta_{2}$), in contrast to large-scale LLMs where $\beta_{2}\approx 0.95$ is preferred. As we use a small number of training steps, we find that increasing the learning rate to the maximum possible value until training instabilities also significantly improves image generation performance.
与LLMs训练超参数的对比。由于扩散架构与用于大型语言模型(LLMs)的变换器非常相似,因此我们对这两个任务的超参数选择进行了比较。类似于常见的LLM训练设置(Touvron et al., 2023; Jiang et al., 2023; Chowdhery et al., 2022),我们发现 SwiGLU 激活(Shazeer, 2020)在前馈层中的表现优于 GELU(Hendrycks & Gimpel, 2016)激活函数。同样地,更高的权重衰减可以带来更好的图像生成性能。然而,我们观察到在使用更高的 AdamW 第二矩运行平均系数( ${\beta }_2\backslash beta\_\{2\}$ )时,图像生成性能更好,这与大规模LLMs中倾向于使用较低的 ${\beta }_2\approx 0.95\backslash beta\_\{2\}\backslash approx\; 0.95$ 相反。由于我们使用较少的训练步骤,我们发现将学习率增加到最大可能值,直到训练出现不稳定性,也能显著提高图像生成性能。

Design choices in masking and patch-mixer. We observe a consistent improvement in performance with a larger patch mixer. However, despite the better performance of larger patch-mixers, we choose to use a small patch mixer to lower the computational budget spent by the patch mixer in processing the unmasked input. We also update the noise distribution $(P_{\text{mean}},P_{\text{std}})$ to $(-0.6,1.2)$ as it improves the alignment between captions and generated images. By default, we mask each patch at random. We ablate the masking strategy to retain more continuous regions of patches using block masking (Figure 5(b)). However, we find that increasing the block size leads to a degradation in the performance of our approach, thus we retain the original strategy of masking each patch at random.
面具和修补程序混合器的设计选择。我们观察到,使用更大的补丁混合器可以一致地提高性能。然而,尽管更大的补丁混合器的性能更好,但我们选择使用一个小的补丁混合器,以降低补丁混合器在处理未屏蔽输入时的计算预算。我们还将噪音分布更新为 $(P_{\text{mean}},P_{\text{std}})(P\_\{\backslash text\{mean\}\},P\_\{\backslash text\{std\}\})$ 到 $(-0.6,1.2)(-0.6,1.2)$ ,因为它可以改善标题和生成图像之间的一致性。默认情况下,我们会随机屏蔽每个补丁。我们对屏蔽策略进行了消融,使用块屏蔽来保留更多连续的补丁区域(图5(b))。然而,我们发现增加块大小会导致我们的方法性能下降,因此我们保留了随机屏蔽每个补丁的原始策略。

We measure the effectiveness of the two modifications in transformer architecture in improving image generation performance.
我们衡量两种修改的转换器架构在提高图像生成性能方面的有效性。

Mixture-of-experts layers. We train a DiT-Tiny/2 transformer with expert-choice routing based mixture-of-experts (MoE) layers in each alternate transformer block. On the small-scale setup, it achieves similar performance to baseline model trained without MoE blocks. While it slightly improves Clip-score from $28.11$ to $28.66$, it degrades FID score from $6.92$ to $6.98$. We hypothesize that the lack of improvement is due to the small number of training steps ($60$K) as using multiple experts effectively reduces the number of samples observed by each router. In top-2 routing for 8-experts, each expert is trained effectively for one fourth number of epochs over the dataset. We observe a significantly improvement in performance with MoE blocks when training on scale for longer training schedules (Section 6.2).
专家混合层。 我们训练一个DiT-Tiny/2变压器 在各个交替的变换器块中基于专家选择的混合专家(MoE)层。在小规模设置中，它的性能与不使用 MoE 块训练的基线模型相似。虽然它略微改善了 Clip-score 从 $28.1128.11$ 到 $28.6628.66$ ，但它使 FID 得分从 $6.926.92$ 降低到 $6.986.98$ 。我们假设性能缺乏改善是由于训练步骤数量较少( $6060$ K)，因为使用多个专家实际上减少了每个路由器观察到的样本数量。在 8 个专家的 top-2 路由中,每个专家在数据集上的有效训练周期数为原来的 1/4。我们观察到,在长时间的训练计划下(第6.2节),使用 MoE 块可以显著提高性能。

Instead of training only on real images, we expand the training data to include 40% additional synthetic images. We compare the performance of two MicroDiT models trained on only real images and combined real and synthetic images, respectively, using identical training processes.
不仅训练真实图像,我们还扩展训练数据包括 40%额外的合成图像。我们比较只使用真实图像和结合真实及合成图像训练的两个MicroDiT模型的性能,采用相同的训练过程。

Under canonical performance metrics, both models apparently achieve similar performance. For example, the model trained on real-only data achieved an FID score of 12.72 and a CLIP score of 26.67, while the model trained on both real and synthetic data achieved an FID score of 12.66 and a CLIP score of 28.14. Even on GenEval (Ghosh et al., 2024), a benchmark that evaluates the ability to generate multiple objects and modelling object dynamics in images, both models achieved an identical score of 0.46. These results seemingly suggest that incorporating a large amount of synthetic samples didn’t yield any meaningful improvement in image generation capabilities.
根据标准性能指标,两种模型的性能似乎都差不多。例如,仅使用实际数据训练的模型 FID 得分为 12.72,CLIP 得分为 26.67,而同时使用实际和合成数据训练的模型 FID 得分为 12.66,CLIP 得分为 28.14。即使在 GenEval (Ghosh et al., 2024)基准测试中,这种评估生成多个对象和建模图像中对象动力学的能力,两种模型也都获得了相同的 0.46 分。这些结果似乎表明,纳入大量合成样本并没有带来图像生成能力的明显改善。

However, we argue that this observation is heavily influenced by the limitations of our existing evaluation metrics. In a qualitative evaluation, we found that the model trained on the combined dataset achieved much better image quality (Figure 9). The real data model often fails to adhere to the prompt, frequently hallucinating key details and often failing to synthesize the correct object. Metrics, such as FID, fail to capture this difference because they predominantly measure distribution similarity (Pernias et al., 2024). Thus, we focus on using human visual preference as an evaluation metric. To automate the process, we use GPT-4o (OpenAI, 2024), a state-of-the-art multimodal model, as a proxy for human preference. We supply the following prompt to the model: Given the prompt ‘{prompt}’, which image do you prefer, Image A or Image B, considering factors like image details, quality, realism, and aesthetics? Respond with ’A’ or ’B’ or ’none’ if neither is preferred. For each comparison, we also shuffle the order of images to remove any ordering bias. We generate samples using DrawBench (Saharia et al., 2022) and PartiPrompts (P2) (Yu et al., 2022), two commonly used prompt databases (Figure 9). On the P2 dataset, images from the combined data model are preferred in 63% of comparisons while images from the real data model are only preferred 21% of the time (16% of comparisons resulted in no preference). For the DrawBench dataset, the combined data model is preferred in 40% of comparisons while the real data model is only preferred in 21% of comparisons. Overall, using a human preference-centric metric clearly demonstrates the benefit of additional synthetic data in improving overall image quality.
然而,我们认为这一观察深受我们现有评估指标局限性的影响。在定性评估中,我们发现在综合数据集上训练的模型能够达到更好的图像质量(图 9)。真实数据模型经常无法恰当地遵循提示,频繁地编造关键细节,通常也无法合成正确的物体。诸如 FID 之类的指标无法捕捉这种差异,因为它们主要衡量分布相似性(Pernias et al., 2024)。因此,我们专注于使用人类视觉偏好作为评估指标。为了自动化这一过程,我们使用了最先进的多模态模型 GPT-4o (OpenAI, 2024)作为人类偏好的代理。我们向模型提供以下提示: 给定提示'{prompt}',您更喜欢图像 A 还是图像 B,考虑诸如图像细节、质量、真实性和美学等因素?请回答'A'、'B'或'none'(如果两者都不是首选)。对于每次比较,我们也会打乱图像的顺序以消除任何排序偏差。我们使用 DrawBench (Saharia et al., 2022)和 PartiPrompts (P2) (Yu et al., 2022)两个常用的提示数据库生成样本(图 9)。在 P2 数据集上,综合数据模型的图像在 63%的比较中被首选,而真实数据模型的图像只在 21%的比较中被首选(16%的比较没有首选)。对于 DrawBench 数据集,综合数据模型在 40%的比较中被首选,而真实数据模型只在 21%的比较中被首选。 总的来说,使用以人类偏好为中心的度量标准清楚地表明了使用额外合成数据提高整体图像质量的好处。

We first examine the effect of using a higher dimensional latent space in micro-budget training. We replace the default four-channel autoencoder with one that has sixteen channels, resulting in a $4\times$ higher dimensional latent space. Recent large-scale models have adopted high dimensional latent space as it provides significant improvements in image generation abilities (Dai et al., 2023; Esser et al., 2024). Note that the autoencoder with higher channels itself has superior image reconstruction capabilities, which further contributes to overall success. Intriguingly, we find that using a higher dimensional latent space in micro-budget training hurts performance. For two MicroDiT models trained with identical computational budgets and training hyperparameters, we find that the model trained in four-channel latent space achieves better FID, Clip-score, and GenEval scores (Table 10). We hypothesize that even though an increase in latent dimensionality allows better modeling of data distribution, it also simultaneously increases the training budget required to train higher quality models. We also ablate the choice of mixture-of-experts (MoE) layers by training another MicroDiT model without them under an identical training setup. We find that using MoE layers improves zero-shot FID-30K from 13.7 to 12.7 on the COCO dataset. Due to the better performance of the sparse transformers with MoE layers, we use MoE layers by default in large-scale training.
我们首先研究在微预算训练中使用更高维的潜在空间的效果。我们将默认的四通道自编码器替换为一个有 16 个通道的自编码器，从而产生 $4\times 4\backslash times$ 更高维的潜在空间。最近的大规模模型采用了高维潜在空间,因为它提供了图像生成能力的显著改善(Dai 等人，2023; Esser 等人，2024)。值得注意的是,具有更多通道的自编码器本身具有更出色的图像重建能力,这进一步有助于整体成功。有趣的是,我们发现在微预算训练中使用更高维的潜在空间会损害性能。对于两个经过相同计算预算和训练超参数训练的MicroDiT模型,我们发现在四通道潜在空间训练的模型实现了更好的 FID、Clip-score 和 GenEval 得分(表10)。我们假设,尽管潜在维度的增加允许更好地建模数据分布,但它同时也增加了训练高质量模型所需的训练预算。我们还剔除了专家组合(MoE)层的选择,在相同的训练设置下训练了另一个MicroDiT模型。我们发现使用 MoE 层可以将 COCO 数据集上的零 shot FID-30K 从 13.7 改善到 12.7。由于稀疏变压器与 MoE 层的性能更好,我们在大规模训练中默认使用 MoE 层。

Comparing zero-shot image generation on COCO dataset (Table 3). We follow the evaluation methodology in previous work (Saharia et al., 2022; Yu et al., 2022; Balaji et al., 2022) and sample 30K random images and corresponding captions from the COCO 2014 validation set. We generate 30K synthetic images using the captions and measure the distribution similarity between real and synthetic samples using FID (referred to as FID-30K in Table 3). Our model achieves a competitive 12.66 FID-30K, while trained with $14\times$ lower computational cost than the state-of-the-art low-cost training method and without any proprietary or humongous dataset to boost performance. Our approach also outperforms Würstchen (Pernias et al., 2024), a recent low-cost training approach for diffusion models, while requiring $19\times$ lower computational cost.
比较 COCO 数据集上的零镜头图像生成效果(表 3)。 我们遵循以前工作 (Saharia 等人, 2022; Yu 等人, 2022; Balaji 等人, 2022) 中的评估方法,从 COCO 2014 验证集中随机抽取 30K 张图像及其对应的标题。我们使用这些标题生成 30K 张合成图像,并使用 FID 来测量真实和合成样本分布的相似度(在表 3 中称为 FID-30K)。我们的模型在 FID-30K 上达到了 12.66 的竞争性成绩,训练成本也低于最先进的低成本训练方法,且无需使用任何专有或庞大的数据集来提升性能。我们的方法还优于最近的低成本训练方法 Würstchen (Pernias 等人, 2024),同时所需的计算成本也更低。

Fine-grained image generation comparison (Table 4). We use the GenEval (Ghosh et al., 2024) framework to evaluate the compositional image generation properties, such as object positions, co-occurrence, count, and color. Similar to most other models, our model achieves near perfect accuracy on single-object generation. Across other tasks, our model performs similar to early Stable-Diffusion variants, notably outperforming Stable-Diffusion-1.5. On the color attribution task, our model also outperforms Stable-Diffusion-XL-turbo and PixArt-$\alpha$ models.
细粒度图像生成比较（表 4）。我们使用 GenEval (Ghosh 等人, 2024)框架来评估图像的组成属性,如对象位置、共现、数量和颜色。与大多数其他模型类似,我们的模型在单个对象生成上达到了近乎完美的准确性。在其他任务中,我们的模型与早期的 Stable-Diffusion 变体表现相似,明显优于 Stable-Diffusion-1.5。在色彩归属任务上,我们的模型也优于 Stable-Diffusion-XL-turbo 和 PixArt- $\alpha \backslash alpha$ 模型。