借文本语义驱动 Token 化:TexTok 突破图像重建与生成,实现 93.5 倍的推理速度提升 !

大模型向量数据库云存储

备注好友:

方向-学校/公司-姓名/昵称

【AIGC 先锋科技】交流群

picture.image

picture.image

图像分割,即将原始图像像素转换为紧凑的低维潜在表示的过程,已被证明对于可扩展和高效的图像生成至关重要。

然而,主流的图像分割方法通常压缩率有限,使得高分辨率图像生成计算成本高昂。为解决这一挑战,作者提出利用语言来实现高效的图像分割,并称之为文本条件图像分割(TexTok)。

TexTok 是一种简单有效的分割框架,通过利用语言提供高层语义信息。

通过在描述性文本说明条件下进行分割过程,TexTok 使分割过程能够专注于将精细视觉细节编码到潜在 Token 中,从而提高了重构质量并增加了压缩率。

与传统的不带文本条件的分割器相比,TexTok 在 ImageNet-256 和 512 基准测试中的平均重构 FID 改进分别为 29.2% 和 48.1%,且这种分割改进始终转化为在不同分割 Token 数量下的平均生成 FID 提高了 6.3% 和 34.3%。

通过仅用 TexTok 替换扩散 Transformer (DiT)中的分割器,作者的系统可以在使用 32 个 Token 时仍比原版 DiT 在 ImageNet-512 上表现更优,同时实现 93.5 倍的推理速度提升。

结合一个基础版 DiT 生成器,TexTok 能够分别在 ImageNet-256 和 -512 上达到最先进的 FID 分数,分别为 1.46 和 1.62。

此外,作者在文本到图像生成任务中展示了 TexTok 的优越性,有效地利用现成的文本说明进行分割。

  1. Introduction

近年来,图像生成取得了显著进展,使其在多种应用中能够合成高质量的图像。这一成功的核心在于图像分词技术的发展,该技术通过训练自编码器将原始图像数据压缩为一种紧凑且富有表达力的潜在表示。分词使得生成模型,如扩散模型和自回归模型可以直接在压缩的潜在空间中操作,而不是在高维像素空间中操作,从而大幅提高计算效率并增强生成质量和保真度。

尽管针对图像 Token 化进行了各种努力以改善训练目标并细化自编码器架构,但当前方法仍然在压缩率和重构质量之间存在根本性的权衡,特别是在高分辨率生成方面。高压缩率虽可降低计算成本,但往往牺牲了重构质量;而优先确保高质量则会导致计算成本显著增加。

为了应对这一限制,需要在 Token 化过程中进行根本性的转变。从根本上讲, Token 化涉及找到一张图像紧凑而有效的表示形式。最简洁且具有意义的图像表示往往来自于其语言描述,即图注。当描述一张图像时,人类通常会从高层次的意义开始,然后详细说明更细微的部分。受到这一洞察的启发,作者引入了一种基于文本条件的图像 Token 化框架

,该框架利用描述性的文本图注来提供高层次的语义内容,从而使 Token 化过程更加专注于细粒度的视觉细节,从而在不牺牲压缩率的情况下提升重建质量。

据作者所知,作者是第一个在分词阶段条件化详细描述的,这种方法通常保留用于生成阶段。文本描述可以从在线图像-文本配对中轻易获得,或者使用视觉-语言模型为图像生成描述。由于文本条件化在图像生成中广泛应用,例如文本到图像生成,作者的方法可以无缝地将这些描述纳入分词过程,而不会增加额外的标注开销。

作者在多样化的任务和场景中展示了TexTok的有效性。与没有文本条件的传统分词器相比,TexTok在重构质量上实现了显著提升,在ImageNet 256×256和512×512分辨率上的平均重构FID分别提高了29.2%和48.1%。这些在分词方面的改进也带来了生成性能的一致提升,在两种分辨率下的生成FID平均改进率分别为16.3%和34.3%。

通过仅仅用TexTok替换Diffusion Transformer (DiT)中的分词器,作者的系统实现了93.5倍的推理速度提升,并且在使用32个Token的情况下仍优于只使用原始DiT的模型。作者最优秀的TexTok变体配合一个普通的DiT生成器,在ImageNet 256×256和512×512分辨率上分别取得了最佳的FID分数1.46和1.62。

由于文本描述词对于此任务而言是固有的,因此 TexTok 能够提升性能而无需增加任何额外的标注开销。

  1. Related Work

图像分词。图像分词器建立了高分辨率像素与低维潜在空间之间的双向映射,显著提高了下游任务(如图像生成[4, 11, 26, 49]和理解[47, 49])的学习效率。图像分词器通常以自编码器(AutoEncoder,AE)框架形式实现,并可选地包含量化器[45],并在变分[23]设置中可能实现。这些自编码器通过最小化输出图像与输入图像在像素空间距离[52]、潜在空间距离或联合训练判别器[11]中的差异来进行训练。编码器和解码器的架构变体包括ResNet[15]和视觉 Transformer [10]。现代分词器设计的一个共同特征是空间对应关系,其中一个Token通常对应一个像素区域的局部邻域。最近,还开发了基于Transformer的模型,以产生更紧凑的全局Token表示[50]。在本文中,作者遵循这一范式,将图像分词为一组全局Token,以便灵活控制Token预算。然而,与以往工作不同的是,作者首次提出将分词过程条件化于图像描述,这大大提高了重建质量和压缩率。

图像生成。在对抗性[3,40]、自回归[6]和扩散[9,18,22]等框架中已经探索了像素的生成学习。对于更高的分辨率,生成学习在压缩的潜在空间中变得流行,这得益于其效率优势。其中,自回归[11,26]和 Mask 预测[4,49]模型通常在类似于GPT[34]和BERT[8]语言模型的离散Token空间中运作。最近的一些变体[28]也可以使用连续的潜在空间,这类似于在隐式扩散模型(LDMs)[37]中所使用的。对于LDMs,架构从基于卷积的U-Net[38]演进到了基于Transformer的DiT[32]。在本文中,作者关注的是使用DiT架构的基于扩散的图像生成,并利用TexTok灵活的Token长度。

利用图像生成和 Token 过程中外部语义信息。许多近期的研究开始利用外部语义信息,例如图像表示和语义地图,来改进图像生成效果 [27, 33, 51]。与这些方法不同,作者的方法侧重于通过文本语义条件化来增强 Token 过程。一些近期的努力 [29, 30, 48, 53] 也在图像 Token 中考虑了将图像 Token 与文本语义对齐,以提高多模态理解能力。它们要么直接将图像映射到固定的大语言模型词汇表中的文本 Token [30, 48, 53],要么使图像 Token 特征与文本特征对齐 [29],以生成语义上有意义的 Token 。

然而,由于视觉和语言表示之间的固有差异,这些工作在强制严格的图像-文本对齐时,导致重建图像质量受限,从而影响生成图像的质量。相比之下,作者的工作采取了一种互补的方法。作者利用文本描述作为外部语义条件化,允许图像 Token 关注捕捉更精细的视觉细节,显著提升了图像重建和生成性能。

  1. Method

3.1. Preliminary

基于潜在表示的形式,图像分词器可以大致分为两类:

1)向量量化(VQ)分词器,如VQ-VAE [45]和VQGAN [11],它们使用一组离散的Token来表示图像;

2)连续潜在分词器[37],它们使用变分自编码器(VAE)[23]将图像嵌入到连续的潜在空间中。在本文中,作者主要关注连续潜在分词器。如附录A所示,TexTok在VQ分词器上也同样表现良好。

标准连续潜在 Token 器通常由编码器( Token 器)

和解码器(反 Token 器)

组成。给定一幅图像

,编码器

将其压缩到一个二维潜在空间

,其中

是空间下采样因子。每个潜在嵌入

被视为一个连续的Token,图像通过总计hw个Token来表示。在解码过程中,这些嵌入

被送入解码器

以重建图像

最近,引入了1维 Token 器[50],允许灵活的 Token 预算,直接将图像I压缩成1维潜在嵌入

,其中有

个 Token 。通过应用重构损失、感知损失[52]和GAN损失[11]来训练 Token 器,使其最小化I与

之间的距离。在这项工作中,作者采用1维 Token 器范式,以允许更灵活的压缩率,并证明TexTok在不同 Token 预算下的有效性和效率。

3.2.TexTok: Text-Conditioned Image Tokenization

作者介绍了Text-ConditionedImageTokenization(TexTok),这是一种简单而有效的分词框架,利用语言提供高级语义信息,并更多地关注细粒度视觉细节的分词。与现有方法将所有视觉信息压缩到潜在分词不同,作者使用描述性文本说明来表示高级语义并指导分词过程。

分词阶段。给定一幅图像的描述,作者使用冻结的T5文本编码器[35]提取文本嵌入。这些嵌入在整个分词过程中被注入到分词器和去分词器中,提供语义指导,并使学习到的分词更侧重于捕获细微的视觉细节。

如图2所示,TexTok 采用Vision Transformer(ViT)作为编码器(Tokenizer)

和解码器(Detokenizer)

的 Backbone 网络,以灵活控制Token的数量。Tokenizer的输入是一个三个部分的拼接:

  1. 来自输入图像wiaprjtion aywhe

,其中

分别表示高度和宽度分割后的部分,

是每个Patch的大小; 2. 一个随机初始化的学习图像Token

,其中

表示输出图像Tokens的数量; 3. 文本Tokens通过文本嵌入线性投影得到的

,其中

表示文本Tokens的数量。在Tokenizer的输出中,仅保留学习到的图像Tokens,并通过线性投影生成输出图像Tokens

Detokenizer同样接收三个拼接的输入:

  1. 可学习的Patch Tokens

, 2. 从输入图像Tokens线性投影得到的图像Tokens

, 3. 来自相同文本Tokens的文本Tokens的线性投影

。在Detokenizer的输出中,仅保留学习到的Patch Tokens,不进行Patch化处理,并通过投影重构图像Patch。

picture.image

作者使用

重构、GAN、感知和LeCAM正则化损失的组合训练分词器和反分词器,遵循[49]中的方法。

通过直接注入包含高阶语义信息的文本Token到分词器和反分词器中,TexTok减轻了分词器和反分词器学习语义的需求,使它们能够更专注于将剩余的细粒度视觉细节编码到图像Token中,从而在不牺牲重建性能的情况下显著提高压缩率。

生成阶段。由于本工作集中于连续潜在Token,作者采用Diffusion Transformer(DiT)[32]作为生成框架,并在TexTok生成的潜在Tokens之上训练DiT。需要注意的是,在生成阶段只需生成潜在图像Token,而文本Token将在去Token化过程中提供。

DiT 被训练以建模 TexTok 潜在 Token 的分布,条件是类别类别(用于条件生成)或文本嵌入(用于文本到图像生成)。

在推理过程中,生成任务的方式有所不同。对于文本到图像生成,作者使用提供的 captions 进行 token 化和生成,将文本嵌入和生成的潜在图像 token 输入到 detokenizer 中以产生输出图像。

而对于类条件生成,DiT 根据指定的类别生成潜在 token;然后从预先生成的列表中抽取一个未见过的 captions,并将其与生成的潜在 token 一起注入到 detokenizer 中以产生最终图像。值得注意的是,在生成过程中仅使用了类别类别,符合标准做法。

  1. Experiments

4.1. Implementation Details

文本标注获取。文本描述对于文本到图像生成任务来说易于获取,可以在分词过程中直接使用这些描述。而对于没有提供描述的其他生成任务,例如作者在使用ImageNet [7]时,作者采用视觉语言模型(VLM),即Gemini v1.5 Flash [42],在离线状态下生成详细的描述。对于训练集,作者为每张给定的图像添加描述。对于评估集,在类别条件生成中,作者预先为每个类别生成未见过的描述,参考的是训练集中该类别的采样描述列表。

默认情况下,每张图像的描述不超过75个单词,并通过T5文本编码器 [35](对于ImageNet-256实验使用XL版本,对于ImageNet512实验使用XXL版本)编码成128个Token的序列。更多细节请参见附录D。

分词与生成。默认情况下,所有TexTok实验均采用ViT-Base作为分词器和反分词器,二者各自包含12层,隐藏层大小为768,注意力头数为12(约176百万参数)。对于GAN损失,作者遵循[47]并使用StyleGAN鉴别器[19](约2400万参数)。除非另有说明,TexTok中的图像 Token 通道维度设置为

。作者由于其有效性和处理一维 Token 的灵活性,默认使用Diffusion Transformer(DiT)[32]作为生成器。

在所有TexTok生成实验中,DiT的块大小均为1。默认情况下,作者训练DiT共350个周期。具体来说,在类条件生成中,作者使用原始的DiT架构。对于文本到图像生成,参考[5],作者在DiT块中的多头自注意力层之后添加了一个额外的多头跨注意力层,以接受文本嵌入。

4.2. Experiment Setup

模型变体。作者比较了两种设置以展示使用文本条件化的有效性:TexTok 在分词器和反分词器中都包含了文本 Token ,对应于图2 所示的架构。相比之下,Baseline(不包含文本)则在分词器和反分词器中都没有条件化文本 Token 。对于每张图像,作者将它分词为“#tokens”数量的潜在 Token ,并训练生成器生成这些 Token 。

评价协议。为了评估分词器的重建性能,作者在ImageNet训练集的50K样本上报告了重建弗雷切尔入射距离(rFID)[17]、重建入射分数(rIS)[39]、峰值信噪比(PSNR)、结构相似性指数测量值(SSIM)以及学习感知图像块相似度(LPIPS)[52]。为了评估条件生成性能,作者按照ADM [9] 提供的评价协议和套件,在50K样本上报告了生成弗雷切尔入射距离(gFID)[17]、生成入射分数(gIS)[39],并遵循其评价准则和工具集,同时报告精度和召回率[24]。

为了评估文本到图像生成性能,作者在ImageNet验证集的50K样本上报告了FID和CLIP得分[16]。

4.3. Effectiveness of Text Conditioning

作者首先评估了文本条件对图像 Token 和生成的有效性。作者将作者的方法TexTok与一个 Baseline 方法(不使用文本)进行了比较,该 Baseline 方法使用相同的设置但不包含文本条件。作者在ImageNet的分辨率为

的情况下进行实验,并通过表1中的定量结果和图3中rFID相对改进的可视化展示了不同的Token数量的效果。

picture.image

picture.image

在 ImageNet

上,无论是在哪种设置下,TexTok 都显著提升了重建和生成性能。具体而言,TexTok 在使用 32、64、128 和 256 个 Tokens 的情况下,分别在 rFID 上取得了 37.2%、25.0%、30.2% 和 24.2% 的提升,这在各次测试中分别对应了 28.6%、12.7%、13.8% 和 7.9% 的 gFID 提升。值得注意的是,使用的 Tokens 越少,文本条件化带来的增益越高。正如图2a 所示,与 Baseline 相比,TexTok 使用的 Token 数量仅为 Baseline 的一半即可达到类似的 rFID(压缩率为 2 倍)。作者注意到,不含文本的 Baseline 表现非常出色。根据表1(a) 的数据显示,即使 Token 数量减少 8 倍,不含文本的 Baseline 也能在重建和生成性能上超越广泛使用的 SD-VAE 分词器 [37]。

在更高分辨率的图像上,比如ImageNet的

大小,TexTok表现出了更强的效果。如表1和图2b所示,在这种高分辨率设置下,TexTok在重构质量上取得了显著改进,并且允许更高的压缩率。具体而言,在32、64、128和256个Token的情况下,TexTok分别在rFID上取得了69.7%、68.4%、30.2%和24.2%的改进,在gFID上取得了60.8%、54.5%、13.2%和8.6%的改进。如图26所示,使用仅为 Baseline 四分之一Token数量的TexTok,就能达到相似的rFID性能。

pression rate).

最后,图1中不同Token数量的定性结果显示,TexTok显著提高了重建质量,特别是在图像中的文本和特定视觉细节(如汽车轮毂和鸟嘴)方面表现尤为明显。这表明TexTok使用相同数量的Token编码了更加精细的视觉细节。

picture.image

4.4. System-level Image Generation Comparison

作者使用TexTok作为分词器进行图像生成实验,并采用经典的DiT图像生成器[

],以研究该系统与其他领先的图像生成系统的性能对比。作者在不同Token数量(压缩率)的类别条件性ImageNet

设置下进行了评估。

在ImageNet

条件图像生成方面,如表2(a)所示,作者的TexTok

DiT-XL 达到了1.46的FID,超越了此前的最佳系统,即使使用了一个更为简单且基础的DiT作为图像生成器。随着Token数量减少和图像压缩率增加,TexTok

DiT仍能保持生成性能。特别地,当扩散 Transformer 仅生成64个图像Token的TexTok

,其性能优于使用扩散 Transformer Patch 化后256个Token的原始DiT-XL/2。

picture.image

在更高分辨率的图像上,例如ImageNet

的图像,如表2(b)所示,TexTok

DiT-XL 仅使用256个图像Token就实现了与先前方法相当的最高1.62 gFID。在最紧凑的一端,TexTok

只使用了32个Token,但在 Patch 化处理后,其生成性能优于使用1024个Token的原始DiT。

作者的系统不仅在生成性能上表现出色,还因其巨大的压缩率而非常高效。作者在图4a中绘制了TexTok

DiT-XL与原始DiT在ImageNet

上的速度与性能折衷关系。简单地将DiT中的分词器替换为TexTok可以实现14.3倍的速度提升,并且FID值更好,或者在相近推理时间的情况下,FID值提高34.3%。这验证了TexTok的有效性和效率。这种改进了的速度/性能折衷关系在ImageNet

(图4b)中进一步得到体现,作者展示了通过将DiT中的分词器替换为TexTok变体,使用32个Token时可以获得更好的FID值,93.5倍的速度提升;使用256个Token时,FID值提高了46.7%,并且推理时间减少了3.7倍。这表明随着图像分辨率的增加,为分词过程提供显式的文本语义能够带来更大的生成性能和推理加速改善。

picture.image

图6中的定性样本表明,TexTok能够生成语义丰富且细节细腻的类条件图像。更多定性样本可以参见附录E。

picture.image

4.5. Text-to-Image Generation

作者现在展示了TexTok在文本到图像生成任务中的优越性。作者使用与ImageNet

上生成的相同VLM(视觉语言模型)描述符,并结合作者修改后的DiT-T2I架构(详见第4.1节)。在训练过程中, Token 器和生成器共享由T5文本编码器提取的相同文本嵌入。在推理过程中,基于ImageNet验证集的描述符生成图像。作者计算这些生成图像与原始ImageNet验证集之间的FID(Frustum)。

如表3所示,与 Baseline (无文本)相比,TexTok在不同数量的图像Token下始终且显著地提高了文本到图像生成的效果。由于文本描述符已经在文本到图像任务中可用,且 Token 器可以直接使用生成器中使用的相同文本嵌入,因此TexTok的性能提升无需额外的成本来获取描述符和文本嵌入。

picture.image

图5中的定性样本表明,TexTok的生成更加真实,并且更好地遵循了 Prompt 。更多定性样本详见附录E。

picture.image

4.6. Tokenization/Generation Inference Efficiency

作者已经证明,TexTok 显著提升了重建、类条件生成和文本到图像生成的质量。在文本到图像任务中,作者的文本条件处理无需额外成本用于提取文本嵌入,因为文本嵌入也作为生成过程中的条件输入。对于其他任务,它仅引入了微小的计算开销来生成文本嵌入并在分词过程中使用它们。

正如表5所示,这种开销可以忽略不计(约0.01秒/张图像)。更为重要的是,生成计算成本的降低弥补了这一小小的增加,这从表5中SD-VAE、无文本 Baseline 和TexTok之间的计算成本比较以及图4中的加速结果中可以得到证实。

picture.image

4.7.AblationStudies

作者通过对TexTok进行消融实验来分析作者设计选择的贡献。作者使用以下默认设置:TexTok128、基础模型规模以及T5-XL文本编码器。Caption长度为75个词,并且在tokenizer和detokenizer中使用上下文条件进行应用。

文本条件量的影响。在表4a中,作者消融了不同类型的类别/文本条件:

(1) 基于类别类别的可学习类别嵌入;

(2) 来源于包含类别名称的简短文本模板的文本嵌入;

(3) 来源于25词描述的文本嵌入;

(4) (作者的方法)来源于75词描述的文本嵌入。作者的结果显示,更具描述性的文本条件有助于提高性能。

picture.image

T5 文本编码器规模。在表4b中,作者研究了文本编码器模型规模的影响。作者发现较大的编码器能够获得更好的重建质量。作者在ImageNet-256上默认使用T5-XL模型以提高效率。

条件生成架构。另一种设计选择是如何将文本注入到分词器和反分词器中。在表4c中,作者发现上下文条件生成(即将文本嵌入与其他输入Token进行连接并送入自我注意层)的效果优于在每个ViT块中添加一个额外的多头跨注意层。

条件化位置。在表4d中,作者消融了文本条件化注入的位置,并发现将其应用于分词器和反分词器可以获得最佳效果。

TexTok 模型规模。在表4e 中,作者研究了 TexTok 模型规模的影响。作者发现使用 TexTok-Base 表现明显优于 TexTok-Small,但进一步增加模型规模仅提供边际改进。因此,作者选择 TexTok-Base 作为默认模型规模。

  1. Conclusion

作者提出了一种基于描述性文本标注的图像Token化框架

,该框架利用描述性的文本标题提供高层次语义,从而使Token化能够专注于将细粒度的视觉细节编码到潜在Token中。

TexTok显著提高了重建和生成性能,在ImageNet上实现了最先进的结果,并具有很高的计算效率。

通过缓解重建质量与压缩率之间的权衡,TexTok使得更高效的图像生成成为可能。

参考

[0]. Language-Guided Image Tokenization for Generation .

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论