这篇论文介绍了Dimba,一个全新的文本到图像扩散模型,它采用了一种独特的混合架构,结合了Transformer和Mamba元素。
具体来说,Dimba逐序堆叠的块在Transformer和Mamba层之间交替,并通过交叉注意力层整合条件信息,从而利用了两种架构范式的优势。
作者研究了几种优化策略,包括质量调整、分辨率适应,并确定了大规模图像生成所需的关键配置。该模型灵活的设计支持针对特定资源限制和目标的场景。
适当地扩展后,Dimba相对于传统的基于纯Transformer的基准提供了显著的高吞吐量和较小的内存占用。广泛的实验表明,Dimba在图像质量、艺术渲染和语义控制方面与基准相比具有可比的性能。
作者还报告了在评估过程中发现的结构的一些有趣特性,并在实验中发布了预训练权重。
作者的发现强调了在大规模扩散模型的基础阶段,混合Transformer-Mamba架构的潜力,为文本到图像生成领域带来了光明的前景。
1 Introduction
扩散模型从噪声中创建数据,并基于将数据路径反转到随机噪声的训练上进行训练。通过利用神经网络固有的近似和泛化特性,这些模型能够生成在训练集中找不到但与其分布一致的全新数据点。这种生成建模技术在处理高维感知数据,如图像方面已被证明非常有效。近年来,诸如DALL-E,Imagen ,Stable Diffusion ,EMU 等文本到图像扩散模型,彻底改变了 photorealistic 图像综合领域的面貌,极大地影响了包括图像编辑,视频生成,3D资产创建等应用。
在反向扩散过程中预测噪声或先前状态通常使用基于卷积神经网络(CNN)的U-Net架构或Transformer,如先前研究所强调的。
另一方面,状态空间模型(SSM)的进步为实现计算效率和模型灵活性之间的平衡铺平了道路。源自经典的卡尔曼滤波模型[36],各种基于SSM的方法已被提出用于处理序列数据,在管理不同任务和模态中的长距离依赖方面被证明非常有效。
它们在处理扩展序列时的效率归功于卷积计算和几乎线性的计算复杂性。最近,Mamba通过将时变参数引入SSM并提出了一种硬件感知算法,促进了高度高效的训练和推理,从而取得了进展。Mamba卓越的扩展性能使其成为语言建模中替代Transformer[71]的有力选择。
此外,一系列基于SSM的主干网络[85; 80; 55]也已被探索用于处理视觉数据,如图像和视频。对Transformer-Mamba混合结构在基于扩散的多模态生成任务中的应用进行深入分析,并识别确保模型容量的关键技术变得至关重要。
在本文中,作者介绍了Dimba,一个新的用于文本到图像合成的扩散模型设计。具体来说,Dimba基于混合架构,结合了Transformer层和Mamba层,Mamba是一种最近先进的态空间模型。尽管CNN或Transformer作为扩散模型的主导架构极为流行,但它们仍然存在重大缺陷。
特别是它们对内存和计算资源的大量需求,限制了在无需特殊设计的情况下处理长上下文的能力。相比之下,基于递归的神经网络模型可以在单个隐藏状态中总结任意长的上下文,从而减少这些问题。Dimba因此结合了两种正交的架构设计,提升了性能和吞吐量,同时保持了可管理的内存占用。
请注意,Dimba架构的灵活性使其能够根据硬件和性能需求容纳各种设计选择。
图1展示了若干生成结果。大量实验表明,在不同参数设置下,Dimba在图像质量和语义对齐方面优于现有主流扩散模型,其在T2I-CompBench [34]上的表现也突显了其在语义控制方面的优势。
作者工作的核心贡献是:
- 作者提出Dimba,一种新型的扩散架构,用于文本到图像生成,结合了注意力和Mamba层。作者展示了Dimba在平衡吞吐量和内存需求方面的灵活性,同时保持了可比较的性能。
- 作者策划了一个大规模高质量图像文本数据集,强调审美上乘的图像。这些图像通过评分网络进行评估,并与先进的自动标签模型生成的密集、精确的标题配对。这得益于文本图像语义对齐的充足性以及图像审美质量的提升。
- 作者采用分阶段逐步策略训练Dimba,先使用大规模数据进行预训练,然后适应高分辨率以及使用精选的小型集合进行质量调整。作者的努力旨在为高效的扩散模型提供有价值的见解,帮助个人研究行人和初创公司以较低的内存成本开发高质量的文本到图像模型。
2 Related Works
基于扩散模型的图像生成。近年来,图像生成领域取得了实质性进展,特别是通过扩散过程。这项技术通常涉及从高斯噪声开始,通过一系列步骤迭代改进,直到与目标分布对齐。它已经显示出卓越的能力,通常超过了基于GAN[22; 17]和VAE[39]的模型的性能。为了根据文本条件生成图像,已经引入了多种方法,例如GLIDE[52],DALLE[60],Imagen[31],LDM[61]等[1; 6; 7; 40; 83]。这些模型使用大规模的文本-图像对训练扩散模型,从而可以根据文本条件创建图像。在此过程中,U-Net负责根据潜在空间中的文本嵌入预测噪声。基于CNN的UNet的特点是一组下采样块,一组上采样块,以及两组之间的长 Short-Cut 。
相比之下,基于Transformer的架构用自注意力替换部分采样块,而其余部分保持不变[56; 79; 2]。基于SSM的UNet[76; 19; 33; 84]在图像序列内也显示出具有计算优化的有希望的结果。在本文中,作者尝试将Transformer与Mamba[23]结合,以提高计算效率以及可比较的性能。
视觉应用的状态空间模型。SSM被用于序列建模,具有一维的输入和输出,并已应用于控制理论、信号处理和自然语言处理[28]。一系列研究集中在结构改进上[25; 66; 21; 49]。
最近,提出了Mamba[23],它具有一个数据相关SSM层,并构建了一个通用的语言模型 Backbone ,在各种规模上优于Transformer,并在大型真实数据上享受序列长度的线性扩展。关于将注意力机制与SSM结合,[86]将S4层[25]与局部注意力层结合,并在小型模型和简单任务上证明了有效性。从SSM层开始,然后是块状Transformer,显示了改进的困惑度指标。
在Transformer层的自注意力之前添加了SSM层,而[64]在自注意力之后添加SSM层,两者在语音识别上都取得了改进。[54]用Mamba层替换了Transformer中的MLP层,并在特定任务上显示了好处。
用更可扩展的基于SSM的 Backbone 替换了注意力机制,以生成高分辨率图像。最相似的工作[41]在因果语言建模中扩展混合Mamba和Transformer结构,而作者的工作专注于双向结构进行文本到图像合成。
3 Methodology
Model Architecture
一般来说,Dimba是一个混合扩散架构,它将一系列Transformer层与Mamba层相结合,并辅以跨注意力模块来处理条件信息。作者将这三者的组合称为Dimba块。
参见图2以获得说明。先前的实验表明,直接将文本信息与前置嵌入 ConCat 会减缓训练的收敛速度。Dimba中Transformer和Mamba元素的结合为平衡有时相互冲突的低内存使用、高吞吐量和高质量目标提供了灵活性。一个重要的考虑因素是Transformer模型在长上下文中的可扩展性,其中内存缓存的二次增长成为一个限制因素。通过将一些注意力层替换为Mamba层,可以减少整体缓存大小。Dimba的架构设计要求比传统Transformer具有更小的内存缓存。
在吞吐量方面,对于短序列,注意力操作在推理和训练FLOPS中占相对较小的部分。然而,当序列不断增加时,注意力操作主导了计算资源。
相反,Mamba层在计算上更有效率,增加它们的比例可以提升吞吐量,特别是对于长序列。
对于实现细节,基本单元Dimba块按顺序重复,前一层的隐藏状态输出作为下一层的输入。每个这样的层包含一个注意力和一个Mamba模块,即作者将注意力-Mamba比例设置为1:,然后是一个多层感知机(MLP)。借鉴[6],作者在AdaLN中包含了全局共享的MLP和逐层嵌入的时间步信息。请注意,每个Dimba块中的Mamba层也包含了几个自适应规范化层(AdaLN)[56],这些层有助于在大规模模型中稳定训练。
Dataset Construction
图像-文本对自动标注与质量估计。主流的图像-文本数据集中的文本描述,通常来源于网络爬取,常常存在诸如文本与图像不对齐、描述不完整、虚构内容以及罕见词汇等问题,如文献所示。此外,值得注意的是,像LAION-400M[65]这样的公共数据集,主要包含来自网站的基础产品预览,这对于训练旨在寻求物体组合多样性的文本到图像生成模型来说并不理想。另外,先前的研究表明,用户通常更偏好具有更高艺术风格的图像[40]。为了解决这些问题并生成高准确性和信息密度的描述,作者首先爬取了一个大规模的内部数据集,然后根据自动评分器筛选高质量图像,以增强生成图像的美学质量,超越真实的照片。在此过程中,作者使用来自2的LAION-Aeshetics-Predictor V2进行图像质量估计,该预测器集成了在具有特定美学评分的数据集上训练的CLIP和MLP模型。最后,作者从开源存储库3中使用了最先进的图像字幕生成模型,即ShareCaptioner[8],来重新标注数据。使用提示“全面而详细地分析图像。”,作者显著提高了描述的质量,如图3(a)所示。
本文介绍了Dimba,这是一种专为文本到图像生成设计的开创性扩散架构,它集成了注意力和Mamba层。
作者还展示了Dimba的一个实现,其性能与精心策划的数据集相当。
此外,作者阐明了Dimba在平衡性能和内存需求的同时保持高吞吐量方面的灵活性。
通过大量实验,作者探索了各种设计考虑因素,包括数据构建、质量调整和分辨率适应,揭示了开发过程中的见解。这些发现有望指导未来在混合注意力-Mamba Backbone 生成方面的努力。
为了促进此类研究,作者打算发布源自较小规模训练迭代的代码和模型预训练权重。
参考
[1].Dimba: Transformer-Mamba Diffusion Models.