AIS-Bench/AIS-Bench huggingface dataset:
Robot-Gated Interactive Imitation Learning with Adaptive Intervention
Mechanism
Haoyuan Cai 1 Zhenghao Peng 1 Bolei Zhou 1
Abstract
Interactive Imitation Learning (IIL) allows agents
to acquire desired behaviors through human in-
terventions, but current methods impose high
cognitive demands on human supervisors. We
propose the Adaptive Intervention Mechanism
(AIM), a novel robot-gated IIL algorithm that
learns an adaptive criterion for requesting hu-
man demonstrations. AIM utilizes a proxy Q-
function to mimic the human intervention rule
and adjusts intervention requests based on the
alignment between agent and human actions. By
assigning high Q-values when the agent deviates
from the expert and decreasing these values as the
agent becomes proficient, the proxy Q-function
enables the agent to assess the real-time align-
ment with the expert and request assistance when
needed. Our expert-in-the-loop experiments re-
veal that AIM significantly reduces expert moni-
toring efforts in both continuous and discrete con-
trol tasks. Compared to the uncertainty-based
baseline Thrifty-DAgger, our method achieves a
40% improvement in terms of human take-over
cost and learning efficiency. Furthermore, AIM
effectively identifies safety-critical states for ex-
pert assistance, thereby collecting higher-quality
expert demonstrations and reducing overall ex-
pert data and environment interactions needed.
Code and demo video are available at https:
//github.com/metadriverse/AIM.
1. Introduction
Reinforcement learning (RL) has been successful in a wide
variety of tasks, from playing Atari games (Mnih et al.,
2013) to autonomous driving (Li et al., 2022a) and robotic
1Computer Science Department, University of Califor-
nia, Los Angeles (UCLA). Correspondence to:
Bolei Zhou
<bolei@cs.ucla.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
I can have a rest.
I can handle it.
Can you help me?
...
Diversity By Design: Leveraging Distribution Matching
for Offline Model-Based Optimization
Michael S. Yao 1 James C. Gee 1 Osbert Bastani 1
Abstract
The goal of offline model-based optimization
(MBO) is to propose new designs that maximize
a reward function given only an offline dataset.
However, an important desiderata is to also pro-
pose a diverse set of final candidates that capture
many optimal and near-optimal design config-
urations. We propose Diversity in Adversarial
Model-based Optimization (DynAMO) as a
novel method to introduce design diversity as an
explicit objective into any MBO problem. Our
key insight is to formulate diversity as a distribu-
tion matching problem where the distribution of
generated designs captures the inherent diversity
contained within the offline dataset. Extensive
experiments spanning multiple scientific domains
show that DynAMO can be used with common
optimization methods to significantly improve the
diversity of proposed designs while still discover-
ing high-quality candidates.
1. Introduction
Discovering designs that optimize certain desirable proper-
ties is a ubiquitous task that spans a wide range of scientific
and engineering domains. For example, we might seek
to design a drug with the most potent therapeutic efficacy
(Brown et al., 2019; Kong et al., 2023; Du et al., 2024); build
a robot that is most capable of navigating complex environ-
ments (Ahn et al., 2020; Trabucco et al., 2021; Wang et al.,
2023); or engineer a material with a certain desirable prop-
erty (Stanev et al., 2018; Pogue et al., 2023; Gashmard et al.,
2024; Ma et al., 2024). However, experimentally validating
every proposed design can be expensive, time-intensive, or
even impossible in many applications. These limitations
can preclude the use of conventional ‘online’ optimization
methods for such generative design tasks.
1University of Pennsylvania, Philadelphia, PA, USA. Corre-
spondence to: Michael Yao <myao2199@seas.upenn.edu>.
Proceedings of the...
BEST-Route: Adaptive LLM Routing with Test-Time Optimal Compute
Dujian Ding 1 Ankur Mallick 2 Shaokun Zhang 3 Chi Wang 4 Daniel Madrigal 2
Mirian Del Carmen Hipolito Garcia 2 Menglin Xia 2 Laks V. S. Lakshmanan 5 Qingyun Wu 3 6 Victor R¨uhle 2
Abstract
Large language models (LLMs) are powerful
tools but are often expensive to deploy at scale.
LLM query routing mitigates this by dynamically
assigning queries to models of varying cost and
quality to obtain a desired tradeoff. Prior query
routing approaches generate only one response
from the selected model and a single response
from a small (inexpensive) model was often not
good enough to beat a response from a large (ex-
pensive) model due to which they end up overus-
ing the large model and missing out on poten-
tial cost savings. However, it is well known that
for small models, generating multiple responses
and selecting the best can enhance quality while
remaining cheaper than a single large-model re-
sponse. We leverage this idea to propose BEST-
Route, a novel routing framework that chooses
a model and the number of responses to sample
from it based on query difficulty and the quality
thresholds. Experiments on real-world datasets
demonstrate that our method reduces costs by up
to 60% with less than 1% performance drop.
1. Introduction
Large language models (LLMs) have revolutionized natural
language processing (NLP) by delivering state-of-the-art
performance across a wide range of tasks, from language
understanding to creative writing, code generation, and be-
yond (Zhao et al., 2023). Their widespread deployment in
applications like ChatGPT (OpenAI, a) and other conversa-
tional agents (Zheng et al., 2023; Zhang et al., 2024a;b) has
made them a cornerstone of modern NLP systems. However,
the superior performance of these models often comes with
1University of British Columbia (work performed during in-
ternship at Microsoft) 2Microsoft 3Pennsylvania State Univer-
sity 4Google (work performed while at Microsoft...
Published as a conference paper at ICLR 2025
OSCILLATORY STATE-SPACE MODELS
T. Konstantin Rusch
MIT
tkrusch@mit.edu
Daniela Rus
MIT
ABSTRACT
We propose Linear Oscillatory State-Space models (LinOSS) for efficiently learn-
ing on long sequences. Inspired by cortical dynamics of biological neural net-
works, we base our proposed LinOSS model on a system of forced harmonic os-
cillators. A stable discretization, integrated over time using fast associative paral-
lel scans, yields the proposed state-space model. We prove that LinOSS produces
stable dynamics only requiring nonnegative diagonal state matrix. This is in stark
contrast to many previous state-space models relying heavily on restrictive param-
eterizations. Moreover, we rigorously show that LinOSS is universal, i.e., it can
approximate any continuous and causal operator mapping between time-varying
functions, to desired accuracy. In addition, we show that an implicit-explicit dis-
cretization of LinOSS perfectly conserves the symmetry of time reversibility of
the underlying dynamics. Together, these properties enable efficient modeling of
long-range interactions, while ensuring stable and accurate long-horizon forecast-
ing. Finally, our empirical results, spanning a wide range of time-series tasks from
mid-range to very long-range classification and regression, as well as long-horizon
forecasting, demonstrate that our proposed LinOSS model consistently outper-
forms state-of-the-art sequence models. Notably, LinOSS outperforms Mamba
and LRU by nearly 2x on a sequence modeling task with sequences of length 50k.
Code: https://github.com/tk-rusch/linoss.
1
INTRODUCTION
State-space models (Gu et al., 2021; Hasani et al., 2022; Smith et al., 2023; Orvieto et al., 2023) have
recently emerged as a powerful tool for learning on long sequences. These models posses the state-
fullness and fast inference capabilities of Recurrent Neural Networks (RNNs) together with many of
the benefits of Transformers (Vaswani, 2017; ...
AttentionPredictor: Temporal Patterns Matter
for KV Cache Compression
Qingyue Yang1*, Jie Wang1†, Xing Li2‡, Zhihai Wang1, Chen Chen2, Lei Chen2,
Xianzhi Yu2, Wulong Liu2, Jianye Hao2, 3, Mingxuan Yuan2, Bin Li1
1MoE Key Laboratory of Brain-inspired Intelligent Perception and Cognition,
University of Science and Technology of China
2Huawei Noah’s Ark Lab
3College of Intelligence and Computing, Tianjin University
{yangqingyue, zhwangx}@mail.ustc.edu.cn,
jiewangx@ustc.edu.cn, li.xing2@huawei.com
Abstract
With the development of large language models (LLMs), efficient inference through
Key-Value (KV) cache compression has attracted considerable attention, especially
for long-context generation. To compress the KV cache, recent methods identify
critical KV tokens through static modeling of attention scores. However, these
methods often struggle to accurately determine critical tokens as they neglect the
temporal patterns in attention scores, resulting in a noticeable degradation in LLM
performance. To address this challenge, we propose AttentionPredictor, which
is the first learning-based method to directly predict attention patterns for
KV cache compression and critical token identification. Specifically, Attention-
Predictor learns a lightweight, unified convolution model to dynamically capture
spatiotemporal patterns and predict the next-token attention scores. An appealing
feature of AttentionPredictor is that it accurately predicts the attention score and
shares the unified prediction model, which consumes negligible memory, among
all transformer layers. Moreover, we propose a cross-token critical cache prefetch-
ing framework that hides the token estimation time overhead to accelerate the
decoding stage. By retaining most of the attention information, AttentionPredictor
achieves 13× KV cache compression and 5.6× speedup in a cache offloading
scenario with comparable LLM performance, significantly outperforming the state-
of-the-arts. The code is available at https...
The Disparate Benefits of Deep Ensembles
Kajetan Schweighofer 1 Adrian Arnaiz-Rodriguez 2 Sepp Hochreiter 1 3 Nuria Oliver 2
Abstract
Ensembles of Deep Neural Networks, Deep En-
sembles, are widely used as a simple way to boost
predictive performance. However, their impact
on algorithmic fairness is not well understood yet.
Algorithmic fairness examines how a model’s per-
formance varies across socially relevant groups
defined by protected attributes such as age, gen-
der, or race. In this work, we explore the interplay
between the performance gains from Deep Ensem-
bles and fairness. Our analysis reveals that they
unevenly favor different groups, a phenomenon
that we term the disparate benefits effect. We
empirically investigate this effect using popular
facial analysis and medical imaging datasets with
protected group attributes and find that it affects
multiple established group fairness metrics, in-
cluding statistical parity and equal opportunity.
Furthermore, we identify that the per-group differ-
ences in predictive diversity of ensemble members
can explain this effect. Finally, we demonstrate
that the classical Hardt post-processing method is
particularly effective at mitigating the disparate
benefits effect of Deep Ensembles by leveraging
their better-calibrated predictive distributions.
1. Introduction
Deep Ensembles (Lakshminarayanan et al., 2017) have
demonstrated their efficacy as a simple and robust method
to improve the performance of individual Deep Neural Net-
works (DNNs). Their superior performance has made them
a popular choice for real-world applications (Bhusal et al.,
2021; Dolezal et al., 2022), including high-stakes scenar-
ios where the impact on people’s lives of Machine Learning
(ML) supported decisions can be profound, such as in health-
care, education, finance or the law. In such applications,
1ELLIS Unit, LIT AI Lab, Institute for Machine Learning,
JKU Linz, Austria 2ELLIS Alicante, Alicante, Spain 3NXAI
GmbH, Linz, Austria. Correspon...
pLSTM: parallelizable Linear Source Transition Mark
networks
Korbinian Pöppel1,2
Richard Freinschlag1
Thomas Schmied1
Wei Lin1
Sepp Hochreiter1,3
1 ELLIS Unit Linz, LIT AI Lab, Institute for Machine Learning
Johannes Kepler University Linz, Austria
2 Zuse School ELIZA
3 NXAI GmbH
poeppel@ml.jku.at
Abstract
Modern recurrent architectures, such as xLSTM and Mamba, have recently chal-
lenged the Transformer in language modeling. However, their structure constrains
their applicability to sequences only or requires processing multi-dimensional
data structures, such as images or molecular graphs, in a pre-defined sequential
order. In contrast, Multi-Dimensional RNNs (MDRNNs) are well suited for data
with a higher level structure, like 2D grids, trees, and directed acyclic graphs
(DAGs). In this work, we extend the notion of multi-dimensionality to linear RNNs.
We introduce parallelizable Linear Source Transition Mark networks (pLSTMs)
using Source, Transition, and Mark gates that act on the linegraph of a general
DAG. This enables parallelization in analogy to parallel associative scans and the
chunkwise-recurrent form of sequential linear RNNs, but for DAGs. For regular
grids (1D and 2D), like images, this scheme can be efficiently implemented using
einsum operations, concatenations, and padding in logarithmic time. pLSTMs
tackle the vanishing/exploding activation/gradient problem for long distances in
DAGs via two distinct modes: a directed propagation mode (P-mode) and a dif-
fusive distribution mode (D-mode). To showcase the long-range capabilities of
pLSTM, we introduce arrow-pointing extrapolation as a synthetic computer vi-
sion task that contains long-distance directional information. We demonstrate
that pLSTMs generalize well to larger image sizes, whereas Transformers strug-
gle to extrapolate. On established molecular graph and computer vision bench-
marks, pLSTMs also show strong performance. The complete code is available at
https://github.com/ml-jku/plstm_ex...
Compressing Tree Ensembles through Level-wise Optimization and Pruning
Laurens Devos * 1 2 Timo Martens * 1 2 Deniz Can Oruc¸ 1 2 Wannes Meert 1 2 Hendrik Blockeel 1 2 Jesse Davis 1 2
Abstract
Tree ensembles (e.g., gradient boosting decision
trees) are often used in practice because they offer
excellent predictive performance while still being
easy and efficient to learn. In some contexts, it
is important to additionally optimize their size:
this is specifically the case when models need
to have verifiable properties (verification of fair-
ness, robustness, etc. is often exponential in the
ensemble’s size), or when models run on battery-
powered devices (smaller ensembles consume less
energy, increasing battery autonomy). For this
reason, compression of tree ensembles is worth
studying. This paper presents LOP, a method for
compressing a given tree ensemble by pruning or
entirely removing trees in it, while updating leaf
predictions in such a way that predictive accuracy
is mostly unaffected. Empirically, LOP achieves
compression factors that are often 10 to 100 times
better than that of competing methods.
1. Introduction
Ensembles of decision trees, a.k.a. decision forests, are
frequently used in practice because of their ease of training,
excellent overall performance, and high efficiency. Their
training requires little or no hyperparameter tuning, the
computational cost of both training and prediction is very
small, and for tabular data, they are usually among the best
performing models (Grinsztajn et al., 2022).
This paper is about compressing decision forests: reducing
their size as measured by, e.g., the total number of leaves
in the forest, while maintaining their accuracy. In the light
of the above successes, one might wonder whether this is
worthwhile: if forests are already so efficient, can we gain
*Equal contribution
1KU Leuven Department of Computer
Science, Leuven, Belgium 2Leuven.AI, KU Leuven Institute for
Artificial Intelligence, Leuven, Belgium.
Co...
Bongard in Wonderland:
Visual Puzzles that Still Make AI Go Mad?
Antonia W¨ust 1 Tim Woydt 1 Lukas Helff 1 2 Inga Ibs 3 4 Wolfgang Stammer 1 2 Devendra S. Dhami 5
Constantin A. Rothkopf 2 3 4 Kristian Kersting 1 2 4 6
Abstract
Recently, newly developed Vision-Language
Models (VLMs), such as OpenAI’s o1, have
emerged, seemingly demonstrating advanced rea-
soning capabilities across text and image modal-
ities.
However, the depth of these advances
in language-guided perception and abstract rea-
soning remains underexplored, and it is unclear
whether these models can truly live up to their
ambitious promises. To assess the progress and
identify shortcomings, we enter the wonderland
of Bongard problems, a set of classic visual rea-
soning puzzles that require human-like abilities of
pattern recognition and abstract reasoning. With
our extensive evaluation setup, we show that while
VLMs occasionally succeed in identifying dis-
criminative concepts and solving some of the
problems, they frequently falter. Surprisingly,
even elementary concepts that may seem trivial to
humans, such as simple spirals, pose significant
challenges. Moreover, when explicitly asked to
recognize ground truth concepts, they continue to
falter, suggesting not only a lack of understanding
of these elementary visual concepts but also an
inability to generalize to unseen concepts. We
compare the results of VLMs to human perfor-
mance and observe that a significant gap remains
between human visual reasoning capabilities and
machine cognition. 1
1AIML Lab, TU Darmstadt 2Hessian Center for AI (hessian.ai)
3Institute of Psychology, TU Darmstadt 4Centre for Cognitive
Science, TU Darmstadt 5Uncertainty in AI Group, TU Eindhoven
6German Center for AI (DFKI). Correspondence to: Antonia W¨ust
<antonia.wuest@cs.tu-darmstadt.de>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1Our
code
is
available
at
https://github.com/...
Diffusion Federated Dataset
Seok-Ju Hahn1
Argonne National Laboratory
hahns@anl.gov
Junghye Lee2
Seoul National University
junghye@snu.ac.kr
Abstract
Diffusion models have demonstrated decent generation quality, yet their deploy-
ment in federated learning scenarios remains challenging. Due to data heterogeneity
and a large number of parameters, conventional parameter averaging schemes of-
ten fail to achieve stable collaborative training of diffusion models. We reframe
collaborative synthetic data generation as a cooperative sampling procedure from
a mixture of decentralized distributions, each encoded by a pre-trained local dif-
fusion model. This leverages the connection between diffusion and energy-based
models, which readily supports compositional generation thereof. Consequently,
we can directly obtain refined synthetic dataset, optionally with differential privacy
guarantee, even without exchanging diffusion model parameters. Our framework
reduces communication overhead while maintaining the generation quality, realized
through an unadjusted Langevin algorithm with a convergence guarantee.
1
Introduction
Federated learning (FL [1]) enables clients (i.e., data owners) to collaboratively train a statistical
model by exchanging locally updated parameters with a central server over iterative communication
rounds, thereby preserving data privacy. While this model-centric FL paradigm is well-established,
sharing public data can substantially enhance FL performance by mitigating statistical heterogeneity
arising from non-independent and identically distributed (non-IID) local data distributions [1–4]. For
instance, public or synthetic datasets can homogenize disparate local distributions and serve as direct
signals for server-side pretraining. This facilitates client-side transfer learning or data augmentation,
both improve the overall utility of FL.
While these data-centric FL scheme offers clear advantages over purely model-centric approaches,
there remain challen...
Published as a conference paper at ICLR 2025
DECOUPLED SUBGRAPH FEDERATED LEARNING
Javad Aliakbari1
Johan ¨Ostman2
Alexandre Graell i Amat1
1Chalmers University of Technology
2AI Sweden
ABSTRACT
We address the challenge of federated learning on graph-structured data distributed
across multiple clients. Specifically, we focus on the prevalent scenario of in-
terconnected subgraphs, where interconnections between different clients play a
critical role. We present a novel framework for this scenario, named FEDSTRUCT,
that harnesses deep structural dependencies. To uphold privacy, unlike existing
methods, FEDSTRUCT eliminates the necessity of sharing or generating sensitive
node features or embeddings among clients. Instead, it leverages explicit global
graph structure information to capture inter-node dependencies. We validate the ef-
fectiveness of FEDSTRUCT through experimental results conducted on six datasets
for semi-supervised node classification, showcasing performance close to the cen-
tralized approach across various scenarios, including different data partitioning
methods, varying levels of label availability, and number of clients.
1
INTRODUCTION
Many real-world data are graph-structured, where nodes represent entities and edges capture their re-
lationships. For example, in anti-money laundering, nodes symbolize accounts and edges correspond
to confidential transactions.
Graph neural networks (GNNs) are specialized neural networks for graph-structured data, showing
success in fields such as drug discovery, social networks, or traffic flow modeling (Stokes et al.,
2020; Fan et al., 2019; Jiang and Luo, 2022). In many practical applications, graph data is inherently
distributed across multiple clients, i.e., the global graph encompasses multiple, non-overlapping
subgraphs. For example, in anti-money laundering, local graphs represent internal transactional
networks of financial institutions.
For this and many other applications, data sharing among clients is ...
Optimal Sensor Scheduling and Selection for Continuous-Discrete Kalman
Filtering with Auxiliary Dynamics
Mohamad Al Ahdab 1 2 John Leth 1 Zheng-Hua Tan 1 2
Abstract
We study the Continuous-Discrete Kalman Filter
(CD-KF) for State-Space Models (SSMs) where
continuous-time dynamics are observed via multi-
ple sensors with discrete, irregularly timed mea-
surements. Our focus extends to scenarios in
which the measurement process is coupled with
the states of an auxiliary SSM. For instance,
higher measurement rates may increase energy
consumption or heat generation, while a sensor’s
accuracy can depend on its own spatial trajec-
tory or that of the measured target. Each sensor
thus carries distinct costs and constraints asso-
ciated with its measurement rate and additional
constraints and costs on the auxiliary state. We
model measurement occurrences as independent
Poisson processes with sensor-specific rates and
derive an upper bound on the mean posterior co-
variance matrix of the CD-KF along the mean
auxiliary state. The bound is continuously differ-
entiable with respect to the measurement rates,
which enables efficient gradient-based optimiza-
tion. Exploiting this bound, we propose a finite-
horizon optimal control framework to optimize
measurement rates and auxiliary-state dynam-
ics jointly. We further introduce a deterministic
method for scheduling measurement times from
the optimized rates. Empirical results in state-
space filtering and dynamic temporal Gaussian
process regression demonstrate that our approach
achieves improved trade-offs between resource
usage and estimation accuracy.
1Department of Electronic Systems, Aalborg University, Aal-
borg Øst 9220, Denmark. 2Pioneer Centre for Artificial Intelli-
gence, Copenhagen 1350, Denmark.. Correspondence to: Mo-
hamad Al Ahdab <maah@es.aau.dk>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1. Introduction
State-space...