订阅
纠错
加入自媒体

扩散模型迎来“终极简化”!何恺明团队新作:像素级一步生成,速度质量双巅峰

2026-02-02 15:51
AI生成未来
关注

作者:Yiyang Lu等

解读:AI生成未来

亮点直击

Pixel MeanFlow (pMF) ,这是一种针对一步生成(one-step generation)的创新图像生成模型。pMF 的核心突破在于成功地在无隐空间(latent-free)的像素级建模中实现了高效的一步生成。

pMF不仅摆脱了对预训练潜在编码器(如 VQ-GAN 或 VAE)的依赖,直接在原始像素空间操作,而且在生成质量上达到了与最先进的多步隐空间扩散模型相媲美的水平。

解决的问题

现代生成模型通常在两个核心维度上进行权衡:

采样效率:多步采样虽然质量高但推理慢。

空间选择:隐空间(Latent Space)通过压缩降低了维度,但引入了复杂的编码器/解码器,且丢失了像素级的直接控制;像素空间(Pixel Space)虽然直观("所见即所得"),但高维数据建模难度极大。

将“一步生成”与“像素空间建模”结合是一个极具挑战性的任务,因为单一神经网络需要同时承担极其复杂的轨迹建模(trajectory modeling)和图像压缩/抽象(manifold learning)任务。现有的方法难以兼顾这两者。

提出的方案

pMF 的核心思想是将网络的预测目标损失函数的计算空间解耦:

预测目标 (Prediction Target) :网络直接预测去噪后的“干净”图像 (即 -prediction)。基于流形假设,干净图像位于低维流形上,更易于神经网络拟合。

损失空间 (Loss Space) :损失函数定义在速度场(velocity space)中,遵循 MeanFlow 的公式,通过最小化瞬时速度误差来学习平均速度场 。

转换机制:引入了一个简单的转换公式,在图像流形  和平均速度场  之间建立联系:。这一转换使得模型能够利用像素空间的流形结构,同时在速度空间进行有效的轨迹匹配。

应用的技术

Pixel-space Prediction:直接在像素空间参数化去噪图像 ,利用低维流形假设降低学习难度,避免直接预测高频噪声或速度场带来的困难。

MeanFlow Formulation:利用 Improved MeanFlow (iMF) 框架,通过瞬时速度  的损失来学习平均速度场 。

Flow Matching:基于流匹配理论,建立从噪声分布到数据分布的概率流。

Perceptual Loss:由于模型直接输出像素,天然适合引入感知损失(LPIPS 和 ConvNeXt 特征),进一步提升生成图像的视觉质量,弥补了像素级 MSE 损失的不足。

达到的效果

pMF 在 ImageNet 数据集上展现了强大的性能,证明了一步无潜在生成的可行性:

ImageNet 256×256: FID 分数达到 2.22,超越了许多多步隐空间模型。

ImageNet 512×512: FID 分数达到 2.48

这表明一步像素级生成模型已经具备了极强的竞争力,且不需要额外的解码器开销(解码器本身在隐空间模型中占据显著计算量)。

背景

本工作的 pMF 建立在 Flow Matching、MeanFlow 以及 JiT的基础之上。

Flow Matching. Flow Matching (FM) 学习一个速度场 ,将先验分布  映射到数据分布 。本文考虑标准的线性插值调度:

其中数据 ,噪声 (例如高斯分布),时间 。在  时,有:。该插值产生一个条件速度 :

FM 通过最小化 -空间中的损失函数(即“-loss”)来优化由  参数化的网络 :

已有研究表明 (Lipman et al., 2023), 的潜在目标是边缘速度 。

在推理阶段,通过求解常微分方程 (ODE): 从  到  生成样本,其中 。这可以通过 Euler 或基于 Heun 的数值求解器来实现。

Flow Matching with x-prediction. 等式 (2) 中的量  是一个带噪声的图像。为了便于使用在像素上操作的 Transformer,JiT 选择通过神经网络参数化数据 ,并通过以下方式将其转换为速度 :

其中  是 Vision Transformer (ViT) 的直接输出。这种公式被称为 -prediction,而在训练中使用等式 (2) 中的 -loss。表 1 列出了这种关系。

Mean Flows. MeanFlow (MF) 框架学习一个平均速度场  用于少步/一步生成。将 FM 的  视为瞬时速度,MF 定义平均速度  为:

其中  和  是两个时间步:。该定义引出了 MeanFlow 恒等式:

该恒等式提供了一种通过网络  定义预测函数的方法:

这里,大写  对应于等式 (6) 的左侧,而在右侧,JVP 表示用于计算  的 Jacobian-vector product,“sg”表示停止梯度(stop-gradient)。本文遵循 iMF的 JVP 计算和实现,这不是本文的重点。根据等式 (7) 的定义,iMF 像等式 (3) 一样最小化 -loss,即 。这种公式可以被视为带有 -loss 的 -prediction(参见表 1)。

Pixel MeanFlow

为了实现一步、无潜在生成,本文提出了 Pixel MeanFlow (pMF)。pMF 的核心设计是在 、 和  的不同场之间建立联系。本文希望网络像 JiT一样直接输出 ,而一步建模则像 MeanFlow一样在  和  空间上进行。

去噪图像场

如前所述,iMF 和 JiT均可视为在最小化瞬时速度  的损失(-loss),区别在于 iMF 执行的是平均速度预测(-prediction),而 JiT 执行的是原始数据预测(-prediction)。基于这一观察,本工作在平均速度  与一种广义形式的  之间建立了一种映射联系。

考虑等式 (5) 中定义的平均速度场 :该场代表了一个由数据分布 、先验分布  以及时间调度决定的底层真实量,它与具体的网络参数  无关。由此,本文推导出一个诱导场(induced field),定义如下:

如下文详述,该场  扮演了类似于“去噪图像”的角色。需要注意的是,本工作定义的  与以往文献中提及的  不同,它是一个受两个时间戳  索引的二元变量:对于给定的观测值 ,本文的  是一个随  变化的二维场,而非仅受  索引的一维轨迹。

广义流形假设

图 1 通过模拟从预训练 FM 模型获得的一条 ODE 轨迹,可视化了  场和  场。如图所示, 由含噪图像组成,因为作为速度场, 包含噪声和数据成分。相比之下, 场具有去噪图像的外观:它们是接近干净的图像,或者是因过度去噪而显得模糊的图像。接下来,本文讨论流形假设如何推广到这个量 。

注意 MF 中的时间步  满足:。本文首先展示在  和  处的边界情况可以近似满足流形假设;然后讨论  的情况。

边界情况 I: . 当  时,平均速度  退化为瞬时速度 ,即 。在这种情况下,等式 (8) 变为:

这本质上是 JiT 中使用的 -prediction 目标。直观地说,这个  是 JiT 要预测的去噪图像。如果噪声水平很高,这个去噪图像可能是模糊的。正如经典图像去噪研究中广泛观察到的那样,可以假设这些去噪图像近似位于低维(或较低维)流形上。

边界情况 II: . 等式 (5) 中  的定义给出:。将其代入等式 (8) 得到:

即,它是 ODE 轨迹的终点。对于真实的 ODE 轨迹,有 ,即它应遵循图像分布。因此,本文可以假设  近似位于图像流形上。

一般情况: . 与边界情况不同,量  不保证对应于来自数据流形的(可能模糊的)图像样本。然而,根据经验,本文的模拟(图 1 右)表明  看起来像去噪图像。这与速度空间量(图 1 中的 )形成鲜明对比,后者噪声明显更多。这种比较表明,通过神经网络对  进行建模可能比对更嘈杂的  进行建模更容易。实验表明,对于像素空间模型,-prediction 表现有效,而 -prediction 则严重退化。

算法

等式 (8) 中的诱导场  提供了 MeanFlow 网络的一种重参数化。具体来说,本文让网络  直接输出 ,并通过等式 (8) 计算相应的速度场 :

这里, 是网络的直接输出,遵循 JiT。这个公式是等式 (4) 的自然扩展。

本文将 (11) 中的  纳入 iMF 公式,即使用带有 -loss 的等式 (7)。具体来说,本文的优化目标是:

其中 。

从概念上讲,这是带有 -prediction 的 -loss,其中  通过  的关系转换为  空间以回归 。表 1 总结了这种关系。相应的伪代码在 Alg. 1 中。

带有感知损失的像素平均流

网络  直接将含噪输入  映射到去噪图像。这使得在训练时能够实现“所见即所得”的行为。因此,除了  损失外,本文还可以进一步结合感知损失。基于潜在的方法在 tokenizer 重建训练期间受益于感知损失,而基于像素的方法尚未能利用这一优势。

形式上,由于  是像素中的去噪图像,本文直接对其应用感知损失(例如 LPIPS)。本文的总体训练目标是 ,其中  表示  和真实干净图像  之间的感知损失, 是权重超参数。在实践中,仅当添加的噪声低于某个阈值(即 )时才应用感知损失,以使去噪图像不会太模糊。本文研究了基于 VGG 分类器的标准 LPIPS 损失和基于 ConvNeXt-V2 的变体。

与前人工作的关系

本文的 pMF 与几种先前的少步/一步方法密切相关,讨论如下。

Consistency Models (CM): 学习从含噪样本  直接到生成图像的映射。在本文的符号中,这对应于固定终点 。此外,CM 通常采用预处理器 (Pre-conditioner),其形式为 。除非  为零,否则网络不执行纯粹的 -prediction。

Consistency Trajectory Models (CTM): 制定了一个双时间量。与基于导数公式的 MeanFlow 不同,CTM 依赖于在训练期间对 ODE 进行积分。

Flow Map Matching (FMM): 也是基于双时间量。在本文符号中,Flow Map 扮演位移的角色,即 。该量通常不位于低维流形上(例如  是含噪图像)。

实验

本文通过 2D 玩具实验(图 2)证明,当底层数据位于低维流形上时,在 MeanFlow 中使用 -prediction 是更可取的。实验设置遵循。

形式上,本文考虑定义在 2D 空间上的底层数据分布(此处为 Swiss roll)。数据使用  列正交矩阵投影到  维观测空间。本文在  维观测空间上训练 MeanFlow 模型,其中 。本文比较了-prediction 与本文的 -prediction。

图 2 显示,-prediction 表现相当不错,而当  增加时,-prediction 迅速退化。本文观察到这种性能差距反映在训练损失的差异上:-prediction 产生的训练损失低于 -prediction 对应物。这表明对于容量有限的网络,预测  更容易。

ImageNet 实验

本文默认在分辨率 256×256 的 ImageNet 上进行消融实验。报告基于 50,000 个生成样本的 FID。所有模型均通过单次函数评估 (1-NFE) 生成原始像素图像。

网络的预测目标

本文的方法基于流形假设,即  位于低维流形上且更容易预测。本文在表 2 中验证了这一假设。

64×64 分辨率: patch 维度为 48 ()。这个维度远低于网络容量。结果显示 pMF 在 -prediction 和 -prediction 下都表现良好。

256×256 分辨率: patch 维度为 768 ()。这导致高维观测空间,神经网络更难建模。在这种情况下,只有 -prediction 表现良好(FID 9.56),这表明  位于较低维流形上,因此更适合学习。相比之下,-prediction 发生灾难性失败(FID 164.89):作为一个含噪量, 在高维空间中具有全支撑,更难建模。

消融研究

优化器 本文发现优化器的选择在 pMF 中起着重要作用。在图 3a 中,本文比较了标准 Adam 优化器与最近提出的 Muon。Muon 表现出更快的收敛速度和大幅提升的 FID(从 11.86 提升至 8.71)。在一步生成设置中,更快的收敛优势被进一步放大,因为更好的网络能提供更准确的停止梯度目标。

感知损失 在图 3b 中,本文进一步结合感知损失。使用标准 VGG-based LPIPS 将 FID 从 9.56 提升至 5.62;结合 ConvNeXt-V2 变体进一步将 FID 提升至 3.53。总体而言,结合感知损失带来了约 6 个 FID 点的提升。

替代方案:预处理器 本文比较了三种预处理器变体:(i) 线性;(ii) EDM 风格;(iii) sCM 风格。表 3a 显示,尽管 EDM 和 sCM 风格优于朴素线性变体,但在本文考虑的极高维输入机制中,简单的 -prediction 更可取且性能更好。这是因为除非 ,否则网络预测会偏离  空间,可能位于更高维流形上。

替代方案:时间采样器 本文研究了限制时间采样的替代设计:仅 (即 Flow Matching),仅 (类似 CM),或两者的组合。表 3b 显示这些受限的时间采样器都不足以解决本文考虑的挑战性场景。这表明 MeanFlow 方法利用  点之间的关系来学习场,限制时间采样可能会破坏这种公式。

高分辨率生成 在表4中,本文研究了分辨率 256、512 和 1024 下的 pMF。通过增加 patch size(例如 )来保持序列长度不变 (),导致极大的 patch 维度(例如 12288)。结果显示 pMF 可以有效处理这种极具挑战性的情况。即使观测空间是高维的,模型始终预测 ,其潜在维度不会成比例增长。

可扩展性 表 5 显示,增加模型大小和训练周期均能提升结果。

系统级比较

ImageNet 256×256. 表 6 显示本文的方法达到了 2.22 FID。据本文所知,该类别中(一步、无潜在扩散/流)唯一的方法是最近提出的 EPG,其 FID 为 8.82。与领先的 GAN 相比,pMF 实现了相当的 FID,但计算量大幅降低(例如 StyleGAN-XL 的计算量是 pMF-H/16 的 5.8 倍)。

ImageNet 512×512. 表 7 显示 pMF 在 512×512 下达到 2.48 FID。值得注意的是,其计算成本(参数量和 Gflops)与 256×256 对应物相当。唯一的开销来自 patch embedding 和预测层。

结论

本质上,图像生成模型是从噪声到图像像素的映射。由于生成建模的固有挑战,该问题通常被分解为更易处理的子问题,涉及多个步骤和阶段。虽然有效,但这些设计偏离了深度学习的端到端精神。

本文关于 pMF 的研究表明,神经网络是具有高度表现力的映射,当设计得当时,能够学习复杂的端到端映射,例如直接从噪声到像素。除了其实际潜力外,本文希望本工作将鼓励未来对直接、端到端生成建模的探索。

参考文献

[1] One-step Latent-free Image Generation with Pixel Mean Flows

       原文标题 : 扩散模型迎来“终极简化”!何恺明团队新作:像素级一步生成,速度质量双巅峰

声明: 本文由入驻维科号的作者撰写,观点仅代表作者本人,不代表OFweek立场。如有侵权或其他问题,请联系举报。

发表评论

0条评论,0人参与

请输入评论内容...

请输入评论/评论长度6~500个字

您提交的评论过于频繁,请输入验证码继续

暂无评论

暂无评论

    人工智能 猎头职位 更多
    扫码关注公众号
    OFweek人工智能网
    获取更多精彩内容
    文章纠错
    x
    *文字标题:
    *纠错内容:
    联系邮箱:
    *验 证 码:

    粤公网安备 44030502002758号