dataset-review/dataset-review huggingface dataset:
Learning Joint Interventional Effects from Single-Variable Interventions in Additive Models
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...
Improved Learning via k-DTW: A Novel Dissimilarity Measure for Curves
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...
EmbedLLM: Learning Compact Representations of Large Language Models
ICLR 2025

Published as a conference paper at ICLR 2025
EMBEDLLM:
LEARNING
COMPACT
REPRESENTA-
TIONS OF LARGE LANGUAGE MODELS
Richard Zhuang∗Tianhao Wu∗
Zhaojin Wen
Andrew Li
Jiantao Jiao
Kannan Ramchandran
University of California, Berkeley
ABSTRACT
With hundreds of thousands of language models available on Huggingface to-
day, efficiently evaluating and utilizing these models across various downstream
tasks has become increasingly critical. Many existing methods repeatedly learn
task-specific representations of Large Language Models (LLMs), which leads to
inefficiencies in both time and computational resources. To address this, we pro-
pose EmbedLLM, a framework designed to learn compact vector representations
of LLMs that facilitate downstream applications involving many models, such as
model routing. We introduce an encoder-decoder approach for learning such em-
beddings, along with a systematic framework to evaluate their effectiveness. Em-
pirical results show that EmbedLLM outperforms prior methods in model routing
both in accuracy and latency. Additionally, we demonstrate that our method can
forecast a model’s performance on multiple benchmarks, without incurring ad-
ditional inference cost. Extensive probing experiments validate that the learned
embeddings capture key model characteristics, e.g. whether the model is special-
ized for coding tasks, even without being explicitly trained on them. We open
source our dataset, code and embedder to facilitate further research and applica-
tion: https://github.com/richardzhuang0412/EmbedLLM.
1
INTRODUCTION
Recent breakthroughs in Large Language Models (LLMs) (Vaswani et al., 2023) have led to the
creation of a vast array of models, each tailored for different use cases. These models, ranging from
small, specialized models to large, general-purpose systems (Hao et al., 2022), differ significantly
in their architecture, size, training data, and performance characteristics. For example, while some
models excel as conversational ...
Generalizable, real-time neural decoding with hybrid state-space models
NeurIPS 2025

Generalizable, real-time neural decoding with hybrid
state-space models
Avery Hee-Woon Ryoo1,2,∗,B
Nanda H Krishna1,2,∗,B
Ximeng Mao1,2,∗,B
Mehdi Azabou3
Eva L Dyer4
Matthew G Perich1,2,†
Guillaume Lajoie1,2,5,†,B
1Mila – Quebec AI Institute
2Université de Montréal
3Columbia University
4University of Pennsylvania
5Canada CIFAR AI Chair
∗Co-first authors
†Co-senior authors
B{hee-woon.ryoo,nanda.harishankar-krishna,ximeng.mao,guillaume.lajoie}@mila.quebec
https://possm-brain.github.io
Abstract
Real-time decoding of neural activity is central to neuroscience and neurotech-
nology applications, from closed-loop experiments to brain-computer interfaces,
where models are subject to strict latency constraints. Traditional methods, includ-
ing simple recurrent neural networks, are fast and lightweight but often struggle
to generalize to unseen data. In contrast, recent Transformer-based approaches
leverage large-scale pretraining for strong generalization performance, but typically
have much larger computational requirements and are not always suitable for low-
resource or real-time settings. To address these shortcomings, we present POSSM, a
novel hybrid architecture that combines individual spike tokenization via a cross-
attention module with a recurrent state-space model (SSM) backbone to enable (1)
fast and causal online prediction on neural activity and (2) efficient generalization
to new sessions, individuals, and tasks through multi-dataset pretraining. We evalu-
ate POSSM’s decoding performance and inference speed on intracortical decoding
of monkey motor tasks, and show that it extends to clinical applications, namely
handwriting and speech decoding in human subjects. Notably, we demonstrate that
pretraining on monkey motor-cortical recordings improves decoding performance
on the human handwriting task, highlighting the exciting potential for cross-species
transfer. In all of these tasks, we find that POSSM achieves decoding accuracy
comparable to state-of-the-art...
Learning Latent Graph Structures and their Uncertainty
ICML 2025

Learning Latent Graph Structures and their Uncertainty
Alessandro Manenti 1 Daniele Zambon 1 Cesare Alippi 1 2
Abstract
Graph neural networks use relational informa-
tion as an inductive bias to enhance prediction
performance. Not rarely, task-relevant relations
are unknown and graph structure learning ap-
proaches have been proposed to learn them from
data. Given their latent nature, no graph obser-
vations are available to provide a direct training
signal to the learnable relations. Therefore, graph
topologies are typically learned on the prediction
task alongside the other graph neural network pa-
rameters. In this paper, we demonstrate that mini-
mizing point-prediction losses does not guarantee
proper learning of the latent relational informa-
tion and its associated uncertainty. Conversely, we
prove that suitable loss functions on the stochastic
model outputs simultaneously grant solving two
tasks: (i) learning the unknown distribution of
the latent graph and (ii) achieving optimal predic-
tions of the target variable. Finally, we propose
a sampling-based method that solves this joint
learning task 1. Empirical results validate our the-
oretical claims and demonstrate the effectiveness
of the proposed approach.
1. Introduction
Relational information processing has provided break-
throughs in the analysis of rich and complex data coming
from, e.g., social networks, natural language, and biology.
This side information takes various forms, from structuring
the data into clusters to defining causal relations and hier-
archies, and enables machine learning models to condition
their predictions on dependency-related observations. In
this context, predictive models take the form y = fψ(x, A),
where the input-output relation x 7→y – modeled by fψ
1Universit`a della Svizzera italiana, IDSIA, Lugano, Switzer-
land 2Politecnico di Milano, Milan, Italy. Correspondence to:
Alessandro Manenti <alessandro.manenti@usi.ch>.
Proceedings of the 42 nd International Conference on ...
AdaSplash: Adaptive Sparse Flash Attention
ICML 2025

ADASPLASH: Adaptive Sparse Flash Attention
Nuno Gonc¸alves 1 Marcos Treviso 2 Andr´e F. T. Martins 1 2 3
Abstract
The computational cost of softmax-based atten-
tion in transformers limits their applicability to
long-context tasks. Adaptive sparsity, of which
α-entmax attention is an example, offers a flexi-
ble data-dependent alternative, but existing imple-
mentations are inefficient and do not leverage the
sparsity to obtain runtime and memory gains. In
this work, we propose ADASPLASH, which com-
bines the efficiency of GPU-optimized algorithms
with the sparsity benefits of α-entmax. We first
introduce a hybrid Halley-bisection algorithm, re-
sulting in a 7-fold reduction in the number of itera-
tions needed to compute the α-entmax transforma-
tion. Then, we implement custom Triton kernels
to efficiently handle adaptive sparsity. Experi-
ments with RoBERTa and ModernBERT for text
classification and single-vector retrieval, along
with GPT-2 for language modeling, show that
our method achieves substantial improvements in
runtime and memory efficiency compared to ex-
isting α-entmax implementations. It approaches—
and in some cases surpasses—the efficiency of
highly optimized softmax implementations like
FlashAttention-2, enabling long-context training
while maintaining strong task performance.1
1. Introduction
Central to the success of transformers (Vaswani et al., 2017)
lies the attention mechanism, where each token in a se-
quence attends directly to every other token. Attention prob-
abilities are computed through the softmax transformation,
which always assigns a nonzero probability to every token.
However, for long context inputs, the accumulation of small
probabilities can lead to dispersion (Veliˇckovi´c et al., 2025).
1Instituto Superior T´ecnico, Universidade de Lisboa, Portu-
gal 2Instituto de Telecomunicac¸˜oes, Lisbon, Portugal 3Unbabel,
Lisbon, Portugal.
Correspondence to:
Nuno Gonc¸alves
<nuno.m.goncalves@tecnico.ulisboa.pt>.
Proceedings of the 42 nd ...
Enforcing convex constraints in Graph Neural Networks
NeurIPS 2025

Enforcing convex constraints in Graph Neural
Networks
Ahmed Rashwan ˚
University of Bath
Keith Briggs
BT Research
Chris Budd
University of Bath
Lisa Kreusser
University of Bath and Monumo
Abstract
Many machine learning applications require outputs that satisfy complex, dy-
namic constraints. This task is particularly challenging in Graph Neural Network
models due to the variable output sizes of graph-structured data. In this paper,
we introduce ProjNet, a Graph Neural Network framework which satisfies input-
dependant constraints. ProjNet combines a sparse vector clipping method with the
Component-Averaged Dykstra (CAD) algorithm, an iterative scheme for solving
the best-approximation problem. We establish a convergence result for CAD and
develop a GPU-accelerated implementation capable of handling large-scale in-
puts efficiently. To enable end-to-end training, we introduce a surrogate gradient
for CAD that is both computationally efficient and better suited for optimization
than the exact gradient. We validate ProjNet on four classes of constrained op-
timisation problems: linear programming, two classes of non-convex quadratic
programs, and radio transmit power optimization, demonstrating its effectiveness
across diverse problem settings.
1
Introduction
In many machine learning (ML) applications, model outputs must satisfy certain constraints to en-
sure feasibility, safety, or interpretability [22, 40]. A common approach to enforcing constraints
is through activation functions (e.g., Sigmoid, Softmax) in the output layer [41], which restrict
outputs to a predefined set. However, while effective for simple linear or box constraints, this ap-
proach is not applicable for more complicated types of constraints. Many real-world applications
require complex, data-dependent constraints. For example, in robotics, constraints can encode fea-
sible movements or collision avoidance [21, 54, 63]. Similarly, in energy systems, power generation
limits vary with demand fluctua...
Relational Conformal Prediction for Correlated Time Series
ICML 2025

Relational Conformal Prediction for Correlated Time Series
Andrea Cini 1 2 Alexander Jenkins 1 3 Danilo Mandic 3 Cesare Alippi 1 4 Filippo Maria Bianchi 5 6
Abstract
We address the problem of uncertainty quantifi-
cation in time series forecasting by exploiting
observations at correlated sequences. Relational
deep learning methods leveraging graph repre-
sentations are among the most effective tools for
obtaining point estimates from spatiotemporal
data and correlated time series. However, the
problem of exploiting relational structures to es-
timate the uncertainty of such predictions has
been largely overlooked in the same context. To
this end, we propose a novel distribution-free ap-
proach based on the conformal prediction frame-
work and quantile regression. Despite the recent
applications of conformal prediction to sequen-
tial data, existing methods operate independently
on each target time series and do not account
for relationships among them when constructing
the prediction interval. We fill this void by in-
troducing a novel conformal prediction method
based on graph deep learning operators.
Our
approach, named Conformal Relational Predic-
tion (COREL), does not require the relational
structure (graph) to be known a priori and can be
applied on top of any pre-trained predictor. Addi-
tionally, COREL includes an adaptive component
to handle non-exchangeable data and changes in
the input time series. Our approach provides ac-
curate coverage and achieves state-of-the-art un-
certainty quantification in relevant benchmarks.
1. Introduction
Many recent advancements in deep learning methods for
time series forecasting rely on learning from large collec-
tions of (related) time series (Benidis et al., 2022; Liang
1IDSIA USI-SUPSI, Universit`a della Svizzera italiana 2Swiss
National Science Foundation Postdoc Fellow 3Imperial College
London 4Politecnico di Milano 5UiT The Arctic University of Nor-
way 6NORCE Norwegian Research Centre AS. Correspondence
to: Andre...
Subgroups Matter for Robust Bias Mitigation
ICML 2025

Subgroups Matter for Robust Bias Mitigation
Anissa Alloula 1 Charles Jones 2 Ben Glocker 2 Bartłomiej W. Papie˙z 1
Abstract
Despite the constant development of new bias mit-
igation methods for machine learning, no method
consistently succeeds, and a fundamental question
remains unanswered: when and why do bias miti-
gation techniques fail? In this paper, we hypothe-
sise that a key factor may be the often-overlooked
but crucial step shared by many bias mitigation
methods: the definition of subgroups. To investi-
gate this, we conduct a comprehensive evaluation
of state-of-the-art bias mitigation methods across
multiple vision and language classification tasks,
systematically varying subgroup definitions, in-
cluding coarse, fine-grained, intersectional, and
noisy subgroups. Our results reveal that subgroup
choice significantly impacts performance, with
certain groupings paradoxically leading to worse
outcomes than no mitigation at all. Our find-
ings suggest that observing a disparity between a
set of subgroups is not a sufficient reason to use
those subgroups for mitigation. Through theoret-
ical analysis, we explain these phenomena and
uncover a counter-intuitive insight that, in some
cases, improving fairness with respect to a par-
ticular set of subgroups is best achieved by using
a different set of subgroups for mitigation. Our
work highlights the importance of careful sub-
group definition in bias mitigation and presents it
as an alternative lever for improving the robust-
ness and fairness of machine learning models.
1. Introduction
A significant barrier to the wider deployment of machine
learning (ML) models is their tendency to fail when tested
on distributions that differ from their training data. One
particularly concerning manifestation of this issue is perfor-
mance degradation for population subgroups, often caused
1University of Oxford, UK 2Imperial College London, UK. Cor-
respondence to: Anissa Alloula <anissa.alloula@dtc.ox.ac.uk>.
Proceedings of ...
Implicit Regularization for Tubal Tensor Factorizations via Gradient Descent
ICML 2025

Implicit Regularization for Tubal Tensor Factorizations via Gradient Descent
Santhosh Karnik * 1 Anna Veselovska * 2 3 Mark Iwen 4 5 Felix Krahmer 2 3
Abstract
We provide a rigorous analysis of implicit regu-
larization in an overparametrized tensor factoriza-
tion problem beyond the lazy training regime. For
matrix factorization problems, this phenomenon
has been studied in a number of works. A par-
ticular challenge has been to design universal ini-
tialization strategies which provably lead to im-
plicit regularization in gradient-descent methods.
At the same time, it has been argued by (Cohen
et al., 2016) that more general classes of neural
networks can be captured by considering tensor
factorizations. However, in the tensor case, im-
plicit regularization has only been rigorously es-
tablished for gradient flow or in the lazy training
regime. In this paper, we prove the first tensor
result of its kind for gradient descent rather than
gradient flow. We focus on the tubal tensor prod-
uct and the associated notion of low tubal rank,
encouraged by the relevance of this model for im-
age data. We establish that gradient descent in an
overparametrized tensor factorization model with
a small random initialization exhibits an implicit
bias towards solutions of low tubal rank. Our the-
oretical findings are illustrated in an extensive set
of numerical simulations show-casing the dynam-
ics predicted by our theory as well as the crucial
role of using a small random initialization.
1. Introduction
Analyzing implicit regularization during Neural Network
(NN) training is considered crucial for understanding why
*Equal contribution 1Department of Mathematics, Northeast-
ern University, Boston, USA 2Department of Mathematics and
Munich Data Science Institute, Technical University of Munich,
Munich, Germany 3Munich Center for Machine Learning, Mu-
nich, Germany 4Department of Mathematics, Michigan State
University, East Lansing, USA 5Department of Computational
Mathematics S...
CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences
ICLR 2025

Published as a conference paper at ICLR 2025
CAKE: CASCADING
AND ADAPTIVE KV CACHE
EVICTION WITH LAYER PREFERENCES
Ziran Qin1,2∗, Yuchen Cao2, Mingbao Lin3†, Wen Hu2, Shixuan Fan2, Ke Cheng2,
Weiyao Lin1†, Jianguo Li2†
1Shanghai Jiao Tong University, 2Ant Group, 3Independent Researcher
ABSTRACT
Large language models (LLMs) excel at processing long sequences, boosting de-
mand for key-value (KV) caching. While recent efforts to evict KV cache have al-
leviated the inference burden, they often fail to allocate resources rationally across
layers with different attention patterns. In this paper, we introduce Cascading and
Adaptive KV cache Eviction (CAKE), a novel approach that frames KV cache
eviction as a “cake-slicing problem.” CAKE assesses layer-specific preferences
by considering attention dynamics in both spatial and temporal dimensions, allo-
cates rational cache size for layers accordingly, and manages memory constraints
in a cascading manner. This approach enables a global view of cache alloca-
tion, adaptively distributing resources across diverse attention mechanisms while
maintaining memory budgets. CAKE also employs a new eviction indicator that
considers the shifting importance of tokens over time, addressing limitations in
existing methods that overlook temporal dynamics. Comprehensive experiments
on LongBench and NeedleBench show that CAKE maintains model performance
with only 3.2% of the KV cache and consistently outperforms current baselines
across various models and memory constraints, particularly in low-memory set-
tings. Additionally, CAKE achieves over 10× speedup in decoding latency com-
pared to full cache when processing contexts of 128K tokens with FlashAttention-
2. Our code is available at https://github.com/antgroup/cakekv.
1
INTRODUCTION
Large language models (LLMs) (Zhao et al., 2023; Achiam et al., 2023; Dubey et al., 2024; An-
thropic, 2024; AI, 2024) have enhanced their long-text processing capabilities, improving perfor-
mance in mul...