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

PRISM: Prior Rectification and Uncertainty-Aware Structure Modeling for Diffusion-Based Text Image Super-Resolution

Zihang Xu, Xiaoyang Liu, Yulun Zhang, Xiaokang Yang

Expand abstract

Text image super-resolution (Text-SR) requires more than visually plausible detail synthesis: slight errors in stroke topology may alter character identity and break readability. Existing methods improve text fidelity with stronger recognitive or generative priors, yet they still face two unresolved challenges under severe degradation: the text condition extracted from low-quality inputs can itself be unreliable, and a plausible global prior does not fully determine fine-grained stroke boundaries. We present PRISM, a single-step diffusion-based Text-SR framework that addresses these two challenges through Flow-Matching Prior Rectification (FMPR) and a Structure-guided Uncertainty-aware Residual Encoder (SURE). FMPR constructs a privileged training-time prior from paired low-quality/high-quality latents and learns a flow matching that transports degraded embeddings toward this restoration-oriented prior space, yielding more accurate and reliable global text guidance. SURE further predicts uncertainty-aware structural residuals to selectively absorb reliable local boundary evidence while suppressing ambiguous stroke cues. Together, these components enable explicit global prior rectification and local structure refinement within a single diffusion restoration pass. Experiments on both synthetic and real-world benchmarks show that PRISM achieves state-of-the-art performance with millisecond-level inference.

Poster
PRECISE: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models

Bo Peng, Tao Huang, Weijie Kong, Junzhe Li, Yue Wu, Qi Tian, Jiangfeng Xiong, Jian-Wei Zhang, Liefeng Bo, Zhao Zhong

Expand abstract

Reinforcement learning (RL) has become an effective way to improve prompt alignment and perceptual quality in diffusion and flow-matching generators. A critical step for applying online RL to flow matching is turning the deterministic sampling trajectory into a stochastic policy, typically by replacing the reverse-time Ordinary Differential Equation (ODE) with a Stochastic Differential Equation (SDE). The stochastic sampler, controlling the exploration behavior and denoising dynamics, is thus part of the policy, and its design can significantly affect the reward optimization performance. We break down the sampler design into two interdependent components: choosing the right amount of stochastic exploration, and discretizing the resulting SDE faithfully at the small step counts used in RL. To address the first component, we analyze the inherent tension between exploration and stability in denoising and derive an SDE schedule that balances the two. Turning to the discretization challenge, we use a toy example to show that existing samplers can deviate from the flow-matching process, either by introducing excessive discretization noise or by relying on heuristic rules that do not guarantee convergence to the data distribution. To address these issues, we propose \textsc{Precise}, a new stochastic sampler that balances effective exploration with stability. Crucially, \textsc{Precise} keeps the denoising trajectory SDE-consistent through a novel approximation that freezes the clean-latent posterior mean, resolving the excess noise issue in standard samplers. Extensive experiments demonstrate that this formulation leads to significantly faster and more stable reward optimization via reinforcement learning, achieving state-of-the-art alignment scores (e.g., PickScore, HPSv2.1) while requiring 13.1--53.2\% less wall-clock training time to match the best in-domain performance of prior samplers.

Poster
PISG: Constraint-Aligned Signal Amplification for Diffusion-Based Combinatorial Optimization

Yuming Zhang, Xianchen Zhou, Hongxia Wang

Expand abstract

Diffusion-based solvers have become a prominent approach to neural combinatorial optimization. At inference time, they assign each decision variable a prediction score. A greedy decoder ranks candidates by these scores and checks hard constraints to ensure feasibility. Top-ranked candidates should be both promising and jointly feasible, making prediction--constraint alignment critical. When this alignment is weak, greedy decoding skips constraint-violating candidates and falls back to lower-ranked alternatives, degrading solution quality. Existing test-time methods improve search, sampling, or decoding, but do not directly strengthen this alignment before decoding. We observe that trained denoisers already carry constraint-aligned structural signals. Building on this observation, we propose Perturbed Instance-Structure Guidance (PISG), a training-free method that amplifies these signals at test time. At each denoising step, PISG constructs a structure-degraded prediction and extrapolates the unperturbed prediction away from it. The resulting guided prediction improves candidate rankings before constrained decoding. Signal analysis shows that this guidance concentrates on constraint-critical variables rather than uniformly rescaling scores. Under greedy decoding and without downstream refinement, PISG improves all 16 TSP solver-scale pairs across DIFUSCO, T2T, FastT2T, and StruDiCO, and all four MIS-ER solvers, with gains up to +8.4\%. These results identify prediction-constraint alignment as an effective test-time enhancement direction for diffusion-based combinatorial optimization.

Poster
PISCO: Precise Video Instance Insertion with Sparse Control

Xiangbo Gao, Renjie Li, Xinghao Chen, Yuheng Wu, Suofei Feng, Jie Yang, Qing Yin, Zhengzhong Tu

Expand abstract

AI video generation is moving beyond general generation, which relies on exhaustive prompt-engineering and "cherry-picking", towards fine-grained, controllable generation and high-fidelity post-processing. In professional AI-assisted filmmaking, the core requirement is the ability to perform precise, targeted modifications. A key task is video instance insertion, which requires precise spatial-temporal placement, physically consistent scene interaction (e.g., shadows and reflections), and the faithful preservation of original dynamics - all achieved under minimal user effort. In this paper, we propose PISCO, a video diffusion model for precise video instance insertion with arbitrary sparse keyframe control. PISCO allows users to specify a single keyframe, start-and-end keyframes, or sparse keyframes at arbitrary timestamps, and automatically propagates object appearance, motion, and interaction. To stabilize generation under sparse conditioning, we introduce Variable-Information Guidance and Distribution-Preserving Temporal Masking, complemented by geometry-aware conditioning. We further construct PISCO-Bench, a benchmark with verified instance annotations and paired clean background videos, and evaluate performance using both reference-based and reference-free perceptual metrics. Experiments demonstrate that PISCO consistently outperforms existing baselines and scales effectively as additional control signals are provided.

Poster
PEEK: One-Step Look-Ahead Exposure Bias Correction for Diffusion Sampling

Yunwei Bai, Ying Kiat Tan, Tsuhan Chen

Expand abstract

Diffusion image generators owe much of their practical success to accurate denoising networks and strong numerical samplers. Yet the two components are not fully matched. During sampling, the denoiser is asked to make predictions based on imperfect noisy states formed by model predictions that differ from the training noisy states formed by real data. This train-test gap, known otherwise as exposure bias in literature, caps image generation quality (i.e., blurry, lacking in details or unnaturally-shaped). To mitigate the problem, we propose PEEK, a training-free deterministic sampler plugin. At each eligible denoising step, PEEK first performs a provisional one-step update, evaluates the denoiser once at the look-ahead state, and either keeps the current clean prediction or replaces it with this one-step-later prediction to form the actual noisy state. By taking the additional step, we generally obtain a higher-signal, more-image-like prediction for the noisy state, mitigating the exposure bias. To optimize the replacement positions, PEEK searches for a binary replacement schedule with a greedy metric-driven search made efficient by the PEEK prefix and baseline noise-schedule-rebasing caches; for a 9-step search, the caches reduce the one-candidate worst-case vanilla denoiser-evaluation count by 69.3\%. Across pixel-space and latent-space diffusion models with unconditional, label-guided, and text-guided settings on CIFAR-10, CelebA-HQ, LSUN Bedroom, ImageNet, and MS-COCO datasets, PEEK improves over NFE-matched DPM-Solver++ and DEIS baselines, with five-run mean FID reductions up to 39.3\% under the same deployment cost. Code will be released.

Poster
PDE-SSM: A Spectral State Space Approach to Spatial Mixing in Diffusion Transformers

Eshed Gal, Moshe Eliasof, Eldad Haber

Expand abstract

The success of vision transformers—especially for generative modeling—is limited by the quadratic cost and weak spatial inductive bias of self-attention. We propose PDE-SSM, a spatial state-space block that replaces attention with a learnable convection–diffusion–reaction partial differential equation. This operator encodes a strong spatial prior by modeling information flow via physically grounded dynamics rather than all-to-all token interactions. Solving the PDE in the Fourier domain yields global coupling with near-linear complexity, delivering a principled and scalable alternative to attention. We integrate PDE-SSM into a flow-matching generative model to obtain the PDE-based Diffusion Transformer PDE-SSM-DiT. Empirically, PDE-SSM-DiT matches or exceeds the performance of state-of-the-art Diffusion Transformers while substantially reducing compute. Our results show that, analogous to 1D settings where SSMs supplant attention, multi-dimensional PDE operators provide an efficient, inductive-bias-rich foundation for next-generation vision models.

Poster
PACE-dLLM: Elastic Block Decoding via Confidence Cliff Estimation for Diffusion Language Models

Xiaocheng Lu, Shuhan Guo, Ziyue Ma, Jie Zhang, Jian Liu, Jingcai Guo, Haoxuan Che, Song Guo

Expand abstract

Diffusion large language models (dLLMs) such as LLaDA and Dream now rival autoregressive LLMs in quality while retaining native parallel decoding. To deploy them efficiently, block-wise decoding partitions generation into blocks of size $B$ and commits a fraction of each block before moving on, but $B$ conflates two roles: \emph{how far the model can look ahead} and \emph{how many tokens get committed per step}. Recent accelerators relax this trade-off with indirect heuristics, yet the underlying difficulty is already exposed at every forward pass by the model's per-step confidence. Specifically, the in-window confidence profile exhibits a context-dependent \emph{cliff} (a high plateau, a sigmoidal transition, and a residual floor) that directly encodes how far ahead is safe to look. We propose \textbf{PACE-dLLM}, a decoder that handles the two roles independently. At each step, PACE-dLLM fits the cliff in closed form on the current window's confidences and reads off the next prediction horizon at the cliff's saturation point; a separate confidence threshold decides which tokens are committed. Both rules read the same signal, add a single hyperparameter, and admit a strict-dominance guarantee over fixed-block decoding. Extensive experiments across four reasoning and code benchmarks on two open-source dLLM backbones demonstrate that PACE-dLLM achieves the best average accuracy at a roughly $\mathbf{4.5\times}$ wall-clock speedup over the unaccelerated semi-AR baseline, pushing the Pareto frontier outward.

Poster
Oversmoothing as Representation Degeneracy in Neural Sheaf Diffusion

Arif Dönmez, Ellen Fritsche, Axel Mosig, Katharina Koch

Expand abstract

Neural Sheaf Diffusion (NSD) generalizes diffusion-based Graph Neural Networks by replacing scalar graph Laplacians with sheaf Laplacians whose learned restriction maps define a task-adapted geometry. While the diffusion limit of NSD is known to be the space of global sections, the representation-theoretic structure of this harmonic space remains largely implicit. In this paper, we develop a quiver-theoretic interpretation of NSD by identifying cellular sheaves on graphs with representations of the associated incidence quiver. Under this correspondence, learned sheaf geometries become points in a finite-dimensional representation space. We prove that direct-sum decompositions of the underlying incidence-quiver representation induce corresponding decompositions of the harmonic space reached in the diffusion limit. This provides an algebraic interpretation of oversmoothing as representation degeneration: a conceptual framing where learned sheaves collapse toward trivial or low-complexity summands whose global sections fail to preserve discriminative information. Building on this viewpoint, we connect sheaf diffusion to stability, moduli, and moment-map principles from Geometric Invariant Theory. We introduce moment-map-inspired regularizers that bias learned restriction maps toward more balanced representation geometries, and we identify a structural obstruction in standard equal-stalk architectures: when $d_v=d_e$, the admissibility condition for learnable stability parameters forces the trivial all-object summand onto a stability wall. We show that non-uniform stalk dimensions remove this obstruction, making adaptive stability meaningful in principle. Empirical evaluations on heterophilic benchmarks are consistent with this mechanism: breaking stalk symmetry can reduce variance or improve validation behavior on some datasets, and adaptive stability regularization becomes more effective in selected rectangular settings. These results support the view that moment-map regularization is a structured but dataset-dependent geometric bias rather than a universal performance booster. Overall, our framework interprets oversmoothing not only as a spectral pathology, but as a degeneration phenomenon in the underlying representation geometry.

Poster
Outlier-robust Diffusion Posterior Sampling for Bayesian Inverse Problems

Yiming Yang, Xiaoyuan Cheng, Yi He, Kaiyu Li, Wenxuan Yuan, Zhuo Sun

Expand abstract

Diffusion models have emerged as powerful learned priors for Bayesian inverse problems (BIPs). Diffusion-based solvers rely on a presumed likelihood for the observations in BIPs to guide the generation process. Likelihood misspecification is common in practical BIPs and is known to degrade recovery performance, particularly under outlier contamination. We investigate this problem by first characterizing the induced posterior deviation and proving the stability of diffusion-based solvers for linear BIPs. Our stability analysis further reveals potential robustness deficiencies of existing diffusion-based solvers under outlier-contaminated measurements. To address this issue, we propose a simple yet effective solution: robust diffusion posterior sampling, which is provably outlier-robust for linear BIPs and compatible with existing gradient-based posterior samplers. Empirical results from scientific inverse problems and natural image tasks demonstrate the effectiveness and robustness of our method, with consistent performance gains in challenging scenarios involving outlier contamination for both linear and nonlinear tasks.

Poster
Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion

Van Chien Nguyen, Chaitra Hegde, Van-Cuong Pham, Ryan Rossi, Franck Dernoncourt, Thien H Nguyen

Expand abstract

We introduce Orthrus, a simple and efficient dual-architecture framework that unifies the exact generation fidelity of autoregressive Large Language Models (LLMs) with the high-speed parallel token generation of diffusion models. The sequential nature of standard autoregressive decoding represents a fundamental bottleneck for high-throughput inference. While diffusion language models attempt to break this barrier via parallel generation, they suffer from significant performance degradation, high training costs, and a lack of rigorous convergence guarantees. Orthrus resolves this dichotomy natively. Designed to seamlessly integrate into existing Transformers, the framework augments a frozen LLM with a lightweight, trainable module to create a parallel diffusion view alongside the standard autoregressive view. In this unified system, both views attend to the exact same high-fidelity Key-Value (KV) cache; the autoregressive head executes context pre-filling to construct accurate KV representations, while the diffusion head executes parallel generation. By employing an exact consensus mechanism between the two views, Orthrus guarantees lossless inference, delivering up to a $7.8\times$ speedup with only an $O(1)$ memory cache overhead and minimal parameter additions.

Spotlight
Order-Optimal Sample Complexity for Distribution Learning via Flow Matching

Hari K Sahoo, Mudit G Gaur, Vaneet Aggarwal

Expand abstract

Flow-based generative models have emerged as an efficient alternative to diffusion models. We study the sample complexity of learning a target distribution in Wasserstein distance via flow matching, which learns a velocity field transporting a base distribution to the target. Under standard assumptions on the network architecture and data distribution, we show that any squared-loss flow-matching variant including rectified flow with bounded derivatives achieves $\widetilde{O}(\varepsilon^{-2})$ sample complexity with the Wasserstein guarantee taking the form $W_2 \le K(\varepsilon + C\sqrt{\varepsilon_{app}})$, where $\varepsilon_{app}$ is the approximation error of the network class. This improves upon existing $\widetilde{O}(\varepsilon^{-4})$ guarantees and does not require the Polyak \L{}ojasiewicz condition assumed in prior work. Instead, we exploit a structural property of the squared-loss objective - it induces an approximate Bernstein condition that directly ties the variance of the excess loss to the excess risk up to additive approximation-error. This condition enables a localized Rademacher complexity analysis yielding fast $O(1/n)$ rates. The flow-matching schedule enters only through the Lipschitz constant of the pointwise loss, which we characterize explicitly for standard schedules.

Poster
Order-Marginalized Scoring for Masked Diffusion Models

Arthur Deng, Sebastian Thrun

Expand abstract

Masked Diffusion Models (MDMs) are trained under an objective that is symmetric over decoding orders, but at sampling time a single order must be committed to. A common task in such models is to evaluate the log-probability of a candidate sequence, either as a downstream score or to rank a finite pool of candidates. We observe that this evaluation depends substantially on the decoding order chosen, even though the training objective averages over orders. For both a model trained on closed-form Probabilistic Context-Free Grammar (PCFG) data and a pre-trained MDM on OpenWebText, the per-order log-probability rankings of fixed candidate pools disagree on most inputs. We characterize this phenomenon and propose Order-Marginalized Scoring (OMS), a simple Monte-Carlo (MC) estimator of the order-marginalized log-likelihood $\tilde \mu(x)=E_\sigma[\log p_\sigma(x)]$. We use it in two complementary algorithms: a reranker over a fixed candidate pool, and a stochastic best-of-$N$ decoder that combines adaptive position selection with $\tilde\mu$-reranking. On Sudoku, the decoder substantially raises exact-match accuracy over the strongest single-order baseline from $65.3$% to $94.7$%. On a zero-shot protein-stability benchmark, the reranker improves Spearman correlation between log-probability scores and experimental folding stabilities over a per-order baseline by $0.07$. These results suggest that OMS is a useful alternative to per-order scoring for order-agnostic problems.

Poster
Optimal Convergence Analysis of DDPM for General Distributions

Yuchen Jiao, Yuchen Zhou, Gen Li

Expand abstract

Score-based diffusion models have achieved remarkable empirical success in generating high-quality samples from target data distributions. Among them, the Denoising Diffusion Probabilistic Model (DDPM) is one of the most widely used samplers, generating samples via estimated score functions. Despite its empirical success, a tight theoretical understanding of DDPM --- especially its convergence properties --- remains limited. In this paper, we provide a refined convergence analysis for the DDPM sampler and establish near-optimal convergence rates under general distributional assumptions. Specifically, we consider a relaxed smoothness condition parameterized by $L$, which is small for many practical distributions (e.g., Gaussian mixture models). We prove that, to approximate a target distribution on $\mathbb{R}^d$ to accuracy $\varepsilon$ in total variation distance and $\varepsilon^2$ in KL-divergence, the DDPM sampler with accurate score estimates requires at most $$T = \widetilde{O}\left(\frac{\sqrt{d}\min\{\sqrt{d},L\}}{\varepsilon}\right)$$ iterations, where $\widetilde{O}$ hides polylogarithmic factors in $d$ and $1/\varepsilon$. This result substantially improves upon the best-known $\widetilde{O}(d/\varepsilon)$ iteration complexity when $L < \sqrt{d}$. By establishing a matching lower bound, we show that our convergence analysis is tight for a wide array of target distributions. Moreover, it reveals that DDPM and Denoising Diffusion Implicit Models (DDIM) share the same dependence on $d$, raising an interesting question of why DDIM often appears empirically faster.

Poster
Opening the Black Box of Classifier-Free Guidance via Information Bottleneck

Jiayang Gao, Tianyi Zheng, Jiayang Zou, Fengxiang Yang, Luyao Fan, Lv Tang, Bo Li, Jia Wang

Expand abstract

Classifier-free guidance (CFG) is a standard technique for improving conditional generation in diffusion models, yet its activation schedule is usually selected heuristically. In practice, most existing CFG-based methods rely on a pre-defined activation policy throughout inference, while recent dynamic variants often depend on hand-designed schedules without a principled theoretical explanation of when conditional information is most useful. In order to understand how the informational relevance between conditioning inputs and the score function evolves over the course of generation, we explore when guidance should be activated from an information-theoretic perspective. In this work, we formulate the diffusion trajectory as a time-dependent information bottleneck and show that the effectiveness of conditional information is not uniform throughout sampling, but becomes significantly stronger after a critical stage characterized by a Fisher-information balance condition. Motivated by this insight, we propose a training-free dynamic CFG strategy for inference. Rather than prescribing a fixed schedule in advance, our method constructs a stochastic process from conditional and unconditional score statistics and activates guidance when the estimated guidance gain crosses a prescribed threshold. This yields an adaptive activation rule that aligns condition injection with the underlying generation dynamics, while also recovering interval-based guidance as a limiting special case. Extensive experiments across diverse architectures and tasks, demonstrate that our method is effective, robust, and broadly applicable. Beyond empirical improvements, our framework provides an interpretation of dynamic guidance methods and a principled foundation for inference-time control in conditional diffusion models.

Poster
OpenVTON-Bench: A Large-Scale High-Resolution Benchmark for Controllable Virtual Try-On Evaluation

Jin Li, Tao Chen, Kai Wen, Siqi Yin, Shuai Jiang, Weijie Wang, Jingwen Luo, Chenhui Wu

Expand abstract

Recent advances in diffusion models have significantly elevated the visual fidelity of Virtual Try-On (VTON) systems, yet reliable evaluation remains a persistent bottleneck. Traditional metrics struggle to quantify fine-grained texture details and semantic consistency, while existing datasets fail to meet commercial standards in scale and diversity. We present OpenVTON-Bench, a large-scale benchmark comprising approximately 100K high-resolution image pairs (up to $1536\times1536$). The dataset combines Gemini-powered dense captioning with stratified sampling across 20 visual clusters obtained by DINOv3 K-Means. To support reliable evaluation, we propose a multi-modal protocol that measures VTON quality along five interpretable dimensions: background consistency, identity fidelity, texture fidelity, shape plausibility, and overall realism. The protocol integrates VLM-based semantic reasoning with a novel Multi-Scale Representation Metric based on SAM3 segmentation and morphological erosion, enabling the separation of boundary alignment errors from internal texture artifacts. Under the specified evaluation rubric, our representation metric shows strong agreement with human model rankings (Kendall's $\tau$ of 0.833 vs. 0.611 for SSIM). Code and dataset are available at https://github.com/RenxingIntelligence/OpenVTON-Bench.

Poster
On the Memorization of Consistency Distillation for Diffusion Models

Bingqing Jiang, Difan Zou

Expand abstract

Diffusion models are central to modern generative modeling, and understanding how they balance memorization and generalization is critical for reliable deployment. Recent work has shown that memorization in diffusion models is shaped by training dynamics, with generalization and memorization emerging at different stages of training. However, deployed diffusion models are often further distilled, introducing an additional training phase whose impact on memorization remains unclear. In this work, we analyze how distillation reshapes memorization behavior in diffusion models, taking consistency distillation as a representative framework. Empirically, we show that when applied to a teacher model that has memorized data, consistency distillation significantly reduces transferred memorization in the student while preserving, and sometimes improving, sample quality. To explain this behavior, we provide a theoretical analysis using a random feature neural network model [Bonnaire et al., 2025], showing that consistency distillation suppresses unstable feature directions associated with memorization while preserving stable, generalizable modes. Our findings suggest that distillation can serve not only as an acceleration tool, but also as a mechanism for improving the memorization-generalization trade-off.

Poster
On the Limits of Latent Reuse in Diffusion Models

Yifeng Yu, Lu Yu

Expand abstract

Diffusion models are often trained in low-dimensional latent spaces, which are then reused for related but shifted datasets. In this work, we study when such latent reuse remains reliable under distribution shift. We consider a source-target setting in which both datasets are approximately low-dimensional but may lie near different subspaces. We show that freezing and reusing a source latent space induces a target-domain score error governed by two quantities: the principal-angle misalignment between the source and target subspaces, and the target ambient noise amplified by the diffusion time scale. Motivated by these limits, we further study mixed source-target training and characterize how the required shared latent dimension depends on the relative geometry of the two distributions. Our results provide theoretical guidance on when latent reuse is reliable and when learning a shared representation may be necessary.

Poster
On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimization

Trevor Chen, Ariel Dai, Jason Yang, Riccardo De Santi, Daniel Khalil, Wenda Chu, Nate Gruver, Pranav Murugan, Alexander Goldberg, Maruan Al-Shedivat, Yisong Yue

Expand abstract

We study how to design online optimization loops for molecular optimization via adaptation of pre-trained generative models. At test time, we aim to utilize limited oracle feedback to steer generation toward top-molecules achieving high rewards. This creates a design problem coupled across several dimensions, including which candidates receive oracle evaluations, how observed rewards are utilized for model update, and how to overcome pre-trained model biases–to effectively explore over complex design spaces. Despite recent algorithmic progress on each individual component, it remains unclear how they interact in practice, within real-world online adaptation loops. For instance, top-$K$ optimization objectives may make adaptation overly greedy and thereby reduce exploration, while model-debiasing techniques may be unnecessary when exploration is already induced by exploratory acquisition functions, e.g., Thompson sampling. To address this, we conduct a controlled study on discrete-diffusion molecular optimization, finding that well-designed components remain beneficial when combined, indicating that they tackle complementary issues. Together, these components yield an online fine-tuning recipe that outperforms offline fine-tuning and search-augmented baselines across several small-molecule binding affinity and protein fitness optimization tasks, under equal oracle-call budgets and GPU-hour accounting.

Poster
On Reparameterizing the Score Function

Mohammad Hossein Amani, Leonardo Bocchieri, Maxime Peyrard, Clement Hongler, Robert West

Expand abstract

Policy gradients for autoregressive sequence models differentiate each token's log-probability while holding the conditioning history fixed. This discards information about how earlier decisions shape later conditional distributions, creating a credit-assignment bottleneck when rewards are sparse. We introduce History-Reparameterized (HR) policy gradients, which replace the frozen prefix in the score function with a differentiable reparameterized or relaxed history, allowing gradients of downstream log-probabilities to propagate backward through the generated trajectory. The resulting history-pathwise term is generally biased when added directly, so we derive Optimal HR (OHR), a centered control-variate estimator with a closed-form coefficient that minimizes the trace of the gradient variance at the population level. The variance reduction is governed by temporal structure in the policy dynamics: random or weakly structured policies provide little benefit, while policies with learned sequential dependencies yield informative history-pathwise corrections. For discrete tokens, the choice of relaxation is crucial: embedding-level relaxation yields substantial variance reduction that grows with horizon, whereas straight-through Gumbel-Softmax yields essentially none in our diagnostics. We validate the method on a Gaussian linear Markov model, a sparse-reward T-maze with a recurrent policy, and an autoregressive Transformer model for small Traveling Salesman Problem (TSP) instances.

Poster
Offline Inverse Reinforcement Learning with Unified Diffusion Planning

Hongmin Zhao, Jiyuan Yin, Kexi Yan, Qinglai Wei, Jie Zhang

Expand abstract

Offline Inverse Reinforcement Learning (IRL) aims to recover a reward function and imitate expert behavior solely from offline demonstrations. While recent Offline IRL approaches employ approximate dynamics to mitigate distribution shift, their performance is constrained by the unstable minimax framework as well as simplified representation and utilization. To address these issues, we propose Offline Inverse Reinforcement Learning with Diffusion Planner (OIDP), which leverages diffusion models to achieve stable Offline IRL. OIDP formulates a model-based conservative Q-optimization through inverse soft Q-learning, which is provably concave in Q-space with a controllable performance gap. Accordingly, we introduce a unified diffusion planner that models dynamics, uncertainty estimator, and policy, it captures multimodal distributions, directly tightens the objective gap bound, and at execution performs trajectory-guided policy sampling using the internalized dynamics knowledge. Experimental results on D4RL benchmarks demonstrate that OIDP outperforms state-of-the-art offline approaches.

Poster