ICML 2026 论文列表

International Conference on Machine Learning

录用类型

6,341篇论文

按官方录用论文列表统计,可能与会议公布的数字略有出入。

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

浏览全部论文

Enhancing Multi-Modal LLMs Reasoning via Difficulty-Aware Group Normalization

Jinghan Li, Junfeng Fang, Jinda Lu, Yuan Wang, Xiaoyan Guo, Tianyu Zhang, Xiang Wang, Xiangnan He

展开摘要

Reinforcement Learning with Verifiable Rewards (RLVR) and Group Relative Policy Optimization (GRPO) have significantly advanced the reasoning capabilities of large language models. Extending these methods to multimodal settings, however, faces a critical challenge: the instability of std-based normalization, which is easily distorted by extreme samples with nearly positive or negative rewards. Unlike pure-text LLMs, multimodal models are particularly sensitive to such distortions, as both perceptual and reasoning errors influence their responses. To address this, we characterize each sample by its difficulty, defined through perceptual complexity (measured via visual entropy) and reasoning uncertainty (captured by model confidence). Building on this characterization, we propose difficulty-aware group normalization (Durian), which re-groups samples by difficulty levels and shares the std within each group. Our approach preserves GRPO's intra-group distinctions while eliminating sensitivity to extreme cases, yielding significant performance gains across multiple multimodal reasoning benchmarks.

Poster
Robust Self-reflective Hashing for Cross-modal Retrieval with Noisy Label

Hao Sun, Qibing Qin, Lei Huang

展开摘要

Cross-modal Hashing (CMH) typically assumes a perfectly complete data annotation, whereas noisy labels are unavoidable in practical scenarios. Existing CMH methods often overlook the uncertainty introduced by noise or semantic ambiguity, making models susceptible to overfitting noisy labels and yielding unreliable similarity judgments during inference. To address this issue, we propose a Robust Self-reflective Hashing (RSH) framework that prudently analyzes semantic discrepancies while accounting for uncertainty, thereby effectively mitigating interference from noisy labels. Specifically, the Double Feature Representation (DFR) method is introduced, employing semantic and uncertainty features to represent the semantic representation and fuzziness of samples. With a double feature, we propose a novel cross-modal similarity metric - the Self-reflective Similarity Metric (SSM), which judges sample similarity by integrating semantic discrepancy and fuzziness, enabling the model to adaptively weaken semantic discrepancy according to uncertainty level. The proposed method is plug-and-play, enabling seamless integration into diverse objective functions to enhance model robustness and reliability. Extensive experiments on benchmark datasets demonstrate that RSH outperforms existing methods. Code is available at https://github.com/QinLab-WFU/RSH.

Poster
Offline Reinforcement Learning of High-Quality Behaviors Under Robust Style Alignment

Mathieu Petitbois, Rémy Portelas, Sylvain Lamprier

展开摘要

We study offline reinforcement learning of style-conditioned policies using explicit style supervision via subtrajectory labeling functions. In this setting, aligning style with high task performance is particularly challenging due to distribution shift and inherent conflicts between style and reward. Existing methods, despite introducing numerous definitions of style, often fail to reconcile these objectives effectively. To address these challenges, we propose a unified definition of behavior style and instantiate it into a practical framework. Building on this, we introduce Style-Conditioned Implicit Q-Learning (SCIQL), which leverages offline goal-conditioned RL techniques, such as hindsight relabeling and value learning, and combine it with a new Gated Advantage Weighted Regression mechanism to efficiently optimize task performance while preserving style alignment. Experiments demonstrate that SCIQL achieves superior performance on both objectives compared to prior offline methods. Code, datasets and visuals are available in: https://mathieu-petitbois.github.io/projects/sciql/.

Spotlight
CrossQ: Task-Aligned Cross-Token Conditional Quantization for Late Interaction Retrieval

Rohit Kumar Salla, Manoj Saravanan, Ramya Manasa Amancherla

展开摘要

Late-interaction retrievers like ColBERT achieve high quality but suffer from large multi-vector indices. Standard compression minimizes token reconstruction error, while ranking depends critically on preserving scores of sparse "winner" tokens. We introduce CrossQ, which adaptively improves effective token fidelity within documents by conditioning token codes on lightweight document context computed at indexing time (but not stored). CrossQ is trained with ranking-aligned objectives that preserve candidate score distributions and protect hard-negative margins. At 2 B/token, CrossQ improves MRR@10 by +0.010 over the strongest strictly footprint-matched quantization baseline and by +0.012 over the strongest candidate-matched system reference. On the nine-dataset BEIR subset reported in Appendix G.1, CrossQ improves average nDCG@10 by +0.009 at 4 B/token over the strongest candidate-matched system reference. At 4 B/token, CrossQ achieves $64\times$ raw token-storage reduction, approximately $61\times$ including metadata and approximately $58\times$ under conservative padding/alignment accounting. At 8 B/token, CrossQ + light fine-tuning retains approximately 98% of full-precision ColBERT MRR@10, improving the footprint-quality tradeoff for memory-constrained late-interaction retrieval.

Poster
Unified Multimodal Autoregressive Modeling with Shared Context—Visual Tokenizer is Key to Unification

Wujian Peng, Lingchen Meng, Yuxuan Cai, Xianwei Zhuang, Yuhuan Yang, Rongyao Fang, Chenfei Wu, Junyang Lin, Zuxuan Wu, Shuai Bai

展开摘要

Unified Multimodal Modeling aims to integrate visual understanding and generation within a single system. However, existing approaches typically rely on two disparate visual tokenizers, which splits the representation space and hinders truly unified modeling. We propose UniAR, a unified autoregressive framework where a single discrete visual tokenizer serves as the key bridge between understanding and generation, enabling a shared context in which the model can directly interpret its own generated visual tokens without additional re-encoding. UniAR adapts a pretrained vision encoder with multi-level feature fusion and a lookup-free bitwise quantization scheme, preserving both high-level semantics and low-level details while scaling the effective visual vocabulary at minimal cost. Building on this, the unified autoregressive model adopts parallel-bitwise-prediction to jointly predict spatially grouped, multi-level visual codes, substantially reducing visual sequence length and accelerating generation. Finally, a diffusion-based visual decoder operates on discrete visual tokens to decode high-fidelity images. Through large-scale pre-training, followed by supervised fine-tuning and reinforcement learning, UniAR achieves state-of-the-art performance on image generation and image editing while remaining competitive on multimodal understanding benchmarks. The project homepage is at https://sharelab-sii.github.io/uniar-web.

Poster
Temporal Difference Learning for Diffusion Models

Qizhen Ying, Yangchen Pan, Victor Adrian Prisacariu, Junfeng Wen

展开摘要

Diffusion models are typically trained with objectives that focus on local denoising targets at individual time steps (or adjacent pairs), which do not enforce consistency between predictions along the denoising trajectory. This lack of cross-time consistency can degrade performance, especially for few-step samplers. We introduce a temporal difference (TD) objective that penalizes inconsistency of the model’s \emph{multi-step} progress along the denoising path. By reformulating the diffusion process as a Markov reward process and casting denoising as a policy evaluation problem in reinforcement learning, we derive a unified TD approach that applies to both discrete- and continuous-time diffusion formulations. We further propose a principled sample-based reweighting method that stabilizes training. Empirically, we show that using our TD training can significantly improve sample quality measured by FID, with stronger advantages when the number of sampling steps is small, highlighting its practical utility under low-computation-budget scenarios. We provide ablation studies to justify our design choices, including pairwise loss reweighting, regularization weight, and one-step stride. Overall, our TD approach can be a general drop-in that enforces cross-time consistency and improves generation quality across different diffusion generative models.

Poster
A Graph Foundation Model with Cross-Modal Alignment and Modality-Aware Expert Fusion for Multi-Modal Graphs

Dongxiao He, Ankang Yang, Jitao Zhao, Di Jin

展开摘要

Graph Foundation Models (GFMs) aim to learn universal patterns through large-scale pretraining on diverse graphs and generalize to open-world scenarios. While GFMs have garnered significant attention, existing works primarily focus on sigle-modal graphs. However, many real-world graphs are multimodal, consisting of structures alongside diverse features derived from modalities such as text and images. To date, exploration into Multimodal Graph Foundation Models (MGFMs) remains limited. Incorporating multimodal data provides a more comprehensive view, allowing models to learn richer semantics, thereby advancing GFMs. We are therefore motivated to explore MGFMs, where the core challenge lies in synergistically encoding structures and multimodal features to achieve effective cross-modal alignment and fusion. To this end, we propose a graph foundation model with Cross-modal Alignment and Modality-aware Expert fusion, CAME. Specifically, CAME first generates graph embeddings for each individual modality. We then introduce a multimodal multi-expert encoding mechanism, which includes a dimension-wise routing strategy to fuse multimodal information. Finally, we employ a cross-modal contrastive loss to train CAME, enabling the adaptive alignment and fusion across different modalities. Extensive experiments demonstrate the effectiveness of CAME across multiple tasks and diverse multimodal graph datasets.

Poster
Language Generation with Feedback: Queries and Mistakes

Steve Hanneke, amin karbasi, Anay Mehrotra, Grigoris Velegkas

展开摘要

We investigate language generation in the limit (Kleinberg & Mullainathan, 2024; Li et al., 2025) in variants where the generator receives some feedback based on its “actions.” We study two such variants. In the first, which is inspired by Littlestone’s model of online learning, the generator observes whether it made a mistake at each iteration. In the second, introduced by Charikar & Pabbaraju (2025a), the generator can query whether a string belongs to the target language. Our main result is a characterization of collections that are generable with mistake feedback. Using similar techniques, we also characterize when generation is possible in the query model with set-based generators; set-based generators have been studied in several works (Charikar & Pabbaraju, 2025a; Kalavasis et al., 2025; Kleinberg & Wei, 2025a; Li et al., 2025). Beyond the characterizations themselves, we derive several implications. First, our results imply new closure properties for generation with mistake and query feedback. Second, our results show that, under feedback, generation is robust to noise: it remains possible with arbitrary contamination in the adversary’s examples and with finite contamination in the feedback. Third, our techniques also yield new sufficient and necessary conditions for generation without feedback among other implications.

Poster
OpenGPT-4o-Image: A Comprehensive Dataset for Advanced Image Generation and Editing

Zhihong Chen, Xuehai Bai, Yang Shi, Chaoyou Fu, Huanyu Zhang, Haotian Wang, Xiaoyan Sun, Zhang Zhang, Liang Wang, Yuanxing Zhang, Pengfei Wan, Yifan Zhang, Yi-Fan Zhang

展开摘要

The performance of unified multimodal models for image generation and editing is fundamentally constrained by the quality and comprehensiveness of their training data. While existing datasets have covered basic tasks like style transfer and simple object manipulation, they often lack the systematic structure and challenging scenarios required for real-world applications. To address this bottleneck, we introduce OpenGPT-4o-Image, a large-scale dataset constructed using a novel methodology that combines hierarchical task taxonomy with automated data generation. Our taxonomy not only includes fundamental capabilities such as text rendering and style control but also introduces highly practical yet challenging categories like scientific imagery for physics/chemistry illustrations and complex instruction editing requiring simultaneous execution of multiple operations. Through an automated pipeline leveraging structured resource pools and GPT-4o, we generate 80k high-quality instruction-image pairs with controlled diversity, covering 11 major domains and 51 subtasks. Extensive experiments show that fine-tuning leading models on our dataset achieves significant performance gains across multiple benchmarks, with improvements of up to 18% on editing tasks UniWorld-V1 on ImgEdit-Bench and 13% on generation tasks Harmon on GenEval. Our work demonstrates that systematic data construction is key to advancing multimodal AI capabilities.

Poster
Video-SVD: Efficient Video Diffusion via Orthogonal Basis Composition

Zhang Wan, Yu Li, Tianze Huang, Haochen Li, Juan Cao, Sheng Tang

展开摘要

Video Diffusion Transformers (VDiTs) represent the state of the art in video generation but remain constrained by the quadratic complexity of dense self-attention. To address this attention bottleneck, we analyze the pre-softmax matrix ($QK^\top$) and reveal two key properties: (1) video attention exhibits an effective low-dimensional structure with rapid singular-value decay, and (2) real motion induces hybrid spatio-temporal patterns rather than rigid ``spatial vs. temporal'' layouts. Guided by these observations, we propose Video-SVD, a training-free and plug-and-play acceleration method that does not modify the original network parameters. Video-SVD learns checkpoint-adaptive orthogonal bases offline and, at inference time, replaces expensive dense attention computation with lightweight online subspace projection and basis composition. To preserve high fidelity, Video-SVD further employs layer-shared dual-stream residual modules to recover fine-grained content details and positional information. Across HunyuanVideo and Wan2.1 backbones, Video-SVD achieves significant end-to-end speedups while maintaining high visual quality, reaching 1.92$\times$ on HunyuanVideo, 1.75$\times$ on Wan2.1-1.3B, and 1.79$\times$ on Wan2.1-14B.

Poster
Doubly Robust Distributionally Robust Offline Contextual Pricing

Min Xu, Xinyi Yin, Yunfan Zhang, Yuxuan Han, Houcai Shen, Caihua Chen

展开摘要

Offline contextual pricing often relies on logged observational data, but faces challenges from distributional shifts between training and deployment environments. Distributionally robust optimization (DRO) provides a principled approach to off-policy evaluation and learning (OPE/L). However, existing methods are mostly limited to discrete actions. Recent work has explored DRO for continuous treatments using inverse propensity weighting (IPW), while such IPW-based estimators can be sensitive to the convergence rate of propensity score estimates, particularly when estimated nonparametrically, which may lead to larger estimation errors and regret. In this work, we develop a doubly robust (DR) framework for distributionally robust OPE/L in continuous pricing settings. For evaluation, we propose a localized DR estimator that addresses the computational challenges of worst-case expectations by fitting only a small number of regressions, comparable to standard non-robust DR, while achieving semiparametric efficiency under mild product rate conditions. For learning, we leverage the inherent smoothness of demand noise to handle pricing-specific discontinuities in revenue outcomes (e.g., threshold-based purchase decisions), establishing a finite-sample regret bound of $\tilde{\mathcal{O}}_p(T^{-s/(2s+1)})$ for smoothness orders $s=1,2$. This bound improves upon existing regret rates in existing DRO-based off-policy learning (OPL) for continuous treatments. Extensive experiments under various levels of distribution shift validate our proposed framework.

Poster
A Theory of How Pretraining Shapes Inductive Bias in Fine-Tuning

Nicolas Anguita, Francesco Locatello, Andrew M Saxe, Marco Mondelli, Flavia Mancini, Samuel Lippl, Clémentine Carla Juliette Dominé

展开摘要

Pretraining and fine-tuning are central stages in modern machine learning systems. In practice, feature learning plays an important role across both stages: deep neural networks learn a broad range of useful features during pretraining and further refine those features during fine-tuning. However, an end-to-end theoretical understanding of how choices of initialization impact the ability to reuse and refine features during fine-tuning has remained elusive. Here we develop an analytical theory of the pretraining–fine-tuning pipeline in diagonal linear networks, deriving exact expressions for the generalization error as a function of initialization parameters and task statistics. We find that different initialization choices place the network into four distinct fine-tuning regimes that are distinguished by their ability to support feature learning and reuse—and therefore by the task statistics for which they are beneficial. In particular, a smaller initialization scale in earlier layers enables the network to both reuse and refine its features, leading to superior generalization on fine-tuning tasks that rely on a subset of pretraining features. We demonstrate empirically that the same initialization parameters impact generalization in ResNets trained on CIFAR-100 and SVHN as well as Transformers trained on modular arithmetic tasks. Overall, our results demonstrate analytically how data and network initialization interact to shape fine-tuning generalization, highlighting an important role for the relative scale of initialization across different layers in enabling continued feature learning during fine-tuning.

Poster
Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients

Dimitris Oikonomou, Nicolas Loizou

展开摘要

The stochastic Polyak step size (SPS) has proven to be a promising choice for stochastic gradient descent (SGD), delivering competitive performance relative to state-of-the-art methods on smooth convex and non-convex optimization problems, including deep neural network training. However, extensions of this approach to non-smooth settings remain in their early stages, often relying on interpolation assumptions or requiring knowledge of the optimal solution. In this work, we propose a novel SPS variant, Safeguarded SPS (SPS$_{safe}$), for the stochastic subgradient method, and provide rigorous convergence guarantees for non-smooth convex optimization with no need for strong assumptions. We further incorporate momentum into the update rule, yielding equally tight theoretical results. Comprehensive experiments on convex benchmarks and deep neural networks corroborate our theory: the proposed step size achieves competitive performance to existing adaptive baselines and exhibits stable behavior across a wide range of problem settings. Finally, in the context of deep neural network training, the gradient norms under our step size do not collapse to (near) zero, indicating robustness to vanishing gradients.

Poster
Exposing Hidden Biases in Text-to-Image Models via Automated Prompt Search

Manos Plitsis, Giorgos Bouritsas, Vassilis Katsouros, Yannis Panagakis

展开摘要

Text-to-image (TTI) diffusion models have achieved remarkable visual quality, yet they have been repeatedly shown to exhibit social biases across sensitive attributes such as gender, race and age. To mitigate these biases, existing approaches frequently depend on curated prompt datasets - either manually constructed or generated with large language models (LLMs) - as part of their training and/or evaluation procedures. Beside the curation cost, this also risks overlooking unanticipated, less obvious prompts that trigger biased generation, even in models that have undergone debiasing. In this work, we introduce Bias-Guided Prompt Search (BGPS), a framework that automatically generates prompts that aim to maximize the presence of biases in the resulting images. BGPS comprises two components: (1) an LLM instructed to produce attribute-neutral prompts and (2) attribute classifiers acting on the TTI’s internal representations that steer the decoding process of the LLM toward regions of the prompt space that amplify the image attributes of interest. We conduct extensive experiments on Stable Diffusion 1.5 and a state-of-the-art debiased model and discover an array of subtle and previously undocumented biases that severely deteriorate fairness metrics. Crucially, the discovered prompts are interpretable, i.e they may be entered by a typical user, quantitatively improving the perplexity metric compared to a prominent hard prompt optimization counterpart. Our findings uncover TTI vulnerabilities, while BGPS expands the bias search space and can act as a new evaluation tool for bias mitigation.

Poster
Decision Transformers As Zero-Shot Learners via Text-Behavior Alignment

Xin Zhang, Jonathan Martinez, Yanhua Li, Yingxue Zhang

展开摘要

Offline meta-reinforcement learning (meta-RL) aims to train agents that can generalize to unseen tasks using pre-collected data from related tasks. Recent approaches leverage the scalability of transformer architectures to model behavior sequences and support task adaptation using target task demonstrations. However, such data is often unavailable in real-world settings, where the task objective may be known but cannot be easily demonstrated. In contrast, humans routinely interpret and perform new tasks based solely on natural language instructions. In this work, we explore the potential of using natural language task descriptions to enable zero-shot task adaptation in offline meta-RL without requiring any data from the target task. We propose the Text-Guided Decision Transformer (TG-DT), a framework that enables zero-shot generalization by grounding policy learning in natural language. TG-DT learns a shared embedding space between task descriptions and behavioral trajectories via a dual contrastive and matching-based objective, ensuring robust alignment. A transformer-based policy is then conditioned on these aligned representations to generate task-appropriate actions. At test time, TG-DT synthesizes policies for unseen tasks using only their text descriptions and can optionally leverage a description-guided data sharing strategy to enhance adaptation. Experiments on standard offline meta-RL benchmarks, including MuJoCo and Meta-World, demonstrate that TG-DT achieves strong generalization to unseen tasks.

Spotlight
GCIB: Graph Contrastive Information Bottleneck for Multi-Behavior Recommendation

Likang Wu, Zihao Chen, Jianxin Zhang, Sangqi Zhu, Yuanyuan Ge, Haipeng Yang, Lei Zhang

展开摘要

With the rapid emergence of multi-behavior learning in recommender systems, leveraging auxiliary user behaviors has proven effective for mitigating target-behavior data sparsity. Yet auxiliary behavior graphs frequently contain noisy or irrelevant interactions that do not align with the target task, impeding the learning of accurate user and item embeddings. Moreover, the scarcity of direct supervised from the target behavior complicates the extraction of informative collaborative signals. In this paper, we introduce GCIB Graph Contrastive Information Bottleneck, a novel framework that denoises auxiliary behavior information and enriches target behavior representations at both the structural and feature levels. At the structural level, GCIB employs a Graph Information Bottleneck (GIB) objective to maximize mutual information between the denoised auxiliary graph and the target-behavior graph while minimizing mutual information with the original auxiliary graph. This formulation preserves task-relevant structural patterns and suppresses spurious interactions. At the feature level, we propose a cross-behavior Graph Contrastive Learning (GCL) scheme in which denoised auxiliary features and target-behavior features serve as complementary views for both users and items. By contrasting these views, GCIB enriches sparse target-behavior representations with semantics distilled from auxiliary behaviors. Extensive experiments on multiple real-world datasets demonstrate that GCIB outperforms state-of-the-art baselines, highlighting its ability to learn noise-resilient and target-aware representations for multi-behavior recommendation.

Poster
Understanding MARS: When Scaling Momentum Provably Helps

Egor Shulgin, Tamaz Gadaev, Sarit Khirirat, Peter Richtárik

展开摘要

MARS (Yuan et al., 2025) has recently emerged as a strong optimizer for large language model (LLM) training by scaling the correction term in momentum-based variance reduction (MVR). However, existing theory does not explain why this modification can improve convergence over the unscaled MVR choice $\gamma=1$. In this paper, we provide a theoretical explanation for this phenomenon. We introduce **$\gamma$-similarity**, a refined similarity condition that captures how the scaling coefficient interacts with the stochastic gradient-difference structure. This condition recovers standard similarity at $\gamma=1$ and smoothness at $\gamma=0$. Using $\gamma$-similarity, we derive convergence guarantees for fixed-$\gamma$ MARS whose complexity depends explicitly on $\gamma$ and the corresponding $\gamma$-similarity constant. The bound reveals why small values of $\gamma$ can be beneficial: they may reduce the similarity term enough to outweigh the penalty from deviating from MVR. We prove that optimizing $\gamma$ gives MARS a lower complexity guarantee than MVR. Experiments with MARS-AdamW on GPT-style LLM pretraining corroborate the theory, showing that properly chosen small values of $\gamma$ improve token efficiency over $\gamma=1$ and AdamW under a fixed training protocol.

Poster
Stabilizing PPO via Latent-Space Regularization and KDE-Driven Exploration

Meiyu Du, Yuqing Gao, Wei Wang

展开摘要

Proximal Policy Optimization (PPO) is widely used in continuous-control tasks, yet its performance is often highly sensitive to training dynamics when neural networks approximate the policy and value functions. This paper introduces SPPO, a drop-in augmentation that preserves PPO’s clipped objective and network architecture while stabilizing actor-critic geometry via three mechanisms: (i) a CKA-based constraint on critic representations, (ii) a no-flip regularizer on actor updates, and (iii) KDE-driven advantage shaping. Theoretical analysis shows that these mechanisms tighten bounds on one-step bootstrapping error, improve expected directional alignment of action updates, and ensure non-decreasing occupancy mass over high-novelty regions. Experiments on standard continuous-control benchmarks demonstrate consistent gains over PPO and recent PPO stabilization methods. Ablation studies further quantify the contribution and complementary effects of each component. Additional training-dynamics analyses indicate that SPPO reduces instability and oscillations in both actor and critic updates, improving training stability and final performance.

Poster
SafeCompass: Dynamic Chain-of-Thought Steering via Inference-Time Safety Signals

Zeyang Zhang, Haotian Xu, Linbao Li, Qi Sun, Xuebo Liu, Yu Li, Cheng Zhuo

展开摘要

Large reasoning models (LRMs) achieve strong performance by explicitly generating chain-of-thought (CoT) reasoning, but this reasoning process can be manipulated by adversarial prompts. Inference-time CoT interventions offer a simple and lightweight approach to improving safety, yet existing methods typically apply static heuristics that ignore the dynamic nature of reasoning, leading to an inherent trade-off between robustness and over-refusal. This paper introduces *SafeCompass*, a plug-and-play framework for dynamically steering chain-of-thought reasoning using inference-time safety signals extracted from internal states. At different reasoning positions, *SafeCompass* derives a latent safety direction through contrastive analysis of internal representations and uses this direction to quantify the model’s current safety state. These signals enable selective intervention, allowing the model’s reasoning trajectory to be modified only when and where it becomes unsafe. Extensive experiments demonstrate that *SafeCompass* significantly improves robustness, reducing the average attack success rate up to $10\times$ compared to the best baseline, while preserving general reasoning performance and minimizing over-refusal rates.

Poster
Learning Permutation Distributions via Reflected Diffusion on Ranks

Sizhuang He, Yangtian Zhang, Shiyang Zhang, David van Dijk

展开摘要

The finite symmetric group $S_n$ provides a natural domain for permutations, yet learning probability distributions on $S_n$ is challenging due to its factorially growing size and discrete, non-Euclidean structure. Recent permutation diffusion methods define forward noising via shuffle-based random walks (e.g., riffle shuffles) and learn reverse transitions with Plackett–Luce (PL) variants, but the resulting trajectories can be abrupt and increasingly hard to denoise as $n$ grows. We propose *Soft-Rank Diffusion*, a discrete diffusion framework that replaces shuffle-based corruption with a structured soft-rank forward process: we lift permutations to a continuous latent representation of order by relaxing discrete ranks into soft ranks, yielding smoother and more tractable trajectories. For the reverse process, we introduce *contextualized generalized Plackett–Luce (cGPL)* denoisers that generalize prior PL-style parameterizations and improve expressivity for sequential decision structures. Experiments on sorting and combinatorial optimization benchmarks show that Soft-Rank Diffusion consistently outperforms prior diffusion baselines, with particularly strong gains in long-sequence and intrinsically sequential settings.

Poster