JSON: Jigsaw Self-play Optimization for Normalizing Flows

Conference: ECCV 2026
Paper: ECCV Official Link
Code: https://media.eventhosts.cc/Conferences/ECCV2026/pdfs/2253.pdf
Area: Image Generation
Keywords: Normalizing Flows, Self-play Optimization, Jigsaw Reassembly, Trajectory Alignment, Spatial Coherence

TL;DR

Addressing the lack of predefined diffusion dynamics in Normalizing Flows (NFs) and the training collapse induced by standard self-play objectives, this paper proposes JSON—a framework integrating an intrinsic jigsaw puzzle reassembly evaluator with bounded self-play loss to achieve state-of-the-art generation performance without external vision priors.

Background & Motivation

Modern generative models have witnessed substantial progress across diffusion models, autoregressive models, and normalizing flows. Normalizing Flows (NFs) establish a bijective mapping between a standard Gaussian prior and real data distributions based on the change of variable formula. Compared to diffusion or continuous-time flow matching models requiring iterative sampling steps, modern latent NFs (such as TarFlow and StarFlow) establish transformations via causal transformers, offering exact likelihood estimation and fast single-pass inference. However, strictly enforcing invertibility and tractable Jacobian determinant computation restricts the architectural expressiveness of NFs. Consequently, their maximum likelihood estimation (MLE) training lacks explicit trajectory constraints, hindering high-level semantic representation and global spatial coherence learning.

In the domain of diffusion models, trajectory fine-tuning (TFT) and self-play paradigms have shown remarkable success by pitting a model against its earlier checkpoint to favor real data dynamics while penalizing synthetic trajectories, all without requiring expensive human feedback. However, directly applying self-play to NFs faces two fundamental obstacles. First, NF models are trained via global density maximization; synthetic samples produced by a preceding checkpoint naturally lie in the current model's high-likelihood region. Forcing the NF model to forget these self-generated samples penalizes valid probability mass, contradicting the core MLE objective and triggering severe training collapse (e.g., NaN loss values). Second, unlike diffusion models equipped with predefined forward-backward schedules, NFs lack explicit physical trajectory anchors. Relying solely on raw likelihood fails to capture high-level spatial logic and structural coherence.

To overcome these barriers, this work investigates whether an intrinsic, verifiable geometric coherence task can be integrated into NFs to distinguish "winning" trajectories without incurring the overhead of external pretrained vision encoders like DINOv2. Core idea: incorporate a self-supervised jigsaw puzzle reassembly task as an intrinsic spatial coherence evaluator alongside a bounded self-play loss and nested-loop training to stably align normalizing flows toward structurally sound generation trajectories.

Method

Overall Architecture

The JSON framework operates across a two-stage training scheme combined with a nested-loop optimization strategy. In Stage 1 (NF pretraining), lightweight convolutional jigsaw heads are attached to intermediate NF blocks, training the feature representations to accurately predict original spatial coordinates from deliberately permuted patches. In Stage 2 (jigsaw self-play), the jigsaw heads are frozen to serve as an intrinsic evaluator of spatial logic, guiding competition between the main model and a frozen competitor snapshot from the previous epoch.

The full pipeline resides in latent space: input images are mapped into latent vectors using a pretrained VAE encoder, and an autoregressive causal transformer calculates the diagonal Jacobian transformations for density estimation. During self-play fine-tuning, both real latents and competitor-generated latents are evaluated by the frozen jigsaw heads, which dynamically determine whether to penalize the synthetic sample via bounded self-play or skip penalization when the synthetic sample exhibits superior spatial logic.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Real Latent x and Class Condition c"] --> B["Stage 1: Spatial Coherence Learning<br/>Pretrain Intermediate Jigsaw Heads"]
    B --> C["Stage 2: Intrinsic Spatial Evaluation<br/>Compute Cumulative Reassembly Error"]
    C --> D{"Evaluate Spatial Coherence<br/>L_jig(x) ≤ L_jig(x') ?"}
    D -->|Yes: Synthetic Trajectory Degraded| E["Bounded Self-play Alignment<br/>Two-Sided Constraint Prevents Collapse"]
    D -->|No: Synthetic Trajectory Superior| F["Bypass Penalization<br/>Retain Coherent Structural Features"]
    E --> G["Nested-Loop Optimization<br/>Inner-Loop Step Play + Outer-Loop Epoch Sync"]
    F --> G

Key Designs

1. Spatial Coherence Learning: Embedding a Verifiable Jigsaw Auxiliary Head

Standard NF training optimizes only MLE objectives via determinant sums and latent norms, lacking fine-grained spatial position awareness across local tokens. To address this, a convolutional jigsaw head \(g_i\) is attached to the output feature \(o_{i+1} = f_i(x_i) \in \mathbb{R}^{HW \times D}\) of each invertible block \(f_i\). During pretraining, the spatial feature map is partitioned into \(p \times p\) patches (with \(p=1\) by default) and indexed with ground-truth coordinates \(\mathcal{Y}_m\). These tokens are randomly permuted along the spatial dimensions into \(\tilde{o}_{i+1}\), and the jigsaw head learns to classify their original locations: $$ \mathcal{L}{\mathrm{jig},i}(x;\theta) = -\sum} \mathcal{Ym \odot \log \mathrm{softmax}(g_i(\tilde{o})) $$ This design establishes an intrinsic, deterministic supervision signal that requires no external reward models or feature extractors. It forces the causal flow transformations to preserve clear spatial structure, and once frozen, it provides a rigorous metric for evaluating generative trajectory quality.

2. Bounded Self-play Loss: Establishing Stable Constraint Boundaries

Directly importing the Bradley-Terry preference model into NFs (basic self-play, B-NSP) minimizes the loss on real data while enforcing unconstrained penalization on synthesized data \(x'\), requiring \(p_\theta(x') < p_{\theta_k}(x')\). Because \(x'\) is sampled from \(\theta_k\)'s high-likelihood region, this penalty conflicts directly with MLE and causes gradient explosion and NaN failures. The proposed bounded self-play (BSP) introduces the real data likelihood under \(\theta_k\) as a lower bound, modifying the relation to \(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} $$ where }}(x';\theta_k) - \mathcal{L}_{\mathrm{nf}\(\sigma(t) = -\log(1/(1 + \exp(-t)))\). This reformulation balances forgetting sub-optimal historical generations with remaining anchored within the model's valid high-likelihood density manifold, guaranteeing optimization stability.

3. Intrinsic Spatial Evaluation: Gated Jigsaw Self-play Objective

Applying bounded self-play blindly to every sample risks penalizing high-quality synthetic data that might exhibit even better semantic completeness than cropped real images. JSON sums the jigsaw losses across all blocks, \(\mathcal{L}_{\mathrm{jig}}(x;\theta) = \sum_i \mathcal{L}_{\mathrm{jig},i}(x;\theta)\), and compares real and generated samples to formulate the final jigsaw self-play objective: $$ \mathcal{L}{\mathrm{jsp}}(x,x',\theta_k;\theta) = \begin{cases} \mathcal{L}}}(x, x', \theta_k; \theta), & \text{if } \mathcal{L{\mathrm{jig}}(x;\theta) \leq \mathcal{L}(x';\theta) \ 0, & \text{otherwise} \end{cases} $$ When synthetic samples demonstrate better spatial coherence than their real-world counterparts (lower jigsaw loss), the model bypasses self-play penalization, preserving advantageous structural logic mined during self-generation.}

4. Nested-Loop Optimization Strategy: Smooth Distribution Tracking

Optimization alternates between an inner loop and an outer loop. In the inner loop, the competitor model \(\theta_k\) is frozen, performing on-the-fly conditional inference to synthesize \(x'\) for each training step and driving the step-level trajectory alignment of \(\theta\). In the outer loop, after completing an entire training epoch, the competitor weights are updated with the latest main model parameters (\(\theta_k \leftarrow \theta\)). This nested strategy allows the competitor to track the evolving distribution smoothly while capping the computational cost of on-the-fly generation.

Loss & Training

The total objective unifies VAE reconstruction, normalizing flow negative log-likelihood, jigsaw classification, and self-play alignment: $$ \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) $$ Training spans }}(x;\theta) + \lambda_2 \mathcal{L}_{\mathrm{jsp}\(E=160\) epochs on 16 NVIDIA H20 GPUs with a global batch size of 160: the first 155 epochs represent Stage 1 pretraining (\(\lambda_1 = 1, \lambda_2 = 0\)), establishing a stable latent manifold and training the jigsaw heads; the final 5 epochs switch to Stage 2 self-play (\(\lambda_1 = 0, \lambda_2 = 1\)) with frozen jigsaw heads. Restricting self-play to the last 5 epochs limits total training overhead while delivering large gains in generation quality.

Key Experimental Results

Main Results

On the ImageNet-1K \(256 \times 256\) benchmark with 50,000 generated samples (CFG scale 1.1), JSON is compared against representative diffusion, autoregressive, and normalizing flow models:

Method Category	Model	#Params	Epochs	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 (Ours)	1.45B	160	2.10	306.2	0.82	0.58

Without using external vision encoders (e.g., DINOv2) for feature supervision, JSON sets a new state of the art for NF architectures with an FID of 2.10 and an Inception Score of 306.2, outperforming previous leading models including SimFlow and StarFlow.

Ablation Study

The ablation investigates the individual impact of the jigsaw head, unconstrained self-play (\(\mathcal{L}_{\mathrm{sp}}\)), bounded self-play (\(\mathcal{L}_{\mathrm{bsp}}\)), and dynamic jigsaw self-play (\(\mathcal{L}_{\mathrm{jsp}}\)):

No.	Jigsaw Head	\(\mathcal{L}_{\mathrm{sp}}\)	\(\mathcal{L}_{\mathrm{bsp}}\)	\(\mathcal{L}_{\mathrm{jsp}}\)	FID (256×256)	Inception Score (IS)	Precision	Recall	Note
No. 1	×	×	×	×	2.15	276.80	0.83	0.57	SimFlow baseline
No. 2	✓	×	×	×	2.37	284.27	0.79	0.58	Jigsaw auxiliary head only
No. 3	✓	✓	×	×	N/A	N/A	-	-	Standard self-play collapses (NaN)
No. 4	×	×	✓	×	2.46	311.64	0.81	0.58	Bounded self-play without jigsaw
No. 5	✓	×	✓	×	2.19	291.85	0.82	0.56	Bounded self-play with jigsaw
No. 6	✓	×	×	✓	2.10	306.20	0.82	0.58	Full JSON framework

Additional design ablations demonstrate: - Jigsaw Head Architecture: A convolutional projection head (Conv Proj., FID 2.10 / IS 306.20) outperforms a 3-layer MLP with spatial pooling (FID 2.28 / IS 294.03) because pooling dilutes localized spatial cues. - Patch Size Sensitivity: Smaller patch sizes provide finer spatial supervision; setting \(p=1\) achieves FID 2.10 and IS 306.20, outperforming \(p=2\) (FID 2.18 / IS 301.38) and \(p=4\) (FID 2.47 / IS 290.54). - Alternative Self-play Strategies: Compared to SPIN (which collapsed) and SPACE (FID 2.25 / IS 286.2), JSON achieves superior generation metrics on normalizing flows.

Key Findings

• Vulnerability of Vanilla Self-Play in NFs: Ablation No. 3 confirms that unconstrained penalization of self-generated samples severely violates MLE probability modeling and causes immediate training divergence; introducing the lower bound in \(\mathcal{L}_{\mathrm{bsp}}\) resolves this collapse.
• Compatibility of Jigsaw and MLE: Reordering accuracy rises steadily from 0% to nearly 90% across training, showing that jigsaw puzzle reassembly is fully compatible with standard likelihood maximization.
• Efficiency of Short Fine-Tuning: Engaging self-play for only the final 5 epochs improves the Inception Score by nearly 30 points and drops the FID by 0.05, keeping training overhead manageable while substantially enhancing fidelity.

Highlights & Insights

• Intrinsic Verification without External Teachers: Unlike methods relying on heavy external vision backbones for trajectory supervision, JSON embeds lightweight self-supervised jigsaw heads into the flow blocks to yield an endogenous spatial coherence metric.
• Bounded Objective Resolving Likelihood-Preference Contradiction: Recognizing that standard self-play objectives penalize valid high-density modes in likelihood-based models, the paper introduces a bounded loss with historical data likelihood anchoring, providing a sound paradigm for NF alignment.
• Quality-Aware Trajectory Gating: Comparing auxiliary task performance between generated and real samples creates a natural filter that protects high-quality synthetic generations from unneeded penalization.

Limitations & Future Work

• Inference Overhead During Self-Play: Online sample synthesis increases single-step iteration time from 580.6ms to 1263.2ms during the final 5 epochs. Scaling to higher resolutions or video synthesis will require more efficient generation pipelines.
• Scaling to High-Dimensional Tokens: Jigsaw reassembly is currently executed on \(16 \times 16\) latent maps with \(1 \times 1\) patches. Extending this to spatio-temporal tokens in video flows may increase combinatorial reassembly complexity.
• Future Directions: Exploring asynchronous sample replay buffers to decouple inference from training, and extending geometric self-play to video generation and fine-grained controllable synthesis.

Related Work & Insights

• vs StarFlow / SimFlow: StarFlow and SimFlow establish scalable deep-shallow causal architectures for latent NFs; SimFlow relies on external DINOv2 encoders for feature alignment. JSON achieves superior generative performance in an entirely self-contained framework without external foundation models.
• vs SPIN / SPACE: In language models, SPIN is prone to collapse and SPACE stabilizes training via noise-contrastive estimation. JSON introduces a bounded loss tailored to density models and pairs it with deterministic spatial evaluation.
• vs Diffusion Trajectory Fine-Tuning: Diffusion models benefit from predefined trajectories for reward assignment, whereas NFs lack explicit trajectory anchors. JSON shows that geometric self-supervised tasks offer an effective alternative for evaluating implicit trajectories.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Pioneering introduction of self-play and jigsaw puzzle reassembly to normalizing flows, creating an endogenous trajectory fine-tuning paradigm.
• Experimental Thoroughness: ⭐⭐⭐⭐☆ Comprehensive evaluation on ImageNet-1K \(256 \times 256\), with rigorous ablations on loss formulations, architectures, and hyperparameters.
• Writing Quality: ⭐⭐⭐⭐⭐ Clear motivation, rigorous mathematical formulation, and well-structured empirical analysis.
• Value: ⭐⭐⭐⭐☆ Demonstrates that normalizing flows can achieve competitive generation fidelity without external teachers, providing valuable insights for likelihood-based generative models.