HybridStitch：像素与时间步级别的模型缝合加速扩散模型

通过像素级区域感知的模型切换策略加速文生图扩散模型推理

前置知识

扩散模型（Diffusion Model）

扩散模型是一类概率生成模型，其推理过程包含一系列去噪步骤。具体而言，模型以高斯噪声 $x_T$ 作为输入，通过迭代预测并消除噪声，经过 $T$ 步去噪后逐步还原为清晰图像 $x_0$。每一步去噪都基于马尔可夫链假设，计算条件概率 $p_{\theta}(x_{t-1}|x_t)$。随着模型参数量从 SD1.5 的 983M 增长到 SD3.5 的 8.1B，甚至商业模型超过 20B 参数，虽然图像质量提升，但推理延迟也随之显著增加。

理解扩散模型的基本推理流程是理解本文提出的加速方法的基础，特别是去噪步骤的时序特性和计算开销问题。

模型融合/混合模型（Mixture of Models）

模型融合是一种加速扩散推理的策略，通过组合大小两个模型来平衡质量与效率。核心思想是在推理过程中定义切换函数，在某些去噪步骤使用大模型，其他步骤使用小模型。例如 MoDM 使用大模型处理初始步骤，T-Stitch 先用小模型再用大模型，SRDiffusion 则先用大模型构建草图再切换到小模型。这些方法的共同局限是都在整张图像粒度上进行切换，忽略了单个时间步内像素间的计算需求差异。

本文正是针对现有模型融合方法的这一局限性提出了像素级区域感知的切换策略，理解已有方法有助于把握本文的创新点。

KV Cache（键值缓存）

KV Cache 是 Transformer 架构中的一种优化技术，缓存前一步计算得到的 Key 和 Value 张量，避免重复计算。在扩散模型中，相邻去噪步骤的 KV 表现出高度相似性。本文利用这一特性，在大模型仅处理部分像素（通过 mask 选择的困难区域）时，用前一步的 KV Cache 填充完整上下文，保证注意力计算的质量。对于 FFN 和归一化等不涉及跨 token 交互的操作，则无需额外处理。

KV Cache 是 HybridStitch 在混合阶段保持生成质量的关键技术组件，理解其工作原理对于理解掩码生成（Masked Generation）机制至关重要。

FID（Fréchet Inception Distance）

FID 是衡量生成图像质量的标准指标，通过计算生成图像与真实图像在 Inception 网络特征空间中的分布差异来评估。FID 越低表示生成图像的分布越接近真实图像。本文在 COCO 数据集上评估，大模型基线 FID 为 27.64，HybridStitch-30% 达到 30.43，在加速方法中表现最优。FID 综合考虑了图像的保真度和多样性，是评价生成模型的主流指标。

FID 是本文评估生成质量的核心指标之一，理解其含义有助于理解实验结果的质量对比。

CLIP Score

CLIP Score 利用 CLIP 模型将输入文本提示和输出图像映射到同一嵌入空间，通过计算余弦相似度来评估生成图像与输入文本的语义对齐程度。分数越高表示图像越符合文本描述。本文中各方法的 CLIP Score 差异较小（31.79-32.01），说明所有加速方法都能较好地保持语义一致性，这也表明质量差异主要体现在视觉细节而非语义层面。

CLIP Score 是评估文本-图像对齐程度的关键指标，与 FID 互补，共同构成质量评估体系。

研究动机

当前扩散模型面临严重的计算开销问题。随着模型参数量从 Stable Diffusion 1.5 的 983M 增长到 SD3.5 的 8.1B，推理延迟显著增加，这对延迟敏感的应用场景构成了重大障碍。现有的混合模型加速方法如 T-Stitch 和 SRDiffusion 虽然通过在不同去噪步骤间切换大小模型来减少计算，但它们都以整张图像为粒度进行切换，忽略了单个时间步内像素间的异质性计算需求。具体而言，一张图像中的某些像素（如背景区域）更容易渲染，可以更早地切换到小模型，而复杂区域（如物体主体部分）则需要大模型持续精炼。整图粒度的切换要么在简单区域就绪时过早切换导致质量下降，要么等到所有像素都充分精炼后才切换导致延迟增加，两者都无法达到最优效率-质量平衡。

本文的目标是本文的具体目标是设计一种像素和时间步级别的模型缝合（Model Stitching）方法，实现区域感知的模型切换，从而在保持甚至提升生成质量的同时，最大化加速效果。作者希望超越现有混合模型方法的加速上限（SRDiffusion 的 1.55x），在 Stable Diffusion 3 上实现显著更高的加速比，同时维持可比或更优的图像质量。

与已有工作不同的是，本文的独特切入角度是将文生图生成视为一种'编辑'过程，观察到大模型和小模型的预测差异在像素级别存在显著的区域多样性——通过分析发现，40% 差异最大的值主要集中在最终图像中的物体部分，表明像素级差异确实存在。基于这一关键观察，作者提出了混合阶段（Hybrid Stage）的概念，在同一个去噪步骤中同时运行大小两个模型：小模型处理整张图像以构建粗略草图并保持全局一致性，大模型仅处理 mask 标记的'困难'区域进行精炼和编辑。这种区域感知的切换策略是对现有方法在空间维度上的根本性创新。

核心方法

HybridStitch 的整体思路可以类比为一位画家的创作过程：先用大号画笔勾勒整体布局，然后切换到小号画笔填充大面积的简单区域，同时用大号画笔专注于细节复杂的部分，最后完全交给小号画笔完成收尾。具体而言，该方法将生成过程分为三个阶段：第一阶段仅使用大模型处理初始高斯噪声，构建图像的整体布局；第二阶段进入混合模式，小模型处理整张 latent 以构建当前步骤的草图，大模型仅处理 mask 标记的困难区域并精炼小模型的输出；第三阶段完全切换到小模型完成剩余去噪步骤。在每个时间步中，HybridStitch 计算相邻步骤的差异张量，根据阈值判断是否进入下一阶段，并动态更新 mask 以确保复杂区域获得足够的精炼。

HybridStitch 的核心创新在于引入了像素级别的区域感知模型切换机制，这是与现有混合模型方法的本质区别。T-Stitch 和 SRDiffusion 都是以整张图像为单位进行模型切换——在某个时间步，要么整张图像由大模型处理，要么整张图像由小模型处理。而 HybridStitch 的混合阶段（第二阶段）中，在同一个去噪步骤内，大模型只处理标记为'困难'的像素子集，小模型处理完整的 latent 状态。通过 Top-K 差异选择策略（选取差异张量中最大的 K 个值作为 mask），系统能够自动识别出图像中需要大模型精炼的区域。这种设计使得大模型的计算量从处理整张图像降低到仅处理 10%-40% 的像素，同时通过 KV Cache 技术保持大模型在处理部分像素时的注意力上下文完整性，实现了效率与质量的双赢。

方法步骤详情

HybridStitch 的完整算法包含三个阶段和两个切换点。首先，给定初始高斯噪声 $x_T$ 和阈值列表，系统初始化阶段标志 stage=0。在第一阶段（步骤 1 到 $T_1$），大模型处理完整的 latent 状态，预测噪声并缓存 KV 值。每个时间步后，系统计算相邻步骤的差异 $D_t = \mathrm{Mean}(\frac{\|x_{t+1}\|_1}{\|x_t\|_1})$，当 $D_t$ 低于第一个阈值时进入第二阶段。在第二阶段（步骤 $T_1$ 到 $T_2$），小模型对完整 latent 进行噪声预测以构建草图，大模型仅对 mask 标记的像素子集进行处理。大模型使用前一步的 KV Cache 填充完整上下文，其输出会替换小模型输出中对应 mask 区域的值，组合后的噪声用于更新 latent。每个时间步后系统继续计算差异并更新 mask（通过 Top-K 选择差异最大的区域），当 $D_t$ 低于第二个阈值时进入第三阶段。在第三阶段（步骤 $T_2$ 到 $T$），完全由小模型处理剩余的所有去噪步骤，直到生成最终图像 $x_0$。此外，大模型的计算遵循约束条件 $M_1 < 1 - \frac{L_s}{L_l}$，确保混合模式确实能带来加速收益。

技术新颖性

HybridStitch 的技术新颖性体现在多个层面。首先，它首次在扩散模型的混合模型加速框架中引入了空间维度（像素级别）的切换策略，打破了此前所有方法只能在时间维度上进行模型切换的范式。其次，混合阶段的设计——同一个去噪步骤中同时运行两个模型并按区域组合输出——是一个全新的计算范式，与传统的'要么全用大模型，要么全用小模型'的二元选择形成鲜明对比。第三，针对大模型仅处理部分像素时丢失完整上下文的问题，作者巧妙地利用相邻步骤 KV 的高相似性，通过 KV Cache 填充机制保证了注意力计算的完整性。第四，动态 mask 更新策略（每步重新计算 Top-K 差异区域）使得系统能够适应去噪过程中困难区域的动态变化，进一步提升了质量。第五，将生成过程类比为'编辑'过程的思想也具有启发性——小模型先'画'出粗略草图，大模型再'编辑'和'精炼'复杂部分。

实验结果

本文在 COCO 数据集上使用 5k 个文本提示生成 768x768 分辨率图像进行了全面评估。核心实验结果表明，HybridStitch 在所有质量指标上都优于 T-Stitch 和 SRDiffusion 基线方法。在 FID 指标上，HybridStitch-30% 达到 30.43，相比 T-Stitch 的 31.87 降低了 4.5%，相比 SRDiffusion 的 31.67 降低了 3.9%。在 LPIPS 指标上，HybridStitch-10% 达到 0.35，远低于 T-Stitch 的 0.72 和 SRDiffusion 的 0.69，表明其生成内容与大模型原始输出最为接近。在效率方面，HybridStitch-30% 实现了 1.83x 加速比（延迟 11.31 秒），相比 SRDiffusion 的 1.55x（延迟 13.37 秒）降低了 18.06% 的延迟。消融实验显示，动态 mask 更新（对比 StaticMask）对质量有显著贡献，而 FullLarge 变体虽然 LPIPS 更好但引入了大模型的偏差导致 FID 恶化。敏感性实验表明，HybridStitch 在 H100 和 A100 GPU 上同样有效，H100 上实现 1.4-1.6x 加速，A100 上实现 1.5-1.7x 加速。多个 mask 配置的实验进一步验证了方法的稳定性，切换步骤分析揭示了 HybridStitch 的优势主要来自第一切换点显著早于 SRDiffusion。

查看结构化数据

任务	指标	本文	基线	提升
COCO 文本到图像生成（768x768）	FID ↓	HybridStitch-30%: 30.43	T-Stitch: 31.87, SRDiffusion: 31.67	比 T-Stitch 降低 4.5%，比 SRDiffusion 降低 3.9%
COCO 文本到图像生成（768x768）	CLIP Score ↑	HybridStitch-30%: 31.84	T-Stitch: 32.01, SRDiffusion: 31.89	与基线基本持平，保持语义对齐
COCO 文本到图像生成（768x768）	LPIPS ↓	HybridStitch-10%: 0.35	T-Stitch: 0.72, SRDiffusion: 0.69	比 T-Stitch 降低 51.4%，比 SRDiffusion 降低 49.3%
COCO 文本到图像生成（768x768）	推理延迟 (s) ↓	HybridStitch-30%: 11.31s	大模型: 20.72s, T-Stitch: 14.74s, SRDiffusion: 13.37s	比大模型加速 1.83x，比 SRDiffusion 降低 18.06% 延迟

局限与改进

尽管 HybridStitch 取得了显著成果，但仍存在若干局限性。首先，该方法假设在推理时两个模型（大模型和小模型）都已加载到 GPU 显存中，这在资源受限的环境下可能不可行，实验是在拥有 48GB VRAM 的 RTX6000 Ada GPU 上进行的。其次，mask 的选择基于 Top-K 差异值，当图像整体都较为复杂时，10%-40% 的 mask 比例可能不足以覆盖所有困难区域，导致部分复杂区域未能获得大模型的精炼。第三，KV Cache 填充机制依赖于相邻步骤 KV 的高相似性假设，但在去噪的早期步骤中（变化剧烈时），这一假设可能不完全成立，可能影响注意力计算的准确性。第四，论文仅在 SD3 系列模型上进行了评估，未验证该方法在其他架构（如 DiT、UNet）或其他任务（如视频生成、图像编辑）上的泛化能力。此外，消融实验中的 StaticMask 变体在某些配置下（如 30% mask）反而获得了更高的加速比（2.08x），暗示动态 mask 更新的开销不容忽视。最后，论文未提供与其他注意力稀疏化、缓存等加速技术的组合实验，未能展示方法的兼容性。

独立分析的弱点

独立分析来看，HybridStitch 存在以下几个值得关注的弱点。第一，双模型显存占用问题：方法要求同时加载大模型和小模型到 GPU，这对于消费级显卡或边缘设备来说是不切实际的。改进方向可以探索模型权重共享、分层加载或 CPU-GPU 异构计算策略。第二，mask 选择策略的局限性：当前基于全局 Top-K 差异值的 mask 选择忽略了空间连续性，可能导致 mask 区域呈现碎片化分布，增加内存访问的不规则性（论文也提到 10% mask 的延迟不严格是全模型的 10%，因为 mask indexing 的开销）。改进方向可以引入基于空间聚类或超像素的 mask 选择策略，提高内存访问的局部性。第三，阈值调优的复杂性：HybridStitch 需要为每个 mask 配置调优两个阈值，且论文中不同 mask 大小对应不同的阈值对，这增加了实际部署的工程负担。改进方向可以设计自适应阈值选择机制。第四，对小模型质量的依赖：第三阶段完全依赖小模型，如果小模型质量较差，最终输出可能在细节上出现明显瑕疵。改进方向可以引入质量感知的自适应切换，在小模型输出质量下降时重新启用大模型。第五，仅评估了文生图任务，未验证在图像编辑、超分辨率等其他扩散模型应用中的有效性。

未来方向

论文作者指出未来可以从以下方向展开研究。首先，将 HybridStitch 扩展到视频生成领域，利用视频帧间的时间冗余性进一步优化 mask 策略——视频中相邻帧的困难区域往往具有时序一致性，可以复用 mask 减少计算。其次，探索与现有注意力稀疏化、特征缓存等加速技术的正交组合，理论上这些技术可以叠加使用以实现更高加速比。第三，研究自适应的 mask 大小选择，根据每张图像的复杂度动态调整 mask 比例而非使用固定配置。基于已有成果，还可以延伸到以下方向：将像素级切换思想应用到模型蒸馏中，指导小模型更好地学习大模型在困难区域的行为；探索跨模态的区域感知加速（如在文本条件生成中对不同语义区域使用不同模型）；将编辑范式推广到图像修复、风格迁移等任务中；研究动态模型加载策略，根据当前阶段的需要按需加载模型权重以降低显存占用。

复现评估

从复现角度来看，本文的复现条件相对明确但有一定门槛。论文使用了公开可用的模型：Stable Diffusion 3.5 Large 作为大模型、Stable Diffusion 3 Medium 作为小模型，这两个模型均可从 HuggingFace 获取。数据集使用 COCO，同样是公开数据集，作者随机采样了 5k 个文本提示。实验在 RTX6000 Ada（48GB VRAM）上进行，使用 50 个去噪步骤，768x768 分辨率。然而，论文未提及是否开源代码，这增加了复现的难度。关键实现细节如 KV Cache 的具体拼接方式、mask 更新的精确实现、差异计算的数值稳定性处理等，虽然在算法伪代码中有描述，但实际实现中可能存在工程细节差异。此外，双模型同时加载需要 48GB VRAM，这对大多数研究者来说是较高的硬件门槛。在 A100（40GB）和 H100（80GB）上的实验也表明结果可在不同 GPU 上复现，但需要相应的硬件资源。整体而言，如果有充足的算力资源和一定的扩散模型工程经验，复现本文结果是可行的。

HybridStitch：像素与时间步级别的模型缝合加速扩散模型 HybridStitch: Pixel and Timestep Level Model Stitching for Diffusion Acceleration