Yancy567/AIS-Bench huggingface dataset:
Conference: ICML 2025

Guarantees of a Preconditioned Subgradient Algorithm for
Overparameterized Asymmetric Low-rank Matrix Recovery
Paris Giampouras 1 HanQin Cai 2 Ren´e Vidal 3
Abstract
In this paper, we focus on a matrix factorization-
based approach to recover low-rank asymmet-
ric matrices from corrupted measurements. We
propose an Overparameterized Preconditioned
Subgradient Algorithm (OPSA) and provide, for
the first time in the literature, linear convergence
rates independent of the rank of the sought asym-
metric matrix in the presence of gross corrup-
tions. Our work goes beyond existing results in
preconditioned-type approaches addressing their
current limitation, i.e., the lack of convergence
guarantees in the case of asymmetric matrices
of unknown rank. By applying our approach to
(robust) matrix sensing, we highlight its merits
when the measurement operator satisfies a mixed-
norm restricted isometry property. Lastly, we
present extensive numerical experiments that val-
idate our theoretical results and demonstrate the
effectiveness of our approach for different levels
of overparameterization and outlier corruptions.
1. Introduction
Low-rank matrix recovery has been a ubiquitous problem
showing up in numerous applications in the fields of sig-
nal/image processing, machine learning, and data science
(Recht et al., 2010; Chen et al., 2013; Davenport & Romberg,
2016; Cai et al., 2021a; Smith et al., 2024; Wang et al., 2024).
For instance, problems such as matrix sensing (Jain et al.,
2013; Wei et al., 2016; Park et al., 2017; Li et al., 2020),
matrix completion (Candes & Recht, 2008; Nie et al., 2012;
1Department of Computer Science, University of Warwick,
Coventry, England, UK 2Department of Statistics and Data Sci-
ence and Department of Computer Science, University of Central
Florida, Orlando, FL, USA 3Center for Innovation in Data Engi-
neering and Science (IDEAS), Departments of ESE and Radiology,
University of Pennsylvania, Philadelphia, PA, USA. Correspon-
dence to: Re...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
LLMOPT: LEARNING TO DEFINE AND SOLVE GEN-
ERAL OPTIMIZATION PROBLEMS FROM SCRATCH
Caigao Jiang♥∗,
Xiang Shu♣∗,
Hong Qian♣†,
Xingyu Lu♥†,
Jun Zhou♥,
Aimin Zhou♣,
Yang Yu♦
♣East China Normal University, China
♥Ant Group, China
♦Nanjing University, China
shux@stu.ecnu.edu.cn,
{hqian,amzhou}@cs.ecnu.edu.cn,
{caigao.jcg,sing.lxy,jun.zhoujun}@antgroup.com,
yuy@nju.edu.cn
ABSTRACT
Optimization problems are prevalent across various scenarios. Formulating and
then solving optimization problems described by natural language often requires
highly specialized human expertise, which could block the widespread applica-
tion of optimization-based decision making. To automate problem formulation
and solving, leveraging large language models (LLMs) has emerged as a potential
way. However, this kind of approach suffers from the issue of optimization gener-
alization. Namely, the accuracy of most current LLM-based methods and the gen-
erality of optimization problem types that they can model are still limited. In this
paper, we propose a unified learning-based framework called LLMOPT to boost
optimization generalization. Starting from the natural language descriptions of op-
timization problems and a pre-trained LLM, LLMOPT constructs the introduced
five-element formulation as a universal model for learning to define diverse op-
timization problem types. Then, LLMOPT employs the multi-instruction tuning
to enhance both problem formalization and solver code generation accuracy and
generality. After that, to prevent hallucinations in LLMs, such as sacrificing solv-
ing accuracy to avoid execution errors, the model alignment and self-correction
mechanism are adopted in LLMOPT. We evaluate the optimization generalization
ability of LLMOPT and compared methods across six real-world datasets cov-
ering roughly 20 fields such as health, environment, energy and manufacturing,
etc. Extensive experiment results show that LLMOPT is able to model vario...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
GOTTACK: UNIVERSAL ADVERSARIAL ATTACKS ON
GRAPH NEURAL NETWORKS VIA GRAPH ORBITS
LEARNING
Zulfikar Alom1,
Tran Gia Bao Ngo1,
Murat Kantarcioglu2,
Cuneyt Gurcan Akcora3
1Department of Computer Science, University of Manitoba, Canada
2Department of Computer Science, Virginia Tech, USA
3AI Initiative - University of Central Florida, USA
md.alom@umanitoba.ca,
ngot1@myumanitoba.ca,
muratk@vt.edu
cuneyt.akcora@ucf.edu
ABSTRACT
Graph Neural Networks (GNNs) have demonstrated superior performance in node
classification tasks across diverse applications. However, their vulnerability to
adversarial attacks, where minor perturbations can mislead model predictions,
poses significant challenges. This study introduces GOttack, a novel adversarial
attack framework that exploits the topological structure of graphs to undermine the
integrity of GNN predictions systematically.
By defining a topology-aware method to manipulate graph orbits, our approach
generates adversarial modifications that are both subtle and effective, posing a
severe test to the robustness of GNNs. We evaluate the efficacy of GOttack
across multiple prominent GNN architectures using standard benchmark datasets.
Our results show that GOttack outperforms existing state-of-the-art adversarial
techniques and completes training in approximately 55% of the time required by
the fastest competing model, achieving the highest average misclassification rate
in 155 tasks. This work not only sheds light on the susceptibility of GNNs to
structured adversarial attacks but also shows that certain topological patterns may
play a significant role in the underlying robustness of the GNNs. Our Python
implementation is shared at https://github.com/cakcora/GOttack.
1
INTRODUCTION
Recent advances in Graph Neural Networks (GNNs) have brought significant progress in node classifi-
cation tasks, utilizing the power of graph topology and node features to generate insightful inferences
across va...
Conference: ICML 2025

Progressive Tempering Sampler with Diffusion
Severi Rissanen * 1 RuiKang OuYang * 2 Jiajun He 2 Wenlin Chen 2 3 Markus Heinonen 1 Arno Solin 1
Jos´e Miguel Hern´andez-Lobato 2
Abstract
Recent research has focused on designing neu-
ral samplers that amortize the process of sam-
pling from unnormalized densities. However, de-
spite significant advancements, they still fall short
of the state-of-the-art MCMC approach, Parallel
Tempering (PT), when it comes to the efficiency
of target evaluations. On the other hand, unlike a
well-trained neural sampler, PT yields only depen-
dent samples and needs to be rerun—at consider-
able computational cost—whenever new samples
are required. To address these weaknesses, we
propose the Progressive Tempering Sampler with
Diffusion (PTSD), which trains diffusion models
sequentially across temperatures, leveraging the
advantages of PT to improve the training of neu-
ral samplers. We also introduce a novel method
to combine high-temperature diffusion models to
generate approximate lower-temperature samples,
which are minimally refined using MCMC and
used to train the next diffusion model. PTSD en-
ables efficient reuse of sample information across
temperature levels while generating well-mixed,
uncorrelated samples. Our method significantly
improves target evaluation efficiency, outperform-
ing diffusion-based neural samplers.
1. Introduction
Sampling from probability distributions is a fundamental
problem in many fields of science, including Bayesian infer-
ence (Gelman et al., 2013; Welling & Teh, 2011), statistical
physics (Von Toussaint, 2011) and molecular simulations
*Equal contribution
1Department of Computer Science,
Aalto University, Finland 2Department of Engineering, Univer-
sity of Cambridge, United Kingdom 3Department of Empir-
ical Inference, Max Planck Institute for Intelligent Systems,
T¨ubingen, Germany. Correspondence to: Severi Rissanen <sev-
eri.rissanen@aalto.fi>, RuiKang OuYang <ro352@cam.ac.uk>,
Jiajun He <jh2383@...
Conference: ICML 2025

Generating Hypotheses of Dynamic Causal Graphs in Neuroscience:
Leveraging Generative Factor Models of Observed Time Series
Zachary C. Brown 1 David Carlson 1 2
Abstract
The field of hypothesis generation promises to re-
duce costs in neuroscience by narrowing the range
of interventional studies needed to study various
phenomena. Existing machine learning methods
can generate scientific hypotheses from complex
datasets, but many approaches assume causal rela-
tionships are static over time, limiting their appli-
cability to systems with dynamic, state-dependent
behavior, such as the brain. While some tech-
niques attempt dynamic causal discovery through
factor models, they often restrict relationships to
linear patterns or impose other simplifying as-
sumptions. We propose a novel method that mod-
els dynamic graphs as a conditionally weighted
superposition of static graphs, where each static
graph can capture nonlinear relationships. This
approach enables the detection of complex, time-
varying interactions between variables beyond lin-
ear limitations. Our method improves f1-scores
of predicted dynamic causal patterns by roughly
22-28% on average over baselines in some of our
experiments, with some improvements reaching
well over 60%. A case study on real brain data
demonstrates our method’s ability to uncover re-
lationships linked to specific behavioral states,
offering valuable insights into neural dynamics.
1. Introduction
Causal discovery methods applied to brain dynamics hold
great promise for designing targeted interventions in var-
ious neurological diseases. A typical approach is to gen-
erate hypotheses about causal relationships between brain
1Department of Electrical and Computer Engineering, Duke
University, Durham, NC 15213, USA 2Departments of Civil and
Environmental Engineering; Biostatistics and Bioinformatics; and
Computer Science, Duke University, Durham, NC 15213, USA.
Correspondence to: Zachary C. Brown <zac.brown@duke.edu>,
David Carlson <dav...
Conference: ICML 2025

World Model Implanting for Test-time Adaptation of Embodied Agents
Minjong Yoo 1 Jinwoo Jang 1 Sihyung Yoon 1 Honguk Woo 1
Abstract
In embodied AI, a persistent challenge is enabling
agents to robustly adapt to novel domains with-
out requiring extensive data collection or retrain-
ing. To address this, we present a world model
implanting framework (WorMI) that combines
the reasoning capabilities of large language mod-
els (LLMs) with independently learned, domain-
specific world models through test-time com-
position.
By allowing seamless implantation
and removal of the world models, the embod-
ied agent’s policy achieves and maintains cross-
domain adaptability. In the WorMI framework,
we employ a prototype-based world model re-
trieval approach, utilizing efficient trajectory-
based abstract representation matching, to incor-
porate relevant models into test-time composition.
We also develop a world-wise compound attention
method that not only integrates the knowledge
from the retrieved world models but also aligns
their intermediate representations with the rea-
soning model’s representation within the agent’s
policy. This framework design effectively fuses
domain-specific knowledge from multiple world
models, ensuring robust adaptation to unseen do-
mains. We evaluate our WorMI on the Virtual-
Home and ALFWorld benchmarks, demonstrating
superior zero-shot and few-shot performance com-
pared to several LLM-based approaches across a
range of unseen domains. These results highlight
the framework’s potential for scalable, real-world
deployment in embodied agent scenarios where
adaptability and data efficiency are essential.
1. Introduction
In recent years, policies powered by large language models
(LLMs) have demonstrated remarkable success in embodied
1Department
of
Computer
Science
and
Engineering,
Sungkyunkwan University. Correspondence to: Honguk Woo
<hwoo@skku.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR ...
Conference: ICML 2025

Branches: Efficiently Seeking Optimal Sparse Decision Trees via AO*
Ayman Chaouki 1 Jesse Read 1 Albert Bifet 2 3
Abstract
Decision Tree (DT) Learning is a fundamental
problem in Interpretable Machine Learning, yet it
poses a formidable optimisation challenge. Prac-
tical algorithms have recently emerged, primarily
leveraging Dynamic Programming and Branch &
Bound. However, most of these approaches rely
on a Depth-First-Search strategy, which is ineffi-
cient when searching for DTs at high depths and
requires the definition of a maximum depth hyper-
parameter. Best-First-Search was also employed
by other methods to circumvent these issues. The
downside of this strategy is its higher memory
consumption, as such, it has to be designed in a
fully efficient manner that takes full advantage of
the problem’s structure. We formulate the prob-
lem within an AND/OR graph search framework
and we solve it with a novel AO*-type algorithm
called BRANCHES. We prove both optimality and
complexity guarantees for BRANCHES and we
show that it is more efficient than the state of
the art theoretically and on a variety of experi-
ments. Furthermore, BRANCHES supports non-
binary features unlike the other methods, we show
that this property can further induce larger gains
in computational efficiency.
1. Introduction
Black-box models are ill-suited for contexts where decisions
carry substantial ramifications. In healthcare for example,
misdiagnoses can delay crucial treatments and lead to severe
outcomes. Likewise, in the criminal justice system, black-
box models can obscure biases and result in discriminatory
rulings. Such risks highlight the necessity of adopting inter-
pretable models in sensitive domains.
Decision Trees (DTs) are highly interpretable due to their
1LIX, Ecole Polytechnique, IP Paris 2AI Institute, University
of Waikato 3LTCI, T´el´ecom Paris, IP Paris. Correspondence to:
Ayman Chaouki <chaoukiayman2@gmail.com>.
Proceedings of the 42 nd International Conference on Mach...
Conference: ICML 2025

Distillation Scaling Laws
Dan Busbridge 1 Amitis Shidani 2 Floris Weers 1 Jason Ramapuram 1 Etai Littwin 1 Russ Webb 1
Abstract
We propose a distillation scaling law that esti-
mates distilled model performance based on a
compute budget and its allocation between the
student and teacher. Our findings mitigate the
risks associated with large-scale distillation by
enabling compute-optimal allocation for both the
teacher and student to maximize student perfor-
mance. We provide compute-optimal distillation
recipes for two key scenarios: when a teacher
already exists, and when a teacher needs train-
ing. In settings involving many students or an
existing teacher, distillation outperforms super-
vised learning up to a compute level that scales
predictably with student size. Conversely, if only
one student is to be distilled and a teacher also re-
quires training, supervised learning is generally
preferable.
Additionally, our large-scale study
of distillation increases our understanding of the
process and helps inform experimental design.
1. Introduction
The study of scaling laws (Hestness et al., 2017; Rosen-
feld et al., 2020; Kaplan et al., 2020; Hoffmann et al.,
2022) revealed that previously trained Language Models
(LMs) could have been more capable if they had followed
a compute optimal training paradigm, which determines
the model size and the number of training tokens that give
the best performing model under a given compute budget.
Many subsequent works have followed compute optimal
training (Dey et al., 2023; Muennighoff et al., 2023b).
The size of compute optimal models grows with compute
(Hoffmann et al., 2022), which makes them challenging
to use due to the growth in inference costs. In practice,
1Apple 2University of Oxford, UK. Work done during an
internship at Apple.
For a full breakdown of contributions
see Appendix J. Correspondence to:
Dan Busbridge <dbus-
bridge@apple.com>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouve...
Conference: ICML 2025

Provable Policy Gradient for Robust Average-Reward MDPs
Beyond Rectangularity
Qiuhao Wang * 1 Yuqi Zha * 2 Chin Pang Ho 2 Marek Petrik 3
Abstract
Robust Markov Decision Processes (MDPs) of-
fer a promising framework for computing reliable
policies under model uncertainty. While policy
gradient methods have gained increasing popular-
ity in robust discounted MDPs, their application
to the average-reward criterion remains largely
unexplored. This paper proposes a Robust Pro-
jected Policy Gradient (RP2G), the first generic
policy gradient method for robust average-reward
MDPs (RAMDPs) that is applicable beyond the
typical rectangularity assumption on transition
ambiguity. In contrast to existing robust policy
gradient algorithms, RP2G incorporates an adap-
tive decreasing tolerance mechanism for efficient
policy updates at each iteration. We also present a
comprehensive convergence analysis of RP2G for
solving ergodic tabular RAMDPs. Furthermore,
we establish the first study of the inner worst-
case transition evaluation problem in RAMDPs,
proposing two gradient-based algorithms tailored
for rectangular and general ambiguity sets, each
with provable convergence guarantees. Numerical
experiments confirm the global convergence of
our new algorithm and demonstrate its superior
performance.
1. Introduction
Markov Decision Processes (MDPs) (Puterman, 2014) pro-
vide a powerful framework for sequential decision-making,
with applications such as game solving (Mnih, 2013), health-
care (Shechter et al., 2008), and finance (Bäuerle & Rieder,
*Equal contribution 1Fintech Innovation Center, Research In-
stitute for Digital Economy and Interdisciplinary Sciences, South-
western University of Finance and Economics 2Department of
Data Science, City University of Hong Kong 3Department of
Computer Science, University of New Hampshire. Correspon-
dence to: Chin Pang Ho <clint.ho@cityu.edu.hk>, Marek Petrik
<mpetrik@cs.unh.edu>.
Proceedings of the 42 nd International Conference on Mach...
Conference: ICML 2025

Causal Discovery from Conditionally Stationary Time Series
Carles Balsells-Rodas 1 Xavier Sumba 1 Tanmayee Narendra 2 Ruibo Tu 3
Gabriele Schweikert 2 Hedvig Kjellstr¨om 3 Yingzhen Li 1
Abstract
Causal discovery, i.e., inferring underlying causal
relationships from observational data, is highly
challenging for AI systems. In a time series mod-
eling context, traditional causal discovery meth-
ods mainly consider constrained scenarios with
fully observed variables and/or data from station-
ary time-series.
We develop a causal discov-
ery approach to handle a wide class of nonsta-
tionary time series that are conditionally station-
ary, where the nonstationary behaviour is mod-
eled as stationarity conditioned on a set of latent
state variables. Named State-Dependent Causal
Inference (SDCI), our approach is able to recover
the underlying causal dependencies, with prov-
able identifiablity for the state-dependent causal
structures. Empirical experiments on nonlinear
particle interaction data and gene regulatory net-
works demonstrate SDCI’s superior performance
over baseline causal discovery methods.
Im-
proved results over non-causal RNNs on mod-
eling NBA player movements demonstrate the
potential of our method and motivate the use of
causality-driven methods for forecasting.
1. Introduction
Deep learning has achieved profound success in vision and
language modelling (Brown et al., 2020; Nichol et al., 2022).
Still, it remains a grand challenge for deep neural networks
to perform causal discovery (Yi et al., 2020; Girdhar & Ra-
manan, 2020; Sauer & Geiger, 2021), which is critical for
interpretability, generalization, and robustness (Lake et al.,
2017; Sch¨olkopf et al., 2021). Theoretically, structure iden-
tifiability is key to ensure a unique correspondence between
observations and the underlying causal structures (Peters
1Imperial College London 2University of Dundee 3KTH Royal
Institute of Technology. Correspondence to: Carles Balsells-Rodas
<cb221@imperial.ac.u...
Conference: ICML 2025

Auditing Prompt Caching in Language Model APIs
Chenchen Gu 1 Xiang Lisa Li 1 Rohith Kuditipudi 1 Percy Liang 1 Tatsunori Hashimoto 1
Abstract
Prompt caching in large language models (LLMs)
results in data-dependent timing variations:
cached prompts are processed faster than non-
cached prompts. These timing differences intro-
duce the risk of side-channel timing attacks. For
example, if the cache is shared across users, an
attacker could identify cached prompts from fast
API response times to learn information about
other users’ prompts. Because prompt caching
may cause privacy leakage, transparency around
the caching policies of API providers is important.
To this end, we develop and conduct statistical au-
dits to detect prompt caching in real-world LLM
API providers. We detect global cache sharing
across users in seven API providers, including
OpenAI, resulting in potential privacy leakage
about users’ prompts. Timing variations due to
prompt caching can also result in leakage of in-
formation about model architecture. Namely, we
find evidence that OpenAI’s embedding model
is a decoder-only Transformer, which was previ-
ously not publicly known.1
1. Introduction
Transformer large language models (LLMs) are computa-
tionally expensive and slow to run. To address this challenge,
recent work has developed optimizations to make LLM in-
ference and serving more efficient, such as prompt caching
(Zheng et al., 2024a; Gim et al., 2024). In prompt caching,
reuse of the attention key-value (KV) cache across requests
results in cache hits and faster response times for prompts
that share a prefix with a cached prompt.
However, prompt caching results in data-dependent timing
1Stanford University.
Correspondence to:
Chenchen
Gu
<cygu@cs.stanford.edu>,
Tatsunori
Hashimoto
<thashim@stanford.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1We release code and data at https://github.c...
Conference: NeurIPS 2025

Teaching Language Models to Reason with Tools
Chengpeng Li∗1,2, Zhengyang Tang∗2,3, Ziniu Li∗3,4, Mingfeng Xue2, Keqin Bao1,2, Tian Ding4,
Ruoyu Sun3,4, Benyou Wang3, Xiang Wang1, Junyang Lin2, and Dayiheng Liu†2
1University of Science and Technology of China
2Qwen Team, Alibaba Inc.
3The Chinese University of Hong Kong, Shenzhen
4Shenzhen International Center for Industrial and Applied Mathematics, Shenzhen Research Institute
of Big Data
Abstract
Large reasoning models (LRMs) like OpenAI-o1 have shown impressive capa-
bilities in natural language reasoning. However, these models frequently demon-
strate inefficiencies or inaccuracies when tackling complex mathematical opera-
tions. While integrating computational tools such as Code Interpreters (CIs) offers
a promising solution, it introduces a critical challenge: a conflict between the
model’s internal, probabilistic reasoning and the external, deterministic knowl-
edge provided by the CI, which often leads models to unproductive deliberation.
To overcome this, we introduce CoRT (Code-Optimized Reasoning Training), a
post-training framework designed to teach LRMs to effectively utilize CIs. We
propose Hint-Engineering, a new data synthesis strategy that strategically injects
diverse hints at optimal points within reasoning paths. This approach generates
high-quality, code-integrated reasoning data specifically tailored to optimize LRM-
CI interaction. Using this method, we have synthesized 30 high-quality samples
to post-train models ranging from 1.5B to 32B parameters through supervised
fine-tuning. CoRT further refines the multi-round interleaving of external CI usage
and internal thinking by employing rejection sampling and reinforcement learn-
ing. Our experimental evaluations demonstrate CoRT’s effectiveness, yielding
absolute improvements of 4% and 8% on DeepSeek-R1-Distill-Qwen-32B and
DeepSeek-R1-Distill-Qwen-1.5B, respectively, across five challenging mathemati-
cal reasoning datasets. Moreover, CoRT sign...