Any-OPD:面向异构流匹配模型的表征空间桥接在策略蒸馏方法 Any-OPD: Heterogeneous On-Policy Distillation for Flow-Matching Models via Representation-Space Bridging
首个在任意异构流匹配模型间实现跨族在策略蒸馏的通用框架
前置知识
在策略蒸馏 (On-Policy Distillation, OPD)
一种蒸馏范式:让学生模型在自己生成的采样轨迹上接受教师反馈,而非在固定的教师输出数据上模仿。它源自语言模型蒸馏,可消除离线模仿的曝光偏差(exposure bias),在扩散/流匹配生成中被认为是最有效的蒸馏形式之一。
本文的核心研究对象。现有 OPD 都默认师生同构,理解这一默认假设如何被打破,是读懂本文动机的钥匙。
流匹配 (Flow Matching) 与隐空间生成器
流匹配学习速度场 $v_\theta(x_t,t)$,沿 ODE $dx_t = v_\theta(x_t,t)\,dt$ 把噪声传输到数据,在最优传输插值 $x_t=(1-t)x_0+t x_1$ 下 $t$ 本身即噪声水平。隐空间生成器是元组 $(E,D,v,\{\sigma_i\})$:编码器 $E$、解码器 $D$、速度场 $v$ 和噪声调度。SD3、FLUX、Z-Image 都属此类。
本文用噪声水平 $\sigma$ 作为对齐变量,且师生是两个作用于不相交隐空间 $\mathcal{Z}_S \not\equiv \mathcal{Z}_T$ 的生成器,必须先掌握这套形式化语言。
同构假设 (Homogeneity Assumption)
现有扩散 OPD(如 Salimans & Ho 2022、DiffusionOPD、Flow-OPD)默认师生共享三项:相同 VAE(从而共享隐坐标,教师隐变量可作回归目标)、相同架构(内部特征可匹配)、相同时间步网格(同一索引意味同一噪声水平)。这个假设把蒸馏限制在同一家族派生出的模型对。
本文要打破的就是这条假设;识别出它在'无共享空间'和'无共享时钟'两个轴上崩溃,是理解所有后续设计的起点。
DINOv2 与 CLS 特征
DINOv2 是自监督视觉表征模型,其 CLS token 聚合所有 patch 的全局信息,判别性强、对结构/语义偏差敏感,但对教师那种保留语义、重绘局部细节的随机再合成保持不变。本文用冻结的 DINOv2-Base CLS 作为师生唯一的耦合点。
选择 DINOv2 CLS 而非像素/隐变量/LPIPS 作为损失空间,是本文方法成立的关键,直接决定损失对局部重合成的鲁棒性。
噪声调度与 shift 函数
推理用 $N$ 步欧拉调度,噪声水平 $1=\sigma_0>\sigma_1>\cdots>\sigma_N=0$,每个 $\sigma_i$ 由模型特定的 shift 函数产生。FLUX 用动态 shifting,SD3.5 用 static shift 3.0,因此同一索引 $i$ 对应不同的真实噪声水平 $\sigma_S^i \neq \sigma_T^i$。
调度不可按索引比较,是'无共享时钟'问题的根源,也是本文用噪声水平而非步索引对齐轨迹的动机所在。
研究动机
在策略蒸馏(OPD)是知识蒸馏中最有效的形式——让学生在自己生成的采样轨迹上接受教师纠正,从而消除离线模仿的曝光偏差。但现有扩散/流匹配 OPD 方法(Salimans & Ho 2022;DiffusionOPD;Flow-OPD)全都依赖一条同构假设:师生共享 VAE(教师隐变量才是合法回归目标)、共享架构(内部特征可匹配)、共享时间步网格('同一步'即同一噪声水平)。这把蒸馏限制在同族派生的模型对。现实里想蒸馏的对几乎全异构:最强教师 FLUX.1-dev(12B)与待部署学生 SD3.5-Medium(2.5B)来自不同团队,VAE、架构、噪声调度全然不同。同构一旦崩塌,OPD 沿两轴失效:其一是'无共享空间'——教师隐变量在学生坐标里非法,像素也无效,因为教师精修是对语义保留、局部细节重绘的随机再合成,像素级 MSE 回归到模糊条件均值,实验中直接训练崩溃;其二是'无共享时钟'——两调度来自不同 shift 函数,$\sigma_S^i \neq \sigma_T^i$,按索引对齐把不可比噪声区间的状态配对。结论是当前最强教师无法训练我们想部署的学生。
本文的目标是本文的目标是提出首个能在任意一对隐空间流匹配生成器之间进行在策略蒸馏的框架,其中师生既不共享 VAE、也不共享架构、更不共享噪声调度,教师只被当作一个黑盒采样器访问。作者希望借此把蒸馏从'同族模型'的牢笼里解放出来,让任何足够强的教师——无论谁造的——都能成为任意待部署学生的可用老师。为证明这一点,他们用 12B 的 FLUX.1-dev 蒸馏 2.5B 的 SD3.5-Medium,目标是在保持学生架构和 40 步采样预算不变的前提下,在偏好类指标(PickScore、HPSv3、ImageReward 等)上逼近甚至超越教师,同时验证当教师换成本身 VAE、调度、架构都不同的 6B Z-Image 时,整套流程一行代码都不用改。
与已有工作不同的是,已有方法总是试图在师生共享的隐变量、特征或像素里建立对应,这些做法无一例外假设同构。Any-OPD 的独特切入角度恰恰相反:既然两个模型说不同的'语言',就不要强行翻译,而是把它们带到唯一一个中立的、外在于两者的第三表征空间(冻结的 DINOv2 CLS)里做判断;同时不再用步索引对齐轨迹,而是直接用噪声水平 $\sigma$ 作为唯一可比的物理量来路由梯度。这种做法彻底回避了对隐变量、特征或架构的任何假设,因此更换教师既不改损失、也不改一行流程。更关键的是,作者把'分布级差距'和'逐样本残差'拆成两段:前者不需要逐样本对应,可离线用数据闭合,留下纯逐样本信息交给在策略阶段修正——这是把'什么时候监督才有用'这个问题纳入方法设计的独到之处。
核心方法
直觉先行:两个模型说不同语言时,别逼它们互译,而是把双方都送到中立第三空间用同一把尺子评判。Any-OPD 把教师当纯采样器,通过'加噪-去噪'把学生的在策略样本投影到教师流形,再把投影结果与学生渲染仅在一个冻结的 DINOv2 表征空间里比较,全程不触及任一模型内部。技术分两阶段。阶段一锚定(离线):分布级差距不需逐样本对应,用数据闭合——教师目标经学生 VAE 编码回学生隐空间,用流匹配速度损失 $\mathcal{L}_{WS}$ 训 400 步($M=10$ 噪声级)。阶段二在策略:学生无梯度滚出轨迹并缓存噪声水平匹配点 $z_j$;教师对终点图像做加噪-去噪投影得 $x_{ref}$;从 $z_j$ 开梯度重放学生、解码 $\hat{x}_S$,用 DINOv2 CLS 余弦损失 $\mathcal{L}_{OPD}=1-\cos(f(\hat{x}_S),f(x_{ref}))$ 监督,梯度穿冻结解码器 $D_S$ 回传到 $v_\theta$。两段均 512×512 训练、1024×1024 评测,共约 1200 步。
三个核心创新把 Any-OPD 与已有方法区分开。其一,表征空间耦合:损失 $\mathcal{L}_{OPD}=1-\cos(f(\hat{x}_S),f(x_{ref}))$ 只比较 DINOv2 CLS 这种全局嵌入,它对教师那种重绘局部细节的随机再合成天然不变,从而躲过了让像素/隐变量回归崩溃的元凶;又因抽取器与模型无关,换教师时损失一行不动。其二,噪声水平匹配取代索引匹配:调度不可按索引比较,于是把校正路由到学生所有满足 $\sigma_j^S \geq \sigma_k^T > \sigma_{j+1}$ 的步骤,即 $j=\max\{i:\sigma_i^S \geq \sigma_k^T\}$,这是对任意一对 shift 函数都良定义的规则。其三,锚定作为前置必要条件而非锦上添花:分布级差距不含对应信息,离线用数据先闭合,让在策略目标只承载纯逐样本残差——消融证明跳过锚定或在锚定未收敛时启动 OPD,学生永远追不上充分锚定的版本。这与所有依赖共享隐状态的先前 OPD 形成本质区别。
方法步骤详情
流程分四步。① 锚定:给定 $c$,教师采样 $x_T$,学生 VAE 重编码得 $z^\star$;抽噪声级 $m$,构造 $z_m=(1-\sigma_m^S)z^\star+\sigma_m^S\epsilon$,按速度损失训 400 步。② rollout:抽半径 $r$ 定教师起始索引 $k=N_T-r$;学生从纯噪声无梯度积分,只缓存满足 $\sigma_j^S\geq\sigma_k^T$ 的匹配点 $z_j$——把该隐变量当常数是保持'在策略'的关键。③ 教师投影:学生终点图像转教师坐标(无损格式转换)后加噪到 $\sigma_k^T$,沿教师调度积分到干净端点 $x_{ref}=\Pi_{\sigma_k^T}^T(x^S)$,半径 $\sigma$ 控制可改动范围。④ 梯度路由:从 $z_j$ 开梯度重放学生到 $\hat{z}_{N_S}$(仅 $N_S-j$ 次带激活、各自梯度检查点),解码后用 DINOv2 CLS 算余弦损失,梯度穿冻结 $D_S$ 回传 $v_\theta$。在策略阶段 800 步,$N_S=N_T=20$,每步只监督一个样本。
技术新颖性
技术新颖性体现在首次形式化了异构在策略蒸馏问题,并精准指出两条阻塞——无共享空间与无共享时钟。它首次把一个冻结的外部表征(DINOv2 CLS)作为师生唯一的耦合点,使损失对教师的局部随机重合成免疫;首次提出用噪声水平而非步索引作为任意调度对之间的对齐变量,并以加噪-去噪投影算子 $\Pi_\sigma^T(x)=D_T\,\text{Euler}_T((1-\sigma)E_T(x)+\sigma\epsilon\to 0,\,c)$ 表达监督,该算子只调用教师的采样接口,于是任何隐空间生成器都自动合格。锚定阶段把蒸馏按粒度拆分、用数据先闭合分布级差距的设计,是从'何时监督才有用'的层面给出的理论洞察。整组组件均按构造与教师的 VAE、架构、调度无关,因此教师天然可替换——这一点由把教师从 FLUX 换成 Z-Image、零改代码即生效得到实证,是与 DiffusionOPD、Flow-OPD 等'假设共享状态'的方法最本质的区别。
实验结果
核心结果 Table 1:把 FLUX.1-dev(12B)蒸馏进 SD3.5-Medium(2.5B)。学生 Aesthetic 5.383→5.788、ImageReward 0.922→1.116(约教师自身优势 +0.049 的 4 倍)、PickScore 0.866→0.884、HPSv3 9.12→10.97、DPG 84.58→84.71,多项偏好指标反超或逼近 12B 教师。学生能超教师,因目标是'教师对学生样本的投影',融合学生构图与教师渲染。消融印证设计:Fig.5 显示隐变量 MSE 训练初期即崩溃、LPIPS 先升后降、DINOv2 稳定最优,排序恰对应各损失假设的局部对应量;Fig.6 证明锚定是 OPD 必要前提;Table 2 显示精修半径 $r\in[15,20]$(HPSv3 10.97)最佳;Fig.7 显示确定性 rollout 全程最优、随机性反而有害。泛化性 Table 3:教师换 6B Z-Image、流程零改动,学生 GenEval 70.70→71.50、DPG 84.58→85.53,提升顺着教师强项走。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| DrawBench PickScore | PickScore↑ | 0.884 | 0.866 (SD3.5-Medium) | +0.018,反超教师 FLUX.1-dev 0.878 |
| DrawBench HPSv3 | HPSv3↑ | 10.97 | 9.12 (SD3.5-Medium) | +1.85,逼近教师 11.19 |
| DrawBench ImageReward | ImageReward↑ | 1.116 | 0.922 (SD3.5-Medium) | +0.194,约为教师自身优势 +0.049 的 4 倍 |
| DrawBench Aesthetic | Aesthetic↑ | 5.788 | 5.383 (SD3.5-Medium) | +0.405,反超教师 5.765 |
| DPG-Bench | Overall↑ | 84.71 | 84.58 (SD3.5-Medium) | +0.13(偏好导向、教师无余量) |
| DPG-Bench(Z-Image 教师) | Overall↑ | 85.53 | 84.58 (SD3.5-Medium) | +0.95,证明教师可替换 |
| GenEval(Z-Image 教师) | Overall↑ | 71.50 | 70.70 (SD3.5-Medium) | +0.80,顺着教师构图强项提升 |
局限与改进
作者自承的局限:方法偏向偏好类指标,组合性基准 DPG 在 FLUX 教师下几乎不变(+0.13),因为该基准上教师本身已无余量(换 Z-Image 教师才能提升);每次迭代只监督一个样本并需重放 $N_S-j$ 份激活,单步计算与显存开销高于离线蒸馏。我的额外观察:(1) 实验仅用 LoRA(rank 32、α=64),未做全量微调或更小学生的对照,结论能否外推到极致小型化学生尚存疑;(2) DINOv2 CLS 是全局嵌入,可能对细粒度局部瑕疵欠惩罚——尽管指标全面上升,但 Fig.5 显示 CLS+Patch 与 CLS 相当,作者并未深究多尺度组合;(3) 锚定目标在 512×512 生成、却用 1024×1024 评测,存在分辨率鸿沟可能压低了上限;(4) 跨族可替换性仅用一位额外教师(Z-Image)验证,更极端的异构(如非最优传输的流变体、跨模态)尚未触及;(5) 精修半径 $r$ 与教师步数 $N_T$ 是较敏感的逐对超参,缺乏自动选择机制。
独立分析的弱点
独立分析的弱点及改进方向:(1) 单样本投影每迭代开销大(12B 教师前向 + 2.5B 学生激活检查点)——可改为一批多半径 $r$ 同时投影,或缓存复用教师 rollout,提升样本效率;(2) 全局 CLS 可能放过局部伪影——CLS+Patch 在 Fig.5 表现相当却未被进一步探索,可设计多尺度或频域感知的表征组合损失;(3) 512 训练/1024 评测的分辨率错配——可直接在 1024 生成锚定目标或加分辨率一致性项;(4) LoRA 配置固定为 rank 32、α=64,缺少适配器容量与全量微调的扫描,无法判断提升是否还有未榨取的余量;(5) 仅两位教师验证可替换性——应纳入 SDXL、非 OT 流变体、甚至视频/音频流模型做压力测试,以坐实'任意教师'的承诺。每条都指向明确、可工程化的改进路线。
未来方向
作者明确提出的方向:用互补的教师组合是自然补救(如 FLUX 主攻偏好、Z-Image 主攻构图);由于整套公式除特征抽取器外都不依赖图像,可把 Any-OPD 扩展到其他隐空间生成模态(视频、音频流模型)。基于成果可延伸的方向:用与人偏好对齐的奖励模型(HPSv3、ImageReward 自身)替换 DINOv2 作为耦合表征,直接把偏好注入监督;做投影不动点收敛性的理论分析;研究连续时间噪声水平对齐与自适应步选择,取代离散半径 $r$ 抽样;把多教师投票或教师课程纳入锚定与在策略两阶段,进一步逼近'谁造的最强模型都能当老师'这一愿景。
复现评估
可复现性中等偏上。超参披露充分:LoRA rank 32、α=64,锚定 400 步($M=10$ 噪声级、512×512、50 步欧拉采样目标),在策略 800 步($N_S=N_T=20$、512×512),学习率 $1\times10^{-4}$,每 GPU batch 4,评测 1024×1024,精修半径采用 $r\in[15,20]$。训练提示采自 Pick-a-Pic 训练集、验证集留出;基准为标准的 DrawBench、DPG-Bench、GenEval。教师 FLUX.1-dev、学生 SD3.5-Medium、DINOv2-Base 均为公开权重。主要障碍有二:论文未提供代码仓库或 GitHub 链接,关键工程细节(如 $\Pi_\sigma^T$ 的梯度截断、CPS 实现细节)需自行补全;推理与训练需同时承载 12B 教师前向 + 2.5B 学生梯度 + 激活检查点,显存与算力门槛较高,普通实验室复现存在资源压力。
论文图表
三联图浓缩了整篇论文。左:异构师生对(FLUX.1-dev→SD3.5-Medium)在隐表征(FLUX latent vs SD3 latent)和噪声调度(FLUX schedule vs SD3.5-M schedule)上同时不同,故无现成在策略方法可监督学生自身轨迹。中:Any-OPD 同时绕过两类失配——教师只精修学生的在策略样本,二者仅在冻结 DINOv2 特征空间比较,从不进入对方隐空间。右:朴素隐变量回归崩溃,而 Any-OPD 把 2.5B 学生的 PickScore 从 0.846 提到 0.884、HPSv3 从 9.12 提到 10.97,在偏好指标上逼近甚至超过 12B 教师。
这一张图同时承载了问题陈述、方法直觉与头条结果,是读者最快抓住论文论点的入口。