JSON: Jigsaw Self-play Optimization for Normalizing Flows¶
会议: ECCV 2026
论文: ECCV 原文
代码: https://media.eventhosts.cc/Conferences/ECCV2026/pdfs/2253.pdf
领域: 图像生成
关键词: 归一化流、自博弈优化、拼图重排、轨迹对齐、空间一致性
一句话总结¶
针对归一化流(NF)缺乏预定义扩散动力学导致胜负轨迹难界定以及标准自博弈易引发训练崩溃的问题,本文提出了结合拼图重排内在空间评测与有界自博弈损失的 JSON 框架,在无需外部特征监督的情况下显著提升了高分辨率图像生成质量。
研究背景与动机¶
现代生成模型在扩散模型、自回归模型与归一化流等方向取得了长足进展。其中,归一化流(Normalizing Flows, NFs)依赖变量替换公式在简单高斯先验与真实数据分布间构建双射映射。相较于需要多步迭代去噪的扩散模型或连续时间流匹配模型,现代归一化流(如 TarFlow、StarFlow)通过因果变换器在隐空间建立动力学,具备一步解析反演与高效采样的先天优势。然而,为维持严格的可逆性与对数行列式计算的解析可行性,NF 在网络结构上受到严苛限制;同时其最大似然估计(MLE)训练缺乏对演化轨迹的显式约束,使得模型在高层语义表征与空间一致性学习上存在天然瓶颈。
在扩散模型领域,轨迹微调(Trajectory Fine-tuning, TFT)与自博弈(Self-play)范式已被证明能够通过让当前模型与历史快照模型博弈竞争、偏好真实数据轨迹并抑制合成数据轨迹,有效改善生成质量且无需昂贵的外部人类标注。然而,将自博弈直接迁移到 NF 会遭遇两大深层矛盾:其一,NF 本质是基于全局概率密度最大化训练的,上一轮历史模型生成的样本天然落在当前模型的高似然区域,强行惩罚这些自生成样本(迫使其欠拟合)直接违背了 MLE 的训练基石,极易引发梯度爆炸与训练崩溃(出现 NaN);其二,与预先定义了正向/反向扩散过程的扩散模型不同,NF 采用隐式动力学映射,缺乏显式的真实轨迹锚点,仅凭似然度这一指标难以敏锐度量图像的全局空间逻辑与结构合理性。
为了打破这一困境,本文探索能否为 NF 注入一种内在可验证的几何一致性度量,从而在无需引入 DINOv2 等外部预训练视觉编码器额外开销的前提下定义“胜出轨迹”。核心 idea:将拼图重排(Jigsaw Puzzle Reassembly)作为自监督内在空间连贯性评测器,并配合松弛边界的有界自博弈损失(Bounded Self-play Loss)构建内外双循环优化,引导归一化流向空间结构更严谨的优质轨迹对齐。
方法详解¶
整体框架¶
JSON 的核心思想在于两阶段协同训练与内外双循环优化:在阶段一(NF 预训练),在各可逆可微块输出后嵌入轻量卷积拼图头,迫使特征学习判别空间拓扑位置的能力;在阶段二(自博弈微调),冻结拼图头作为固有的空间逻辑判别器,引导主模型与上一轮历史快照模型开展对抗博弈。
整个训练流程由潜空间因果流架构承载,输入图像首先经过 VAE 编码器映射为潜变量,在因果变换器流块中完成仿射变换与密度估计。在自博弈微调阶段,主模型与历史快照模型对合成数据与真实数据分别计算拼图重排重构误差,从而动态判定该合成轨迹是否值得学习或抑制。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入真实潜变量 x 与标签 c"] --> B["阶段1:空间连贯性表征学习<br/>多尺度卷积拼图头预训练"]
B --> C["阶段2:内在空间逻辑自评估<br/>冻结拼图头计算累计重排误差"]
C --> D{"比较真实样本与合成样本误差<br/>L_jig(x) ≤ L_jig(x') ?"}
D -->|是:合成样本存在空间缺陷| E["有界自博弈轨迹对齐<br/>双向边界约束防崩溃"]
D -->|否:合成样本语义更优| F["跳过自博弈惩罚<br/>保留高质量合成几何特征"]
E --> G["内外双循环参数更新<br/>内循环步级博弈 + 外循环代际迭代"]
F --> G
关键设计¶
1. 空间连贯性表征学习:引入可验证的拼图重排辅助头
传统 NF 模型仅通过 MLE 损失计算对数雅可比行列式与二次范数,对局部 patch 的宏观空间排布不具备显式敏感性。为此,本文在每个可逆可微块 \(f_i\) 的输出特征 \(o_{i+1} = f_i(x_i) \in \mathbb{R}^{HW \times D}\) 处接入轻量卷积拼图头 \(g_i\)。在预训练阶段,将空间特征沿二维网格切分为大小为 \(p \times p\) 的局部块(默认 \(p=1\)),赋予其在原始网格上的空间坐标索引标签 \(\mathcal{Y}_m\)。随后将这组空间 token 沿空间维度进行随机乱序打乱得到 \(\tilde{o}_{i+1}\),并由拼图头预测其原始空间位置: $$ \mathcal{L}{\mathrm{jig},i}(x;\theta) = -\sum} \mathcal{Ym \odot \log \mathrm{softmax}(g_i(\tilde{o})) $$ 这一设计巧妙之处在于,拼图重排是一个具备确定性、可自验证真值标签的任务,无需人工奖励打分或外部网络。它不仅迫使因果变换器流块在前向变换过程中维持严格的空间临近连贯性,而且在后续微调中构成了天然的轨迹打分标准。
2. 规避训练崩溃的有界自博弈损失:构建松弛约束区间
在将经典 Bradley-Terry 偏好模型直接引入 NF 时,标准自博弈(B-NSP)要求最大化主模型相对历史快照在真实数据上的对数似然比,同时无节制地压低主模型在合成数据 \(x'\) 上的似然度,即强制让 \(p_\theta(x') < p_{\theta_k}(x')\)。然而 \(x'\) 本身是由 \(\theta_k\) 在高置信区间生成的,这种对抗性惩罚直接破坏了 NF 训练的似然下界,造成损失函数出现 NaN 并导致崩溃。本文提出的有界自博弈(Bounded Self-play, BSP)通过引入真实数据作为下界锚点,将似然约束调整为 \(p_{\theta_k}(x) < p_\theta(x') < p_{\theta_k}(x')\),构建具有保护边界的优化目标: $$ \begin{aligned} \mathcal{L}{\mathrm{bsp}}(x, x', \theta_k; \theta) = \;& \sigma\Big(\mathcal{L}}}(x;\theta_k) - \mathcal{L{\mathrm{nf}}(x;\theta)\Big) \ & + \sigma\Big(\mathcal{L}}}(x';\theta) - \mathcal{L{\mathrm{nf}}(x;\theta_k)\Big) \ & - \sigma\Big(\mathcal{L}(x';\theta)\Big) \end{aligned} $$ 其中 }}(x';\theta_k) - \mathcal{L}_{\mathrm{nf}\(\sigma(t) = -\log(1/(1 + \exp(-t)))\)。该公式巧妙地在“遗忘历史模型的不良合成轨迹”与“停留在历史模型的高似然有效流形”之间取得数学平衡,彻底消除了训练不稳定性。
3. 内在空间逻辑自评估:动态门控轨迹判定
单纯采用有界自博弈依然无法区分合成样本中是否存在甚至优于真实样本的语义结构(例如真实图像可能存在不完整截断,而生成图像主体完整清晰)。JSON 借助预训练中冻结的拼图头,定义跨所有块的累积拼图误差 \(\mathcal{L}_{\mathrm{jig}}(x;\theta) = \sum_i \mathcal{L}_{\mathrm{jig},i}(x;\theta)\)。在自博弈微调时,对真实样本 \(x\) 与合成样本 \(x'\) 进行空间合理性比对,形成自适应分段自博弈损失: $$ \mathcal{L}{\mathrm{jsp}}(x,x',\theta_k;\theta) = \begin{cases} \mathcal{L}}}(x, x', \theta_k; \theta), & \text{若 } \mathcal{L{\mathrm{jig}}(x;\theta) \leq \mathcal{L}(x';\theta) \ 0, & \text{其他} \end{cases} $$ 当合成样本的重排误差小于真实样本时,表明该生成轨迹展现了高度自洽的空间几何逻辑,模型拒绝盲目惩罚该合成轨迹,从而有效沉淀了生成过程中的高价值结构特征。}
4. 内外双循环优化策略:平滑跟踪演化分布
在整个优化体系中,模型由内循环(Inner-loop)与外循环(Outer-loop)交替驱动。在每个训练 epoch 内部的步骤迭代中,历史快照模型 \(\theta_k\) 参数冻结,用于实时条件推断合成 \(x'\) 并参与主模型 \(\theta\) 的步级对抗微调;在完成一个完整 epoch 之后,外循环将最新的主模型权重同步覆盖至竞争者 \(\theta_k \leftarrow \theta\)。这种渐进推进机制使得竞争分布平滑演进,既控制了在线生成带来的时间开销,又保证了自博弈轨迹判别依据的时效性。
损失函数 / 训练策略¶
JSON 框架采用联合优化策略,将 VAE 重构约束、NF 负对数似然损失、拼图重排损失与自博弈损失融为一体: $$ \mathcal{L}{\mathrm{total}}(\mathbf{I}; \psi, \phi, \theta) = \mathcal{L}}}(\mathbf{I};\psi,\phi) + \mathcal{L{\mathrm{nf}}(x;\theta) + \lambda_1 \mathcal{L}(x, x', \theta_k; \theta) $$ 在训练调度上,总训练轮数设为 }}(x;\theta) + \lambda_2 \mathcal{L}_{\mathrm{jsp}\(E=160\) 轮:前 155 轮为阶段一预训练,激活 \(\lambda_1 = 1, \lambda_2 = 0\),专注于建立平稳的潜变量流形并使拼图头达到极高的位置重排精度;最后 5 轮进入阶段二自博弈微调,切换为 \(\lambda_1 = 0, \lambda_2 = 1\) 并冻结拼图头权重。通过仅在末尾 5 轮引入高强度的动态轨迹对比,在极大抑制额外显存和计算时间消耗的同时,实现了生成保真度的跃升。
实验关键数据¶
主实验¶
在 ImageNet-1K \(256 \times 256\) 基准上,基于 50,000 张采样样本与代表性 Diffusion、AR、NF 模型进行综合对比(CFG 指数取 1.1):
| 方法类别 | 模型 | 参数量 | 训练轮数 | FID (w/ Guidance) | IS (w/ Guidance) | Precision | Recall |
|---|---|---|---|---|---|---|---|
| Pixel-Level NF | TarFlow | 1.4B | 320 | 4.69 | - | - | - |
| Pixel-Level NF | JetFormer | 2.8B | 500 | 6.64 | - | 0.69 | 0.56 |
| Pixel-Level NF | FARMER | 1.9B | 320 | 3.60 | 269.2 | 0.81 | 0.51 |
| Latent-Level Diffusion | DiT-XL/2 | 675M | 1400 | 2.27 | 278.2 | 0.83 | 0.57 |
| Latent-Level Diffusion | SiT-XL/2 | 675M | 1400 | 2.06 | 270.3 | 0.82 | 0.59 |
| Latent-Level AR | MAR-L | 943M | 800 | 1.55 | 303.7 | 0.81 | 0.62 |
| Latent-Level NF | StarFlow | 1.4B | 320 | 2.40 | - | - | - |
| Latent-Level NF | BiFlow | 133M | 160 | 2.39 | 303.0 | - | - |
| Latent-Level NF | SimFlow | 1.4B | 160 | 2.15 | 276.8 | 0.83 | 0.57 |
| Latent-Level NF | JSON (本文) | 1.45B | 160 | 2.10 | 306.2 | 0.82 | 0.58 |
在无需借助任何外部视觉基础模型(如 DINOv2)进行特征监督的纯内生训练模式下,JSON 取得了 NF 架构中的当前最优性能,FID 达到 2.10,IS 达到 306.2,在保持高保真样本分布的同时显著超越了前代 SOTA 方案 SimFlow 与 StarFlow。
消融实验¶
消融实验聚焦于探究拼图头(Jigsaw Head)、原始自博弈损失(\(\mathcal{L}_{\mathrm{sp}}\))、有界自博弈损失(\(\mathcal{L}_{\mathrm{bsp}}\))以及动态门控拼图自博弈损失(\(\mathcal{L}_{\mathrm{jsp}}\))对生成质量与训练稳定性的影响:
| 实验编号 | 拼图头 (Jigsaw) | 原始自博弈 \(\mathcal{L}_{\mathrm{sp}}\) | 有界自博弈 \(\mathcal{L}_{\mathrm{bsp}}\) | 动态拼图自博弈 \(\mathcal{L}_{\mathrm{jsp}}\) | FID (256×256) | Inception Score (IS) | Precision | Recall | 说明与现象 |
|---|---|---|---|---|---|---|---|---|---|
| No. 1 | × | × | × | × | 2.15 | 276.80 | 0.83 | 0.57 | SimFlow 基础基线 |
| No. 2 | ✓ | × | × | × | 2.37 | 284.27 | 0.79 | 0.58 | 仅加入拼图辅助任务 |
| No. 3 | ✓ | ✓ | × | × | N/A | N/A | - | - | 原始自博弈直接崩溃 (NaN) |
| No. 4 | × | × | ✓ | × | 2.46 | 311.64 | 0.81 | 0.58 | 无拼图头,仅有界自博弈 |
| No. 5 | ✓ | × | ✓ | × | 2.19 | 291.85 | 0.82 | 0.56 | 拼图头 + 无条件有界自博弈 |
| No. 6 | ✓ | × | × | ✓ | 2.10 | 306.20 | 0.82 | 0.58 | 完整 JSON 框架 |
此外,针对拼图头投影层架构、特征切片尺度与不同自博弈策略的对比实验如下: - 拼图头结构选择:卷积投影(Conv Proj.,FID 2.10 / IS 306.20)明显优于基于池化的 MLP 映射(FID 2.28 / IS 294.03),印证了保留局部空间拓扑是驱动有效自博弈的前提。 - Patch 大小敏感性:当划分大小 \(p=1\) 时取得最佳指标(FID 2.10 / IS 306.20),增大至 \(p=2\)(FID 2.18 / IS 301.38)和 \(p=4\)(FID 2.47 / IS 290.54)时性能平滑下滑,说明更细粒度的空间重排任务能施加更强的结构约束。 - 博弈策略对比:相较于在 LLM 中用于稳定自博弈的 SPACE 策略(FID 2.25 / IS 286.2)以及发生崩溃的 SPIN 策略,JSON 在 NF 域上表现最为优越。
关键发现¶
- 原始自博弈在 NF 上的脆弱性:消融实验 No.3 证明,盲目惩罚自身高概率样本直接导致可逆流矩阵参数梯度发散;而引入有界下限后(No.4、No.5),训练全程稳定,彻底解决了崩溃难题。
- 空间一致性与似然度的互补性:拼图损失的准确率随训练进程从 0% 稳步提升至近 90%,表明重排任务与最大似然估计高度兼容;在自博弈微调期间,BSP 触发比例由最初的较高水平逐渐回落至收敛稳态,充分印证了合成样本通过拼图筛选对修正生成盲区的关键贡献。
- 少轮次微调的高效性:仅在最后 5 轮引入拼图自博弈即可带来近 30 点的 IS 暴增与 0.05 的 FID 优化,兼顾了端到端训练的计算开销。
亮点与洞察¶
- 将拼图重排内化为无需外挂模型的判别标尺:不同于以往方法依赖昂贵且庞大的预训练大模型进行表征对齐,本文通过在可逆层间植入自监督拼图头,使模型获得了内生的高精度空间结构评估能力。
- 基于理论对立推导的有界自博弈边界:洞察到 NF 最大似然优化与标准自博弈遗忘惩罚的内在冲突,设计了以历史数据似然为下界锚点的三项有界损失函数,为生成模型的偏好优化提供了普适范式。
- 可复用的轨迹质量筛选机制:通过对比生成样本与真实样本在内在辅助任务上的表现来决定微调梯度的启闭,该思路完全可推广至扩散模型、连续流匹配等其他生成范式中,用以规避低质量或异常合成轨迹的负面干扰。
局限与展望¶
- 在线采样的时间开销:自博弈阶段需要在每步前向迭代中由竞争者模型在线反演合成样本,虽然仅在最后 5 轮微调中开启,但单步时间开销仍从 580ms 增加到 1263ms,大规模扩展至超长序列或高分辨率视频时仍面临吞吐瓶颈。
- 拼图粒度与感受野的潜在约束:当前拼图头主要在潜变量空间(16×16)以 \(1 \times 1\) patch 执行,若推广到更高维的潜空间或者视频生成场景,token 置换维度的剧增可能会提高拼图头的学习难度。
- 未来方向:探索异步缓冲池采样策略以重叠推理与训练时间,并将内在几何逻辑自博弈拓展至视频生成(如时空块立体拼图)以及条件控制生成任务中。
相关工作与启发¶
- vs StarFlow / SimFlow: StarFlow 与 SimFlow 奠定了现代隐式归一化流的深浅层设计与因果变换架构,但 SimFlow 需要借助 DINOv2 做 REPA 特征对齐以弥补高层语义不足。本文的 JSON 则在纯内生框架下,利用自监督拼图自博弈在无需外置模型的前提下刷新了 SOTA。
- vs SPIN / SPACE: 语言模型中的 SPIN 易引发自生成分布崩溃,SPACE 采用噪声对比估计进行缓解。本文针对连续生成流设计了有界自博弈损失(\(\mathcal{L}_{\mathrm{bsp}}\)),并首次引入辅助重排任务动态选优,提供了更稳健的轨迹对齐途径。
- vs Diffusion Trajectory Fine-tuning: 扩散模型拥有预设的加噪去噪动力学,易于定义轨迹微调奖励;归一化流缺乏显式轨迹标签,本文证明了空间几何自监督是解决此类隐式动力学模型轨迹评估的有效途径。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将自博弈与空间拼图重排任务引入归一化流,开创了无外部先验下的 NF 轨迹微调新范式。
- 实验充分度: ⭐⭐⭐⭐☆ 在 ImageNet-1K 256×256 上对比详实,提供了全面的组件消融、超参敏感性分析及训练动态可视化。
- 写作质量: ⭐⭐⭐⭐⭐ 动机阐述清晰深刻,数学公式推导自洽严密,问题定义与解决方案高度呼应。
- 价值: ⭐⭐⭐⭐☆ 突破了现代可逆流模型受制于表征约束的瓶颈,为轻量、高效的非扩散生成模型演进提供了重要借鉴。