如您有工作需要分享,欢迎联系:aigc_to_future
作者:Yanzuo Lu等
解读:AI生成未来
文章链接:https://arxiv.org/pdf/2509.18824
项目链接:https://hyper-bagel.github.io/
亮点直击
Hyper-Bagel,一个统一的加速框架,旨在同时提升多模态理解和生成任务的速度。方法采用分治策略,利用推测解码进行下一token预测,并采用多阶段蒸馏过程进行扩散去噪,确保一个领域的加速不会影响另一个领域。 该框架带来了显著的性能提升,在多模态理解上实现了超过2倍的加速。对于生成任务,它实现了更显著的加速,将文本到图像生成速度提升16.67倍,图像编辑速度提升22倍,同时保持了原始模型的高质量输出。 进一步开发了一个高效的1-NFE模型,能够实现近乎实时的交互式编辑和生成。通过结合先进的对抗蒸馏与人类反馈学习,该模型实现了极致的成本效益和响应能力,使复杂的多模态交互变得无缝且即时。
总结速览
解决的问题
计算开销大: 随着多模态内容(交织的图像和文本token)越来越复杂,统一多模态模型在理解和生成任务中的迭代过程(如扩散去噪和自回归解码)带来了巨大的计算负担,导致速度缓慢。 现有加速方法在多模态场景下效果不佳: 现有的推测解码技术在处理纯文本的大型语言模型时效果显著,但在处理多模态模型(如BAGEL)时,由于图像和文本嵌入空间的复杂性,加速效果大打折扣。 保持模型能力与质量: 在加速扩散过程时,需要同时保持模型的关键能力,包括通过分类器无关引导实现的控制能力(指令遵循、图像编辑一致性),以及生成图像的结构完整性和视觉保真度。
提出的方案
Hyper-Bagel框架: 一个统一的加速框架,采用"分而治之"的策略,同时针对多模态理解(自回归解码)和生成(扩散去噪)任务进行加速。 针对理解任务的加速: 采用改进的推测解码。训练一个轻量级的"草稿模型"来预测多个连续标记,然后由目标模型并行验证,将内存访问瓶颈转化为计算瓶颈。 针对生成任务的加速: 采用多阶段蒸馏过程。将目标分解为三个关键维度(CFG控制、结构完整性、图像保真度)并分别优化,逐步减少采样所需的步骤。
应用的技术
推测解码: 基于EAGLE-3范式,但针对多模态挑战进行了关键改进,设计了一个高效的中间层架构来桥接目标模型和草稿模型,提升草稿模型在复杂多模态序列中的预测准确性。 扩散蒸馏: CFG蒸馏: 将文本/图像引导尺度与时间步一起作为控制条件嵌入到单次前向传播中。 对抗性蒸馏: 设计了多头判别器,在多尺度上判别潜在表示的真伪,以增强生成图像的结构完整性。 分数蒸馏: 提出了DMDO方法,采用基于常微分方程的欧拉离散采样器生成图像,避免了基于SDE的方法可能产生的图像过度平滑问题,且无需额外的正则化器。 极致加速模型(1-NFE): 在6-NFE模型基础上,通过基于修正流的对抗性训练和基于人类反馈的奖励学习(ReFL) 进行进一步优化,使用具备更强视觉理解能力的VLM-based HPSv3作为奖励模型。
达到的效果
多模态理解: 实现了超过2倍(2.16倍)的加速(Tokens Per Second从98.3提升至212.4)。 多模态生成(6-NFE模型): 文本到图像生成: 实现16.67倍加速。 图像编辑: 实现22倍加速。 质量保持: 在上述加速下,在GenEval和GEdit-Bench指标上保持了与原始模型相当的输出质量,是无损加速。 极致效率(1-NFE模型): 实现了接近实时的交互式编辑和生成,在保证成本效益的同时,使复杂的多模态交互变得无缝和即时。
数据
本文中用于草稿模型训练和扩散蒸馏的VLM图像-文本配对数据、文本到图像生成数据、图像到图像编辑数据以及交织数据均源自线上可获取的开源数据集。
VLM图像-文本配对数据:采用LLaVA-OneVision中的单图像阶段数据混合集作为VLM任务的训练数据,其中包含约400万个图像-文本对。为了与目标模型的预测分布对齐并训练草稿模型,我们也调用目标模型为每个问题生成新的答案(如EAGLE3中的做法),而不是使用数据集中的答案。
文本到图像生成数据:引入JourneyDB作为文本到图像扩散蒸馏的训练数据,这是一个包含超过400万张由Midjourney生成的图像以及由VLM标注的详细标题的合成数据集。由于图像的原始提示均由用户提供,这种合成分布与实际使用情况高度吻合,其丰富的多样性也有利于蒸馏训练。
图像到图像编辑与交织数据:为了保留BAGEL的新兴特性,我们在蒸馏训练过程中也加入了交织数据。具体来说,使用SEED-Data-Edit 数据集的Part-2和Part-3中的编辑数据。Part-2包含52K个真实场景中的编辑图像对,而Part-3包含21K个人工标注的多轮对话(每轮最多5个回合),总计95K个编辑图像对。
方法
推测解码
在采用推测解码时,遵循EAGLE-3的训练范式。为了解决不同模态标记之间的差异使得能力有限的草稿模型难以编码目标特征这一挑战,特别在中间层、初始化策略和损失函数方面实施了一系列改进以增强EAGLE-3。所提出的中间层架构如下图3所示,其中还将草稿解码器层数扩展到N=2以增强表示能力。
使用元查询的目标特征聚合
EAGLE的中间层在为草稿模型执行下一token预测提供必要信息方面起着至关重要的作用,但正如我们在引言中讨论的,这对BAGEL来说挑战更大,因为包括文本、ViT标记、干净潜在表示和噪声潜在表示在内的更多多模态标记都交织在一起。本文第一个改进是通过注意力机制而非简单的全连接层来聚合来自目标模型的更多特征,这限制了集成更多特征层的可能性。
初始化几个可学习的嵌入作为元查询,这些查询与序列中每个token的所有目标层特征(作为键和值)通过交叉注意力进行交互,其中是输入序列的长度,是目标模型层数,是特征维度。将展平的注意力输出向量记为。其中的每一个计算如下:
其中 在序列中的不同标记之间共享,在下文中使用 表示交叉注意力模块。序列 随后继续前向通过一个与目标模型架构相同的 Transformer 解码器层 ,以进一步聚合信息,最终形成草稿模型的输入 。请注意,在这个中间阶段,输入特征并未与标记嵌入进行拼接,并且该解码器层独立于草稿模型中参与迭代解码的那些层。
带残差的零初始化全连接层
关于初始化策略,实验表明,利用预训练目标模型的最后几层和语言建模头有助于最终性能。这基于一个基本直觉,即我们的目标是降低草稿模型的训练难度。因此,希望预训练的权重能够通过梯度反向传播来指导两个最关键的全连接层的更新,即中间层中负责降维的 ,以及在每个迭代解码步骤中通过拼接融入标记嵌入的 。
为了实现这一点,除了利用预训练权重外,本文进行了以下改进:(1) 对中间层中 和 的最后一个投影层进行零初始化;(2) 为 和 添加跳跃连接,其中我们指定目标模型前向传播中的低、中、高层特征序列作为残差,类似于 EAGLE-3;(3) 对 和 进行零初始化。
通过这种方式,在零输入的训练初始阶段,从损失函数反向传播到 的梯度与目标模型保持一致。这意味着 需要聚合现有信息来模拟目标模型在相应位置的特征序列。当草稿模型仅包含一个解码器层时, 需要模拟倒数第二层的输出;如果草稿模型有两个解码器层,则 模拟倒数第三层的输出。而 最初通过跳跃连接获得的输入特征与 EAGLE-3 中的三个不同层次的特征相同。
减弱前向 KL 散度监督
一个潜在的问题是,虽然来自目标模型概率分布的软标签包含了丰富的知识,但使用前向 KL 散度作为损失函数来覆盖所有模式对于能力非常有限的草稿模型来说可能过于困难。为了放宽这一约束,额外引入了由目标模型输出的 one-hot 硬标签监督的交叉熵损失。
其中 和 分别表示草稿模型和目标模型的输出分布概率向量, 表示概率分布 中第 个类别的概率。在实践中,设置 ,并发现其性能优于 EAGLE-3 中使用的原始前向 KL 散度。
扩散蒸馏
关于扩散蒸馏,本文的主要目标是保留模型的全部能力,包括控制和质量两个方面。本文训练了一个在图像生成和编辑基准测试中无损的 6-NFE 模型,以及一个高成本效益的 1-NFE 模型。6-NFE 模型通过包含 CFG 蒸馏、TSCD和 DMDO的三阶段训练获得。1-NFE 模型则从 6-NFE 模型出发,通过两个额外的阶段(即 AD和 ReFL)进行微调得到。
阶段一:CFG 蒸馏
为了保留模型通过 CFG 控制指令遵循程度以及在编辑场景中保持与原始图像一致性的能力,在第一阶段的训练中将 CFG 嵌入蒸馏到其单次前向传播中。本文精心设计了两个额外的时间步编码层,分别用于注入文本引导尺度和图像引导尺度。文本引导尺度在图像生成和编辑中通用,而图像引导尺度仅用于编辑场景。它们的架构和注入位置与时间步的编码完全一致,这与 FLUX中的最佳实践类似,确保了控制信号能够精确传播到 DiT的每一层。
在蒸馏训练过程中,对于文本到图像数据,随机选择一个范围在 1 到 5 之间的文本引导尺度值;而对于编辑样本,额外采样一个在 1.0 到 2.5 之间的随机图像引导尺度值作为输入。关于训练时间步调度,发现将扩散时间步偏移设置为至少 3.0 对于实现更高的结构完整性和图像保真度是必要的。
阶段二:轨迹分段一致性蒸馏
在 CFG 蒸馏之后,采用与之前工作 Hyper-SD 类似的一致性蒸馏方法。尽管我们的目标是得到一个 6-NFE 模型,但我们不像 Hyper-SD 那样采用渐进式蒸馏(例如,将分段数从 8→4→2→1 减少),而是直接在一个阶段中实现 3 段配置。并且我们弃用了均方误差(MSE)损失函数,完全转向对抗性损失。在判别器设计上,采用与 DMDX 相同的多头架构以增强判别能力,并且预训练骨干网络的参数也被设置为可训练。
所有这些改进背后的直觉是为了在第二阶段增强结构完整性,同时将提升图像保真度的任务委托给第三阶段的分数蒸馏。消除多阶段渐进式蒸馏简化了训练流程并减少了训练时间,而纯对抗性损失和多头判别器设计使模型能够更专注于跨多个尺度的整体图像构图。
阶段三:通过 ODE 的分布匹配蒸馏
DMD 系列方法的一个主要缺点是在少步生成器中使用了基于 SDE 的一致性采样器,这导致生成的图像过于平滑且缺乏细节。这与在第三阶段提升图像保真度的目标相悖。为此,本文提出了通过 ODE 的分布匹配蒸馏(DMDO),旨在保持原始采样器不变,并尽可能保持学生模型和教师模型之间 ODE 轨迹的对齐。
在伪造模型更新步骤中,随机初始化一个噪声 并利用少步生成器的 ODE 来获得完整的轨迹。以 6-NFE 为例,保存完整的轨迹 ,并在 和 之间进行线性插值以获得 作为伪造模型的输入。因此,伪造模型的目标也转变为预测速度场 ()。这是为了使伪造模型能更好地捕捉少步生成器在每个时间步的实际分布,消除 DMD 原始实践中添加随机噪声可能引入的分布偏移。
在少步生成器更新步骤中,重用伪造模型更新步骤中已有的轨迹,以帮助降低训练成本并缩短每次迭代的时间。以采样到一个落在 和 之间的时间步 为例。从轨迹中提取带噪潜在表示 ,然后再次通过启用梯度的少步生成器的 ODE 进行前向传播。在时间步 获得预测的 后,重新采样一个新的时间步 ,并与 进行线性插值以获得 ,作为伪造模型和真实模型的分数函数输入。后续的损失函数计算遵循与 DMD 相同的方法,如下图 4 所示。
通过这种伪造模型和少步生成器的交替优化,最终获得了一个在所有基准维度上都无损的 6-NFE 模型。在这个分数蒸馏阶段,模型显著提高了生成图像的保真度,在色彩鲜艳度和细节丰富度上都实现了与原始模型的高度近似。
阶段四:对抗性扩散预训练
为了进一步实现 1-NFE 图像生成和编辑以达到极致的成本效益,我们在 6-NFE 模型的基础上额外引入了第四和第五阶段的微调。我们在构建 1-NFE 模型时采用了与训练 6-NFE 模型类似的方法:首先增强结构完整性以建立图像的整体构图,然后细化生成内容的颜色和细节。然而,考虑到模型容量在 1-NFE 时显著下降,要求其与教师模型的分布对齐的原始训练范式是不现实的。本方法是首先尝试与 6-NFE 模型的分布进行结构级别的对齐,然后通过人类反馈来弥补保真度方面的不足。
关于 1-NFE 模型的结构完整性,采用了基于 DMDX 中修正流的对抗性方法,即对抗性扩散预训练(ADP)。该方法利用 6-NFE 模型采样一条 ODE 轨迹,然后在起点 和终点 之间进行线性插值以获得一个带噪潜在表示 ,随后将其输入到 1-NFE 生成器进行预测。1-NFE 生成器预测的 被两个分别在隐空间和像素空间的判别器评估为伪造样本,而真实目标样本则是 ODE 轨迹终点 本身。这种方法与我们的训练目标完美契合,因为我们只旨在实现与 6-NFE 模型的对齐,并且使用 6-NFE 模型采样 ODE 轨迹的成本相对较低。
阶段五:奖励反馈学习
关于 1-NFE 模型的图像保真度,遵循 Hyper-SD 中的常见做法,引入人类反馈。与先前的方法不同,本文不采用多个奖励模型的训练范式。相反,我们仅使用一个更全面的基于 VLM 的奖励模型来提供指导。这是考虑到 BAGEL 也是一个多模态模型,其能力侧重于对提示词的语义理解。因此,基于 VLM 的奖励模型非常合适,因为 VLM 可以扩展到更大的容量,从而容纳更丰富的知识,并且奖励模型的缩放已在 RewardDance 中被证明是非常有效的。在 ReFL 中使用的损失函数与 Hyper-SD 中的审美监督损失函数一致,其中 ReLU 函数的奖励激活阈值 设置为 6.0。
实验
设置
在扩散蒸馏过程中,冻结了理解分支的参数。因此,由于通过推测解码训练的草稿模型在预测后需经过目标模型的验证,理解任务的性能指标不受影响。遵循 EAGLE-3的报告方式,汇报草稿模型的平均接受长度。对于生成基准测试,本文遵循 BAGEL 的做法,分别报告图像生成和编辑任务在 GenEval 和 GEdit-Bench 上的性能。
有些人可能担心扩散蒸馏是否会影响文本生成,尤其是在交织场景中的思维能力,因为 VAE 标记被纳入了序列上下文。然而,在 BAGEL 的上下文管理中,每次扩散去噪后的干净潜在表示是通过预填充(prefill)方式纳入上下文的,并不保留带噪潜在表示的 KV 缓存。这意味着在部署时,可以分别部署蒸馏后的模型的生成分支权重和原始模型的权重。只在去噪过程中调用蒸馏后的模型,并在预填充阶段切换回原始模型,从而确保理解性能没有损失。
定量结果
推测解码。在下表 1 中,展示了包括本文方法在内的多种 EAGLE-3 变体在使用链式解码进行 10 个外推步时的平均接受长度 和接受率 。完整的 Hyper-Bagel 框架实现了最佳性能,优于原始的 Vanilla EAGLE-3 基线。结果表明,移除零初始化策略会导致最显著的性能下降,突显了其在桥接目标模型和草稿模型中的关键作用。虽然交叉熵损失也被证明是有益的,但其缺失导致的下降幅度较小。有趣的是,模型在同时缺少这两个组件时的表现比仅移除零初始化时更好,这表明在没有零初始化提供的基础对齐的情况下,严格的交叉熵损失约束可能会产生反效果,从而阻碍训练过程。
图像生成。如下表 2 所示,我们的 6-NFE Hyper-BAGEL 模型在 GenEval 基准测试中展示了无损的性能。它取得了 0.8647 的综合得分,略微超过了 100-NFE BAGEL 基线的 0.8640 分。这一结果证实了我们的蒸馏过程在将采样步数从 100 NFE 减少到 6 NFE 从而实现 16.67 倍加速的同时,并未损害生成质量。此外,高效的 1-NFE 模型与其他最先进的统一模型相比仍具有竞争力,其 0.7962 的综合得分与 Janus-Pro-7B 和 MetaQuery-XL(两者均为 0.80 分)等领先模型相当。
图像编辑。本文加速模型的无损特性在 GEdit-Bench(下表 3)上评估的图像编辑任务中表现明显。6-NFE Hyper-BAGEL 在英文和中文数据集上均一致地优于 132-NFE 基线,分别取得了 6.612 和 6.671 的更高综合得分。尽管计算成本大幅降低,相当于实现了约 22 倍的推理加速,但仍保持了这种卓越的一致性。值得注意的是,为追求极致效率而设计的 1-NFE 模型依然表现出强大的性能,其在英文和中文数据集上分别取得的 5.975 和 5.966 的综合得分显著超过了如 OmniGen 等现有方法。
定性结果
图像生成。下图 5 展示了本文的加速模型与 100-NFE 基线在图像生成任务上的定性比较。从所有四个案例中可以明显看出,我们的 6-NFE 模型的结果在视觉上与基线无法区分,忠实地再现了复杂的细节,例如猴子衬衫上的数字"619"、黑暗骑士场景中复杂的构图和光照,以及动物毛皮的纹理。相比之下,为追求极致效率而设计的 1-NFE 模型则表现出明显的权衡。虽然它能快速生成与提示词高度相关的图像,但这是以降低细节保真度为代价的。例如,它有时可能会忽略提示词中的关键元素,如小猫的眼镜,或者在细节上出现微小偏差,如猴子衬衫上的数字。尽管如此,其输出的核心语义和整体质量仍然具有很强的竞争力。这些视觉比较有力地验证了6-NFE 模型实现了无损加速,而 1-NFE 模型则为优先考虑实时交互的应用提供了一个高效可靠的选择。
图像编辑。下图 6 展示了本文的加速模型与 132-NFE 基线在图像编辑任务中的定性性能表现。对于每个案例,6-NFE Hyper-BAGEL 模型都展示了卓越的保真度,执行了精确的编辑,例如移除花生、替换文本或消除人物,其结果与高 NFE 基线几乎无法区分。这证实了6-NFE 加速的无损特性,在提供显著加速的同时,没有任何可感知的编辑质量下降。1-NFE 模型在编辑场景中展现出一个显著优势。它能够利用源图像的结构和上下文信息,使其能够保持强大的视觉连贯性并成功应用所请求的编辑。虽然与更高 NFE 版本的模型相比,细微的细节或完美的照片真实感可能略有不足,但 1-NFE 模型仍然能提供高度可用且上下文准确的编辑,使其成为近乎实时的交互式编辑的强大工具,在这种场景下,快速模型提供的即时视觉反馈是无价的。
结论
Hyper-Bagel,一个旨在成功减轻先进多模态模型中显著计算开销的统一框架。采用分治策略,对理解任务使用推测解码,对生成任务使用多阶段蒸馏,通过全面的实验证明该策略非常有效。已经证明,无损 6-NFE 模型将文本到图像生成和编辑任务分别加速了超过 16.67 倍和 22 倍,实现了与高 NFE 基线相当甚至更优的性能,同时将多模态理解的速度提高了一倍。高效的 1-NFE 模型为近实时应用提供了一个稳健且实用的解决方案,尤其在交互式编辑中证明特别有效。最终,Hyper-Bagel 提供了一个整体解决方案,弥合了统一多模态模型的强大能力与实际部署需求之间的差距,实现了无缝且即时的创造性交互,而无须妥协。
参考文献
[1] Hyper-Bagel: A Unified Acceleration Framework for Multimodal Understanding and Generation
致谢
如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。同时,欢迎给个星标⭐,以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进,共同探索未知,见证一个充满希望和伟大的未来!
技术交流
加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!
没有评论:
发表评论