AI绘图王炸！PRX+T5Gemma狂提速38%内存砍半！

在数字技术飞速发展的今天，人工智能生成图像技术（Text-to-Image）正以前所未有的速度改变着我们认识和创造世界的方式。从简单的草图到逼真的艺术作品，AI正在逐步实现将文字描述转化为生动视觉的奇迹。这背后，是无数科学家和工程师在模型架构、训练方法上不懈探索的成果。今天，我们将深入探讨其中核心的“骨架”——模型架构，看看不同的设计理念如何影响着图像生成的质量、效率与可扩展性。

Figure: Evolution of generated images accross the first 100K steps.

深度解析：流匹配与潜在扩散技术的奥秘

在深入探讨具体的模型架构之前，我们有必要先了解一下支撑这些模型运行的核心机制——流匹配（Flow Matching）与潜在扩散（Latent Diffusion）框架。

传统上，许多生成模型，尤其是扩散模型（Diffusion Models），通过模拟一个从噪音到数据、充满随机性的渐进过程来生成图像。它们先给一张图片添加噪音，再学习如何逐步去除噪音，最终得到清晰图像。这就像给一张照片蒙上白雾，然后教AI如何一点点拨开迷雾，还原真实面貌。这个过程虽然有效，但往往涉及复杂的随机微分方程，计算量较大。

而流匹配框架则提供了一种更为直接、确定性的生成路径。它不像扩散模型那样模拟随机的噪声扰动和去噪过程，而是学习一个连续的时间变换，直接描述样本如何从简单的初始分布（比如纯噪音）演变为目标数据分布（即真实图像）。你可以把它想象成一条平滑的河流，AI学习的不是河水随机的流动，而是水流从源头到入海口最直接、最有效率的“流速场”，引导噪音样本沿着这条确定的路径，最终转化为我们想要的图像。

具体来说，研究人员定义了一个在时间t \in [0,1]内连续变化的中间分布家族p_t，它像一座桥梁，连接着初始的噪音分布p_0和最终的图像数据分布p_1。我们的目标就是找到一个“速度场”v_t^\star(z_t)，它能精确地描述样本z_t如何在时间t下演变，使其沿着常微分方程\\frac{d z_t}{d t} = v_\\theta(z_t, t)的路径，从噪音z_0平稳地“流动”到图像z_1。

在实际操作中，研究人员通常采用一种简化的流匹配变体，称为“修正流”（Rectified Flow）方法。这种方法假设样本在噪音和数据之间沿着一条线性的、直来直去的路径传输。在这种理想假设下，最优的流场变得与时间无关，可以简化表示为v_t^\star(z_t) = z_1 - z_0。这意味着，AI不再需要模拟复杂的时间动态，只需学习如何直接从噪音样本“跳跃”到真实样本的“位移向量”。通过\\ell_2损失函数来训练网络v_\\theta，使其预测的这个“位移向量”尽可能接近真实值，从而实现稳定且高效的训练。这种方法避免了随机微分方程的复杂性，却依然保持了强大的生成能力。

为了进一步提高计算效率，流匹配操作通常在“潜在空间”（Latent Space）而非像素空间进行。想象一下，原始的图像数据非常庞大，直接处理就像在一个巨大的画板上操作。而潜在空间，就像把这个画板上的主要信息“压缩”成一个更小、更紧凑的抽象形式。通过一个编码器\\mathcal{E}，可以将高维图像x映射为低维的潜在表示z = \\mathcal{E}(x)；再通过一个解码器\\mathcal{D}，可以将潜在表示z重建回原始图像，确保\\mathcal{D}(z) \\approx x。这种在潜在空间进行操作的方式，在保持图像感知质量的同时，显著降低了数据维度，使得模型训练更快，内存消耗更少，极大地提升了效率。

当然，我们的最终目标是根据文字描述生成图像，所以模型还需要一个“文本条件”（Text Conditioning）。一个文本编码器f_\\phi（比如强大的T5或T5Gemma模型）会将输入的文字描述y（如“一只在雪地里泡温泉的猴子”）转化为一系列具有语义信息的嵌入向量c = f_\\phi(y)。这些嵌入向量就像给AI提供了一份详细的“蓝图”，指导生成过程，确保最终的图像内容与文字描述高度契合。最终，整个网络学习的是一个条件速度场v_\\theta(z_t,t,c)，巧妙地结合了潜在空间建模的高效、修正流的确定性以及文本条件指导的强大表现力。

AI的“大脑”：多种Transformer架构的探索

在文本到图像生成领域，模型架构的选择至关重要。它决定了模型如何处理图像和文本信息，从而影响生成质量、训练速度和资源消耗。研究人员对一系列基于Transformer的架构进行了深入评估，包括经典的DiT、多模态的MMDiT、高效的DiT-Air、U形结构的U-ViT，以及他们自己设计的PRX模型。这些探索的目的不是追求最大的模型，而是为了找到在效率、稳定性和文本-图像对齐效果之间取得最佳平衡的架构原则。

1. DiT (Diffusion Transformer)：开山之作的演进

Diffusion Transformer（DiT）是率先将Transformer模块引入图像生成扩散模型的先驱。它最初是为类别条件生成而设计，后来逐渐扩展到文本到图像的合成任务，为后续众多模型的诞生奠定了坚实基础。

在实验中，研究人员采纳了DiT的PixArt-α变体。这一改进在原有的自注意力层和前馈网络层之间巧妙地插入了一个“交叉注意力”（cross-attention）机制。你可以想象，传统的DiT是让AI自己看图像的各个部分（自注意力），而PixArt-α则是在AI看图的同时，提供一份文字“说明书”，让AI能够对照着说明书来理解图像（交叉注意力）。这种设计使得视觉特征和文本特征能够更直接地融合，显著提升了生成图像与文字描述的契合度。

PixArt-α还引入了一种更精细的归一化策略——共享自适应层归一化（Adaptive LayerNorm，简称AdaLN）。与原始DiT在每个模块中独立维护自适应归一化参数不同，PixArt-α从时间步嵌入中派生出一组全局的缩放和平移参数，并在所有层之间共享。这就像为AI的所有层提供了一把统一的“刻度尺”，减少了参数冗余，同时通过轻量级的每模块嵌入保持了灵活性。尽管后来出现了更高效、表现力更强的架构，DiT凭借其简洁性和可扩展性，至今仍是稳健且被广泛采用的基线模型，许多现代文本到图像系统仍在沿用DiT启发的核心思想。

Figure: PixArt-α DiT block

2. MMDiT (Multimodal Diffusion Transformer)：文本与图像的深度交织

Multimodal Diffusion Transformer（MMDiT）是作为Stable Diffusion 3的一部分推出的架构，它将DiT家族进一步扩展，引入了一种“双流”（dual-stream）架构。与PixArt-α通过交叉注意力将文本条件注入到纯图像骨干网络不同，MMDiT在整个网络中并行维护两条独立的“信息流”——一条处理文本特征，一条处理图像特征。

你可以把MMDiT想象成AI同时拥有“左脑”和“右脑”，左脑处理文字信息，右脑处理图像信息。每个脑区都有自己的归一化、调制和前馈层，但它们共享一个核心的注意力机制。这意味着在注意力计算时，文本和图像的“注意力焦点”（queries, keys, values）可以互相参考，实现模态间的全面双向交流。这就像左脑和右脑在共同思考一个问题时，可以随时互相“询问”和“提供信息”。

每个信息流都保留了自己的AdaLN参数，并通过时间步和模态嵌入进行调节，确保在不同领域内扩散条件的一致性。这种设计让MMDiT能够比单流架构更明确地捕捉跨模态的深层依赖关系，生成更符合语义的图像。当然，这种深度交织也伴随着更高的内存消耗和计算成本。

Figure: MMDiT block

3. DiT-Air：轻量与效率的平衡之道

DiT-Air是一种混合架构，它巧妙地在DiT和MMDiT之间架起了一座桥梁，既保留了单流Transformer的简洁高效，又汲取了双流设计中富有表现力的多模态交互优势。与MMDiT为文本和图像信息流维护独立通道不同，DiT-Air在一个统一的“信息序列”中操作，将两种模态的数据并存处理。

想象一下，DiT-Air不是让AI同时用两个脑区独立思考，而是把所有相关信息（包括图像和文本）都放到一个统一的工作区，让AI在这个工作区里进行一次全面的“自查”。它沿用了DiT中的AdaLN机制，确保时间和条件信息能够始终如一地整合到整个网络中。这种设计在MMDiT结构化多模态推理能力和原始DiT高效性之间找到了实用的平衡点。通过避免双路径带来的计算和内存开销，DiT-Air在联合注意力机制下实现了强大的文本-图像对齐效果，同时保持了轻量化和良好的可扩展性。在大规模应用中，DiT-Air能够达到甚至超越大型架构的性能，但其参数量却显著减少——比MMDiT减少了约66%，比PixArt-α减少了约25%，使其成为高效文本到图像扩散模型的一个强大基线。

Figure: DiT-Air block

4. U-ViT (U-shaped Vision Transformer)：U形结构的创新融合

U-shaped Vision Transformer（U-ViT）的拓扑结构让人联想到经典的U-Net架构，但其内部却完全由Transformer模块实现。它像U-Net一样，拥有一个编码器堆栈和一个解码器堆栈，并通过“长跳跃连接”（long skip connections）将它们连接起来。这些跳跃连接允许来自浅层模块的低级空间特征被直接传递并融合到更深层的模块中，从而显著提高重建图像的质量。这就像在图像处理过程中，不仅关注整体，还能记住很多局部细节，让最终生成的图片更加精细。

与DiT-Air类似，U-ViT也在一个统一的信息序列上运行，将视觉和条件信息（如时间步和文本）通过自注意力机制共同处理。然而，U-ViT设计上的一大特点是它完全取消了自适应归一化机制（如AdaLN或每层调制）。相反，时间步和文本嵌入等条件信息被直接拼接（concatenated）到输入的信息序列中。这意味着Transformer可以在一个单一的注意力空间中，对图像块、时间令牌和文本令牌进行联合推理。这种设计使得U-ViT在概念上更为简洁优雅，它将Transformer的全局上下文建模能力与编码器-解码器架构的层次结构巧妙地融合在一起。

Figure: U-ViT architecture

5. PRX (Photoroom eXperimental)：效率与性能的巧妙平衡

为了探索更多创新的设计方案，研究团队开发了自己的专属架构——PRX（Photoroom eXperimental）。这是一种混合设计，融合了单流和双流Transformer的特点。PRX同时接收图像和文本信息作为输入，但其设计目标是仅输出图像信息，从而将大部分计算资源集中在图像生成路径上。

每个PRX模块直接从文本编码器接收文本信息，这与PixArt-α的设计有相似之处。然而，与典型的交叉注意力或双流设置不同，PRX首先独立处理图像和文本信息，然后将它们拼接起来进行自注意力操作。但关键在于，注意力计算只针对图像信息进行，这极大地降低了计算和内存成本。这种设计与此前一些研究中提出的“自注意力DiT浅层融合”基线模型密切相关。

PRX的巧妙之处在于，它避免了显式的文本信息更新，只需执行一次注意力操作（而非标准DiT中的两次），并且与MMDiT相比，其注意力矩阵也更小（MMDiT的跨模态注意力计算量与文本和图像信息数量的乘积成正比）。此外，研究人员观察到文本信息在扩散的时间步长中是静态不变的。基于这一发现，PRX在文本信息流中省略了时间步调制。由于文本信息无需修改，它们可以在推理时一次性投影并缓存起来，消除了每个步骤中重复的计算，从而显著加速了图像生成过程。

新媒网跨境了解到，这种简单而有效的设计，在速度和内存效率方面，相较于DiT和MMDiT都有显著提升，同时保持了强大的文本-图像对齐能力和具有竞争力的生成质量。尤其是在推理阶段可以缓存文本流的能力，进一步巩固了PRX作为一个实用选择的地位：它大幅减少了计算量和延迟，为实际部署提供了明显优势，即便其CMMD分数略高于MMDiT，也足以证明其价值。

Figure: PRX block.

性能大比拼：实战评估揭示最佳选择

为了科学地评估这些架构的优劣，研究人员设计了一套严格的基准测试。外媒相关研究表明，模型架构的效率和相对性能趋势在小规模实验中观察到的结果，通常能够可靠地预测其在大规模应用时的表现。

Figure: Schematic of the PRX block architecture.

基于这一重要发现，研究团队也遵循了相同的原则来设计他们的基准测试，以便快速迭代：在较低的分辨率和模型规模下进行操作，并期望这些小规模结果能够代表其大规模对应模型的性能。与之前DiT-AIR的研究不同，本次比较并未严格控制参数数量。相反，研究人员统一固定了所有模型的Transformer模块数量、注意力头数量和隐藏维度，以确保模型之间的公平比较。这种方法使得他们能够更清晰地分离出架构结构（如流配置、条件策略和归一化设计）对模型性能的贡献，而不将其与整体模型容量或规模混淆。

所有模型都在一个包含100万张图像的定制数据集上，以256x256的分辨率进行训练，实验设置如下：

• 批处理大小（Batch size）：256
• Transformer模块数量：16
• 注意力头数量：28
• 令牌嵌入维度：1792
• 潜在空间：采用Flux VAE，具有16个潜在通道和8倍压缩因子
• 文本编码器：GemmaT5
• 位置编码：除U-ViT使用学习到的1D位置编码外，所有架构均采用旋转位置编码（Rotary Positional Encoding，RoPE）

评估标准则涵盖了多个维度，力求全面衡量模型的性能：

• 重建损失（Reconstruction loss）： 衡量在独立评估集上重建样本与目标样本之间的均方误差（MSE）。数值越低，表示重建质量越好。
• Frechet Inception距离（FID）： 通过Inception v3特征统计量，衡量生成图像分布与真实图像分布之间的相似性。FID值越低，视觉保真度越高，图像越逼真。
• Clip-最大均值差异（CMMD）： 利用CLIP嵌入和最大均值差异（MMD）指标，评估真实图像和生成图像分布之间的距离。它被认为是比FID更稳健且样本效率更高的评估方法，与人类感知判断的契合度更高。
• 内存使用（Memory usage）： 训练期间的GPU峰值内存消耗。数值越低，表示资源效率越高。
• 网络吞吐量（Network throughput）： 每秒处理的平均样本数量，衡量整体的计算效率。数值越高，表示处理速度越快。

从评估结果来看，MMDiT在重建损失（MSE）和CMMD分数上表现出色，展示了强大的生成性能，但它也是所有模型中最“重”的一个，需要最多的参数和GPU内存，并且吞吐量最低。这说明它虽然强大，但资源消耗巨大。DiT、DiT-Air和U-ViT变体在各项指标上都取得了有竞争力的结果，但在图像质量方面（尤其是FID和CMMD）略微落后于MMDiT，不过它们的整体效率更高。

值得注意的是，研究团队提出的PRX架构在综合表现上达到了最佳平衡。它在重建质量上与MMDiT不相上下，但在FID、吞吐量和内存效率方面均超越了MMDiT。特别是PRX在推理过程中能够缓存文本流的特性，进一步提升了其实用价值：它显著减少了计算量和延迟，即便其CMMD分数略高于MMDiT，也为实际部署带来了明显的优势。新媒网跨境获悉，这充分证明了PRX在寻求高性能与高效率之间平衡的成功。

文本编码器：迈向更智能的T5Gemma

在文本到图像模型中，文本编码器扮演着核心角色，它是连接自然语言理解与视觉生成之间的桥梁。文本表示的质量和结构直接影响着模型在生成图像时捕捉语义和构图的能力。

过去，大多数基于扩散技术的文本到图像架构，例如Flux或Stable Diffusion 3，普遍依赖于T5模型。T5是一种经典的编码器-解码器Transformer模型，以文本到文本的范式进行训练。其中，广泛使用的T5-XXL版本拥有约110亿参数，并生成维度为4096的嵌入向量。凭借其强大的上下文理解能力和双向注意力机制，T5长期以来一直是扩散模型中文本条件化的标准骨干网络。

然而，近期的研究趋势正逐渐转向基于大型语言模型（LLM）的编码器。这些新型编码器能够提供更丰富、语义基础更强的嵌入向量。在这其中，T5Gemma脱颖而出，它是一个源自Gemma 2家族的现代编码器-解码器模型。T5Gemma通过一种巧妙的改编过程构建，将预训练的仅解码器LLM转换为编码器-解码器架构，这使得T5Gemma能够继承Gemma模型深厚的表征能力，同时保留T5的双向推理优势。

研究团队将T5Gemma作为T5的替代品，在他们的文本到图像管道中进行了评估，并观察到多项显著优势：

• 更少的参数： T5Gemma的参数量约为20亿，远低于T5-XXL的110亿，大幅降低了模型规模和部署成本。
• 更小的嵌入维度： 嵌入维度为2304，相较于T5的4096更小，减少了内存使用和计算成本，提升了效率。
• 更低的评估损失： 表明T5Gemma能够产生信息更丰富、对齐更好的文本嵌入，从而提升了模型的整体生成质量。
• 多语言能力： 继承自Gemma 2基础模型，使得T5Gemma无需额外适应训练，即可支持从多种语言的提示词生成图像。这为全球用户带来了极大的便利，拓宽了应用的边界。

鉴于这些显著优势，研究团队已决定采用T5Gemma 2B作为未来模型的文本编码器。这一选择将进一步提升模型在效率、可扩展性和多语言支持方面的表现，为未来的迭代发展奠定坚实基础。

Figure: Multilingual generations with PRX + T5Gemma. The same prompt rendered in English, French, Spanish, and Italian demonstrates T5Gemma’s multilingual understanding without retraining.
English: A professional close-up photograph of a monkey bathing in a hot spring during a snowstorm, steam rising gently from the water as snowflakes melt on its fur.
Français: Une photo professionnelle en gros plan d’un singe se baignant dans une source chaude pendant une tempête de neige, la vapeur s’élevant doucement de l’eau tandis que les flocons fondent sur son pelage.
Español: Una fotografía profesional en primer plano de un mono bañándose en una fuente termal durante una tormenta de nieve, con vapor elevándose suavemente del agua mientras los copos se derriten sobre su pelaje.
Italiano: Una fotografia professionale in primo piano di una scimmia che si bagna in una sorgente termale durante una tempête de neige, con il vapore che si solleva dolcemente dall’acqua mentre i fiocchi si sciolgono sul suo pelo.

潜在空间与自编码器：压缩与效率的艺术

潜在表示的选择对模型的训练效率和生成质量有着重大影响。回顾整个实验过程，FluxVAE在当时平衡了重建质量与计算速度，因此成为了研究早期迭代的自然首选。它的稳定性以及与基于Transformer的扩散架构的兼容性，使其成为可靠的起点。

与此同时，研究团队也积极探索了另一种前沿技术：麻省理工学院Han Lab开发的深度压缩自编码器（Deep-Compression Autoencoders，DC-AE）。DC-AE旨在通过在通道和空间维度上应用结构化压缩，学习紧凑而富有表现力的潜在空间。你可以将其理解为一种更高级的“图像压缩技术”，它不仅能将图像编码成更小的潜在张量，还能在极大程度上保留图像的感知质量，确保压缩后的信息依然能很好地重建出原图的精髓。

传统的FluxVAE通常以8倍的空间压缩比运行，而DC-AE则能实现高达32倍的压缩比，且不显著损失表现力或视觉保真度。这种突破性进展极大地降低了潜在空间的空间分辨率，从而显著提升了训练吞吐量和内存效率，这对于开发高分辨率的扩散模型尤为关键。考虑到这些显著优势，研究团队发布了使用DC-AE潜在空间训练的PRX模型检查点，这将有助于社区探索更快、更轻量级的设置，推动大规模文本到图像训练的进一步发展。

Figure: Images generated with the PRX and Deep-Compression Auto-encoder.

展望未来：持续创新，共创价值

此次分享标志着研究团队在从零开始构建新的文本到图像基础模型征程中，首次深度揭示其核心技术细节。从Transformer骨干网络到文本编码器，再到自编码器和潜在表示，每一个环节的设计与优化都为今日发布的模型奠定了坚实的基础。

新媒网认为，当前，研究团队仍在积极地迭代和完善其方法。目前发布的检查点对应的是小规模的12亿参数PRX模型，旨在验证其架构实验的有效性。更大规模的版本已在规划中，但尚未启动训练，因为团队仍在不断优化设计选择和训练流程。

在接下来的系列文章中，我们将继续深入探讨训练技术——如何优化大规模训练以实现速度和稳定性，加速收敛的方法，以及在此过程中汲取的宝贵经验。我们满怀期待地看到，全球的开发者和研究者如何利用PRX进行创新实验，并在其基础上创造出更多令人惊叹的应用。这些技术的进步，无疑将为数字内容创作、设计、教育、娱乐等多个领域带来深远影响，让文字与图像的结合变得前所未有的自由与高效，共同推动人工智能技术服务社会、造福人民，迈向更加美好的智能时代。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/ai-img-prx-t5gemma-38-speed-half-mem.html