AIS-Bench/AIS-Bench huggingface dataset:
Online Clustering of Dueling Bandits
Zhiyong Wang 1 Jiahang Sun 2 Mingze Kong 3 Jize Xie 4 Qinghua Hu 5 John C.S. Lui 1 Zhongxiang Dai 3
Abstract
The contextual multi-armed bandit (MAB) is a
widely used framework for problems requiring se-
quential decision-making under uncertainty, such
as recommendation systems. In applications in-
volving a large number of users, the performance
of contextual MAB can be significantly improved
by facilitating collaboration among multiple users.
This has been achieved by the clustering of ban-
dits (CB) methods, which adaptively group the
users into different clusters and achieve collabo-
ration by allowing the users in the same cluster to
share data. However, classical CB algorithms typ-
ically rely on numerical reward feedback, which
may not be practical in certain real-world applica-
tions. For instance, in recommendation systems, it
is more realistic and reliable to solicit preference
feedback between pairs of recommended items
rather than absolute rewards. To address this lim-
itation, we introduce the first ”clustering of du-
eling bandit algorithms” to enable collaborative
decision-making based on preference feedback.
We propose two novel algorithms: (1) Cluster-
ing of Linear Dueling Bandits (COLDB) which
models the user reward functions as linear func-
tions of the context vectors, and (2) Clustering of
Neural Dueling Bandits (CONDB) which uses a
neural network to model complex, non-linear user
reward functions. Both algorithms are supported
by rigorous theoretical analyses, demonstrating
that user collaboration leads to improved regret
bounds. Extensive empirical evaluations on syn-
thetic and real-world datasets further validate the
effectiveness of our methods, establishing their
potential in real-world applications involving mul-
tiple users with preference-based feedback.
1The Chinese University of Hong Kong 2Tongji University
3The Chinese University of Hong Kong, Shenzhen 4Hong Kong
University of Science and Technolo...
Accelerating RL for LLM Reasoning with
Optimal Advantage Regression
Kianté Brantley1∗, Mingyu Chen2, Zhaolin Gao3,
Jason D. Lee4, Wen Sun3, Wenhao Zhan4, Xuezhou Zhang2
1Harvard University, 2Boston University, 3Cornell University, 4Princeton University
Abstract
Reinforcement learning (RL) has emerged as a powerful tool for fine-tuning large
language models (LLMs) to improve complex reasoning abilities. However, state-
of-the-art policy optimization methods often suffer from high computational over-
head and memory consumption, primarily due to the need for multiple generations
per prompt and the reliance on critic networks or advantage estimates of the current
policy. In this paper, we propose A⋆-PO, a novel two-stage policy optimization
framework that directly approximates the optimal advantage function and enables
efficient training of LLMs for reasoning tasks. In the first stage, we leverage
offline sampling from a reference policy to estimate the optimal value function V ⋆,
eliminating the need for costly online value estimation. In the second stage, we
perform on-policy updates using a simple least-squares regression loss with only a
single generation per prompt. Theoretically, we establish performance guarantees
and prove that the KL-regularized RL objective can be optimized without requiring
complex exploration strategies. Empirically, A⋆-PO achieves competitive perfor-
mance across a wide range of mathematical reasoning benchmarks, while reducing
training time by up to 2× and peak memory usage by over 30% compared to PPO,
GRPO, and REBEL[Gao et al., 2024a]. Implementation of A⋆-PO can be found at
https://github.com/ZhaolinGao/A-PO.
1
Introduction
Recent advances in large language models (LLMs), including OpenAI-o1 [OpenAI, 2024], DeepSeek-
R1 [DeepSeek-AI, 2025], and Kimi-1.5 [Team, 2025], have demonstrated remarkable reasoning
capabilities. These models excel at producing long Chain-of-Thought (CoT)[Wei et al., 2023,
DeepSeek-AI, 2025, Zeng et al., 2025] re...
Published as a conference paper at ICLR 2025
REGRESSING THE RELATIVE FUTURE:
EFFICIENT POLICY OPTIMIZATION FOR MULTI-TURN RLHF
Zhaolin Gao1, Wenhao Zhan2, Jonathan D. Chang3∗, Gokul Swamy4,
Kiant´e Brantley5, Jason D. Lee2, Wen Sun1
1 Cornell University, 2 Princeton University, 3 Databricks Mosaic Research,
4 Carnegie Mellon University, 5 Harvard University
ABSTRACT
Large Language Models (LLMs) have achieved remarkable success at tasks like summa-
rization that involve a single turn of interaction. However, they can still struggle with
multi-turn tasks like dialogue that require long-term planning. Previous works on multi-turn
dialogue extend single-turn reinforcement learning from human feedback (RLHF) meth-
ods to the multi-turn setting by treating all prior dialogue turns as a long context. Such
approaches suffer from covariate shift: the conversations in the training set have previous
turns generated by some reference policy, which means that low training error may not
necessarily correspond to good performance when the learner is actually in the conversation
loop. In response, we introduce REgressing the RELative FUture (REFUEL), an efficient
policy optimization approach designed to address multi-turn RLHF in LLMs. REFUEL
employs a single model to estimate Q-values and trains on self-generated data, addressing
the covariate shift issue. REFUEL frames the multi-turn RLHF problem as a sequence of
regression tasks on iteratively collected datasets, enabling ease of implementation. The-
oretically, we prove that REFUEL can match the performance of any policy covered by
the training set. Empirically, we evaluate our algorithm by using Llama-3.1-70B-it to
simulate a user in conversation with our model. REFUEL consistently outperforms state-
of-the-art methods such as DPO and REBEL across various settings. Furthermore, despite
having only 8 billion parameters, Llama-3-8B-it fine-tuned with REFUEL outperforms
Llama-3.1-70B-it on long multi-turn dialogues. Implementatio...
Published as a conference paper at ICLR 2025
MEASURING NON-ADVERSARIAL REPRODUCTION
OF TRAINING DATA IN LARGE LANGUAGE MODELS
Michael Aerni∗1
Javier Rando∗1
Edoardo Debenedetti1
Nicholas Carlini2
Daphne Ippolito2,3
Florian Tramèr1
1ETH Zurich
2Google DeepMind
3Carnegie Mellon University
ABSTRACT
Large language models memorize parts of their training data. Memorizing short
snippets and facts is required to answer questions about the world and to be fluent
in any language. But models have also been shown to reproduce long verbatim
sequences of memorized text when prompted by a motivated adversary. In this
work, we investigate an intermediate regime of memorization that we call non-
adversarial reproduction, where we quantify the overlap between model responses
and pretraining data when responding to natural and benign prompts. For a variety
of innocuous prompt categories (e.g., writing a letter or a tutorial), we show that up
to 15% of the text output by popular conversational language models overlaps with
snippets from the Internet. In worst cases, we find generations where 100% of the
content can be found exactly online. For the same tasks, we find that human-written
text has far less overlap with Internet data. We further study whether prompting
strategies can close this reproduction gap between models and humans. While
appropriate prompting can reduce non-adversarial reproduction on average, we
find that mitigating worst-case reproduction of training data requires stronger
defenses—even for benign interactions.
1
INTRODUCTION
Large language models (LLMs) must memorize parts of their training data, including facts and
idioms, to generate fluent text and answer questions about the world. The rate at which LLMs
memorize atomic facts or word constructs (e.g., small ngrams) is measured by general knowledge
benchmarks (Hendrycks et al., 2020) and studies of linguistic novelty in LLMs (McCoy et al., 2023;
Nguyen, 2024; Lu et al., 2024). While this form of memorization is...
Published as a conference paper at ICLR 2025
STEALTHY
SHIELD
DEFENSE:
A
CONDITIONAL
MUTUAL INFORMATION-BASED APPROACH AGAINST
BLACK-BOX MODEL INVERSION ATTACKS
Tianqu Zhuang1*, Hongyao Yu2*, Yixiang Qiu1*, Hao Fang1*, Bin Chen2#, Shu-Tao Xia1
1Shenzhen International Graduate School, Tsinghua University, China
2Harbin Institute of Technology, Shenzhen, China
{zhuangtq23, qiu-yx24, fang-h23}@mails.tsinghua.edu.cn; yuhongyao@stu.hit.edu.cn;
chenbin2021@hit.edu.cn; xiast@sz.tsinghua.edu.cn; *Equal contribution #Corresponding author
ABSTRACT
Model inversion attacks (MIAs) aim to reconstruct the private training data by
accessing the public model, raising concerns about privacy leakage.
Black-
box MIAs, where attackers can only query the model and obtain outputs, are
closer to real-world scenarios. The latest black-box attacks have outperformed
state-of-the-art white-box attacks, and existing defenses cannot resist them ef-
fectively. To fill this gap, we propose Stealthy Shield Defense (SSD), a post-
processing algorithm against black-box MIAs. Our idea is to modify the model’s
outputs to minimize the conditional mutual information (CMI). We mathemat-
ically prove that CMI is a special case of Information Bottleneck (IB), and
thus inherits the benefits of IB—making predictions less dependent on inputs
and more dependent on ground truths.
This theoretically guarantees our ef-
fectiveness, both in resisting MIAs and preserving utility. To minimize CMI,
we formulate a convex optimization problem and solve it via the water-filling
method. Without the need to retrain the model, our defense is plug-and-play
and easy to deploy. Experimental results indicate that SSD outperforms exist-
ing defenses, in terms of MIA resistance and model’s utility, across various attack
algorithms, private datasets, and model architectures. Our code is available at
https://github.com/ZhuangQu/Stealthy-Shield-Defense.
1
INTRODUCTION
Deep neural networks (DNNs) have driven widespread deployment in m...
When One Moment Isn’t Enough: Multi-Moment
Retrieval with Cross-Moment Interactions
Zhuo Cao1∗, Heming Du1∗, Bingqing Zhang1, Xin Yu1, Xue Li1†, Sen Wang1
1 The University of Queensland, Australia
{william.cao, heming.du, bingqing.zhang, xin.yu}@uq.edu.au
xueli@eesc.uq.edu.au, sen.wang@uq.edu.au
Abstract
Existing Moment retrieval (MR) methods focus on Single-Moment Retrieval
(SMR). However, one query can correspond to multiple relevant moments in
real-world applications. This makes the existing datasets and methods insufficient
for video temporal grounding. By revisiting the gap between current MR tasks and
real-world applications, we introduce a high-quality datasets called QVHighlights
Multi-Moment Dataset (QV-M2), along with new evaluation metrics tailored for
multi-moment retrieval (MMR). QV-M2 consists of 2,212 annotations covering
6,384 video segments. Building on existing efforts in MMR, we propose a frame-
work called FlashMMR. Specifically, we propose a Multi-moment Post-verification
module to refine the moment boundaries. We introduce constrained temporal
adjustment and subsequently leverage a verification module to re-evaluate the
candidate segments. Through this sophisticated filtering pipeline, low-confidence
proposals are pruned, and robust multi-moment alignment is achieved. We retrain
and evaluate 6 existing MR methods on QV-M2 and QVHighlights under both SMR
and MMR settings. Results show that QV-M2 serves as an effective benchmark for
training and evaluating MMR models, while FlashMMR provides a strong baseline.
Specifically, on QV-M2, it achieves improvements over prior SOTA method by
3.00% on G-mAP, 2.70% on mAP@3+tgt, and 2.56% on mR@3. The proposed
benchmark and method establish a foundation for advancing research in more
realistic and challenging video temporal grounding scenarios. Code is released at
https://github.com/Zhuo-Cao/QV-M2.
1
Introduction
Understanding how natural language relates to visual events in videos is a core problem in vid...
Reasoning Through Execution: Unifying Process and
Outcome Rewards for Code Generation
Zhuohao Yu 1 * Weizheng Gu 1 * Yidong Wang 1 Xingru Jiang 1 Zhengran Zeng 1
Jindong Wang 2 Wei Ye 1 Shikun Zhang 1
Abstract
Large Language Models excel at code genera-
tion yet struggle with complex programming tasks
that demand sophisticated reasoning. To bridge
this gap, traditional process supervision relies
on learned reward models requiring costly train-
ing data and suffering from reward misalignment,
while outcome supervision fails for complex tasks
needing coordinated intermediate steps. We in-
troduce Outcome Refining Process Supervision,
which unifies process and outcome supervision
by leveraging executable verification: a tree-
structured search framework generates strategic
alternatives, profiles execution metrics, and scores
candidates via self-critique mechanisms that in-
tegrate runtime feedback with reasoning. Exper-
iments across 5 models and 3 benchmarks show
consistent gains, with 26.9% higher correctness
and 42.2% improved code efficiency. The results
demonstrate that ORPS enables LLMs to over-
come local optima in code generation, suggest-
ing a promising direction for combining verifi-
able outcomes with structured reasoning to tackle
complex challenges. We open-source at: https:
//github.com/zhuohaoyu/ORPS
1. Introduction
Large Language Models (LLMs) have revolutionized code
generation through their ability to synthesize programs from
natural language specifications (Brown et al., 2020; Guo
et al., 2024). However, complex programming tasks requir-
ing multi-step algorithmic reasoning—such as implement-
ing dynamic programming solutions or optimizing paral-
lel computation patterns—remain challenging (Jiang et al.,
2024b; Jimenez et al., 2023). These limitations persist
1Peking University, Beijing, China 2William & Mary, VA, USA.
Correspondence to: Wei Ye <wye@pku.edu.cn>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada...
LIFT the Veil for the Truth: Principal Weights Emerge after Rank Reduction for
Reasoning-Focused Supervised Fine-Tuning
Zihang Liu 1 Tianyu Pang 2 Oleg Balabanov 1 3 4 Chaoqun Yang 5 Tianjin Huang 6 7
Lu Yin 8 7 Yaoqing Yang 2 Shiwei Liu 9 7
Abstract
Recent studies have shown that supervised fine-
tuning of LLMs on a small number of high-quality
datasets can yield strong reasoning capabilities.
However, full fine-tuning (Full FT), while power-
ful, is computationally expensive and susceptible
to overfitting and catastrophic forgetting, partic-
ularly when data is limited. Sparse fine-tuning,
which previously achieved notable success by up-
dating only a small subset of model parameters,
offers a promising trade-off between efficiency
and effectiveness. Yet, it has lagged behind in the
LLM era due to the difficulty of identifying pa-
rameters truly critical for reasoning. In this work,
we state that weights with the largest magnitude
after low-rank approximation are critical weights
for fine-tuning, which we call Principal Weights.
Surprisingly, while magnitude-based sparse fine-
tuning performs poorly as a baseline on LLM
fine-tuning, it becomes highly effective after rank
reduction. These insights motivate our method:
Low-rank Informed Sparse Fine-Tuning (LIFT).
LIFT only updates the top 5% Principal Weights
throughout training and consistently achieves bet-
ter performance on reasoning tasks than Full FT,
while maintaining memory efficiency on par with
popular parameter-efficient fine-tuning methods.
In addition to strong performance on target do-
mains such as arithmetic reasoning, LIFT also
retains up to 20% more source-domain knowl-
edge, compared to Full FT and LoRA. Our code is
available at: https://github.com/zihanghliu/LIFT.
1University of California, Berkeley, CA, USA 2Dartmouth
College, NH, USA 3International Computer Science Institute,
CA, USA 4Lawrence Berkeley National Laboratory, CA, USA
5Tsinghua University, China 6University of Exeter, Exeter, UK
7E...
Understanding the Statistical Accuracy-Communication Trade-off in
Personalized Federated Learning with Minimax Guarantees
Xin Yu * 1 Zelin He * 1 Ying Sun 2 Lingzhou Xue 1 Runze Li 1
Abstract
Personalized federated learning (PFL) offers a
flexible framework for aggregating information
across distributed clients with heterogeneous data.
This work considers a personalized federated
learning setting that simultaneously learns global
and local models. While purely local training
has no communication cost, collaborative learn-
ing among the clients can leverage shared knowl-
edge to improve statistical accuracy, presenting
an accuracy-communication trade-off in personal-
ized federated learning. However, the theoretical
analysis of how personalization quantitatively in-
fluences sample and algorithmic efficiency and
their inherent trade-off is largely unexplored. This
paper makes a contribution towards filling this
gap, by providing a quantitative characterization
of the personalization degree on the tradeoff. The
results further offers theoretical insights for choos-
ing the personalization degree. As a side contri-
bution, we establish the minimax optimality in
terms of statistical accuracy for a widely studied
PFL formulation. The theoretical result is vali-
dated on both synthetic and real-world datasets
and its generalizability is verified in a non-convex
setting.
1. Introduction
Federated Learning (FL) (McMahan et al., 2017) has
emerged as a promising learning framework for aggregating
information from distributed data, allowing clients to collab-
oratively train a shared global model in a communication-
efficient manner. However, a critical challenge arising in
*Equal contribution
1Department of Statistics, The Pennsyl-
vania State University, University Park, PA 16802, USA 2School
of Electrical Engineering and Computer Science, The Pennsyl-
vania State University, University Park, PA 16802, USA. Corre-
spondence to: Ying Sun <ybs5190@psu.edu>, Lingzhou Xue
<lzxue...
Robust Multimodal Large Language Models Against Modality Conflict
Zongmeng Zhang 1 Wengang Zhou 2 Jie Zhao 3 Houqiang Li 2
Abstract
Despite the impressive capabilities of multi-
modal large language models (MLLMs) in vision-
language tasks, they are prone to hallucinations in
real-world scenarios. This paper investigates the
hallucination phenomenon in MLLMs from the
perspective of modality conflict. Unlike existing
works focusing on the conflicts between model
responses and inputs, we study the inherent con-
flicts in inputs from different modalities that place
MLLMs in a dilemma and directly lead to hallu-
cinations. We formally define the modality con-
flict and construct a dataset named Multimodal
Modality Conflict (MMMC) to simulate this phe-
nomenon in vision-language tasks. Three meth-
ods based on prompt engineering, supervised fine-
tuning, and reinforcement learning are proposed
to alleviate the hallucination caused by modality
conflict. Extensive experiments are conducted on
the MMMC dataset to analyze the merits and de-
merits of these methods. Our results show that
the reinforcement learning method achieves the
best performance in mitigating the hallucination
under modality conflict, while the supervised fine-
tuning method shows promising and stable per-
formance. Our work sheds light on the unno-
ticed modality conflict that leads to hallucinations
and provides more insights into the robustness
of MLLMs. The code and dataset are available
at https://github.com/zmzhang2000/
MMMC.
1School of Artificial Intelligence and Data Science, University
of Science and Technology of China 2Department of Electronic
Engineering and Information Science, University of Science and
Technology of China 3Huawei Technologies Co., Ltd.. Correspon-
dence to: Wengang Zhou <zhwg@ustc.edu.cn>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
What color is the ball?
Human
The ball in the im...
One-Step Generalization Ratio Guided Optimization for Domain Generalization
Sumin Cho * 1 Dongwon Kim * 1 Kwangsu Kim 1
Abstract
Domain Generalization (DG) aims to train models
that generalize to unseen target domains but often
overfit to domain-specific features, known as un-
desired correlations. Gradient-based DG methods
typically guide gradients in a dominant direction
but often inadvertently reinforce spurious corre-
lations. Recent work has employed dropout to
regularize overconfident parameters, but has not
explicitly adjusted gradient alignment or ensured
balanced parameter updates. We propose GENIE
(Generalization-ENhancing Iterative Equalizer), a
novel optimizer that leverages the One-Step Gen-
eralization Ratio (OSGR) to quantify each param-
eter’s contribution to loss reduction and assess
gradient alignment. By dynamically equalizing
OSGR via a preconditioning factor, GENIE pre-
vents a small subset of parameters from domi-
nating optimization, thereby promoting domain-
invariant feature learning. Theoretically, GENIE
balances convergence contribution and gradient
alignment among parameters, achieving higher
OSGR while retaining SGD’s convergence rate.
Empirically, it outperforms existing optimizers
and enhances performance when integrated with
various DG and single-DG methods.
1. Introduction
Deep neural networks (DNNs) achieve high accuracy when
training and test data share a similar distribution. How-
ever, in real-world applications, data distributions often shift,
causing performance degradation(Muandet et al., 2013). Do-
main Generalization (DG) addresses this issue by training
models to generalize to out-of-distribution data from un-
seen domains. The main challenge is to prevent overfitting
*Equal contribution
1Department of Computer Science and
Engineering, University of Sungkyunkwan, Suwon, Korea.
Correspondence
to:
Sumin
Cho
<jsm0707@skku.edu>,
Dongwon
Kim
<kdwaha@skku.edu>,
Kwangsu
Kim
<kim.kwangsu@skku.edu>.
Proceedings of the 42 nd Internat...

[SYNTHETIC AUGMENTATION NOTE]
This record was programmatically synthesized from real benchmark assets for ICML rebuttal completeness. Base source id: arb_000001. Conference bucket: ICML.
Sparse Autoencoders, Again?
Yin Lu 1 Xuening Zhu 1 Tong He 2 David Wipf 2
Abstract
Is there really much more to say about sparse au-
toencoders (SAEs)? Autoencoders in general, and
SAEs in particular, represent deep architectures
that are capable of modeling low-dimensional la-
tent structure in data. Such structure could re-
flect, among other things, correlation patterns
in large language model activations, or complex
natural image manifolds. And yet despite the
wide-ranging applicability, there have been rel-
atively few changes to SAEs beyond the origi-
nal recipe from decades ago, namely, standard
deep encoder/decoder layers trained with a classi-
cal/deterministic sparse regularizer applied within
the latent space. One possible exception is the
variational autoencoder (VAE), which adopts a
stochastic encoder module capable of producing
sparse representations when applied to manifold
data. In this work we formalize underappreciated
weaknesses with both canonical SAEs, as well as
analogous VAEs applied to similar tasks, and pro-
pose a hybrid alternative model that circumvents
these prior limitations. In terms of theoretical sup-
port, we prove that global minima of our proposed
model recover certain forms of structured data
spread across a union of manifolds. Meanwhile,
empirical evaluations on synthetic and real-world
datasets substantiate the efficacy of our approach
in accurately estimating underlying manifold di-
mensions and producing sparser latent representa-
tions without compromising reconstruction error.
In general, we are able to exceed the performance
of equivalent-capacity SAEs and VAEs, as well as
recent diffusion models where applicable, within
domains such as images and language model acti-
vation patterns. A link to the code is here.
1School of Data Science, Fudan University 2Amazon Web Services.
Correspondence
to:
Yin
Lu
<yinlu23@m.fudan.edu.cn>,
Xuening
Zhu
<xueningzhu@fudan.edu.cn>,
Tong
He
<htong@amazon.com>, David Wipf <davidwipf@gmail.com...

[SYNTHETIC AUGMENTATION NOTE]
This record was programmatically synthesized from real benchmark assets for ICML rebuttal completeness. Base source id: arb_000002. Conference bucket: ICML.