JSON: Jigsaw Self-play Optimization for Normalizing Flows

Normalizing Flows (NFs) is a principle generative framework,which learns to generate images by establishing a bijective mapping tra-jectory between noise and data with likelihood maximization. In thispaper, we leverage the self-play paradigm, a general trajectory fine-tuningmethod to improve generative capability of NFs without auxiliary fine-tuning data. The general self-play mechanism formulates competitionbetween model and its historical counterpart in favoring real data trajec-tory (i.e., “wining trajectory”) while forgetting trajectory of self-generatedsynthesized data (i.e., “losing trajectory”), improving model capabilitythrough preference alignment. However, NF-based self-play faces twochallenges. (1) Enforcing NF to forget data synthesized with historicalmodel contradicts with NF’s original training goal, leading to potentialtraining collapse. (2) Self-play is naturally plausible for generative modelwith pre-defined data trajectory (e.g., diffusion models), but faces com-patibility problem with NFs due to its implicit wining data trajectorydefinition. We thus propose jigsaw self-play optimization for NFs (JSON)to overcome these challenges. Our core contribution is the integration ofa jigsaw puzzle reassembly task to serve as an intrinsic spatial coherenceevaluator, identifying “winning” trajectories that exhibit superior spatiallogic for self-play fine-tuning. We further propose a bounded self-play lossand a nested-loop optimization strategy to mitigate training collapse andensure stable fine-tuning. By aligning the model with high-quality struc-tural anchors mined during self-play, JSON achieves promising results onImageNet-1K, providing a robust framework for high-fidelity NF-basedgeneration.