
多模态进入“实时交互”时代!字节Hyper-Bagel驯服巨模型:三大任务无损加速最高22倍!
文章链接:https://arxiv.org/pdf/2509.18824
项目链接:https://hyper-bagel.github.io/
经过6-NFE加速的BAGEL模型生成的图像生成样例
经过6-NFE加速的BAGEL模型生成的图像编辑样例
亮点直击
- 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标记、干净潜在表示和噪声潜在表示在内的更多多模态标记都交织在一起。本文第一个改进是通过注意力机制而非简单的全连接层来聚合来自目标模型的更多特征,这限制了集成更多特征层的可能性。
带残差的零初始化全连接层
减弱前向 KL 散度监督
一个潜在的问题是,虽然来自目标模型概率分布的软标签包含了丰富的知识,但使用前向 KL 散度作为损失函数来覆盖所有模式对于能力非常有限的草稿模型来说可能过于困难。为了放宽这一约束,额外引入了由目标模型输出的 one-hot 硬标签监督的交叉熵损失。
扩散蒸馏
关于扩散蒸馏,本文的主要目标是保留模型的全部能力,包括控制和质量两个方面。本文训练了一个在图像生成和编辑基准测试中无损的 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 轨迹的对齐。
通过这种伪造模型和少步生成器的交替优化,最终获得了一个在所有基准维度上都无损的 6-NFE 模型。在这个分数蒸馏阶段,模型显著提高了生成图像的保真度,在色彩鲜艳度和细节丰富度上都实现了与原始模型的高度近似。
阶段四:对抗性扩散预训练
为了进一步实现 1-NFE 图像生成和编辑以达到极致的成本效益,我们在 6-NFE 模型的基础上额外引入了第四和第五阶段的微调。我们在构建 1-NFE 模型时采用了与训练 6-NFE 模型类似的方法:首先增强结构完整性以建立图像的整体构图,然后细化生成内容的颜色和细节。然而,考虑到模型容量在 1-NFE 时显著下降,要求其与教师模型的分布对齐的原始训练范式是不现实的。本方法是首先尝试与 6-NFE 模型的分布进行结构级别的对齐,然后通过人类反馈来弥补保真度方面的不足。
阶段五:奖励反馈学习
实验
设置
在扩散蒸馏过程中,冻结了理解分支的参数。因此,由于通过推测解码训练的草稿模型在预测后需经过目标模型的验证,理解任务的性能指标不受影响。遵循 EAGLE-3的报告方式,汇报草稿模型的平均接受长度。对于生成基准测试,本文遵循 BAGEL 的做法,分别报告图像生成和编辑任务在 GenEval 和 GEdit-Bench 上的性能。
有些人可能担心扩散蒸馏是否会影响文本生成,尤其是在交织场景中的思维能力,因为 VAE 标记被纳入了序列上下文。然而,在 BAGEL 的上下文管理中,每次扩散去噪后的干净潜在表示是通过预填充(prefill)方式纳入上下文的,并不保留带噪潜在表示的 KV 缓存。这意味着在部署时,可以分别部署蒸馏后的模型的生成分支权重和原始模型的权重。只在去噪过程中调用蒸馏后的模型,并在预填充阶段切换回原始模型,从而确保理解性能没有损失。
定量结果
图像生成。如下表 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 提供了一个整体解决方案,弥合了统一多模态模型的强大能力与实际部署需求之间的差距,实现了无缝且即时的创造性交互,而无须妥协。
本文转自AI生成未来 ,作者:AI生成未来
