NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

Topic: Diffusion Models · 661 papers · Reset filters

Every Measurement, Every Direction, All at Once: Multimodal Flow Matching for Molecules and Spectra

Thorben Prein, Elton Pan, Rafael Gomez-Bombarelli, Santiago Miret

Expand abstract

Generative modeling for science is data-scarce: each sample is expensive but arrives with several coupled measurements of the same system. Standard conditional models often use single-direction generation, discarding supervision from the rest. We argue that this asymmetry should not be a constraint during training: Under random target/conditioning masking, every measurement can act as a training target, hence one expensive sample yields multiple supervised denoising tasks. We present MOSAIC, a multimodal flow matching transformer over molecular 3D structure and UV, IR, and Raman spectra. On QM9S, MOSAIC achieves state-of-the-art performance with 88.8% Acc@1 for spectra-to-structure task, outperforming the strongest baseline by over 20% with 20$\times$ fewer integration steps. MOSAIC enables any-to-any generation from any combination of input to any combination of output modalities common in chemical sciences. We provide ablations to demonstrate the effectiveness of MOSAIC and show consistent performance scaling with increasing model and data size.

Poster
EverAnimate: Minute-Scale Human Animation via Latent Flow Restoration

Wuyang Li, Yang Gao, Mariam Hassan, Lan Feng, Wentao Pan, Po-Chien Luan, Alexandre Alahi

Expand abstract

We propose **EverAnimate**, an efficient post-training method for long-horizon animated video generation that preserves visual quality and character identity. Long-form animation remains challenging because highly dynamic human motion must be synthesized against relatively static environments, making chunk-based generation prone to accumulated drift: (i) low-level quality drift, such as progressive degradation of static backgrounds, and (ii) high-level semantic drift, such as inconsistent character identity and view-dependent attributes. To address this issue, **EverAnimate** restores drifted flow trajectories by anchoring generation to a persistent latent context memory, consisting of two complementary mechanisms. *(i) Persistent Latent Propagation* maintains a context memory across chunks to propagate identity and motion in latent space while mitigating temporal forgetting. *(ii) Restorative Flow Matching* introduces an implicit restoration objective during sampling through velocity adjustment, improving within-chunk fidelity. With only lightweight LoRA tuning, **EverAnimate** outperforms state-of-the-art long-animation methods in both short- and long-horizon settings: at 10 seconds, it improves PSNR/SSIM by 8%/7% and reduces LPIPS/FID by 22%/11%; at 90 seconds, the gains increase to 15%/15% and 32%/27%, respectively. Anonymous project page: [https://anonymous-xxx-lab.github.io/](https://anonymous-xxx-lab.github.io/).

Poster
Escaping the Curse of Dimensionality in One‑Step Flow-Based Generative Models

Konstantin Yakovlev

Expand abstract

Flow and diffusion-based generative models achieve state-of-the-art quality, but require multiple forward passes at inference time, making them computationally costly. Consistency models overcome this by learning a flow map directly, enabling few‑step generation and dramatically reducing the computation needed at test time. However, formal statistical guarantees for such one‑step generators are scarce. Existing results either suffer from the curse of dimensionality or rely on impractical loss functions. We investigate whether one‑step flow-based models can escape the curse of dimensionality under the assumption that the data distribution is generated by a low-dimensional surface perturbed with Gaussian noise. In this setup, we derive approximation and generalization error bounds for the flow map estimate that circumvent the curse of dimensionality. Our findings provide the first principled explanation for practical success of one-step flow-based generative models.

Poster
Ergodic Trajectory Design by Learned Pushforward Maps: Provable Coverage via Conditional Flow Matching

Ehsan Aghazadeh, Masoud Malekzadeh, Ahmad Ghasemi, Hossein Pishro-Nik

Expand abstract

Designing continuous trajectories whose time-averaged occupancy provably matches a prescribed spatial density (the *ergodic coverage* problem) is central to UAV-assisted data collection and sensing, robotic exploration, and mobile monitoring. For flying agents in particular, this challenge is acute: trajectories must balance coverage fidelity against tight energy budgets, no-fly zones, and acceleration limits. Existing methods either re-optimize each trajectory online (with cost growing in the horizon and re-running for every target, agent, and realization) or rely on bespoke analytical constructions that must be re-derived for each new constraint. We propose a *pushforward* framework that decouples ergodicity from density matching: an analytic latent trajectory provides exact uniform ergodicity on a simple annular domain, and a single map, learned offline by optimal-transport conditional flow matching, transports this latent occupancy onto the prescribed target density. The composed trajectory is then asymptotically ergodic with respect to the learned pushforward distribution, with deviation from the target controlled by the flow-matching training loss. Once trained for a given target density and constraint set, the map serves an unbounded number of trajectories and a multi-agent fleet without per-agent retraining, and many differentiable operational constraints (no-fly zones, acceleration ceilings, or fairness penalties) enter as additive soft penalties in the training loss without re-deriving the design. We prove three results (an acceleration-energy bound, an $O(1/\sqrt{K})$ ergodic convergence rate in the number of trajectory cycles $K$, and an approximation-error bound) that combine into an end-to-end coverage bound estimable from CFM training diagnostics (certified given an architectural Lipschitz bound on $v_\theta$). Experiments on synthetic targets empirically support each theoretical envelope; on a real UAV-coverage dataset, the proposed method gives the strongest observed coverage--energy and coverage--constraint tradeoffs among the evaluated time-warping, optimization, and concurrent flow-matching baselines.

Poster
Epigenomics-Guided Flow Matching for 3D Genome Super-Resolution

Yubiao Zhao, Shicheng Song, Yusen Ye, Huan Liu, Juan Liu, Lihua Zhang

Expand abstract

Mapping fine-scale 3D chromatin landscapes is critical for understanding gene regulation, yet the resolution of Hi-C data remains fundamentally constrained by both physical bottlenecks and sequencing costs. Although deep learning has shown promise in super-resolving Hi-C data, existing methods suffer from absolute coordinate loss during localized patching and a lack of structural and epigenomic priors in unconditioned generation. We propose a novel epigenomics-guided flow matching framework EpiFlow for 3D genome super-resolution. EpiFlow introduces two key innovations: (1) an EpiCond module that deterministically maps 1D epigenomic features to 2D spatial constraints via orthogonal expansion with built-in positional grounding, avoiding quadratic attention complexity; and (2) a Toeplitz-regularized classifier-free guidance mechanism that enforces distance-decay priors using a learnable symmetric prototype matrix. Extensive evaluation shows that our method achieves state-of-the-art cross-cell-type loop detection, improving recall by $2.8\times$ over raw Hi-C on unseen H1-hESC (Loop $F_1$ of $0.325$ vs. $0.199$) and outperforming all baselines on same-cell HFFc6 ($F_1=0.520$). Aggregate Peak Analysis (APA) confirms biological validity, and strong generalization to a held-out cell line demonstrates that the model learns transferable principles of 3D genome folding rather than cell-specific biases.

Poster
EpicWorldModel: Exploration-driven Planning with Latent World Models

Bowen Feng, Julian Ost, Zhiting Mei, Anirudha Majumdar, Felix Heide

Expand abstract

Latent world models based on Joint-Embedding Predictive Architecture (JEPA) are deterministic by design. While successful in fully observable scenarios, this paradigm breaks down when past observations and actions lead to multiple plausible future possibilities, e.g., due to occlusion. We introduce EpicWorldModel, a framework to train stochastic JEPAs for environments and tasks with inherent uncertainty under partially observability. We jointly train the EpicWorldModel predictor with its latent representation space to directly predict multiple potential future states using a flow-matching objective, when the goal-relevant scene content is absent from the conditioning history. We show that flow predictive variance approximates an upper bound on Expected Information Gain and serves as a useful exploration guidance for planning. By incorporating this uncertainty signal into Cross-Entropy Method (CEM)-based planning, our approach balances goal-reaching with exploration of uncertain regions where occluded or unseen goals are most likely to be located. We demonstrate the effectiveness of EpicWorldModel through a series of latent planning experiments with the best performance among all baseline methods, showing up to $22\%$ empirical improvement in success rate over LeWorldModel.

Poster
Entangled Schrödinger Bridge Matching

Sophia Tang, Yinuo Zhang, Pranam Chatterjee

Expand abstract

Simulating trajectories of multi-particle systems on complex energy landscapes is a central task in molecular dynamics (MD) and drug discovery, but remains challenging at scale due to computationally expensive and long simulations. Previous approaches leverage techniques such as flow or Schrödinger bridge matching to implicitly learn joint trajectories through data snapshots. However, many systems, including biomolecular systems and heterogeneous cell populations, undergo *dynamic* interactions that evolve over their trajectory and cannot be captured through static snapshots. To close this gap, we introduce **Entangled Schrödinger Bridge Matching (EntangledSBM)**, a framework that learns the first- and second-order stochastic dynamics of interacting, multi-particle systems where the direction and magnitude of each particle's path depend dynamically on the paths of the other particles. We define the Entangled Schrödinger Bridge (EntangledSB) problem as solving a coupled system of bias forces that *entangle* particle velocities. We show that our framework accurately simulates heterogeneous cell populations under perturbations and rare transitions in high-dimensional biomolecular systems. Anonymous code is provided at https://anonymous.4open.science/r/EntangledSBManon.

Poster
Enhancing Novel View Synthesis via Geometry Grounded Set Diffusion

Farhad Zanjani, Herbert Cai, Amirhossein Habibian

Expand abstract

We present SetDiff, a geometry-grounded multi-view diffusion framework that enhances novel-view renderings produced by 3D Gaussian Splatting. Our method integrates explicit 3D priors, pixel-aligned coordinate maps and pose-aware Plucker ray embeddings, into a set-based diffusion model capable of jointly processing variable numbers of reference and target views. This formulation enables robust occlusion handling, reduces hallucinations under low-signal conditions, and improves photometric fidelity in visual content restoration. A unified set mixer performs global token-level attention across all input views, supporting scalable multi-camera enhancement while maintaining computational efficiency through latent-space supervision and selective decoding. Extensive experiments on EUVS, Para-Lane, nuScenes, and DL3DV demonstrate significant gains in perceptual fidelity, structural similarity, and robustness under severe extrapolation. SetDiff establishes a state-of-the-art diffusion-based solution for realistic and reliable novel-view synthesis in autonomous driving scenarios.

Poster
Enhanced convergence guarantees of score-based generative models in $\mathcal{W}_2$-distance beyond log-concavity

Zhang Xiaoyan, Chenxu Pang, Xiaojie Wang

Expand abstract

Although score-based generative models (SGMs) have achieved remarkable success in real-world sample generation tasks, it is still far from sufficient to understand them from a mathematical perspective. In this paper, we aim to provide enhanced convergence guarantees for two existing SGMs in $\mathcal{W}_2$-distance beyond log-concavity. More precisely, we develop a novel framework of error analysis for Euler-Maruyama (EM) and Poisson midpoint (PM) time discretization schemes for SGMs. Under a non-log-concavity condition, we show that $\tilde{\mathcal{O}}(\sqrt{d}/\epsilon)$ iterations suffice to approximate the target distributions in $\epsilon$-accuracy for the classical EM-based SGM, considerably improving upon the existing iteration complexity $\tilde{\mathcal{O}}(d/\epsilon^2)$. Notably, the novel framework of error analysis enables us to establish an $\tilde{\mathcal{O}}(\sqrt{d}/\epsilon^{2/3})$ iteration complexity for the PM-based SGM in $\epsilon$-accuracy, significantly outperforming current state-of-the-art Wasserstein convergence guarantees for SDE-based diffusion samplers.

Poster
Energy-Tweedie: Score meets Score, Energy meets Energy

Andrej Leban

Expand abstract

Denoising and score estimation are classically linked through Tweedie’s formula, which relates the posterior mean under Gaussian corruption to the Stein score of the noisy marginal. In this work, we extend this perspective beyond Gaussian noise to a broad class of elliptical, energy-based noise distributions, with particular emphasis on generalized Gaussian corruptions. We derive the Energy–Tweedie identity: when the denoising posterior is viewed through the lens of proper scoring rules, the path derivative of a matched, possibly non-Euclidean energy score recovers the Stein score of the noisy marginal. Thus, the familiar correspondence between Gaussian noise, posterior means, squared loss, and Tweedie’s formula is lifted to a distributional correspondence between generalized Gaussian noise, full posterior laws, Mahalanobis energy scores, and the Energy–Tweedie identity. Among its consequences, this identity gives a posterior-sample-based route to score estimation, yields a principled criterion for noise-parameter calibration, and supplies a score-based perspective on recent diffusion-style generative methods trained with scoring rules.

Poster
Energy-Guided Transport for Projection-Free Physics-Informed Flow Matching

Giuliano Galadini, Aurelio Uncini, Danilo Comminiello

Expand abstract

Deep generative models, particularly Flow Matching frameworks, have emerged as powerful tools for synthesizing complex physical systems. However, standard architectures remain purely statistically driven and agnostic to the underlying governing equations, often violating fundamental physical constraints during inference. Existing zero-shot approaches designed to enforce these constraints typically rely on rigid algebraic projections. While effective for sparse boundary conditions, applying such discrete geometric corrections to dense, full-domain PDEs disrupts the probability path, pulling samples off the learned manifold and introducing intractable computational overhead. To address these challenges, we introduce Energy-Guided Transport, a novel training-free framework that formulates non-linear physical constraints as a smooth potential energy landscape. Rather than interrupting the ODE integration, our method applies a continuous dissipative force directly to the intermediate generative states, seamlessly steering them toward the physical manifold. By evaluating this force implicitly via reverse-mode automatic differentiation, our framework preserves the flow's natural dynamics while maintaining an efficient $\mathcal{O}(N)$ algorithmic complexity. We extensively evaluate our framework across various 1D and 2D physical systems, including the chaotic Navier-Stokes equations. Our method achieves stable generation without structural divergence, significantly improves the stability-accuracy trade-off compared to state-of-the-art methods, and substantially accelerates inference with respect to exact projection strategies. Code available at: \url{https://anonymous.4open.science/r/Projection-free-energy-guided-flow-matching-A302}.

Poster
End-to-End Differentiable Diffusion Conditioning for Physics-Informed Optimization

Ricardo Luna Gutierrez, Vineet Gundecha, Rahman Ejaz, Varchas Gopalaswamy, Riccardo Betti, Sahand Ghorbanpour, Aarne Lees, Soumyendu Sarkar

Expand abstract

Optimizing high-dimensional physical systems under expensive evaluations and complex constraints is a central challenge in science and engineering. In offline model-based optimization (MBO), methods must improve designs using only a fixed dataset and no test-time access to the true objective. However, existing methods typically suffer from conservatism, unconstrained surrogate exploitation, or function as static inference-time samplers that cannot refine candidates against precise physical objectives. In this paper, we propose Differentiable Diffusion Conditioning (D2C), and show that deterministic sampling makes the reverse denoising chain differentiable with respect to the conditioning signal. D2C turns a frozen conditional diffusion model into an end-to-end gradient-based optimizer by directly optimizing conditioning variables at test time via gradients from frozen surrogate objectives and physics-informed penalties backpropagated through the reverse chain. We implement this idea in a composed proposer-evaluator framework and evaluate it on laser pulse optimization for inertial confinement fusion, sustainable data-center workload scheduling, and four Design-Bench tasks. Our method achieves the strongest results on the two real-world tasks and the best average rank across four Design-Bench tasks.

Poster
Enabling VLA Action Self-Verification via VLM Token Probability Bucketing

Chen Zhao, Zhuoran Wang, Haoyang Li, Guanlin Li, Shifeng Bao, Youhe Feng, Yang Li, Jie Tang, Jing Zhang

Expand abstract

Test-time sampling can improve Vision-Language-Action (VLA) policies, but only if the model can efficiently select a good action from multiple candidates. Prior work often relies on separately trained external verifiers, adding extra models and inference overhead. We propose \textbf{Token Probability Bucketing (TokenPB)}, a self-verification approach that discretizes continuous action chunks into bucket tokens and co-trains a single VLA checkpoint with a joint objective: a flow-matching loss for action generation and an auxiliary bucket loss $\mathcal{L}_{\text{bucket}}$ that teaches the VLM backbone to model observation-conditioned distributions over bucket-token sequences. At inference, TokenPB samples $M$ candidates, scores each via teacher-forced evaluation under the learned distribution, and selects the best---enabling efficient batched scoring with observation-prefix KV-cache reuse. We provide theory connecting the training objective to a monotonic ranking property and bounds on expected selection regret that separate candidate-set effects from calibration. Across multiple backbones and benchmarks, TokenPB yields consistent success-rate improvements over single-sample inference from the same checkpoint (+2.7 in simulation, +17.1 in real-world, +16.4 under out-of-distribution shifts), and further improves over external-verifier baselines under the same sampling budget.

Poster
Enabling Preference-driven Unlearning in Few-step Distilled Text-to-Image Diffusion Models

Gaurav Patel, Jun Fang, Greg Ver Steeg, Qiang Qiu, Sravan Sripada

Expand abstract

Text-to-image (T2I) diffusion models are increasingly distilled into few-step variants and being deployed to enable fast inference. However, their ability to generate harmful or undesired content poses significant safety risks. Data-driven unlearning methods suppress targeted generations by fine-tuning model weights using specialized unlearning objectives. Crucially, these objectives implicitly rely on multi-step denoising dynamics, an assumption that breaks down for few-step distilled (FSD) models, resulting in ineffective forgetting. Furthermore, performing unlearning on the non-distilled base model and subsequently re-distilling it to obtain an unlearned FSD model incurs substantial computational and time overhead, making it impractical in many settings. Hence, we address this limitation with a preference-driven unlearning framework that revisits Direct Preference Optimization (DPO) for diffusion models. We show that standard DPO and its unlearning derivatives, formulated around noise-prediction error, transfer poorly to FSD models due to their altered generation dynamics. To overcome this, we introduce a modified preference optimization formulation explicitly aligned with the few-step generation properties, $\textit{enabling direct concept removal in FSD models}$ while preserving few-step efficiency and maintaining strong retention of desirable (non-targeted) capabilities. We evaluate our framework primarily on identity and NSFW (nudity) removal tasks and also extend our method to object-level unlearning. Extensive experiments demonstrate consistent and effective forgetting, and strong retention performance, establishing our method as a practical and principled solution for unlearning in FSD models.

Poster
Enabling Denoising Score Matching Type Training for Manifold Diffusion Model via Momentum and Splitting

Lingkai Kong, Yuchen Zhu, Molei Tao

Expand abstract

Denoising score matching (DSM) is the cornerstone of scalable diffusion model training. However, DSM becomes inapplicable on curved manifolds: the curved geometry makes the forward diffusion process nonlinear and the transition kernel is no longer tractable. Existing approaches therefore revert to implicit score matching (ISM) with expensive Brownian simulation, approximate the intractable heat kernel, or bypass the transition kernel altogether via flow matching with the logarithm map, each incurring significant computational cost or approximation error on general manifolds. We propose \emph{splitting diffusion}, a new framework that restores DSM-like training on general Riemannian manifolds by replacing the intractable Brownian forward process with a splitting scheme that alternates two steps: (1) an Ornstein-Uhlenbeck velocity update in the tangent space, whose transition kernel has a closed-form solution; (2) a deterministic geodesic transport step. This splitting yields a fully tractable transition kernel locally, enabling DSM-style training while requiring only the exponential map as a geometric oracle. The effectiveness of splitting diffusion is demonstrated on complicated manifolds and high-dimensional data.

Poster
Empowering Masked Diffusion Models to Self-Correct with Leave-One-Out Transformers

Sofian Zalouk, Vincent Counathe, Paul Jünger, Daniel Cao, Daniel Zlotnick, Kilian Weinberger, Christopher De Sa

Expand abstract

Diffusion models promise iterative refinement, yet masked diffusion language models (MDMs) suffer from two shortcomings: (1) inability to revise tokens, and (2) sparse training signal. The root cause is structural: MDMs control information flow only through input masking, so unmasked positions cannot be supervised, and self-correction is never learned. We propose the Leave-One-Out Transformer (LOOT), which addresses both issues. With LOOT, MDMs predict the distribution over potential replacements conditioned on the rest of the sequence at each position, in a single forward pass. Supervision can then be applied at masked and unmasked positions alike, yielding a lower-variance training objective with a dense loss signal at every token. Our method empowers MDMs to learn self-correction during training and supports continuous refinement at inference time via Gibbs sampling. LOOT achieves a new best validation perplexity on LM1B in a parameter-matched setting. Finetuned from a MDM checkpoint on OpenWebText, LOOT establishes a generative quality-vs-diversity frontier that surpasses prior remasking methods across NFE budgets.

Poster
Emergence of Distortions in High-Dimensional Guided Diffusion Models

Enrico Ventura, Beatrice Achilli, Luca Ambrogioni, Carlo Lucibello

Expand abstract

Classifier-free guidance (CFG) is the de facto standard for conditional sampling in diffusion models, yet it often reduces sample diversity. Using tools from statistical physics, we analyze the emergence of generative distortions induced by CFG, namely the mismatch between the CFG sampling distribution and the true conditional distribution. We study this phenomenon in analytically tractable settings with exact score functions, characterizing its dependence on data dimensionality and the number of classes. For high-dimensional Gaussian mixtures, we use dynamic mean-field theory to show that distortions arise when the number of classes scales exponentially with the data dimension, whereas they vanish in the sub-exponential regime due to a dynamical phase transition. We further prove that, in the infinite-class limit, distortions remain unavoidable regardless of dimensionality because of the increasing density of classes. Finally, we show that standard CFG schedules cannot prevent variance shrinkage, and we propose a theoretically grounded guidance schedule incorporating a negative-guidance window that improves both class separability and sample diversity in real-world latent diffusion models.

Poster
EgoHMP: Achieving Precise Human Motion Prediction in 3D Scenes via Egocentric Cues

Xin Zhao, Pengzhan Zhou, Ziyi Li, Riheng Jia, Zhida Qin, Yu Liu

Expand abstract

The complex interactive dependencies between human behavior and 3D scenes make accurately predicting human intent a major challenge for human motion prediction (HMP). While recent studies use human gaze to infer intent, gaze signals suffer from ambiguity, drift, and a reliance on eye-tracking hardware. Research in cognitive science indicates that egocentric views inherently encode potential interactions, while head motion serves as a crucial precursor in motion planning. Motivated by this, we propose EgoHMP, a novel framework that uses egocentric views and head motions as robust carriers of interaction intent. By translating this intent into an envisioning of future interactions, EgoHMP achieves precise HMP in 3D scenes. EgoHMP first employs a vision-motion alignment encoder to align the visual and motion feature spaces. Based on these aligned features, an egocentric intent predictor with a modality-aware modulation mechanism adaptively suppresses invalid interaction cues to predict an accurate interaction probability map. Furthermore, to enhance the stability and realism of the generated motions, we propose a contact-aware motion decoder. Specifically, inspired by the human cognitive process of planning trajectories prior to action execution, we first predict a global trajectory to serve as a self-prompt, effectively mitigating the optimization conflict in diffusion models. Subsequently, we introduce a multi-stage denoiser, which relieves motion artifacts by incorporating human contact modeling through a multi-stage optimization strategy. Finally, a GCN-based motion decoder is employed to synthesize physically plausible and semantically consistent motions. Extensive experiments demonstrate that EgoHMP achieves state-of-the-art performance.

Poster
Efficient and Robust Physical 3DGS-MPM Simulation via Interior Filling and Text-Physics Optimization

Yikun Ma, Yiqing Li, Jingwen Ye, Zhongkai Wu, Weidong Zhang, Lin Gao, Zhi Jin

Expand abstract

Extending 3D Gaussian Splatting (3DGS) to 4D physical simulation remains challenging. Based on the Material Point Method (MPM), existing methods either rely on manual parameter tuning or distill dynamics from video diffusion models, limiting the generalization and optimization efficiency. Recent attempts using LLMs/VLMs suffer from a text/image-to-3D perceptual gap, yielding unstable physics behavior. In addition, they often ignore the surface structure of 3DGS, leading to implausible motion. We propose FastPhysGS, a fast and robust framework for physics-based dynamic 3DGS simulation: (1) Instance-aware Particle Filling (IPF) with Monte Carlo Importance Sampling (MCIS) to efficiently populate interior particles while preserving geometric fidelity; (2) Bidirectional Graph Decoupling Optimization (BGDO), an adaptive strategy that rapidly optimizes material parameters predicted from a VLM. Experiments show FastPhysGS achieves high-fidelity physical simulation in 1 minute using only 7 GB runtime memory, outperforming prior works with broad potential applications.

Poster
Efficient Retrosynthesis Prediction with Integral Flow Matching and Latent Inversion

Tao Yin, Xiaohong Zhang, Yinjie Zhu, JiaCheng Zhang, Haotian Zou, Li Huang, Jiajun Cai, Zhibin Zhang, Meng Yan

Expand abstract

Fast and accurate retrosynthesis prediction is desired for downstream drug discovery and synthetic planning tasks. Currently, training and sampling state-of-the-art discrete diffusion or flow-based models for retrosynthesis requires significant computational resources. In this work, we propose FlashRetro, an Integral Flow Matching (IFM) framework with Latent Inversion for efficient retrosynthesis prediction. FlashRetro learns finite-interval latent transport rather than local instantaneous velocity, directly matching the transport quantity required for one-step latent inversion and avoiding test-time numerical integration. For inference, FlashRetro uses a principled 1-NFE Latent Inversion rule that maps a noise latent conditioned on the product latent to the reactant latent with a single learned integral-transport update. Experimental results on standard benchmarks (USPTO-50K) show that FlashRetro achieves superior top@1 accuracy, outperforming prior discrete flow-matching methods by 13.5\% and diffusion-based baselines by 14.2\%, while its average inference time is only 5.1 ms per reactant---38$\times$ faster than multi-step diffusion-based models and nearly 22$\times$ faster than recent discrete flow models. These results show that FlashRetro provides an effective path toward highly efficient 1-NFE retrosynthesis prediction with flow-based models.

Poster