AiPaper
论文状态:已完成

Mean Flows for One-step Generative Modeling

发表:2025/05/20
原文链接PDF 下载
价格:0.10
价格:0.10
已有 9 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出MeanFlow模型,引入平均速度概念替代瞬时速度,推导恒等式指导神经网络训练,简化了生成过程。该模型无需预训练或蒸馏,单次函数评估在ImageNet 256x256上实现3.43 FID,显著优于现有一步生成模型,缩小一步与多步模型性能差距。

摘要

We propose a principled and effective framework for one-step generative modeling. We introduce the notion of average velocity to characterize flow fields, in contrast to instantaneous velocity modeled by Flow Matching methods. A well-defined identity between average and instantaneous velocities is derived and used to guide neural network training. Our method, termed the MeanFlow model, is self-contained and requires no pre-training, distillation, or curriculum learning. MeanFlow demonstrates strong empirical performance: it achieves an FID of 3.43 with a single function evaluation (1-NFE) on ImageNet 256x256 trained from scratch, significantly outperforming previous state-of-the-art one-step diffusion/flow models. Our study substantially narrows the gap between one-step diffusion/flow models and their multi-step predecessors, and we hope it will motivate future research to revisit the foundations of these powerful models.

思维导图

论文精读

中文精读

论文基本信息 (Bibliographic Information)

  • 标题 (Title): Mean Flows for One-step Generative Modeling (用于一步生成建模的平均流)
  • 作者 (Authors): Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter, Kaiming He
  • 隶属机构 (Affiliations):
    • Zhengyang Geng, J. Zico Kolter: CMU (卡内基梅隆大学)
    • Mingyang Deng, Xingjian Bai, Kaiming He: MIT (麻省理工学院)
  • 发表期刊/会议 (Journal/Conference): 预印本 (arXiv preprint),但根据上下文可能已被接受或将在 ICLR 2025 等顶级会议上发表(例如,引用中的 Kevin Frans et al. 2025 ICLR)。在生成建模领域,ICLR 是一个具有极高声誉和影响力的会议。
  • 发表年份 (Publication Year): 2025
  • 摘要 (Abstract): 我们提出了一个用于一步生成建模的原则性且有效的框架。该框架引入了平均速度 (average velocity) 的概念来表征流场,这与流匹配 (Flow Matching) 方法中建模的瞬时速度 (instantaneous velocity) 形成对比。论文推导了一个定义明确的平均速度与瞬时速度之间的恒等式 (identity),并利用该恒等式指导神经网络训练。我们提出的方法称为 MeanFlow 模型 (MeanFlow model),它是自洽的 (self-contained),不需要预训练 (pre-training)、蒸馏 (distillation) 或课程学习 (curriculum learning)。MeanFlow 展示了强大的经验性能:在 ImageNet 256x256 上从零开始训练,使用单次函数评估 (1-NFE) 实现了 3.43 的 FID (Fréchet Inception Distance),显著优于此前最先进的一步扩散/流模型。我们的研究大大缩小了一步扩散/流模型与其多步前身之间的差距,我们希望它能激励未来的研究重新审视这些强大模型的基础。
  • 原文链接 (Source Link): https://arxiv.org/abs/2505.13447
  • PDF 链接 (PDF Link): https://arxiv.org/pdf/2505.13447v1.pdf
  • 发布状态 (Publication Status): 预印本 (arXiv preprint),于 2025-05-19T17:59:42.000Z 发布。

整体概括 (Executive Summary)

研究背景与动机 (Background & Motivation - Why)

生成模型 (generative modeling) 的目标是将一个先验分布 (prior distribution) 转换为数据分布 (data distribution)。流匹配 (Flow Matching) 和扩散模型 (Diffusion Models) 是当前最先进的生成范式,它们通过学习一个速度场 (velocity field) 来定义从噪声到数据的连续轨迹。然而,这些模型在采样 (sampling) 阶段通常需要多次迭代(即多次函数评估,NFE),这导致生成过程缓慢且计算成本高昂。

为了解决这一问题,研究者们提出了少步 (few-step) 乃至一步生成 (one-step generation) 的方法。早期的工作如一致性模型 (Consistency Models) 尝试通过在网络输出上施加一致性约束来实现一步生成。但这些方法往往存在训练不稳定、需要精心设计的离散化课程 (discretization curriculum) 以及依赖预训练 (pre-training)蒸馏 (distillation) 等问题,限制了它们的普适性和从零开始训练的能力。

本论文的动机在于寻找一个更具原则性 (principled)自洽 (self-contained)高效 (effective) 的一步生成框架,能够避免上述限制,并显著缩小一步生成模型与多步生成模型之间的性能差距。

核心贡献/主要发现 (Main Contribution/Findings - What)

论文的核心贡献在于提出了 MeanFlow 模型 (MeanFlow model),一个用于一步生成建模的创新框架:

  • 引入平均速度概念 (Introduction of Average Velocity): MeanFlow 首次引入了平均速度 (average velocity) 的概念来表征流场,这与传统流匹配方法中建模的瞬时速度 (instantaneous velocity) 形成鲜明对比。平均速度被定义为位移 (displacement) 与时间间隔之比,位移是瞬时速度在时间上的积分。
  • 推导 MeanFlow 恒等式 (Derivation of MeanFlow Identity): 论文从平均速度的定义出发,推导出了一个定义明确的、内在的平均速度与瞬时速度之间的关系——MeanFlow 恒等式 (MeanFlow Identity)。这个恒等式为神经网络训练提供了一个坚实的、原则性的基础,无需引入额外的启发式一致性约束。
  • 自洽的训练框架 (Self-Contained Training Framework): MeanFlow 模型是自洽的 (self-contained),它直接从零开始训练,无需任何预训练 (pre-training)、蒸馏 (distillation) 或课程学习 (curriculum learning)。这极大地简化了训练流程,并提高了模型的通用性。
  • 支持分类器无关引导 (Classifier-Free Guidance, CFG) (Support for CFG): MeanFlow 框架能够自然地集成 (naturally incorporate) CFG,并且在采样时不增加额外的计算成本(即保持 1-NFE)。这是通过将 CFG 视为底层真实场 (ground-truth fields) 的属性来实现的。
  • 显著超越现有技术水平 (Outperforming SOTA): 在 ImageNet 256x256 数据集上,MeanFlow 模型使用 1-NFE 实现了 3.43 的 FID,显著优于此前所有最先进的一步扩散/流模型,相对性能提升达 50% 至 70%。
  • 缩小与多步模型的差距 (Narrowing Gap with Multi-step Models): 在 2-NFE 条件下,MeanFlow 实现了 2.20 的 FID,与顶级多步扩散/流模型(如 DiT 和 SiT 在 250x2 NFE 下的性能)不相上下。这表明少步生成模型完全有潜力与多步模型竞争。

预备知识与相关工作 (Prerequisite Knowledge & Related Work)

本节旨在为理解 MeanFlow 模型的创新之处和技术细节奠定基础。

基础概念 (Foundational Concepts)

  • 生成模型 (Generative Modeling):
    • 概念定义: 生成模型是一类机器学习模型,其目标是学习训练数据的内在分布,并能够生成与训练数据具有相似特征的新样本。例如,给定一组猫的图片,一个生成模型可以学习猫的特征,然后生成从未见过的猫的图片。
    • 在本文语境中的目标: 将一个简单的先验分布 (prior distribution)(通常是高斯噪声)转换成复杂的数据分布 (data distribution)(例如图像数据)。
  • 扩散模型 (Diffusion Models):
    • 概念定义: 扩散模型是一类基于马尔可夫链 (Markov chain) 的生成模型,它们通过两个过程工作:前向扩散过程 (forward diffusion process)反向去噪过程 (reverse denoising process)。在前向过程中,模型逐渐向数据中添加噪声,直到数据完全变成噪声。在反向过程中,模型学习从噪声中逐步去除噪声,从而恢复出原始数据。这个去噪过程通常通过求解随机微分方程 (Stochastic Differential Equation, SDE)常微分方程 (Ordinary Differential Equation, ODE) 来实现。
    • 关键特点: 扩散模型通常在采样时需要多次迭代 (iterative sampling),每次迭代都需要进行一次或多次神经网络的函数评估 (Function Evaluation, NFE)
  • 流匹配 (Flow Matching, FM):
    • 概念定义: 流匹配是一种与扩散模型密切相关的生成建模框架。它专注于学习一个速度场 (velocity field),该速度场定义了从先验分布到数据分布的流路径 (flow paths)。流匹配可以被看作是一种连续时间 (continuous-time) 的归一化流 (Normalizing Flows)
    • 瞬时速度 (Instantaneous Velocity): 在流匹配中,通常建模的是物体在某一时刻的瞬间运动方向和速率,即瞬时速度 v(zt,t)v(z_t, t)
    • 条件流 (Conditional Flow) 与边际流 (Marginal Flow):
      • 条件流: 如图 3 左图所示,给定一个中间状态 ztz_t,它可能由不同的初始数据 xx 和噪声 ϵ\epsilon 对产生。因此,条件速度 vt(ztx)v_t(z_t|x) 是在给定 xx 的条件下,从 xxztz_t 的瞬时速度。
      • 边际流: 如图 3 右图所示,由于 ztz_t 可以有多种来源,流匹配通过对所有可能的条件速度进行期望 (expectation) 或边际化 (marginalizing),得到边际速度场 v(zt,t)v(z_t, t)。这个边际速度场是流匹配中神经网络学习的真实场 (ground-truth field)
    • 训练目标: 流匹配训练一个神经网络 vθv_\theta 来拟合边际速度场。虽然计算边际速度场的损失很困难,但研究表明,通过最小化条件流匹配损失 (conditional Flow Matching loss, LCFM\mathcal{L}_{\mathrm{CFM}}),即拟合条件速度,可以等价地最小化边际速度场的损失。
    • 采样: 采样过程通过数值求解速度场定义的 ODE (Ordinary Differential Equation) 来实现,例如使用欧拉法 (Euler method),这仍然需要多次函数评估。
  • 函数评估次数 (Number of Function Evaluations, NFE):
    • 概念定义: NFE 是指在生成一个样本时,神经网络进行前向传播计算的次数。在扩散模型和流匹配中,由于采样是迭代的,每一步迭代都会增加 NFE。
    • 重要性: NFE 越少,生成速度越快,计算成本越低。一步生成 (one-step generation) 意味着 NFE 为 1,即仅需一次神经网络前向传播即可生成最终样本。
  • 一致性模型 (Consistency Models, CM):
    • 概念定义: 一致性模型是一类旨在实现一步生成的新型生成模型。它们的核心思想是引入一致性约束 (consistency constraint),强制在同一流路径上不同时间步的网络输出都收敛到相同的最终数据点。
    • 局限性: 尽管取得了初步成功,但一致性约束通常是作为神经网络行为的属性施加的,而非从底层真实场的数学定义中推导而来。这可能导致训练不稳定,并且通常需要离散化课程 (discretization curriculum) 来逐步训练模型以适应不同的时间范围。
  • 分类器无关引导 (Classifier-Free Guidance, CFG):
    • 概念定义: CFG 是一种在条件生成模型中提高生成质量和控制性的技术。它通过结合类别条件 (class-conditional) 预测和类别无关 (class-unconditional) 预测来“放大”条件的影响。
    • 传统实现: 通常在采样时,对条件和无条件模型进行两次函数评估,然后进行加权组合,导致 NFE 翻倍(例如,1-NFE 变成 2-NFE)。
  • FID (Fréchet Inception Distance):
    • 概念定义: FID 是衡量生成图像质量和多样性的一个重要指标。它计算生成图像分布和真实图像分布在 Inception-v3 网络特征空间中的 Fréchet 距离。FID 值越低,表示生成图像的质量越高,与真实图像的分布越接近。
  • VAE 编码器/解码器 (VAE tokenizer):
    • 概念定义: 变分自编码器 (Variational AutoEncoder, VAE) 可以学习将高维数据(如图像)压缩到低维的潜在空间 (latent space) 中(编码器),并能从潜在空间重构回高维数据(解码器)。在大型图像生成任务中,通常会在 VAE 学习到的潜在空间中进行扩散或流模型的训练,以降低计算复杂度。这里的 tokenizer 指的是 VAE 的编码器部分,将图像转换为潜在表示。
  • 雅可比向量积 (Jacobian-Vector Product, JVP):
    • 概念定义: JVP 是计算函数梯度的一种高效方法。给定一个函数 f:RnRmf: \mathbb{R}^n \to \mathbb{R}^m 及其输入 xRnx \in \mathbb{R}^n 和一个切向量 vRnv \in \mathbb{R}^n,JVP 计算雅可比矩阵 Jf(x)J_f(x)vv 的乘积 Jf(x)vJ_f(x)v。它在反向传播 (backpropagation) 中被广泛使用,尤其是在需要计算高阶导数或优化某些特定表达式时非常高效。
    • 在本文中的作用: 用于计算平均速度 uuztz_ttt 的偏导数与瞬时速度 vv 的乘积,这是 MeanFlow 恒等式中关键的项。

前人工作与技术演进 (Previous Works & Technological Evolution)

  • 扩散模型与流匹配的迭代采样 (Iterative Sampling in Diffusion/Flow Models):
    • 早期和经典的扩散模型(如 DDPM [19]、ADM [8])和流匹配模型(如 Flow Matching [28])都需要多步迭代采样,NFE 通常在 50-1000 之间。这是因为它们学习的是瞬时变化,需要逐步积分才能得到最终结果。
    • 技术挑战: 高 NFE 导致生成速度慢,限制了实际应用。
  • 蒸馏方法 (Distillation-based Methods):
    • 为了减少采样步数,一种主流方法是蒸馏 (distillation)。例如,[39, 14, 41] 和基于分数蒸馏 (score distillation) 的方法 [32, 50, 53] 尝试将一个预训练好的多步扩散模型“压缩”成一个少步模型。
    • 局限性: 这些方法通常需要一个高质量的预训练 (pre-trained) 大模型作为“教师”,且蒸馏过程本身可能复杂。
  • 自洽训练的少步模型 (Standalone Few-step Models):
    • 一致性模型 (Consistency Models, CM) [46]: 作为这一方向的先驱,CM 不依赖于蒸馏,而是通过施加一致性约束 (consistency constraints) 来直接训练生成模型以实现少步(尤其是 1-NFE)采样。后续工作如 iCT [43]、ECT [15]、sCT [31] 等对此进行了改进。
    • 核心思路: 路径上任何一点的输出都应该“一致”地指向同一个最终数据点。
    • 局限性: 这种约束是施加在神经网络行为上的,而非从底层物理/数学原理出发。训练可能不稳定,通常需要离散化课程 (discretization curriculum) 来逐步引导模型适应不同时间范围。
  • 基于两时间变量 (two-time variables) 的方法:
    • Flow Map [3]: 定义为两时间步之间流的积分,与本文的平均速度概念有一定相似性,但 Flow Map 对应的是位移 (displacement),而本文聚焦于速度。
    • Shortcut Models [13]: 在流匹配基础上引入了一个额外的自洽损失函数 (self-consistency loss function),捕获不同离散时间间隔内流之间的关系。
    • Inductive Moment Matching (IMM) [52]: 建模了随机插值 (stochastic interpolants) 在不同时间步的自洽性 (self-consistency)
    • 共同点与差异: 这些方法都引入了对两个时间变量的考虑来解决少步生成问题。然而,它们通常通过额外增加启发式损失或约束来实现,而 MeanFlow 则是从平均速度的定义出发,推导出内在的数学恒等式

差异化分析 (Differentiation)

MeanFlow 与现有工作的核心区别和创新点在于:

  1. 原则性 (Principled Foundation): 与一致性模型等依赖启发式一致性约束的方法不同,MeanFlow 从平均速度的数学定义出发,严格推导出了 MeanFlow 恒等式 (MeanFlow Identity)。这个恒等式是瞬时速度和平均速度之间内在的、必然的数学关系,与神经网络的参数化无关。这提供了一个更坚实、更具原则性的训练目标。
  2. 自洽与无依赖 (Self-Contained & Independent): MeanFlow 完全从零开始训练,不依赖于任何预训练模型、蒸馏过程或复杂的课程学习策略。这使其成为一个真正独立的生成模型,简化了研究和应用的门槛。而许多少步方法(尤其是蒸馏方法)都高度依赖于高质量的多步教师模型。
  3. 对两时间变量的统一处理 (Unified Treatment of Two Time Variables): 虽然 Shortcut 和 IMM 也考虑了两个时间变量,但它们通常通过额外的自洽损失来建立关系。MeanFlow 的恒等式直接将这两个时间变量(以及它们之间的平均速度)与瞬时速度联系起来,形成了一个统一且自然的训练目标。
  4. 单次函数评估的 CFG (1-NFE CFG): MeanFlow 能够将分类器无关引导 (CFG) 直接融入到其目标场 (target field) 中,从而在采样时保持单次函数评估 (1-NFE) 的效率,避免了传统 CFG 导致 NFE 翻倍的问题。

4. 方法论 (Methodology - Core Technology & Implementation Details)

MeanFlow 模型的核心思想是引入平均速度 (average velocity) 的概念来表征流场,并从其定义推导出MeanFlow 恒等式 (MeanFlow Identity) 作为神经网络训练的指导。

4.1 方法原理 (Methodology Principles)

传统的流匹配 (Flow Matching) 方法主要建模瞬时速度 (instantaneous velocity) v(zt,t)v(z_t, t),即在某一时刻 tt 和状态 ztz_t 下的瞬时变化率。通过对瞬时速度进行时间积分,可以得到从一个时间点到另一个时间点的轨迹。然而,当需要一步到位地从初始噪声 ϵ\epsilon 采样到最终数据 xx 时,直接利用瞬时速度进行单步积分会导致较大的数值误差,因为轨迹通常是弯曲的。

MeanFlow 提出了一种新的视角:直接建模平均速度 (average velocity) uu。平均速度可以直观地理解为在一段时间间隔 [r, t] 内,物体总位移 (displacement) 除以时间间隔 (t-r)。如果能够准确地预测这段时间内的平均速度,那么只需一步就可以从 ztz_t 得到 zrz_r,从而实现高效的单步生成。

核心直觉:

  • 瞬时速度 vv 决定了路径的切线方向 (tangent direction)
  • 平均速度 uu 决定了在一段时间内从起点到终点的直线方向 (straight-line direction) 和平均速率。
  • 目标: 训练一个神经网络直接预测这个“直线方向”和“平均速率”,使得一步采样 zr=zt(tr)u(zt,r,t)z_r = z_t - (t-r)u(z_t, r, t) 尽可能准确。

4.2 方法步骤与流程 (Steps & Procedures)

4.2.1 平均流的定义 (Mean Flows Definition)

首先,我们形式化平均速度 (average velocity) 的概念。给定瞬时速度场 v(zτ,τ)v(z_\tau, \tau),连接两个时间步 ttrr 之间的平均速度 uu 定义为:

u(zt,r,t)1trrtv(zτ,τ)dτ(3) \boxed{u(z_t, r, t) \triangleq \frac{1}{t-r} \int_{r}^{t} v(z_{\tau}, \tau) d\tau} \quad (3)

  • 符号解释:
    • u(zt,r,t)u(z_t, r, t): 在时间 tt 的状态 ztz_t 下,从时间 rr 到时间 tt 的平均速度场。它是一个依赖于当前状态 ztz_t 和两个时间点 r, t 的向量场。
    • v(zτ,τ)v(z_\tau, \tau): 在时间 τ\tau 的状态 zτz_\tau 下的瞬时速度场。
    • rtv(zτ,τ)dτ\int_{r}^{t} v(z_{\tau}, \tau) d\tau: 从时间 rr 到时间 tt 沿流路径的瞬时速度的积分,代表了这段时间内的总位移。
    • (t-r): 时间间隔。

图示 (Figure 3):

  • 左侧子图展示了瞬时速度 vv 决定了路径的切线方向,而平均速度 uu 通常不与 vv 对齐,它指向的是位移 (t-r)u(z,r,t) 的方向。
  • 右侧三个子图展示了平均速度场 u(z, r, t) 是如何依赖于 rrtt 的,随着 tt 的变化,平均速度场的分布也会改变。

特性:

  • rtr \to t 时,平均速度 uu 的极限是瞬时速度 vv (即 limrtu=v\lim_{r \to t} u = v)。
  • 平均速度满足一种“一致性”:在 [r, t] 上迈一大步等价于在 [r, s][s, t] 上迈两小步之和,这直接来源于积分的加性。 (tr)u(zt,r,t)=(sr)u(zs,r,s)+(ts)u(zt,s,t) (t-r)u(z_t, r, t) = (s-r)u(z_s, r, s) + (t-s)u(z_t, s, t) 这种内在的一致性是 MeanFlow 模型稳定性的基础。

4.2.2 MeanFlow 恒等式的推导 (Derivation of the MeanFlow Identity)

直接使用 Eq. (3) 作为神经网络的训练目标是不可行的,因为它涉及到积分的计算。关键在于找到一个可训练的、等价的表达形式。 我们将 Eq. (3) 重写为:

(tr)u(zt,r,t)=rtv(zτ,τ)dτ(4) (t-r)u(z_t, r, t) = \int_{r}^{t} v(z_{\tau}, \tau) d\tau \quad (4)

接下来,我们对等式两边关于 tt 求导,将 rr 视为独立于 tt 的变量。

左侧求导 (LHS): 运用乘法法则 (product rule) 对 (tr)u(zt,r,t)(t-r)u(z_t, r, t) 求导: ddt[(tr)u(zt,r,t)]=1u(zt,r,t)+(tr)ddtu(zt,r,t) \frac{d}{dt} [(t-r)u(z_t, r, t)] = 1 \cdot u(z_t, r, t) + (t-r) \frac{d}{dt} u(z_t, r, t)

右侧求导 (RHS): 运用微积分基本定理 (Fundamental Theorem of Calculus) 对 rtv(zτ,τ)dτ\int_{r}^{t} v(z_{\tau}, \tau) d\tau 求导: ddtrtv(zτ,τ)dτ=v(zt,t) \frac{d}{dt} \int_{r}^{t} v(z_{\tau}, \tau) d\tau = v(z_t, t) (这里假设 rr 不依赖于 tt。如果 rr 依赖于 tt,则需要使用莱布尼茨积分法则的更一般形式,但本文的设定是 rr 独立于 tt。)

将左右两边的结果等同起来,得到: u(zt,r,t)+(tr)ddtu(zt,r,t)=v(zt,t)(5) u(z_t, r, t) + (t-r) \frac{d}{dt} u(z_t, r, t) = v(z_t, t) \quad (5)

重新排列项,我们得到了 MeanFlow 恒等式 (MeanFlow Identity)

u(zt,r,t)average vel.=v(zt,t)instant. vel.(tr)ddtu(zt,r,t)time derivative(6) \boxed{ \underbrace{u(z_t, r, t)}_{\mathrm{average\ vel.}} = \underbrace{v(z_t, t)}_{\mathrm{instant.\ vel.}} - (t-r) \underbrace{\frac{d}{dt} u(z_t, r, t)}_{\mathrm{time\ derivative}} } \quad (6)

这个恒等式揭示了平均速度 uu 与瞬时速度 vv 之间的内在联系,以及 uu 对时间 tt全导数 (total derivative)

全导数 ddtu(zt,r,t)\frac{d}{dt} u(z_t, r, t) 的分解: ddtu(zt,r,t)\frac{d}{dt} u(z_t, r, t) 是一个全导数,需要考虑 uuztz_ttt 的依赖。根据链式法则 (chain rule): ddtu(zt,r,t)=uztdztdt+urdrdt+utdtdt(7) \frac{d}{dt} u(z_t, r, t) = \frac{\partial u}{\partial z_t} \frac{dz_t}{dt} + \frac{\partial u}{\partial r} \frac{dr}{dt} + \frac{\partial u}{\partial t} \frac{dt}{dt} \quad (7) 由于 ztz_t 遵循 ODE dztdt=v(zt,t) \frac{dz_t}{dt} = v(z_t, t) (Eq. (2)),且 rr 被视为独立于 tt,所以 drdt=0\frac{dr}{dt} = 0,同时 dtdt=1\frac{dt}{dt} = 1。代入这些项,得到: ddtu(zt,r,t)=v(zt,t)zu+tu(8) \boxed{ \frac{d}{dt} u(z_t, r, t) = v(z_t, t) \partial_z u + \partial_t u } \quad (8) 其中 v(zt,t)zuv(z_t, t) \partial_z u 是一个雅可比向量积 (Jacobian-Vector Product, JVP),表示瞬时速度 vv 与平均速度场 uu 对状态 zz 的雅可比矩阵的乘积。这一项可以通过现代深度学习框架(如 PyTorch 的 jvp 或 JAX 的 jax.jvp)高效计算。

4.2.3 神经网络训练 (Neural Network Training)

我们的目标是训练一个神经网络 uθu_\theta 来近似平均速度场 uu。训练过程通过最小化一个损失函数来实现,该损失函数鼓励 uθu_\theta 满足 MeanFlow 恒等式 (Eq. (6))。

损失函数 (Loss Function): L(θ)=Euθ(zt,r,t)sg(utgt)22(9) \mathcal{L}(\boldsymbol{\theta}) = \mathbb{E} \big\| u_{\boldsymbol{\theta}}(z_t, r, t) - \mathbf{sg}(u_{\mathrm{tgt}}) \big\|_2^2 \quad (9) 其中,目标 utgtu_{\mathrm{tgt}} (target) 由 MeanFlow 恒等式 (Eq. (6)) 和全导数分解 (Eq. (8)) 给出,并用神经网络参数化的 uθu_\theta 替换 uu 的导数项: utgt=v(zt,t)(tr)(v(zt,t)zuθ+tuθ)(10) u_{\mathrm{tgt}} = v(z_t, t) - (t-r) \big( v(z_t, t) \partial_z u_{\boldsymbol{\theta}} + \partial_t u_{\boldsymbol{\theta}} \big) \quad (10)

  • 关键点:
    • sg()\mathbf{sg}(\cdot) (stop-gradient): 在目标 utgtu_{\mathrm{tgt}} 上应用了停止梯度 (stop-gradient) 操作。这意味着在计算损失并进行反向传播时,utgtu_{\mathrm{tgt}} 被视为常数,其梯度不会回传到 uθu_\theta 的导数项。这避免了双重反向传播 (double backpropagation),从而降低了计算复杂性和内存消耗,并提高了训练稳定性。这与一致性模型 [46] 等常用做法一致。
    • 瞬时速度 v(zt,t)v(z_t, t): 在 Eq. (10) 中,v(zt,t)v(z_t, t) 是流匹配中的边际瞬时速度 (marginal instantaneous velocity)。根据流匹配的惯例 [28],我们用条件瞬时速度 (conditional instantaneous velocity) vtv_t 来替换它,因为最小化条件流匹配损失等价于最小化边际流匹配损失。 vt=atx+btϵv_t = a_t' x + b_t' \epsilon 在本文默认的线性时间表 at=1t,bt=ta_t = 1-t, b_t = t 下,条件瞬时速度简化为 vt=ϵxv_t = \epsilon - x
    • 最终目标 utgtu_{\mathrm{tgt}}:v(zt,t)v(z_t, t) 替换为 vtv_t,得到实际训练中使用的目标: utgt=vt(tr)(vtzuθ+tuθ)(11) u_{\mathrm{tgt}} = v_t - (t-r) \big( v_t \partial_z u_{\boldsymbol{\theta}} + \partial_t u_{\boldsymbol{\theta}} \big) \quad (11)

算法 1:MeanFlow 训练伪代码 (Algorithm 1: MeanFlow Training)

Algorithm 1 MeanFlow: Training.
Input: Neural network u_theta, data x, prior distribution p_prior, time scheduler (a_t, b_t)
Output: Trained parameters theta
 
1: Loop for many training iterations:
2:   Sample a mini-batch of data x from p_data.
3:   Sample a mini-batch of noise epsilon from p_prior.
4:   Sample time steps (r, t) (e.g., from U(0,1), lognorm).
     (Note: Ensure t > r for most samples, optionally include r=t samples)
5:
6:   # Compute z_t and conditional instantaneous velocity v_t
7:   z_t = a_t * x + b_t * epsilon
8:   v_t = a_t_prime * x + b_t_prime * epsilon  (e.g., epsilon - x for a_t=1-t, b_t=t)
9:
10:  # Compute u_theta(z_t, r, t) and its JVP
11:  u_theta_output, jvp_output = jvp(u_theta, (z_t, r, t), (v_t, 0, 1))
     # jvp(function, (inputs), (tangent_vectors))
     # tangent_vectors: v_t for z_t, 0 for r (as dr/dt=0), 1 for t (as dt/dt=1)
     # jvp_output corresponds to the total derivative term: (v_t * d_z u_theta + d_t u_theta)
12:
13:  # Construct the target u_tgt using MeanFlow Identity
14:  u_tgt = v_t - (t - r) * jvp_output
15:
16:  # Compute loss (with stop-gradient on target)
17:  loss = mean_squared_error(u_theta_output, stop_gradient(u_tgt))
18:
19:  # Backpropagate and update theta
20:  loss.backward()
21:  optimizer.step()
22:  optimizer.zero_grad()

JVP 计算的效率 (Efficiency of JVP Computation):

  • jvp\mathsf{jvp} 操作本质上是一个反向传播 (backward pass)。在我们的框架中,它被用于计算目标 utgtu_{\mathrm{tgt}} 的一部分,而 utgtu_{\mathrm{tgt}} 又被施加了停止梯度。因此,在对神经网络参数 θ\theta 进行优化时,jvp\mathsf{jvp} 的结果被视为常数。
  • 这意味着 jvp\mathsf{jvp} 引入的额外计算成本仅相当于一次额外的反向传播,但这次反向传播只需要计算到输入变量(zt,r,tz_t, r, t),而不需要计算到网络参数 θ\theta。因此,它的开销远低于标准的 θ\theta-反向传播,通常只占总训练时间的少部分(本文提到小于 20%)。

4.2.4 采样 (Sampling)

MeanFlow 模型的采样过程非常直接。一旦神经网络 uθu_\theta 被训练好,它就直接近似平均速度。

采样方程: zr=zt(tr)uθ(zt,r,t)(12) z_r = z_t - (t-r) u_\theta(z_t, r, t) \quad (12)

  • 符号解释:
    • zrz_r: 在时间 rr 的状态(目标输出)。
    • ztz_t: 在时间 tt 的状态(当前输入)。
    • (tr)uθ(zt,r,t)(t-r) u_\theta(z_t, r, t): 神经网络预测的从 ttrr 的总位移。

算法 2:MeanFlow 一步采样伪代码 (Algorithm 2: MeanFlow 1-step Sampling)

Algorithm 2 MeanFlow: 1-step Sampling
Input: Trained neural network u_theta, prior distribution p_prior
Output: Generated sample x_0
 
1: Sample z_1 = epsilon from p_prior.  (This corresponds to t=1)
2:
3: # Perform 1-step generation to t=0
4: x_0 = z_1 - u_theta(z_1, r=0, t=1)  (Since t-r = 1-0 = 1)
5:
6: Return x_0

虽然本文主要关注一步采样 (1-NFE),但 Eq. (12) 同样适用于少步采样 (few-step sampling),只需将时间间隔 [r, t] 划分为多个小步即可。

4.2.5 与先前工作的关系 (Relation to Prior Work)

  • 一致性模型 (Consistency Models): CM 专注于锚定在数据侧的路径(即 r0r \equiv 0),并且只关注单个时间变量。它们通过对神经网络输出施加启发式一致性约束。而 MeanFlow 是基于平均速度的定义,从数学上推导出了恒等式,且考虑了两个任意时间变量 rrtt
  • Shortcut [13] 和 IMM [52] 模型: 这些模型也考虑了两个时间变量,但通过引入额外的自洽损失来实现。MeanFlow 的训练是完全由平均速度的定义及其导出的恒等式驱动的,没有额外的假设。

4.3 带引导的平均流 (Mean Flows with Guidance)

MeanFlow 框架能够自然地集成 (naturally supports) 分类器无关引导 (Classifier-Free Guidance, CFG),且在采样时保持 1-NFE,避免了传统 CFG 导致 NFE 翻倍的问题。这是通过将 CFG 的概念直接融入到真实场 (ground-truth fields) 中来实现的。

4.3.1 真实场 (Ground-truth Fields)

首先,我们构建一个新的引导瞬时速度场 (guided instantaneous velocity field) vcfgv^{\mathrm{cfg}}

vcfg(zt,tc)ωv(zt,tc)+(1ω)v(zt,t)(13) v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}) \triangleq \omega v(z_t, t \mid \mathbf{c}) + (1-\omega) v(z_t, t) \quad (13)

  • 符号解释:
    • vcfg(zt,tc)v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}): 带类别条件 c\mathbf{c} 的引导瞬时速度场。

    • v(zt,tc)v(z_t, t \mid \mathbf{c}): 类别条件瞬时速度场,它是对给定类别 c\mathbf{c} 下瞬时速度 vtv_t 的期望。

    • v(zt,t)v(z_t, t): 类别无关瞬时速度场,它是对所有类别 c\mathbf{c} 进行平均后的瞬时速度场(即 Ec[v(zt,tc)]\mathbb{E}_{\mathbf{c}} [v(z_t, t \mid \mathbf{c})])。

    • ω\omega: 引导尺度 (guidance scale),控制条件影响的强度。

      与 MeanFlow 的核心思想一致,我们引入与 vcfgv^{\mathrm{cfg}} 相对应的引导平均速度场 (guided average velocity field) ucfgu^{\mathrm{cfg}}。根据 MeanFlow 恒等式 (Eq. (6)), ucfgu^{\mathrm{cfg}} 必须满足:

ucfg(zt,r,tc)=vcfg(zt,tc)(tr)ddtucfg(zt,r,tc)(15) u^{\mathrm{cfg}}(z_t, r, t \mid \mathbf{c}) = v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}) - (t-r) \frac{d}{dt} u^{\mathrm{cfg}}(z_t, r, t \mid \mathbf{c}) \quad (15) 这里,vcfg(zt,tc)v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}) 可以进一步分解为: vcfg(zt,tc)=ωv(zt,tc)+(1ω)ucfg(zt,t,t)(16) v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}) = \omega v(z_t, t \mid \mathbf{c}) + (1-\omega) u^{\mathrm{cfg}}(z_t, t, t) \quad (16) 注意,当 r=tr=t 时,ucfg(zt,t,t)u^{\mathrm{cfg}}(z_t, t, t) 退化为瞬时速度 vcfg(zt,t)v^{\mathrm{cfg}}(z_t, t)

4.3.2 带引导的训练 (Training with Guidance)

我们训练一个神经网络 uθcfgu_\theta^{\mathrm{cfg}} 来直接建模 ucfgu^{\mathrm{cfg}}。损失函数类似于 Eq. (9),但使用一个修改后的目标 v~t\tilde{v}_t:

L(θ)=Euθcfg(zt,r,tc)sg(utgt)22(17) \mathcal{L}(\theta) = \mathbb{E} \big\| u_{\theta}^{\mathrm{cfg}}(z_t, r, t \mid \mathbf{c}) - \mathbf{sg}(u_{\mathrm{tgt}}) \big\|_2^2 \quad (17) 其中 utgtu_{\mathrm{tgt}} 是: utgt=v~t(tr)(v~tzuθcfg+tuθcfg)(18) u_{\mathrm{tgt}} = \tilde{v}_t - (t-r) \big( \tilde{v}_t \partial_z u_{\theta}^{\mathrm{cfg}} + \partial_t u_{\theta}^{\mathrm{cfg}} \big) \quad (18) 而修改后的 v~t\tilde{v}_t 来源于 Eq. (16),其中将边际速度 v(zt,tc)v(z_t, t \mid \mathbf{c}) 替换为条件瞬时速度 vtv_t: v~tωvt+(1ω)uθcfg(zt,t,t)(19) \tilde{v}_t \triangleq \omega v_t + (1-\omega) u_{\theta}^{\mathrm{cfg}}(z_t, t, t) \quad (19)

  • 注意: 如果 ω=1\omega = 1,这个损失函数就退化回无 CFG 的情况 (Eq. (9))。
  • 单次函数评估与 CFG (Single-NFE with CFG): 由于 uθcfgu_\theta^{\mathrm{cfg}} 直接建模了引导后的平均速度 ucfgu^{\mathrm{cfg}},因此在采样时,我们只需要对 uθcfgu_\theta^{\mathrm{cfg}} 进行一次前向传播,无需进行线性组合,从而保持了 1-NFE 的采样行为。

4.3.3 改进的 CFG (Improved CFG, Appendix B.1)

为了进一步提升性能,可以引入一个混合尺度 κ\kappa 来更好地融合类别条件和类别无关的信息,类似于传统 CFG 中随机丢弃类别条件 (random dropping) 的做法:

vcfg(zt,tc)=ωv(zt,tc)+κucfg(zt,t,tc)+(1ωκ)ucfg(zt,t,t)(20) v^{\mathrm{cfg}}(z_t, t \mid \mathbf{c}) = \omega v(z_t, t \mid \mathbf{c}) + \kappa u^{\mathrm{cfg}}(z_t, t, t \mid \mathbf{c}) + (1-\omega-\kappa) u^{\mathrm{cfg}}(z_t, t, t) \quad (20) 对应的训练目标中的 v~t\tilde{v}_t 变为: v~tω(ϵx)sample vt+κ uθcfg(zt,t,tc)clscond output+(1ωκ)uθcfg(zt,t,t)clsuncond output(21) \tilde{v}_t \triangleq \omega \underbrace{(\epsilon - x)}_{\mathrm{sample}\ v_t} + \underbrace{\kappa \ u_{\theta}^{\mathrm{cfg}}(z_t, t, t \mid \mathbf{c})}_{\mathrm{cls-cond}\ \mathrm{output}} + \underbrace{(1-\omega-\kappa) u_{\theta}^{\mathrm{cfg}}(z_t, t, t)}_{\mathrm{cls-uncond}\ \mathrm{output}} \quad (21) 通过调整 κ\kappa 参数,可以在目标中混合类别条件和类别无关的 uθcfgu_{\theta}^{\mathrm{cfg}} 输出,从而进一步改善生成质量。

4.4 设计决策 (Design Decisions)

4.4.1 损失度量 (Loss Metrics)

  • 问题: 传统的平方 L2 损失 L=Δ22\mathcal{L} = ||\Delta||_2^2 (其中 Δ=uθutgt\Delta = u_\theta - u_{\mathrm{tgt}}) 在少步/一步生成中可能表现不佳。
  • 解决方案: 采用自适应加权损失 (adaptive loss weighting),类似于功率 L2 损失 Lγ=Δ22γ\mathcal{L}_\gamma = ||\Delta||_2^{2\gamma}。这等价于最小化加权平方 L2 损失 sg(w)L\mathbf{sg}(w) \cdot \mathcal{L},其中权重 ww 定义为: w=1/(Δ22+c)p(22) w = 1 / ( ||\Delta||_2^2 + c )^p \quad (22)
    • 符号解释:
      • pp: 一个超参数,与 γ\gamma 相关 (p=1γp = 1-\gamma)。
      • cc: 一个小的正数 (例如 10310^{-3}),用于避免除以零。
    • 与 Pseudo-Huber 损失的关系:p=0.5p=0.5 时,这个损失函数类似于 Pseudo-Huber 损失 [43]。
    • 实验发现: 在消融实验中,发现 p=1p=1 时效果最佳。

4.4.2 采样时间步 (r, t) 的选择 (Sampling Time Steps (r, t))

  • 策略: 在训练时,需要从一个分布中采样两个时间步 rrtt
  • 分布类型:
    • 均匀分布 (Uniform distribution): U(0,1)\mathcal{U}(0, 1)
    • Logit-Normal 分布 (Logit-normal distribution): 从正态分布 N(μ,σ)\mathcal{N}(\mu, \sigma) 中采样,然后通过逻辑函数 (logistic function) 映射到 (0,1)(0, 1) 区间。
  • 后处理: 采样后,将较大的值赋给 tt,较小的值赋给 rr,以确保 t>rt > r。此外,会设置一定比例的样本,使其 r=tr=t,以保证模型也能学习瞬时速度的行为。
  • 实验发现: Logit-Normal 采样器通常表现更好,这与流匹配的观察结果一致 [11]。

4.4.3 对 (r, t) 的条件化 (Conditioning on (r, t))

  • 方法: 使用位置嵌入 (positional embedding) [48] 来编码时间变量 rrtt,然后将嵌入后的信息作为神经网络的条件输入。
  • 不同的条件化形式:
    • 直接条件化在 (t, r) 上。
    • 条件化在 (t,Δt)(t, \Delta t) 上,其中 Δt=tr\Delta t = t-r
    • 仅条件化在 Δt\Delta t 上。
  • JVP 计算: JVP 总是针对函数 uθ(,r,t)u_\theta(\cdot, r, t) 来计算,无论网络内部如何处理 rrtt 的嵌入。
  • 实验发现: 条件化在 (t,tr)(t, t-r) 上效果最好,但直接条件化在 (t, r) 上也表现出色。即使仅条件化在时间间隔 t-r 上也能获得合理的结果,这表明时间间隔本身是 MeanFlow 的一个重要信号。

5. 实验设置 (Experimental Setup)

数据集 (Datasets)

  • ImageNet 256x256 [7]:

    • 来源与特点: 一个大型、多样化的图像数据集,包含数百万张图像和上千个类别。本文在 256×256256 \times 256 分辨率上进行条件生成。
    • 使用方式: 遵循 [34, 13, 52] 的做法,在预训练的 VAE 编码器 (pre-trained VAE tokenizer) [37] 的潜在空间 (latent space) 中进行模型训练。潜在空间的大小为 32×32×432 \times 32 \times 4
    • 目的: 验证模型在复杂、高分辨率图像生成任务中的性能和可扩展性。
    • 样本示例: 由于系统无法直接展示 ImageNet 图片,请读者参考 Fig. 6 中的生成样本,以直观理解 ImageNet 数据的复杂性和多样性。
  • CIFAR-10 [25]:

    • 来源与特点: 一个较小规模的图像分类数据集,包含 10 个类别的 32×3232 \times 32 彩色图像。
    • 使用方式: 进行无条件生成 (unconditional generation),模型直接在像素空间 (pixel space) 上训练,输入为 32×32×332 \times 32 \times 3
    • 目的: 验证模型在较简单数据集上的基本性能,并与其他基线进行对比,尤其是在像素空间直接生成的能力。

评估指标 (Evaluation Metrics)

  • Fréchet Inception Distance (FID) [17]:
    • 概念定义 (Conceptual Definition): FID 是一种广泛用于评估生成模型质量的指标。它通过计算生成图像和真实图像在预训练的 Inception-v3 网络特征空间中的统计量(均值和协方差)之间的 Fréchet 距离来衡量两组图像分布的相似度。FID 值越低,表示生成图像的质量越高,与真实图像的分布越接近,同时通常也反映了更好的多样性。
    • 数学公式 (Mathematical Formula): FID=μxμg22+Tr(Σx+Σg2(ΣxΣg)1/2) \mathrm{FID} = ||\mu_x - \mu_g||^2_2 + \mathrm{Tr}(\Sigma_x + \Sigma_g - 2(\Sigma_x \Sigma_g)^{1/2})
    • 符号解释 (Symbol Explanation):
      • 22||\cdot||_2^2: 向量的 L2 范数的平方。
      • Tr()\mathrm{Tr}(\cdot): 矩阵的迹 (trace)。
      • μx\mu_x: 真实图像(xx)在 Inception-v3 网络某个中间层(通常是全局平均池化层之前的 2048 维特征)的特征均值向量。
      • μg\mu_g: 生成图像(gg)在 Inception-v3 网络相同中间层的特征均值向量。
      • Σx\Sigma_x: 真实图像特征的协方差矩阵。
      • Σg\Sigma_g: 生成图像特征的协方差矩阵。
      • (ΣxΣg)1/2(\Sigma_x \Sigma_g)^{1/2}: 矩阵 ΣxΣg\Sigma_x \Sigma_g 的平方根,这需要进行矩阵代数运算(通常通过特征分解实现)。
    • 计算: 本文在 50K 张生成图像上计算 FID-50K。

对比基线 (Baselines)

论文将 MeanFlow 与多种类型的生成模型进行比较,主要分为以下几类:

  • 一步/少步扩散/流模型 (1-NFE/Few-step Diffusion/Flow Models from scratch):
    • iCT-XL/2 [43]: 一致性模型 (Consistency Models) 的改进版本。
    • Shortcut-XL/2 [13]: 引入自洽损失的流模型。
    • iMM-XL/2 [52]: 归纳矩匹配模型。
    • MeanFlow 的优势: MeanFlow 在此类别中显著超越所有基线。
  • 其他生成模型家族 (Other Families of Generative Models):
    • GANs (Generative Adversarial Networks): BigGAN [5], GigaGAN [21], StyleGAN-XL [40]。这些模型通常一步生成,但面临训练不稳定、模式崩溃等挑战。
    • 自回归/掩码模型 (Autoregressive/Masking Models): AR w/ VQGAN [10], MaskGIT [6], VAR-d30 [47], MAR-H [27]。这类模型通常通过逐个预测或掩码预测来生成,NFE 较高或以迭代精华为代价。
    • 多步扩散/流模型 (Multi-step Diffusion/Flow Models): ADM [8], LDM-4-G [37], SimDiff [20], DiT-XL/2 [34], SiT-XL/2 [33], SiT-XL/2+REPA [51]。这些是当前性能最优异的生成模型,但通常需要数百次甚至上千次 NFE。
    • MeanFlow 的目标: 在 1-NFE 或 2-NFE 条件下,MeanFlow 旨在缩小与这些高性能多步模型之间的差距。

实验配置 (Experiment Setting)

  • ImageNet 256x256:
    • 模型架构: 采用基于 Transformer 的架构,遵循 DiT [34] 的设计,使用 ViT [9] 作为骨干网络,并结合 adaLN-Zero [34] 进行条件化。
    • 时间条件化:rrtt 两个时间变量使用位置嵌入 (positional embedding),然后通过 2 层 MLP 结合并作为网络的条件输入。
    • 训练: 所有模型都从零开始训练 (trained from scratch)。
    • 优化器: Adam [24]。
    • 学习率: 0.0001 (常数)。
    • EMA 衰减: 0.9999。
    • 批大小 (Batch Size): 通常为 2048。
    • 训练轮数 (Epochs): 80 到 1000 轮不等,具体取决于模型大小和配置(详见 Table 4)。
    • CFG: 应用分类器无关引导 (CFG),其中 ω\omegaκ\kappa 参数根据实验进行调整,并采用 0.1 的类别条件丢弃率。
  • CIFAR-10:
    • 模型架构: 使用 U-net [38] 架构,类似于 [44] 中的实现 (约 55M 参数)。
    • 时间条件化: 使用位置嵌入编码 (t,tr)(t, t-r),然后拼接作为条件。
    • 预处理器 (Preconditioner): 不使用 EDM 风格的预处理器 [22]。
    • 优化器: Adam,学习率 0.0006,批大小 1024,dropout 0.2,EMA 衰减 0.99995。
    • 训练迭代次数: 800K 迭代 (包含 10K warm-up)。
    • 数据增强: 遵循 [22],但禁用垂直翻转和旋转。

Table 4: ImageNet 256x256 的配置详情 (转录)

configs B/4 B/2 M/2 L/2 XL/2 XL/2+
params (M) 131 131 497.8 459 676 676
FLOPs (G) 5.6 23.1 54.0 119.0 119.0 119.0
depth 12 12 16 24 28 28
hidden dim 768 768 1024 1024 1152 1152
heads 12 12 16 16 16 16
patch size 4x4 2x2 2x2 2x2 2x2 2x2
epochs 80 240 240 240 240 1000
batch size \multicolumn{6}{ c }{2048}
dropout \multicolumn{6}{ c }{0.0}
optimizer \multicolumn{6}{ c }{Adam [24]}
lr schedule \multicolumn{6}{ c }{constant}
lr \multicolumn{6}{ c }{0.0001}
Adam (β1,β2\beta_1, \beta_2) \multicolumn{6}{ c }{(0.9, 0.95)}
weight decay \multicolumn{6}{ c }{0.0}
ema decay \multicolumn{6}{ c }{0.9999}
ratio of r!=t \multicolumn{6}{ c }{25%}
(r, t) cond \multicolumn{6}{ c }{(t, t - r)}
(r, t) sampler \multicolumn{6}{ c }{lognorm(0.4, 1.0)}
p for adaptive weight \multicolumn{6}{ c }{1.0}
CFG effective scale ω\omega' 2.0 2.0 2.5 2.5 2.5 2.5
CFG ω\omega, Eq. (21) \multicolumn{6}{ c }{ω=ω(1κ)\omega=\omega'(1-\kappa)}
CFG κ\kappa, Eq. (21) \multicolumn{6}{ c }{κ=1ω/ω\kappa=1-\omega/\omega'}
CFG cls-cond drop [18] \multicolumn{6}{ c }{0.1 [18]}
CFG triggered if t is in: [0.0, 1.0] [0.0, 1.0] [0.0, 1.0] [0.0, 0.8] [0.0, 0.75] [0.3, 0.8]

6. 实验结果与分析 (Results & Analysis)

本节将详细分析 MeanFlow 模型在不同设置下的实验结果,并与现有技术进行比较。

核心结果分析 (Core Results Analysis)

ImageNet 256×256256 \times 256 上的性能 (ImageNet 256×256256 \times 256 Performance)

  • 1-NFE 性能显著提升 (Significant 1-NFE Performance Improvement):
    • MeanFlow 在 ImageNet 256×256256 \times 256 上使用 1-NFE 实现了 3.43 的 FID,这是迄今为止一步生成模型中的最佳结果。
    • 与之前最先进的一步扩散/流模型(如 Shortcut [13] 的 10.60 FID)相比,MeanFlow 实现了高达 50% 到 70% 的相对性能提升。
    • 值得注意的是,MeanFlow 的训练是从零开始 (from scratch)不依赖任何预训练、蒸馏或课程学习。这凸显了其框架的原则性和鲁棒性。
  • 2-NFE 性能媲美多步模型 (2-NFE Performance Rivals Multi-step Models):
    • 在 2-NFE 条件下,MeanFlow 模型的 FID 达到了 2.20。
    • 这一结果与领先的多步扩散/流模型(如 DiT [34] 的 2.27 FID 和 SiT [33] 的 2.15 FID,两者均使用 250×2250 \times 2 NFE)在相同的 XL/2 骨干网络下处于同等水平。
    • 这有力地证明了少步扩散/流模型有能力媲美甚至超越其多步前身,极大地缩小了两者之间的差距。
  • CFG 的有效性 (Effectiveness of CFG):
    • 在 1-NFE 场景下,CFG 显著提高了生成质量(如表 1f 所示,从 61.06 FID 提升到 15.53 FID),这与多步生成模型的观察结果一致。
    • MeanFlow 独特的 CFG 整合方式 (Sec. 4.2) 使得引导过程不增加采样时的 NFE,保持了单步采样的效率。

CIFAR-10 上的性能 (CIFAR-10 Performance)

  • 在 CIFAR-10 无条件生成任务中,MeanFlow 实现了 2.92 的 FID。
  • 尽管 MeanFlow 未使用 EDM 风格的预处理器 (preconditioner),但其性能仍与使用了预处理器的 iCT [43] (2.83 FID)、sCT [31] (2.97 FID) 和 IMM [52] (3.20 FID) 等竞争方法相当。这表明 MeanFlow 在像素空间直接生成时也具有竞争力。

数据呈现 (表格)

Table 1: ImageNet 256×256256 \times 256 1-NFE 生成的消融研究 (转录) 默认配置用灰色背景标记:B/4 骨干网络,80 epoch 从零开始训练。

(a) r ≠ t 采样比率 (b) JVP 计算 (c) 位置嵌入
% of r≠t FID, 1-NFE jvp tangent FID, 1-NFE pos. embed FID, 1-NFE
0% (= FM) 328.91 (v, 0, 1) 61.06 (t, r) 61.75
25% 61.06 (v, 0, 0) 268.06 (t, t−r) 61.06
50% 63.14 (v, 1, 0) 329.22 (t, r, −r) 63.98
100% 67.32 (v, 1, 1) 137.96 t−r only 63.13
(d) 时间采样器 (e) 损失度量 (f) CFG 引导尺度
t, r sampler FID, 1-NFE p FID, 1-NFE ω FID, 1-NFE
uniform(0, 1) 65.90 0.0 79.75 1.0 (w/o cfg) 61.06
lognorm(0.2, 1.0) 63.83 0.5 63.98 1.5 33.33
lognorm(0.2, 1.2) 64.72 1.0 61.06 2.0 20.15
lognorm(0.4, 1.0) 61.06 1.5 66.57 3.0 15.53
lognorm(0.4, 1.2) 61.79 2.0 69.19 5.0 20.75

Table 2: ImageNet 256×256256 \times 256 上的类别条件生成 (转录) 所有条目均在适用时报告了 CFG。左侧:从零开始训练的 1-NFE 和 2-NFE 扩散/流模型。右侧:其他生成模型家族作为参考。两表中,“×2\times 2”表示 CFG 会导致每次采样步的 NFE 为 2。我们的 MeanFlow 模型都训练了 240 轮,除了“MeanFlow-XL/2+”训练了更多轮并采用了附录中指定的更长训练配置。†: iCT [43] 结果由 [52] 报告。

1-NFE diffusion/flow from scratch
method params NFE FID method params NFE FID
iCT-XL/2 [43]† 675M 1 34.24 GANs
Shortcut-XL/2 [13] 675M 1 10.60 BigGAN [5] 112M 1 6.95
MeanFlow-B/2 131M 1 6.17 GigaGAN [21] 569M 1 3.45
MeanFlow-M/2 308M 1 5.01 StyleGAN-XL [40] 166M 1 2.30
MeanFlow-L/2 459M 1 3.84 autoregressive/masking
MeanFlow-XL/2 676M 1 3.43 AR w/ VQGAN [10] 227M 1024 26.52
2-NFE diffusion/flow from scratch MaskGIT [6] 227M 8 6.18
iCT-XL/2 [43]† 675M 2 20.30 VAR-d30 [47] 2B 10×2 1.92
iMM-XL/2 [52] 675M 1×2 7.77 MAR-H [27] 943M 256×2 1.55
MeanFlow-XL/2 676M 2 2.93 diffusion/flow
MeanFlow-XL/2+ 676M 2 2.20 ADM [8] 554M 250×2 10.94
LDM-4-G [37] 400M 250×2 3.60
SimDiff [20] 2B 512×2 2.77
DiT-XL/2 [34] 675M 250×2 2.27
SiT-XL/2 [33] 675M 250×2 2.06
SiT-XL/2+REPA [51] 675M 250×2 1.42

Table 3: 无条件 CIFAR-10 (转录)

method precond NFE FID
iCT [43] EDM 1 2.83
ECT [15] EDM 1 3.60
sCT [31] EDM 1 2.97
IMM [52] EDM 1 3.20
MeanFlow none 1 2.92

Table 5: MeanFlow 的改进 CFG (转录) κ\kappa 如 Eq. (20) 所定义,其目标是使类别条件 ucfg(c)u^{\mathrm{cfg}}(\cdot | \mathbf{c}) 和类别无关 uˉcfg()\bar{u}^{\mathrm{cfg}}(\cdot) 都出现在目标中。在本表中,我们固定有效引导尺度 ω\omega' 为 2.0。相应地,对于不同的 κ\kappa 值,我们通过 ω=(1κ)ω\omega = (1-\kappa) \cdot \omega' 设置 ω\omega。如果 κ=0\kappa = 0,则退回到 Eq. (19) 中的 CFG 情况(也参见表 1f)。类似于标准 CFG 随机丢弃类别条件 [18] 的做法,我们观察到在目标中混合类别条件和类别无关的 ucfgu^{\mathrm{cfg}} 能够提高生成质量。

κ FID, 1-NFE
0.0 20.15
0.5 19.15
0.8 19.10
0.9 18.63
0.95 19.17

消融实验/参数分析 (Ablation Studies / Parameter Analysis)

从流匹配到平均流 (From Flow Matching to Mean Flows) (Table 1a)

  • r!=tr != t 的采样比率为 0% 时(即退化为标准流匹配),1-NFE 的 FID 高达 328.91,表明传统流匹配难以进行一步生成。
  • r!=tr != t 的比率增加时,MeanFlow 开始发挥作用。在 25% 的比率下,FID 显著下降到 61.06,在 50% 和 100% 的比率下也能产生合理结果(63.14 和 67.32)。
  • 分析: 这表明为了实现有效的一步生成,模型必须学习到平均速度的概念,而不仅仅是瞬时速度。模型需要在学习瞬时速度和通过修改后的目标学习 rtr \ne t 时的平均速度之间取得平衡。

JVP 计算 (JVP Computation) (Table 1b)

  • 正确的 JVP 导数向量 (tangent) (v,0,1)(v, 0, 1)(对应于 (zu,ru,tu)(\partial_z u, \partial_r u, \partial_t u))实现了 61.06 的 FID。
  • 当 JVP 的导数向量设置错误时(例如 (v,0,0)(v, 0, 0)(v,1,0)(v, 1, 0)(v,1,1)(v, 1, 1)),FID 急剧恶化(268.06 到 329.22)。
  • 分析: 这验证了 MeanFlow 恒等式中 JVP 项的数学正确性和关键作用。即使 rrtt 只是单维变量,它们对平均速度场的决定作用也是至关重要的。错误的 JVP 计算会导致模型无法正确学习平均速度与瞬时速度之间的关系。

(r, t) 的条件化 (Conditioning on (r, t)) (Table 1c)

  • 不同的位置嵌入方式都产生了有意义的 1-NFE 结果。
  • 将时间 tt 和时间间隔 t-r 作为条件 (embedding (t,tr)(t, t-r)) 取得了最佳效果 (61.06 FID)。
  • 直接使用 (t, r) 作为条件 (embedding (t, r)) 表现也几乎一样好 (61.75 FID)。
  • 即使仅使用时间间隔 t-r 作为条件 (embedding t-r only) 也能产生合理结果 (63.13 FID)。
  • 分析: 这表明 MeanFlow 框架对于时间变量的条件化方式具有一定的鲁棒性,并且 tt 和时间间隔 t-r 都是重要的信息。

时间采样器 (Time Samplers) (Table 1d)

  • Logit-normal 采样器 (lognorm) (例如 lognorm(0.4, 1.0)) 相比均匀采样器 (uniform(0, 1)) 表现更好 (61.06 vs 65.90)。
  • 分析: 这与流匹配 [11] 中的观察结果一致,即时间步的分布会影响生成质量。Logit-normal 采样器可能在时间域上提供了更有效的采样密度,尤其是在接近 0 和 1 的边界。

损失度量 (Loss Metrics) (Table 1e)

  • 采用自适应加权损失中的参数 p=1.0p=1.0 实现了最佳性能 (61.06 FID)。
  • p=0.5p=0.5 时(类似于 Pseudo-Huber 损失),也取得了有竞争力的结果 (63.98 FID)。
  • 标准的平方 L2 损失 (p=0.0p=0.0) 表现不佳 (79.75 FID),但仍能产生有意义的结果。
  • 分析: 这表明自适应加权损失在少步/一步生成任务中非常重要,它可以帮助模型更好地聚焦于误差较大的样本或阶段,从而改善训练。

CFG 引导尺度 (Guidance Scale) (Table 1f)

  • CFG 显著提升了生成质量。随着引导尺度 ω\omega 的增加,FID 显著下降。
  • ω=1.0\omega=1.0 (无 CFG,61.06 FID) 到 ω=3.0\omega=3.0 (15.53 FID) 达到了最佳效果。过高的 ω\omega (如 ω=5.0\omega=5.0) 可能导致 FID 略微上升 (20.75 FID),这可能意味着过强的引导会牺牲多样性或导致模式崩溃。
  • 分析: 这证实了 CFG 在 MeanFlow 的 1-NFE 生成中同样有效,并且 MeanFlow 的设计使其能够保持 1-NFE 的采样效率。

改进的 CFG (Improved CFG) (Table 5)

  • 通过引入混合尺度 κ\kappa 来在目标中混合类别条件和类别无关的 uθcfgu_{\theta}^{\mathrm{cfg}} 输出,可以进一步改善生成质量。
  • κ=0\kappa=0 时(即不混合,退化到表 1f 的 CFG 情况),FID 为 20.15。
  • κ\kappa 增加到 0.9 时,FID 降至 18.63,达到了最佳性能。
  • 分析: 这类似于传统 CFG 中随机丢弃类别条件以暴露模型于无条件输入。在 MeanFlow 的目标中显式混合条件和无条件项,有助于模型学习更鲁棒的引导机制。

可扩展性 (Scalability)

  • 图 5 (Figure 4) 展示了 MeanFlow 模型在 ImageNet 256×256256 \times 256 上的可扩展性。随着模型规模(B/2, M/2, L/2, XL/2)的增加,1-NFE 的 FID 值持续下降,表明更大的模型可以带来更好的生成质量。
  • 分析: 这与基于 Transformer 的扩散/流模型(如 DiT [34] 和 SiT [33])的行为一致,MeanFlow 在 1-NFE 生成方面也展现出良好的可扩展性。

定性结果 (Qualitative Results)

  • 图 6 (Figure 5) 展示了 MeanFlow-XL/2 模型在 ImageNet 256×256256 \times 256 上使用 1-NFE 生成的精选样本。这些样本展示了高质量和逼真度,进一步支持了定量指标的优越性。
  • 分析: 视觉质量是生成模型性能的直观体现。高保真度的生成图像表明 MeanFlow 模型能够有效地捕捉数据分布的复杂特征。

7. 总结与思考 (Conclusion & Personal Thoughts)

结论总结 (Conclusion Summary)

本论文提出了 MeanFlow,一个用于一步生成建模 (one-step generative modeling)原则性 (principled)高效 (effective) 的框架。其核心创新在于:

  1. 引入平均速度概念 (Average Velocity): 与传统的瞬时速度建模不同,MeanFlow 聚焦于建模平均速度 (average velocity),将其定义为位移与时间间隔之比。

  2. 推导 MeanFlow 恒等式 (MeanFlow Identity): 从平均速度的定义出发,论文严格推导出了一个连接平均速度和瞬时速度的数学恒等式,为神经网络训练提供了坚实的理论基础,避免了启发式一致性约束。

  3. 自洽的训练范式 (Self-Contained Training Paradigm): MeanFlow 实现了完全从零开始 (from scratch) 训练,无需预训练、蒸馏或课程学习,大大简化了模型的开发和应用。

  4. 高效的 CFG 集成 (Efficient CFG Integration): 该框架能够自然地将分类器无关引导 (CFG) 融入到目标场中,从而在采样时保持单次函数评估 (1-NFE),避免了传统 CFG 带来的计算开销。

  5. 卓越的经验性能 (Superior Empirical Performance): 在 ImageNet 256x256 上,MeanFlow 使用 1-NFE 取得了 3.43 的 FID,显著超越了所有现有的一步生成模型,并在 2-NFE 条件下达到了与顶级多步模型(如 DiT 和 SiT)相当的 2.20 FID。

    MeanFlow 的研究成果极大地缩小了一步生成模型与多步生成模型之间的性能差距,并提供了一个全新的、更具原则性的视角来理解和构建连续时间生成模型。

局限性与未来工作 (Limitations & Future Work)

论文中并未明确列出“局限性”,但从其展望中可以推断出一些方向:

  • 多尺度模拟的关联 (Connection to Multi-scale Simulation): 论文在结论中提及,其工作与物理学中的多尺度模拟问题相关,这些问题涉及空间或时间上的尺度、长度和分辨率范围。这意味着 MeanFlow 的概念可能更适合于在不同时间粒度下统一描述动态系统,但其在生成建模领域之外的应用尚待探索。
  • 理论基础的进一步深化 (Further Theoretical Deepening): 尽管 MeanFlow 恒等式是原则性的,但其在复杂数据分布和高维空间中的优化特性、收敛性保证以及与现有连续时间模型的更深层理论联系仍值得进一步研究。
  • 与其他正交改进的结合 (Integration with Orthogonal Improvements): 论文指出,其他正交的改进方法(如 REPA [51])可以应用于 MeanFlow,但尚未在本文中进行探索。未来的工作可以研究如何将 MeanFlow 与最新的骨干网络、优化技术、数据增强策略等结合,以进一步提升性能。
  • 更广泛的应用领域 (Broader Application Domains): 目前主要在图像生成领域进行了验证。MeanFlow 的概念是否能有效推广到其他模态(如音频、视频、文本)或更复杂的动态系统建模任务,是一个开放的问题。

个人启发与批判 (Personal Insights & Critique)

个人启发 (Personal Insights)

  1. 从第一性原理出发的创新 (Innovation from First Principles): MeanFlow 最引人注目之处在于其从平均速度的数学定义出发,通过严谨的推导得出训练目标。这种从“第一性原理”出发的创新,比单纯的启发式设计更具 Elegance 和鲁棒性。它揭示了瞬时量和平均量之间内在的数学美感,为生成模型提供了新的理论支点。对于初学者而言,这种思路非常值得学习:在遇到问题时,不仅仅是修补现有方法,更要深入挖掘其背后的基本定义和数学关系。
  2. 高效与原则性的统一 (Unifying Efficiency with Principled Design): 长期以来,高效(一步生成)和原则性(理论基础)在生成模型中似乎是一对矛盾。MeanFlow 巧妙地将两者结合,通过平均速度的概念,在单次函数评估下实现了与多步模型相当的性能,这为生成模型的未来发展指明了方向。特别是其将 CFG 自然融入目标场,解决了传统 CFG 效率低下的问题,是工程实践上的一个亮点。
  3. 对连续时间模型的重新思考 (Rethinking Continuous-Time Models): MeanFlow 挑战了我们对连续时间生成模型中“速度”的传统理解。它提醒我们,除了瞬时变化率,对一段时间内的平均行为进行建模可能更适用于某些任务需求(如一步生成)。这种思维方式的转变,可能会激励研究者重新审视现有连续时间模型的基础构建块。

批判与潜在改进 (Critique & Potential Improvements)

  1. JVP 的平滑性假设 (Smoothness Assumption for JVP): 尽管论文强调 JVP 计算的高效性,并指出其是目标的一部分且被 stop-gradient。但 JVP 仍然依赖于神经网络 uθu_\theta 对其输入 ztz_ttt可导性 (differentiability)平滑性 (smoothness)。如果 uθu_\theta 学习到的函数不够平滑,其导数可能不稳定,从而影响训练。虽然 Transformer 结构通常具有良好的平滑性,但这种依赖性仍然是需要考虑的。是否存在隐式或显式的正则化来保证这种平滑性?
  2. 时间步采样分布的重要性 (Importance of Time Step Sampling Distribution): 论文通过消融实验证实了时间采样器(如 lognorm)的选择对性能有显著影响。虽然论文声明无需“课程学习”,但选择最佳的 (r, t) 采样分布和其比例(例如 25% 的 r!=tr != t 比例)仍然是一个需要仔细调优的超参数,这在某种程度上可以看作是一种隐式的课程设计。未来的工作可以探索如何自适应地学习或优化这个采样分布。
  3. “自洽”的边界 (Boundary of "Self-Contained"): 论文强调 MeanFlow 是“自洽的”,无需预训练或蒸馏。然而,在 ImageNet 实验中,MeanFlow 仍然依赖于一个预训练的 VAE 编码器 (pre-trained VAE tokenizer) 来将图像映射到潜在空间。虽然这是当前高分辨率图像生成任务的普遍做法,但从最严格的意义上讲,模型并非完全从像素到像素的“自洽”。澄清这种依赖性有助于更准确地评估其“自洽”的程度。
  4. 计算成本的更详细对比 (More Detailed Computational Cost Comparison): 论文提供了 JVP 带来的额外训练时间开销(约 16%),这很棒。但如果能进一步对比 MeanFlow 整体训练时间与同等性能的多步或蒸馏模型的训练时间,将更有助于评估其端到端效率。例如,一个 1-NFE 的 MeanFlow 训练 240 轮,与一个 250-NFE 的 DiT 模型训练 240 轮,其总计算量和收敛速度如何?
  5. 理论收敛性分析 (Theoretical Convergence Analysis): 尽管提出了原则性恒等式,但对于神经网络在何种条件下能够收敛到这个真实场,以及收敛速度和稳定性如何,论文没有给出详细的理论分析。对于初学者,更深入的理论探讨将有助于理解其鲁棒性来源。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。