dataset-review/dataset-review huggingface dataset:
World Model Implanting for Test-time Adaptation of Embodied Agents
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 ...
Branches: Efficiently Seeking Optimal Sparse Decision Trees via AO*
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...
Distillation Scaling Laws
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...
Provable Policy Gradient for Robust Average-Reward MDPs  Beyond Rectangularity
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...
Causal Discovery from Conditionally Stationary Time Series
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...
Auditing Prompt Caching in Language Model APIs
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...
Teaching Language Models to Reason with Tools
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...
Information-Theoretic Discrete Diffusion
NeurIPS 2025

Information-Theoretic Discrete Diffusion
Moongyu Jeon1
Sangwoo Shin1
Dongjae Jeon2
Albert No1∗
1Department of Artificial Intelligence, Yonsei University
2Department of Computer Science, Yonsei University
Abstract
We present an information-theoretic framework for discrete diffusion models that
yields principled estimators of log-likelihood using score-matching losses. Inspired
by the I-MMSE identity for the Gaussian setup, we derive analogous results for the
discrete setting. Specifically, we introduce the Information–Minimum Denoising
Score Entropy (I-MDSE) relation, which links mutual information between data
and its diffused version to the minimum denoising score entropy (DSE) loss. We
extend this theory to masked diffusion and establish the Information–Minimum
Denoising Cross-Entropy (I-MDCE) relation, connecting cross-entropy losses to
mutual information in discrete masked processes. These results provide a time-
integral decomposition of the log-likelihood of the data in terms of optimal score-
based losses, showing that commonly used losses such as DSE and DCE are not
merely variational bounds but tight and principled estimators of log-likelihood. The
I-MDCE decomposition further enables practical extensions, including time-free
formula, conditional likelihood estimation in prompt–response tasks, and coupled
Monte Carlo estimation of likelihood ratios. Experiments on synthetic and real-
world data confirm the accuracy, variance stability, and utility of our estimators. The
code is publicly available at https://github.com/Dongjae0324/infodis.
1
Introduction
Diffusion models have emerged as a powerful framework for generative modeling, enabling state-
of-the-art performance in continuous domains such as image and audio generation (Sohl-Dickstein
et al., 2015; Ho et al., 2020; Chen et al., 2021; Kong et al., 2021; Saharia et al., 2022). Central to
these models is the idea of gradually corrupting data through a forward noising process, and learning
to reverse this...
Learning the Electronic Hamiltonian of Large Atomic Structures
ICML 2025

Learning the Electronic Hamiltonian of Large Atomic Structures
Chen Hao Xia * 1 Manasa Kaniselvan * 1 Alexandros Nikolaos Ziogas 1 Marko Mladenovi´c 1 Rayen Mahjoub 1
Alexander Maeder 1 Mathieu Luisier 1
Abstract
Graph neural networks (GNNs) have shown
promise in learning the ground-state electronic
properties of materials, subverting ab initio den-
sity functional theory (DFT) calculations when
the underlying lattices can be represented as small
and/or repeatable unit cells (i.e., molecules and
periodic crystals). Realistic systems are, how-
ever, non-ideal and generally characterized by
higher structural complexity. As such, they re-
quire large (10+ ˚A) unit cells and thousands of
atoms to be accurately described. At these scales,
DFT becomes computationally prohibitive, mak-
ing GNNs especially attractive. In this work, we
present a strictly local equivariant GNN capable
of learning the electronic Hamiltonian (H) of
realistically extended materials. It incorporates
an augmented partitioning approach that enables
training on arbitrarily large structures while pre-
serving local atomic environments beyond bound-
aries. We demonstrate its capabilities by predict-
ing the electronic Hamiltonian of various systems
with up to 3,000 nodes (atoms), 500,000+ edges,
∼28 million orbital interactions (nonzero entries
of H), and ≤0.53% error in the eigenvalue spec-
tra. Our work expands the applicability of current
electronic property prediction methods to some of
the most challenging cases encountered in compu-
tational materials science, namely systems with
disorder, interfaces, and defects.
1. Introduction
Graph neural networks (GNNs) have proven capable of
learning properties that are functions of molecular and
*Equal contribution
1Integrated Systems Laboratory, Depart-
ment of Information Technology and Electrical Engineering, ETH
Zurich, Zurich, Switzerland.
Correspondence to: Chen Hao
Xia <chexia@iis.ee.ethz.ch>, Manasa Kaniselvan <mkanisel-
van@iis.ee.ethz.ch>.
Proce...
PENCIL: Long Thoughts with Short Memory
ICML 2025

PENCIL: Long Thoughts with Short Memory
Chenxiao Yang 1 Nathan Srebro 1 David McAllester 1 Zhiyuan Li 1
Abstract
While state-of-the-art LLMs have demonstrated
great promise of using long Chains-of-Thought
(CoT) to boost reasoning, scaling it up to more
challenging problems is fundamentally limited by
suboptimal memory usage — intermediate com-
putations accumulate indefinitely in context even
no longer needed for future thoughts. We intro-
duce PENCIL, which incorporates a novel reduc-
tion mechanism into the autoregressive genera-
tion process that recursively clean up intermedi-
ate thoughts based on patterns learned from train-
ing. By alternately generating and erasing, PEN-
CIL can think deeper to solve harder problems
using shorter context and less computes. Empiri-
cally, for example, we demonstrate PENCIL with
a small 25M-parameter transformer and 2048 con-
text length solves Einstein’s puzzle — a task that
challenges much larger models like GPT-4. Theo-
retically, we prove PENCIL can perform universal
efficient computation by simulating any Turing
machines with optimal time and space complexity,
and thus can solve arbitrary computable tasks that
are otherwise intractable for vanilla CoT.
1. Introduction
Recently, there has been a surge of interest in reasoning
with Chain-of-Thought (CoT) (Wei et al., 2022) and gen-
erating longer thoughts at test-time to tackle larger-scale
and more complicated problems (OpenAI, 2024; Guo et al.,
2025; Snell et al., 2024; Muennighoff et al., 2025). CoT is
an iterative generation process: each intermediate reason-
ing step is appended to the current context and treated as
the input in subsequent reasoning. The context grows until
reaching a final answer. While such an iterative model is
theoretically powerful – capable, in principle, of tackling
many intricate problems given unlimited length (Merrill &
1Toyota Technological Institute at Chicago. Correspondence
to: Chenxiao Yang <chenxiao@ttic.edu>.
Proceedings of the 42 nd I...
Feedforward Few-shot Species Range Estimation
ICML 2025

Feedforward Few-shot Species Range Estimation
Christian Lange 1 Max Hamilton 2 Elijah Cole 3 Alexander Shepard 4 Samuel Heinrich 5 Angela Zhu 2
Subhransu Maji 2 Grant Van Horn 2 Oisin Mac Aodha 1
Abstract
Knowing where a particular species can or can-
not be found on Earth is crucial for ecological
research and conservation efforts. By mapping
the spatial ranges of all species, we would ob-
tain deeper insights into how global biodiversity
is affected by climate change and habitat loss.
However, accurate range estimates are only avail-
able for a relatively small proportion of all known
species. For the majority of the remaining species,
we typically only have a small number of records
denoting the spatial locations where they have
previously been observed. We outline a new ap-
proach for few-shot species range estimation to
address the challenge of accurately estimating the
range of a species from limited data. During infer-
ence, our model takes a set of spatial locations as
input, along with optional metadata such as text or
an image, and outputs a species encoding that can
be used to predict the range of a previously un-
seen species in a feedforward manner. We evalu-
ate our approach on two challenging benchmarks,
where we obtain state-of-the-art range estimation
performance, in a fraction of the compute time,
compared to recent alternative approaches.
1. Introduction
Understanding the spatial distribution of plant and ani-
mal species is essential to mitigate the ongoing decline
in global biodiversity (Jetz et al., 2019). Monitoring these
distributions over time allows us to quantify the impacts
of climate change, habitat loss, and conservation interven-
tions (Mantyka-pringle et al., 2012). Estimating a species’
spatial distribution typically starts with collecting a set of
observations that denote the locations where the species has
been confirmed to be present or absent. Traditionally, this
1University of Edinburgh
2UMass Amherst
3GenBio AI
4iNaturalist
5Co...