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

Graph Neural Network Generalization
with Gaussian Mixture Model Based Augmentation
Yassine Abbahaddou 1 * Fragkiskos D. Malliaros 2 Johannes F. Lutzeyer 1
Amine M. Aboussalah 3 Michalis Vazirgiannis 1 4
Abstract
Graph Neural Networks (GNNs) have shown
great promise in tasks like node and graph classi-
fication, but they often struggle to generalize, par-
ticularly to unseen or out-of-distribution (OOD)
data.
These challenges are exacerbated when
training data is limited in size or diversity. To
address these issues, we introduce a theoretical
framework using Rademacher complexity to com-
pute a regret bound on the generalization error
and then characterize the effect of data augmen-
tation. This framework informs the design of
GRATIN, an efficient graph data augmentation
algorithm leveraging the capability of Gaussian
Mixture Models (GMMs) to approximate any dis-
tribution. Our approach not only outperforms
existing augmentation techniques in terms of gen-
eralization but also offers improved time com-
plexity, making it highly suitable for real-world
applications. Our code is publicly available at:
https://github.com/abbahaddou/GRATIN.
1. Introduction
Graphs are a fundamental and ubiquitous structure for mod-
eling complex relationships and interactions. In biology,
graphs are employed to represent complex networks of pro-
tein interactions and in drug discovery by modeling molecu-
lar relationships (Gaudelet et al., 2021; Jagtap et al., 2022).
Similarly, social networks capture relationships and commu-
nity interactions (Aboussalah et al., 2023a; Zeng et al., 2022;
Malliaros & Vazirgiannis, 2013; Newman et al., 2002). To
address the unique challenges posed by graph-structured
*Work partially completed during a research visit to NYU.
1LIX, ´Ecole Polytechnique, IP Paris 2Universit´e Paris-Saclay,
CentraleSup´elec, Inria 3NYU Tandon School of Engineering
4MBZUAI. Correspondence to:
Yassine Abbahaddou <yas-
sine.abbahaddou@polytechnique.edu>.
Proceedings of the 42 nd...
Conference: ICML 2025

Loss Functions and Operators Generated by f-Divergences
Vincent Roulet 1 Tianlin Liu 1 Nino Vieillard 1 Micha¨el E. Sander 1 Mathieu Blondel 1
Abstract
The logistic loss (a.k.a. cross-entropy loss) is one
of the most popular loss functions used for mul-
ticlass classification. It is also the loss function
of choice for next-token prediction in language
modeling. It is associated with the Kullback–
Leibler (KL) divergence and the softargmax oper-
ator. In this work, we build upon Fenchel-Young
losses to construct convex loss functions gener-
ated from f-divergences. Our loss functions gen-
eralize the logistic loss in two directions: i) by
replacing the KL divergence with f-divergences
and ii) by allowing non-uniform reference mea-
sures. We instantiate our framework for numerous
f-divergences, recovering existing losses and cre-
ating new ones. By analogy with the logistic loss,
the loss function generated by an f-divergence
is associated with an operator, that we dub f-
softargmax. We derive a novel parallelizable bi-
section algorithm for computing the f-softargmax
associated with any f-divergence. On the empiri-
cal side, one of the goals of this paper is to deter-
mine the effectiveness of loss functions beyond
the classical cross-entropy in a language model
setting, including on pre-training, post-training
(SFT) and distillation. We show that the loss
function generated by the α-divergence (which is
equivalent to Tsallis α-negentropy in the case of
unit reference measures) with α = 1.5 performs
well across several tasks.
1. Introduction
The logistic loss, a.k.a. cross-entropy loss, is widely used
for multiclass classification. It is associated with the sof-
targmax (a.k.a. softmax) operator, which turns logits into
class probabilities. The logistic loss and the softargmax
operators are also frequently used in the space of tokens for
language model pre-training, post-training and distillation.
1Google DeepMind. Correspondence to: Mathieu Blondel,
Vincent Roulet ...
Conference: ICML 2025

Zero-Inflated Bandits
Haoyu Wei * 1 Runzhe Wan * 2 Lei Shi 2 Rui Song 2
Abstract
Many real-world bandit applications are charac-
terized by sparse rewards, which can significantly
hinder learning efficiency. Leveraging problem-
specific structures for careful distribution model-
ing is recognized as essential for improving esti-
mation efficiency in statistics. However, this ap-
proach remains under-explored in the context of
bandits. To address this gap, we initiate the study
of zero-inflated bandits, where the reward is mod-
eled using a classic semi-parametric distribution
known as the zero-inflated distribution. We de-
velop algorithms based on the Upper Confidence
Bound and Thompson Sampling frameworks for
this specific structure. The superior empirical
performance of these methods is demonstrated
through extensive numerical studies.
1. Introduction
Bandit algorithms have been widely applied in areas such
as clinical trials (Durand et al., 2018), finance (Shen et al.,
2015), recommendation systems (Zhou et al., 2017), among
others. Accurate uncertainty quantification is key to ad-
dressing the exploration-exploitation trade-off and typically
requires on certain reward distribution assumptions. Exist-
ing assumptions can be roughly classified into two groups:
• Parametric: the reward distribution is assumed to be-
long to a parameterized family, such as Gaussian or
Bernoulli distributions (Audibert et al., 2010; Krause
& Ong, 2011; Agrawal & Goyal, 2012). The strong
assumption typically ensures clean algorithmic design
and nice theoretical results. However, when misspeci-
fied, it may lead to over- or under-exploration.
*Equal contribution 1Department of Economics, University of
California San Diego, La Jolla, USA 2Amazon, Seattle, USA.
This work does not relate to the positions at Amazon.
Cor-
respondence to: Haoyu Wei <h8wei@ucsd.edu>, Rui Song
<songray@gmail.com>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 2...
Conference: NeurIPS 2025

Activation-Informed Merging of
Large Language Models
Amin Heyrani Nobari∗1, Kaveh Alim∗1, Ali ArjomandBigdeli2
Akash Srivastava3, Faez Ahmed1, Navid Azizan1
1Massachusetts Institute of Technology
2Stony Brook University
3MIT-IBM Watson AI Lab & Red Hat AI Innovation
Abstract
Model merging, a method that combines the parameters and embeddings of mul-
tiple fine-tuned large language models (LLMs), offers a promising approach to
enhance model performance across various tasks while maintaining computational
efficiency. This paper introduces Activation-Informed Merging (AIM), a technique
that integrates the information from the activation space of LLMs into the merging
process to improve performance and robustness. AIM is designed as a flexible,
complementary solution that is applicable to any existing merging method. It aims
to preserve critical weights from the base model, drawing on principles from con-
tinual learning (CL) and model compression. Utilizing a task-agnostic calibration
set, AIM selectively prioritizes essential weights during merging. We empirically
demonstrate that AIM significantly enhances the performance of merged models
across multiple benchmarks. Our findings suggest that considering the activation-
space information can provide substantial advancements in the model merging
strategies for LLMs with up to 40% increase in benchmark performance. Our code
is publicly available at https://github.com/ahnobari/ActivationInformedMerging.
1
Introduction
Foundation models are rapidly becoming the dominant force for building Artificial Intelligence (AI)
systems. In many cases, researchers build their machine learning models by starting from pre-trained
foundation models and fine-tuning (FT) these pre-trained models for some desired target task [27]. In
such a paradigm, numerous fine-tuned models are developed for various tasks. However, an important
opportunity is missed, as these fine-tuned task-specialized models typically operate in isolation
without leve...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
OPTIONZERO: PLANNING WITH LEARNED OPTIONS
Po-Wei Huang1,2, Pei-Chiun Peng1,2, Hung Guei1, Ti-Rong Wu1†
1Institute of Information Science, Academia Sinica, Taiwan
2Department of Computer Science, National Yang Ming Chiao Tung University, Taiwan
ABSTRACT
Planning with options – a sequence of primitive actions – has been shown ef-
fective in reinforcement learning within complex environments. Previous studies
have focused on planning with predefined options or learned options through ex-
pert demonstration data. Inspired by MuZero, which learns superhuman heuristics
without any human knowledge, we propose a novel approach, named OptionZero.
OptionZero incorporates an option network into MuZero, providing autonomous
discovery of options through self-play games. Furthermore, we modify the dy-
namics network to provide environment transitions when using options, allowing
searching deeper under the same simulation constraints. Empirical experiments
conducted in 26 Atari games demonstrate that OptionZero outperforms MuZero,
achieving a 131.58% improvement in mean human-normalized score. Our behav-
ior analysis shows that OptionZero not only learns options but also acquires strate-
gic skills tailored to different game characteristics. Our findings show promising
directions for discovering and using options in planning. Our code is available at
https://rlg.iis.sinica.edu.tw/papers/optionzero.
1
INTRODUCTION
Reinforcement learning is a decision-making process in which an agent interacts with environments
by selecting actions at each step to maximize long-term rewards. Actions, commonly referred to as
primitive action, advance the state by one step each. While this granularity allows precise control
at each time step, it can lead to inefficiencies in scenarios where predictable sequences of actions
are beneficial. For example, in a maze navigation task, it is more efficient to choose a sequence
of actions – such as following a strai...
Conference: NeurIPS 2025

Ascent Fails to Forget
Ioannis Mavrothalassitis∗
Pol Puigdemont∗
Noam Itzhak Levi∗
Volkan Cevher
LIONS, École Polytechnique Fédérale de Lausanne (EPFL), Lausanne, Switzerland
{ioannis.mavrothalassitis, pol.puigdemontplana, noam.levi}@epfl.ch
Abstract
Contrary to common belief, we show that gradient ascent-based unconstrained
optimization methods frequently fail to perform machine unlearning, a phenomenon
we attribute to the inherent statistical dependence between the forget and retain
data sets. This dependence, which can manifest itself even as simple correlations,
undermines the misconception that these sets can be independently manipulated
during unlearning. We provide empirical and theoretical evidence showing these
methods often fail precisely due to this overlooked relationship. For random forget
sets, this dependence means that degrading forget set metrics (which, for the oracle,
should mirror test set metrics) inevitably harms overall test performance. Going
beyond random sets, we consider logistic regression as an instructive example
where a critical failure mode emerges: inter-set dependence causes gradient descent-
ascent iterations to progressively diverge from the oracle. Strikingly, these methods
can converge to solutions that are not only far from the oracle but are potentially
even further from it than the original model itself, rendering the unlearning process
actively detrimental. A toy example further illustrates how this dependence can
trap models in inferior local minima, inescapable via finetuning. Our findings
highlight that the presence of such statistical dependencies, even when manifest
only as correlations, can be sufficient for ascent-based unlearning to fail. Our
theoretical insights are corroborated by experiments on complex neural networks,
demonstrating that these methods do not perform as expected in practice due to
this unaddressed statistical interplay.
1
Introduction
Machine learning models have become an integral part of modern r...
Conference: NeurIPS 2025

Low-Precision Streaming PCA
Sanjoy Dasgupta
University of California San Diego
sadasgupta@ucsd.edu
Syamantak Kumar
University of Texas at Austin
syamantak@utexas.edu
Shourya Pandey
University of Texas at Austin
shouryap@utexas.edu
Purnamrita Sarkar
University of Texas at Austin
purna.sarkar@utexas.edu
Abstract
Low-precision Streaming PCA estimates the top principal component in a stream-
ing setting under limited precision. We establish an information-theoretic lower
bound on the quantization resolution required to achieve a target accuracy for the
leading eigenvector. We study Oja’s algorithm for streaming PCA under linear and
nonlinear stochastic quantization. The quantized variants use unbiased stochas-
tic quantization of the weight vector and the updates. Under mild moment and
spectral-gap assumptions on the data distribution, we show that a batched version
achieves the lower bound up to logarithmic factors under both schemes. This
leads to a nearly dimension-free quantization error in the nonlinear quantization
setting. Empirical evaluations on synthetic streams validate our theoretical findings
and demonstrate that our low-precision methods closely track the performance of
standard Oja’s algorithm.
1
Introduction
Quantization (or discretization) is the mapping of a continuous set of values to a small, finite set
of outputs close to the original values; standard methods for quantization include rounding and
truncation. The current popularity of training large-scale Machine Learning models has brought a
renewed focus on quantization, though its origins go back to the 1800s. Some early examples include
least-squares methods applied to large-scale data analysis in the early nineteenth century [Sti86].
In 1867, discretization was introduced for the approximate calculation of integrals [Rie67], and the
effects of rounding errors in integration were examined in 1897 [She97]. For an excellent survey and
history of quantization, see [GKD`22].
In the context of efficie...
Conference: ICML 2025

DUNIA: Pixel-Sized Embeddings via Cross-Modal Alignment for Earth
Observation Applications
Ibrahim Fayad 1 2 Max Zimmer 3 Martin Schwartz 1 Fabian Gieseke 4 Philippe Ciais 1 Gabriel Belouze 1
Sarah Brood 5 Aurelien De Truchis 2 Alexandre d’Aspremont 5 2
Abstract
Significant efforts have been directed towards
adapting self-supervised multimodal learning for
Earth observation applications. However, most
current methods produce coarse patch-sized em-
beddings, limiting their effectiveness and inte-
gration with other modalities like LiDAR. To
close this gap, we present DUNIA, an ap-
proach to learn pixel-sized embeddings through
cross-modal alignment between images and full-
waveform LiDAR data. As the model is trained in
a contrastive manner, the embeddings can be di-
rectly leveraged in the context of a variety of envi-
ronmental monitoring tasks in a zero-shot setting.
In our experiments, we demonstrate the effec-
tiveness of the embeddings for seven such tasks:
canopy height mapping, fractional canopy cover,
land cover mapping, tree species identification,
plant area index, crop type classification, and per-
pixel waveform-based vertical structure mapping.
The results show that the embeddings, along with
zero-shot classifiers, often outperform specialized
supervised models, even in low-data regimes. In
the fine-tuning setting, we show strong perfor-
mances near or better than the state-of-the-art on
five out of six tasks.
1. Introduction
With the rapid expansion of Earth Observation (EO) satellite
missions, deep learning has emerged as a powerful solution
1Laboratoire des Sciences du Climat et de l’Environnement,
LSCE/IPSL, France
2Kayrros SAS, Paris 75009,
France
3Department for AI in Society, Science, and Technology, Zuse In-
stitute Berlin, Germany 4Department of Information Systems, Uni-
versity of M¨unster, Germany 5Department of Computer Science,
CNRS, INRIA & ´Ecole Normale Sup´erieure, Paris 75230, France.
Correspondence to: Ibrahim Fayad <ifayad@lsce.ipsl....
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
CORESET SPECTRAL CLUSTERING
Ben Jourdan1, Gregory Schwartzman2, Peter Macgregor3, and He Sun1
1University of Edinburgh, UK
2Japan Advanced Institute of Science and Technology (JAIST), Japan
3University of St Andrews, UK
ben.jourdan@ed.ac.uk, greg@jaist.ac.jp, prm4@st-andrews.ac.uk,
h.sun@ed.ac.uk
ABSTRACT
Coresets have become an invaluable tool for solving k-means and kernel k-means clustering
problems on large datasets with small numbers of clusters. On the other hand, spectral clustering
works well on sparse graphs and has recently been extended to scale efficiently to large numbers
of clusters. We exploit the connection between kernel k-means and the normalised cut problem
to combine the benefits of both. Our main result is a coreset spectral clustering algorithm
for graphs that clusters a coreset graph to infer a good labelling of the original graph. We
prove that an α-approximation for the normalised cut problem on the coreset graph is an O(α)-
approximation on the original. We also improve the running time of the state-of-the-art coreset
algorithm for kernel k-means on sparse kernels, from ˜O(nk) to ˜O(n · min{k, davg}), where
davg is the average number of non-zero entries in each row of the n × n kernel matrix. Our
experiments confirm our coreset algorithm is asymptotically faster on large real-world graphs
with many clusters, and show that our clustering algorithm overcomes the main challenge faced
by coreset kernel k-means on sparse kernels which is getting stuck in local optima.
1
INTRODUCTION
Kernel k-means and spectral clustering are two popular algorithms which are capable of learning
non-linear decision boundaries between clusters. For this reason, they have been applied to many
practical problems in machine learning, including in the fields of medical research and network
science (Gönen & Margolin, 2014; Kuo et al., 2014; White & Smyth, 2005).
Given a data set X, both kernel k-means and spectral clustering ...
Conference: NeurIPS 2025

Embodied Cognition Augmented End2End
Autonomous Driving
Ling Niu
Tsinghua University
nl23@mails.tsinghua.edu.cn
Xiaoji Zheng
Tsinghua University
zhengxj24@mails.tsinghua.edu.cn
Han Wang
Tsinghua University
kelsulhgod@gmail.com
Chen Zheng
Tsinghua University
zhengchen@air.tsinghua.edu.cn
Ziyuan Yang
Tsinghua University
ziyuan86@uw.edu
Bokui Chen*
Tsinghua University
chenbokui@tsinghua.edu.cn
Jiangtao Gong∗
Tsinghua University
gongjiangtao2@gmail.com
Abstract
In recent years, vision-based end-to-end autonomous driving has emerged as a new
paradigm. However, popular end-to-end approaches typically rely on visual feature
extraction networks trained under label supervision. This limited supervision frame-
work restricts the generality and applicability of driving models. In this paper, we
propose a novel paradigm termed E3AD, which advocates for comparative learn-
ing between visual feature extraction networks and the general EEG large model,
in order to learn latent human driving cognition for enhancing end-to-end planning.
In this work, we collected a cognitive dataset for the mentioned contrastive learning
process. Subsequently, we investigated the methods and potential mechanisms
for enhancing end-to-end planning with human driving cognition, using popular
driving models as baselines on publicly available autonomous driving datasets.
Both open-loop and closed-loop tests are conducted for a comprehensive evalua-
tion of planning performance. Experimental results demonstrate that the E3AD
paradigm significantly enhances the end-to-end planning performance of baseline
models. Ablation studies further validate the contribution of driving cognition and
the effectiveness of comparative learning process. To the best of our knowledge,
this is the first work to integrate human driving cognition for improving end-to-end
autonomous driving planning. It represents an initial attempt to incorporate embod-
ied cognitive data into end-to-end autonomous driving, providing valuable i...
Conference: ICML 2025

Learning Joint Interventional Effects from Single-Variable Interventions
in Additive Models
Armin Keki´c 1 Sergio Hernan Garrido Mejia 1 2 Bernhard Sch¨olkopf 1 3 4
Abstract
Estimating causal effects of joint interventions on
multiple variables is crucial in many domains, but
obtaining data from such simultaneous interven-
tions can be challenging. Our study explores how
to learn joint interventional effects using only ob-
servational data and single-variable interventions.
We present an identifiability result for this prob-
lem, showing that for a class of nonlinear additive
outcome mechanisms, joint effects can be inferred
without access to joint interventional data. We
propose a practical estimator that decomposes the
causal effect into confounded and unconfounded
contributions for each intervention variable. Ex-
periments on synthetic data demonstrate that our
method achieves performance comparable to mod-
els trained directly on joint interventional data,
outperforming a purely observational estimator.
1. Introduction
Understanding the effects of interventions is fundamental
across many domains, from designing public health policies
to optimizing business operations and administering medical
treatments. Particularly challenging and important are joint
interventional effects, where simultaneous interventions on
multiple action variables influence a target outcome. Such
scenarios are common in fields like epidemiology (Keki´c
et al., 2023a), e-commerce (Kunz et al., 2023; Schultz et al.,
2023), and medicine (Prosperi et al., 2020).
Consider a company trying to optimize their marketing strat-
egy across multiple channels like social media, email cam-
paigns, and display advertising. While the ultimate goal is
to understand how these channels work together to drive
sales, running experiments for every possible combination
1Max Planck Institute for Intelligent Systems, T¨ubingen, Ger-
many 2Amazon Research, T¨ubingen, Germany 3T¨ubingen AI Cen-
ter, T¨ubingen, Germa...
Conference: ICML 2025

Improved Learning via k-DTW: A Novel Dissimilarity Measure for Curves
Amer Krivoˇsija 1 Alexander Munteanu 1 2 Andr´e Nusser 3 Chris Schwiegelshohn 4
Abstract
This paper introduces k-Dynamic Time Warping
(k-DTW), a novel dissimilarity measure for polyg-
onal curves. k-DTW has stronger metric proper-
ties than Dynamic Time Warping (DTW) and is
more robust to outliers than the Fr´echet distance,
which are the two gold standards of dissimilarity
measures for polygonal curves. We show inter-
esting properties of k-DTW and give an exact
algorithm as well as a (1 + ε)-approximation al-
gorithm for k-DTW by a parametric search for
the k-th largest matched distance. We prove the
first dimension-free learning bounds for curves
and further learning theoretic results. k-DTW not
only admits smaller sample size than DTW for the
problem of learning the median of curves, where
some factors depending on the curves’ complexity
m are replaced by k, but we also show a surpris-
ing separation on the associated Rademacher and
Gaussian complexities: k-DTW admits strictly
smaller bounds than DTW, by a factor ˜Ω(√m)
when k ≪m. We complement our theoretical
findings with an experimental illustration of the
benefits of using k-DTW for clustering and near-
est neighbor classification.
1. Introduction
Handwriting, panel data, time series, sensor-generated tra-
jectories, and many more types of data are instances of
polygonal curves in Euclidean space. These curves are
input for learning processes, both supervised and unsuper-
vised. To compare and quantify the distance of two curves,
a suitable dissimilarity measure is needed that reflects the
fact that similar curves that represent, for instance, the same
1TU Dortmund, Germany 2University of Cologne, Germany
3Universit´e Cˆote d’Azur, CNRS, Inria, France 4Aarhus University,
Denmark. Correspondence to: A. Krivoˇsija <amer.krivosija@tu-
dortmund.de>,
A.
Munteanu
<alexander.munteanu@tu-
dortmund.de>,
A.
Nusser
<andre.nusser@cnrs.fr>,
C.
Schwiege...