NeurIPS 2026 Papers

In this paper, we present gfnx, a fast and scalable package for training and evaluating Generative Flow Networks (GFlowNets) written in JAX. gfnx provides an extensive set of environments and metrics for benchmarking, accompanied with single-file implementations of core objectives for training GFlowNets. We include synthetic hypergrids, multiple sequence generation environments with various editing regimes and particular reward designs for molecular generation, phylogenetic tree construction, Bayesian structure learning, and sampling from the Ising model energy. Across different tasks, gfnx achieves significant wall-clock speedups compared to Pytorch-based benchmarks (such as torchgfn library) and author implementations. For example, gfnx achieves up to 55 times speedup on CPU-based sequence generation environments, and up to 80 times speedup with the GPU-based Bayesian network structure learning setup. Our package provides a diverse set of benchmarks and aims to standardize empirical evaluation and accelerate research and applications of GFlowNets.

Masked diffusion language models generate text by iteratively unmasking positions, with practical samplers often selecting reveal positions from the model's own logits. The reveal order is therefore an endogenous latent variable of the rollout process with an enormous discrete latent space. Policy optimization objectives relying on pre-defined schedules or sampled trajectories introduce either bias or high variance. In this paper, we propose \texttt{FEROM}, \emph{Frontier Endogenous Reveal-Order Marginal Policy Optimization}, which targets the rollout-induced marginal response policy of masked diffusion LMs. \texttt{FEROM} derives a Rao--Blackwellized policy-gradient identity over latent reveal orders and expresses the resulting estimator as posterior edge occupancy on a reveal-state DAG. To make marginalization practical, we introduce Frontier Reveal Marginalization, a budgeted estimator that combines scorelaw local reveal estimation with frontier expansion of high-mass partial states. Integrated into a GRPO-style objective, \texttt{FEROM} replaces single-path log-scores with a locally marginalized edge-based surrogate. Experiments on math and coding tasks show comparable or improved results over existing methods under matched compute budgets. Offline proxy study further shows potential in gains with increased budgets.

Concept unlearning in text-to-image diffusion models aims to suppress a target concept, e.g., \texttt{horse}, while preserving the model's ability to generate semantically related but distinct content, like \texttt{donkey}. Yet existing methods either leak under indirect prompts or visibly degrade remaining concepts. {\em We show that these failure modes arise naturally from overlapping concept representations.} Our main contribution is demonstrating that the observed trade-offs stem from the geometry of concept representations, rather than from weaknesses in any particular algorithm. We support this claim through both theoretical analysis and empirical observations. By formalizing concepts as activation-space regions, we show that overlap between a target and other concepts lower-bounds the unavoidable degradation on those other concepts when the target is erased. Our analysis suggests that strong erasure and preservation become fundamentally coupled when concepts occupy overlapping activation regions, with the trade-off scaling linearly in the degree of overlap. We verify this trade-off on seven unlearning methods spanning fine-tuning, adversarially-robust fine-tuning, and inference-time interventions. Averaged across target concepts, STEREO almost completely suppresses the target concept under indirect prompts but cuts the model's ability to generate semantically related concepts by more than 75\%. On the other hand, sparse inference-time methods (SAeUron, SEOT) better preserve utility but leave substantial target leakage. The trade-off is steep for concepts whose internal representations heavily overlap with their evaluated semantic neighborhoods, (e.g., \texttt{horse} with \texttt{cat}, \texttt{dog}, and \texttt{bear}) and milder for \texttt{castle}, which is comparatively less entangled within our concept pool and probe layer in activation space. Across targets, damage to related concepts scales monotonically with our activation-overlap measure. These results suggest that perfect unlearning is the wrong target for entangled concepts. The field should evaluate on the Pareto frontier our theorem establishes; current benchmarks, which decouple unlearning accuracy from utility preservation, hide this trade-off and need to be revised.

High-dynamic-range (HDR) formats and displays are becoming increasingly prevalent, yet state-of-the-art image generators (e.g., Stable Diffusion and FLUX) typically remain limited to low-dynamic-range (LDR) output due to the lack of large-scale HDR training data. In this work, we show that existing pretrained diffusion models can be easily adapted to HDR generation without retraining from scratch. A key challenge is that HDR images are natively represented in linear RGB, whose intensity and color statistics differ substantially from those of sRGB-encoded LDR images. This gap, however, can be effectively bridged by converting HDR inputs into perceptually uniform encodings (e.g., using PU21 or PQ). Empirically, we find that LDR-pretrained variational autoencoders (VAEs) reconstruct PU21-encoded HDR inputs with fidelity comparable to LDR data, whereas linear RGB inputs cause severe degradations. Motivated by this finding, we describe an efficient adaptation strategy that freezes the VAE and finetunes only the denoiser via low-rank adaptation in a perceptually uniform space. This results in a unified computational method that supports both text-to-HDR synthesis and single-image RAW-to-HDR reconstruction. Experiments demonstrate that our perceptually encoded adaptation consistently improves perceptual fidelity, text-image alignment, and effective dynamic range, relative to previous techniques.

Semi-dense local feature matching commonly aggregates contextual information before coarse-to-fine correspondence estimation. However, intra-image and inter-image aggregation serve different roles: the former should propagate spatial evidence within each image while retaining structured two-dimensional dependencies, whereas the latter should gather information from plausible corresponding regions while limiting noise from unrelated locations. We propose WovenAnchor Matcher (WAM), a semi-dense matching framework based on role-specialized context aggregation. WAM introduces two complementary operators. WovenMamba performs intra-image aggregation through horizontal-then-vertical state-space propagation, allowing vertical updates to operate on horizontally contextualized features and thereby encouraging a structured two-dimensional receptive-field bias. AnchorCrossAttention performs inter-image aggregation by using cross-attention to gather context from plausible corresponding regions across the image pair. To make this retrieval robust, it operates on local representative anchors rather than dense point-wise tokens, reducing sensitivity to noisy affinities that can otherwise lead to incorrect correspondence propagation. On MegaDepth, WAM achieves 66.1 pose AUC@5$^\circ$ with a runtime of 35.4 ms on a single H100 GPU, improving over JamMa, which achieves 64.1 pose AUC@5$^\circ$ at 43.8 ms under the same setup. Ablations show that removing or replacing either specialized component reduces accuracy, providing evidence that intra-image propagation and inter-image retrieval benefit from role-specific aggregation designs.

Single-image 3D methods often trade off pixel alignment and completeness: monocular reconstructors stay registered to visible pixels but stop at the first surface, while image-to-3D generators complete shapes in a canonical frame that is no longer aligned to the input. We introduce methodlong (method), a pixel-aligned multilayer geometry representation that keeps generated 3D in the input camera frame while completing geometry beyond the visible surface. For each input ray, method predicts an ordered stack of camera-space 3D points, where the first entry explains the observed surface and subsequent entries complete occluded surfaces along the same pixel ray. methodnet, a flow-matching diffusion transformer with a frozen 2D foundation encoder, learns this representation from depth-filled dense targets, removing the need for a separate per-layer visibility-mask head. The same objective covers static objects, static scenes, and dynamic clips, with temporal attention added for video. Because the output preserves the input camera frame, intrinsics, and pixel correspondences, it provides a direct 2D-to-3D interface for downstream systems, including text-driven 3D scene editing, geometry-conditioned novel-camera video generation, and training-free integration with existing textured-mesh generators. Across object, scene, and dynamic benchmarks, method improves visible-surface accuracy and complete-geometry quality over strong reconstruction and generation baselines.

Equipping artificial agents with spatial intelligence requires a comprehensive generative prior over the dynamic 3D world. We propose World Motion Models (WMMs) that capture "what was, is, and will be where across time" via sparse SE(3) pose trajectories. WMMs are built on the observation that elements of dynamic scenes can be well-approximated by a set of rigid SE(3) trajectories, a minimal yet expressive primitive for 4D modeling. This representation unifies articulated objects, human bodies, hand-object interactions, piecewise-rigid scene dynamics, camera motion, and even robot states and actions into a single shared space. Given this representation, we cast the joint distribution of these entities as a flexible sequence modeling problem, utilizing flow-matching with per-token noise levels. Coupled with a context token mechanism with non-sequential conditioning, this formulation supports any-to-any marginal conditioning across an arbitrary number of entities and time steps. Tasks such as future prediction, motion infilling, model-predictive control, inverse kinematics, cross-embodiment retargeting, and policy learning all reduce to the application of different masks over the same network. Experiments on diverse 6 applications of 3D vision and robotics demonstrate the versatility and flexibility of WMMs with strong performance.

Token merging has emerged as an effective training-free strategy for accelerating diffusion models by reducing redundant token computation while minimizing generation degradation. A key challenge is determining which tokens can be safely merged and which should be preserved for subsequent computation. Existing method addresses this by preserving prompt-aligned foreground regions, thereby protecting semantically relevant structures. However, this text-relevance criterion overlooks background objects, local details that are not explicitly mentioned in the prompt, leading to degraded structural coherence. In this work, we propose Dynamics-aware Token Merging (DaTMe), a training-free framework that grounds token merging in the underlying diffusion dynamics. Our key insight is that token importance should reflect whether a token remains unresolved during denoising and whether it carries structural information. To this end, DaTMe constructs a dynamics-aware importance map from the Tweedie denoised estimate. The map combines a temporal signal, measuring inter-step changes in the predicted clean image, with a structural signal, capturing local high-frequency content such as boundaries and fine details. Using this map, DaTMe adaptively assigns tokens into merging roles under a given compression ratio: important tokens are preserved as anchors, while tokens that are both temporally stabilized and structurally homogeneous are selected as safe merge candidates. Experiments on PixArt-$\alpha$ and FLUX demonstrate that DaTMe maintains visual fidelity across the entire image, highlighting diffusion dynamics as an effective token-level merging decision in efficient generation.

Text-to-image diffusion models fail on semantically rare attribute--object compositions, but it is unclear whether this brittleness arises during denoising or is already present in conditioning representations. We study this question at three levels of granularity. First, using a controlled stress protocol across SD1.5, SDXL, and SD3, we show that a Compositional Stress Index derived from text-side non-factorization residuals separates anchor from stress prompts with perfect discrimination (AUC${=}1.0$) cross-model, establishing a robust pre-denoising risk axis. Second, embedding-level low-rank repair reduces representation residuals without any positive semantic improvement, while a downstream cross-attention intervention succeeds on the same subset---a site-sensitive boundary between diagnosis and control. Third, we decompose this boundary at block-group resolution under three intervention modalities. Diagnostic accessibility peaks at the deep encoder (D2) across all conditions, but token-level intervention of either sign produces no measurable improvement at any block. Only broad cross-attention ablation moves semantics at D2, while decoder blocks respond to selective boost and collapse under ablation. Together, these results establish that compositional defects are diagnosable before denoising, but \textbf{where to look is not how to fix}: the diagnostic site and the control site require different intervention modalities, and even at the correct site, the wrong modality is futile.

Diffusion-based generative attacks have emerged as a promising paradigm for targeted adversarial example generation, leveraging generative priors to improve transferability and visual fidelity. However, existing methods typically entangle image reconstruction and target manipulation within the same reverse denoising trajectory, where target guidance can disrupt early structural recovery. This coupling creates a persistent trade-off among attack effectiveness, transformation robustness, and perceptual fidelity. In this paper, we propose SDTI (Stage Decoupled Target Injection), a stage-decoupled reverse diffusion framework that separates source-structure preservation from adversarial target injection along the denoising timeline. SDTI first maps the input image to an intermediate latent through null-text DDIM inversion, then recovers and anchors the source structure with a frozen null-text denoising prefix, and finally activates target-conditioned LoRA guidance only in the late denoising suffix to inject adversarial target semantics. To confine adversarial optimization to this suffix phase, we further introduce suffix-wise supervision with inter-step gradient truncation. Extensive experiments on ImageNet-NeurIPS, robust victim models, common input transformations, and MS-COCO cross-domain transfer show that SDTI achieves a favorable attack--fidelity--robustness trade-off, improving transformation robustness and perceptual quality while maintaining competitive targeted transferability. The code is available at: https://anonymous.4open.science/r/AY_code-1431.

Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on non-Euclidean domains. Existing generative methods largely assume Euclidean geometry and fail to capture this structure. We introduce Riemannian Wasserstein Entropic Flow Matching (RWEFM), a generative framework on the Wasserstein space $\mathcal{P}_2(\mathcal{M})$ of a Riemannian manifold $(\mathcal{M},g)$. RWEFM is trained by regressing a neural vector field onto Riemannian optimal transport velocities, using McCann displacement interpolations as conditional paths. We confirm theoretically that this construction leads to a valid flow matching approach on $\mathcal{P}_2(\mathcal{M})$ and introduce the Riemannian Entropic Map, a GPU-efficient approximation of the optimal transport map on manifolds. Our experiments show that by respecting the intrinsic geometry of the data, RWEFM can generate whole single-cell samples in hyperspherical latent spaces and protein conformational ensembles on the torus. Code and tutorials are available at \href{https://github.com/AnonRWEFM/RWEFM}{RWEFM}.

RLHF is widely used to align flow-matching text-to-image models with human preferences, but often leads to severe diversity collapse after fine-tuning. In RL, diversity is often assumed to correlate with policy entropy, motivating entropy regularization. However, we show this intuition breaks in flow models: policy entropy remains constant, even while perceptual diversity collapses. We explain this mismatch both theoretically and empirically: the constant entropy arises from the fixed, pre-defined noise schedule, while the diversity collapse is driven by the mode-seeking nature of policy gradients. As a result, policy entropy fails to prevent the model from converging to a narrow high-reward region in the perceptual space. To this end, we introduce perceptual entropy that captures diversity in a perceptual space and maintains the property of standard entropy. Building upon this insight, we propose two entropy-regularized strategies, Perceptual Entropy Constraint and Perceptual Constraints on Generation Space, to preserve perceptual diversity and improve the quality. Experiments across two base models, neural and rule-based rewards, and three perceptual spaces demonstrate consistent gains in the quality-diversity trade-off; PEC achieves the best overall score of 0.734 (vs.\ baseline's 0.366); a complementary setting of PEC further reaches a diversity average of 0.989 (vs.\ baseline's 0.047). Code will be released.

Discrete graph diffusion models corrupt graph structure through independent edge updates with a Markovian noise model. Combined with the exchangeability requirement on the output distribution, the terminal prior is necessarily an Erd\H{o}s--R\'enyi (E-R) random graph. This has a fundamental consequence for unconditional graph generation without node features: in the large graph limit, any GNN denoiser applied to an E-R graph approximately outputs an E-R graph, making effective denoising difficult. In practice, generated graphs fail to recover meaningful node heterogeneity and higher-order structure. Previous models introduced specialized node features to compensate, framing this as an expressivity problem, while the role of the prior remained unexamined. We instead address the root cause by introducing the \emph{Latent Sociability Prior} (LSP) and a \emph{joint} diffusion process that co-evolves latent node sociability variables and graph structure, removing edge independence while preserving exchangeability. Unlike existing latent graph diffusion methods, the sociability mechanism requires no complex encoder: its sole role is to guide edge updates and preserve node heterogeneity throughout the trajectory. We validate our approach on multiple graph generation datasets and show that the proposed framework more accurately captures graph distributions than baseline methods without node features.

Latent Diffusion Models dominate generative modeling but face two-stage training complexity and the inherent reconstruction limits of pre-trained autoencoders. Pixel-space Flow Matching offers a simplified end-to-end alternative yet often encounters training instabilities when using the standard $v$-prediction objective. Recent research attributes this difficulty to the off-manifold nature of noised targets, prompting a widespread shift toward $x$-prediction. In this paper, we carefully investigate this prevailing assumption through patch-wise Principal Component Analysis (PCA), and we reveal that the failure of $v$-prediction is primarily a bandwidth issue. Specifically, the velocity field in minor components is mostly determined by the raw noisy input, which differs from the dynamics in the principal component space. Deep backbones must propagate this minor information to the final layer to predict the velocity field, which consumes significant model capacity that should be dedicated to principal components. This is particularly critical in pixel space where noise dimensionality is comparable to the hidden dimension. To resolve this, we simply incorporate a term containing the raw noisy input into the final layer. This architectural bypass enables the backbone to focus on high-level semantics while the final layer efficiently reconstructs the minor velocity field. Our method successfully revives $v$-prediction and achieves competitive performance on ImageNet $256\times256$ dataset compared to $x$-prediction baselines.

Recent work has shown that flow matching models can be trained without explicit time conditioning, challenging the standard view that the interpolation time is needed to disambiguate velocity targets. But why should a time-blind model work at all? Decomposing the time-blind flow matching loss, we identify two sources of irreducible error: a *coupling variance*, which arises from ambiguous velocity targets induced by how noise and data points are paired, and the *time-blindness gap*, which is the additional error caused by ignoring time. This gap shows that time-blind training is strictly harder than conventional training, reinforcing the puzzle that time-blind models work so well in practice. We resolve this tension by showing that the geometry of high-dimensional data makes time identifiable directly from noisy observations. When data concentrates near a $k$-dimensional subspace, time can be recovered from the statistical structure of noisy interpolants in directions orthogonal to the data; under a spiked-covariance model, this yields a closed-form estimator that recovers $t$ from a single observation $z$ at rate $O(1/\sqrt{d-k})$ for ambient dimension $d$. As a consequence, we prove that the time-blindness gap is asymptotically negligible relative to the coupling variance. We empirically demonstrate our identifiability result on real-world data and show that changing the coupling has a much larger effect on loss and sample quality than removing time conditioning across CIFAR-10, CelebA-HQ, and FFHQ. These results explain why time-blind flow matching works and show that the main practical lever is the choice of coupling, not explicit time conditioning.

Tokenizers are a crucial component of latent diffusion models, as they define the latent space in which diffusion models operate. However, existing tokenizers are primarily designed to improve reconstruction fidelity or inherit pretrained representations, leaving unclear what kind of latent space is truly friendly for generative modeling. In this paper, we study this question from the perspective of latent manifold organization. By constructing controlled tokenizer variants, we identify three key properties of a diffusion-friendly latent manifold: coherent spatial structure, local manifold continuity, and global manifold semantics. We find that these properties are more consistent with downstream generation quality than reconstruction fidelity. Motivated by this finding, we propose the $\textit{\textbf{P}rior-\textbf{A}ligned Auto\textbf{E}ncoder} (\textbf{PAE})$, which explicitly shapes the latent manifold instead of leaving diffusion-friendly manifold to emerge indirectly from reconstruction or inheritance. Specifically, PAE leverages refined VFM-derived priors and perturbation-based regularization to turn spatial structure, local continuity, and global semantic organization into explicit training objectives. On ImageNet $256{\times}256$, PAE improves both training efficiency and generation quality over existing tokenizers, reaching comparable performance up to $13\times$ faster than RAE under the same LightningDiT setup and achieving a new state-of-the-art gFID of $\textbf{1.03}$. These results highlight the importance of organizing the latent manifold for latent diffusion models.

The manifold hypothesis suggests that real-world data concentrates on low-dimensional structures, yet navigating meaningful trajectories on latent manifolds remains an open challenge. A common approach is to construct density-aware metrics, but such constructions bundle two distinct questions: (i) where valid data lies and (ii) how we choose to traverse on the manifold. Since different tasks may demand different traversal behaviors even on the same data manifold, specifying these two aspects independently offers additional modeling flexibility. We propose to factor the path objective into a kinetic term governed by a freely chosen geometric prior and a score-based potential, derived from the negative log-density of a pretrained diffusion model, that acts as a soft constraint for manifold adherence. The potential encourages dynamics to evolve along the data support, while the geometric prior remains free to encode task-specific notions of path optimality. To solve the resulting optimization problem, we introduce \textit{geodesic force matching} (GFM), a direct-collocation algorithm that discretizes the trajectory and optimizes all waypoints jointly to satisfy the Euler--Lagrange force balance in a least-squares sense. On synthetic manifolds and 3D shape interpolation, our method produces smooth, on-manifold paths and performs favorably against recent density-aware baselines, most notably in the low-noise regime where those baselines become numerically unstable.

Accurate emulation of multi-scale physical systems governed by PDEs demands models that remain stable over long autoregressive rollouts while preserving fine-scale structures. Deterministic emulators produce overly-smoothed predictions, while generative approaches better capture details but are costly. Latent-space generative models have emerged as a compromise but with the additional cost of separately pre-trained autoencoders. We propose **Wavelet Flow Matching** (WFM), a novel generative emulator that overcomes current trade-offs between cost and skill by performing optimal-transport directly in the multi-scale wavelet space. Rather than learning a latent compression, WFM leverages the hierarchical structure of a U-Net to jointly predict transport velocities of a prescribed wavelet representation. On three challenging systems of chaotic fluid dynamics, WFM achieves superior long-horizon stability, accuracy and spectral coherence compared to state-of-the-art models. Our results clearly position the wavelet space as an effective training-free representation for generative emulation of complex physical dynamics.

Recent advances in latent diffusion models have enabled high-quality image generation, but also raise critical concerns for intellectual property protection in model distribution scenarios, where downstream users have unrestricted access to models, allowing arbitrary modifications. Existing watermarking methods either rely on inference-time control of inputs (e.g., specific prompts or noise initialization) or embed watermark signals in auxiliary components, making them easily removable in such settings. In this paper, we propose INMARK, which treats watermarking as an intrinsic property learned by the model. Instead of relying on input control or auxiliary watermarking components, INMARK enables the core denoising network to internalize and reproduce watermark patterns. Extensive experiments demonstrate that INMARK achieves strong generation fidelity, high watermark detectability, and robustness against attacks, while remaining fully compatible with standard diffusion training pipelines. Our results highlight a new perspective on diffusion model watermarking: the denoising network can learn a reliable and persistent watermarking capability, which is crucial in practical model distribution scenarios.