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

CACTI: Leveraging Copy Masking and Contextual Information to Improve
Tabular Data Imputation
Aditya Gorla 1 2 Ryan Wang 3 Zhengtong Liu 3 Ulzee An 1 3 Sriram Sankararaman 1 3 4
Abstract
We present CACTI, a masked autoencoding ap-
proach for imputing tabular data that leverages
the structure in missingness patterns and contex-
tual information. Our approach employs a novel
median truncated copy masking training strategy
that encourages the model to learn from empirical
patterns of missingness while incorporating se-
mantic relationships between features – captured
by column names and text descriptions – to better
represent feature dependence. These dual sources
of inductive bias enable CACTI to outperform
state-of-the-art methods – an average R2 gain of
7.8% over the next best method (13.4%, 6.1%,
and 5.3% under missing not at random, at ran-
dom and completely at random, respectively) –
across a diverse range of datasets and missing-
ness conditions. Our results highlight the value
of leveraging dataset-specific contextual informa-
tion and missingness patterns to enhance imputa-
tion performance. Code is publicly available at
github.com/sriramlab/CACTI
1. Introduction
Missingness is a pervasive problem in real-world tabular
datasets with the potential to adversely affect downstream
inferential tasks (Rubin, 1987; Schafer & Graham, 2002).
While many techniques to estimate or impute missing en-
tries have been proposed (see Section 2.1), missing data
imputation remains a challenging problem.
A primary reason underlying this challenge is that missing-
1Department of Computational Medicine, David Geffen School
of Medicine, UCLA, Los Angeles, CA, USA 2Bioinformatics
Interdepartmental Program, UCLA, Los Angeles, CA, USA
3Department of Computer Science, UCLA, Los Angeles, CA,
USA 4Department of Human Genetics, UCLA, Los Angeles, CA,
USA. Correspondence to: Aditya Gorla <adityagorla@ucla.edu>,
Sriram Sankararaman <sriram@cs.ucla.edu>.
Proceedings of the 42 nd International...
Conference: NeurIPS 2025

KGGen: Extracting Knowledge Graphs from Plain
Text with Language Models
Belinda Mo∗1, Kyssen Yu∗2, Joshua Kazdan∗1, Joan Cabezas, Proud Mpala1, Lisa Yu2, Chris Cundy3,
Charilaos Kanatsoulis1, Sanmi Koyejo1
1Stanford University
2University of Toronto
3 FAR AI
Abstract
Recent interest in building foundation models for knowledge graphs has highlighted
a fundamental challenge: knowledge graph data is scarce. The best-known knowl-
edge graphs are primarily human-labeled, created by pattern-matching, or extracted
using early NLP techniques. While human-generated knowledge graphs are in
short supply, automatically extracted ones are of questionable quality. We present
KGGen, a novel text-to-knowledge-graph generator that uses language models to
extract high-quality graphs from plain text with a novel entity resolution approach
that clusters related entities, significantly reducing the sparsity problem that plagues
existing extractors. Unlike other KG generators, KGGen clusters and de-duplicates
related entities to reduce sparsity in extracted KGs. Along with KGGen, we release
Measure of Information in Nodes and Edges (MINE), the first benchmark to test an
extractor’s ability to produce a useful KG from plain text. We benchmark our new
tool against leading existing generators such as Microsoft’s GraphRAG; we achieve
comparable retrieval accuracy on the generated graphs and better information re-
tention. Moreover, our graphs exhibit more concise and generalizable entities and
relations. Our code is open-sourced at https://github.com/stair-lab/kg-gen/.
1
Introduction
Knowledge graph (KG) applications and graph-based Retrieval-Augmented Generation (RAG)
systems are increasingly bottlenecked by the scarcity and incompleteness of available KGs. KGs
consist of a set of subject-predicate-object triples, and have become a fundamental data structure
for information retrieval [Schneider, 1973]. Most real-world KGs, including Wikidata [contributors,
2024], DBpedia [Lehmann et al., 20...
Conference: ICML 2025

Logits are All We Need to Adapt Closed Models
Gaurush Hiranandani 1 Haolun Wu* 2 3 Subhojyoti Mukherjee† 4 Sanmi Koyejo 2
Abstract
Many commercial Large Language Models
(LLMs) are often closed-source, limiting devel-
opers to prompt tuning for aligning content gen-
eration with specific applications. While these
models currently do not provide access to token
logits, we argue that if such access were available,
it would enable more powerful adaptation tech-
niques beyond prompt engineering. In this paper,
we propose a token-level probability reweighting
framework that, given access to logits and a small
amount of task-specific data, can effectively steer
black-box LLMs toward application-specific con-
tent generation. Our approach views next-token
prediction through the lens of supervised classi-
fication. We show that aligning black-box LLMs
with task-specific data can be formulated as a la-
bel noise correction problem, leading to Plugin
model – an autoregressive probability reweighting
model that operates solely on logits. We provide
theoretical justification for why reweighting logits
alone is sufficient for task adaptation. Extensive
experiments with multiple datasets, LLMs, and
reweighting models demonstrate the effectiveness
of our method, advocating for broader access to
token logits in closed-source models. We provide
our code at this https URL.
1. Introduction
The rise of Large Language Models (LLMs) has revolution-
ized generative Artificial Intelligence, yet the most capable
models are often closed-source or black-box (Achiam et al.,
2023; Bai et al., 2022a). These models generate text based
on input prompts but keep their internal weights and training
data undisclosed, limiting transparency and customization.
*Work done while visiting Stanford University. †Work done
while PhD Candidate at UW Madison. 1Typeface 2Stanford Univer-
sity 3Mila - Quebec AI Institute 4Adobe Research. Correspondence
to: Gaurush Hiranandani <gaurush@typeface.ai>, Haolun Wu
<haol...
Conference: NeurIPS 2025

TROVE: Discovering Error-Inducing Static Feature
Biases in Temporal Vision-Language Models
Maya Varma1
Jean-Benoit Delbrouck1,2
Sophie Ostmeier1
Akshay Chaudhari1
Curtis Langlotz1
1Stanford University
2HOPPR
mayavarma@cs.stanford.edu
Abstract
Vision-language models (VLMs) have made great strides in addressing temporal
understanding tasks, which involve characterizing visual changes across a sequence
of images. However, recent works have suggested that when making predictions,
VLMs may rely on static feature biases, such as background or object features,
rather than dynamic visual changes. Static feature biases are a type of shortcut
and can contribute to systematic prediction errors on downstream tasks; as a
result, identifying and characterizing error-inducing static feature biases is critical
prior to real-world model deployment. Existing approaches for identifying such
systematic failure modes in trained models (i) are typically designed for non-
temporal settings and (ii) are challenging to evaluate in temporal settings due to
the lack of quantitative evaluation frameworks. In this work, we address these
challenges by introducing TROVE, an automated approach for discovering error-
inducing static feature biases learned by temporal VLMs. Given a trained VLM
and an annotated validation dataset associated with a downstream classification
task, TROVE extracts candidate static features from the dataset and scores each
feature by (i) the effect of the feature on classification errors as well as (ii) the
extent to which the VLM relies on the feature when making predictions. In
order to quantitatively evaluate TROVE, we introduce an evaluation framework
consisting of 101 trained temporal VLMs paired with ground-truth annotations for
learned static feature biases. We use this framework to demonstrate that TROVE
can accurately identify error-inducing static feature biases in VLMs, achieving
a 28.6% improvement over the closest baseline. Finally, we apply TROVE to
7 off-th...
Conference: ICML 2025

Dynamic Similarity Graph Construction
with Kernel Density Estimation
Steinar Laenen 1 Peter Macgregor 2 He Sun 1
Abstract
In the kernel density estimation (KDE) problem,
we are given a set X of data points in Rd, a kernel
function k : Rd × Rd →R, and a query point
q ∈Rd, and the objective is to quickly output an
estimate of P
x∈X k(q, x). In this paper, we con-
sider KDE in the dynamic setting, and introduce
a data structure that efficiently maintains the esti-
mates for a set of query points as data points are
added to X over time. Based on this, we design
a dynamic data structure that maintains a sparse
approximation of the fully connected similarity
graph on X, and develop a fast dynamic spectral
clustering algorithm. We further evaluate the ef-
fectiveness of our algorithms on both synthetic
and real-world datasets.
1. Introduction
Given a set X
= {x1, . . . , xn} of data points, a set
Q = {q1, . . . , qm} of query points, and a kernel func-
tion k : Rd × Rd →R⩾0, the KDE problem is to quickly
approximate µq ≜P
xi∈X k(q, xi) for every q ∈Q. As a
basic question in computer science, this problem has been
actively studied since the 1990s (Greengard & Strain, 1991)
and has comprehensive applications in machine learning
and statistics (Bakshi et al., 2023; Genovese et al., 2014;
Scholkopf & Smola, 2018; Schubert et al., 2014).
In this paper we first study the KDE problem in the dynamic
setting, where both the sets of data and query points change
over time. The objective is to dynamically update all the
KDE estimates µq for every q ∈Q as data points are added
to X. Building on the framework for static KDE developed
by Charikar et al. (2020), our algorithm processes: (i) in-
sertions and deletions of query points, and (ii) insertions of
1University of Edinburgh, United Kingdom 2University of
St Andrews, United Kingdom.
Correspondence to: He Sun
<h.sun@ed.ac.uk>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Co...
Conference: ICML 2025

Counterfactual Effect Decomposition in
Multi-Agent Sequential Decision Making
Stelios Triantafyllou 1 Aleksa Sukovic 1 Yasaman Zolfimoselo 1 Goran Radanovic 1
Abstract
We address the challenge of explaining counter-
factual outcomes in multi-agent Markov decision
processes. In particular, we aim to explain the to-
tal counterfactual effect of an agent’s action on the
outcome of a realized scenario through its influ-
ence on the environment dynamics and the agents’
behavior. To achieve this, we introduce a novel
causal explanation formula that decomposes the
counterfactual effect by attributing to each agent
and state variable a score reflecting their respec-
tive contributions to the effect. First, we show that
the total counterfactual effect of an agent’s action
can be decomposed into two components: one
measuring the effect that propagates through all
subsequent agents’ actions and another related
to the effect that propagates through the state
transitions. Building on recent advancements in
causal contribution analysis, we further decom-
pose these two effects as follows. For the former,
we consider agent-specific effects – a causal con-
cept that quantifies the counterfactual effect of an
agent’s action that propagates through a subset
of agents. Based on this notion, we use Shapley
value to attribute the effect to individual agents.
For the latter, we consider the concept of structure-
preserving interventions and attribute the effect
to state variables based on their “intrinsic” contri-
butions. Through extensive experimentation, we
demonstrate the interpretability of our approach
in a Gridworld environment with LLM-assisted
agents and a sepsis management simulator.
1. Introduction
Applying counterfactual reasoning to retrospectively an-
alyze the impact of different actions in decision making
1Max Planck Institute for Software Systems, Germany. Corre-
spondence to: Stelios Triantafyllou <strianta@mpi-sws.org>.
Proceedings of the 42 nd International Conference...
Conference: NeurIPS 2025

Not All Data are Good Labels:
On the Self-supervised Labeling for Time Series
Forecasting
Yuxuan Yang1,2, Dalin Zhang3, Yuxuan Liang4, Hua Lu5, Gang Chen1, 2, Huan Li1,2 B
1The State Key Laboratory of Blockchain and Data Security, Zhejiang University
2Hangzhou High-Tech Zone (Binjiang) Institute of Blockchain and Data Security
3Space Information Research Institute, Hangzhou Dianzi University
4The Hong Kong University of Science and Technology (Guangzhou)
5Department of Computer Science, Aalborg University, Denmark
{yangyuxuan, cg, lihuan.cs}@zju.edu.cn
zhangdalin@hdu.edu.cn, yuxliang@outlook.com, luhua@cs.aau.dk
Abstract
Time Series Forecasting (TSF) is a crucial task in various domains, yet existing TSF
models rely heavily on high-quality data and insufficiently exploit all available data.
This paper explores a novel self-supervised approach to re-label time series datasets
by inherently constructing candidate datasets. During the optimization of a simple
reconstruction network, intermediates are used as pseudo labels in a self-supervised
paradigm, improving generalization for any predictor. We introduce the Self-
Correction with Adaptive Mask (SCAM), which discards overfitted components
and selectively replaces them with pseudo labels generated from reconstructions.
Additionally, we incorporate Spectral Norm Regularization (SNR) to further sup-
press overfitting from a loss landscape perspective. Our experiments on eleven
real-world datasets demonstrate that SCAM consistently improves the performance
of various backbone models. This work offers a new perspective on constructing
datasets and enhancing the generalization of TSF models through self-supervised
learning. The code is available at https://github.com/SuDIS-ZJU/SCAM.
1
Introduction
Time Series Forecasting (TSF) is a crucial task with extensive applications in energy, finance,
engineering, and many other domains. Recent advances in deep learning have resulted in TSF
methods that outperform traditional metho...
Conference: NeurIPS 2025

An Evidence-Based Post-Hoc Adjustment Framework
for Anomaly Detection Under Data Contamination
Sukanya Patra
Department of Computer Science
University of Mons
Mons, Belgium
sukanya.patra@umons.ac.be
Souhaib Ben Taieb∗
Department of Statistics and Data Science
Mohamed bin Zayed University of Artificial Intelligence
Abu Dhabi, United Arab Emirates
souhaib.bentaieb@mbzuai.ac.ae
Abstract
Unsupervised anomaly detection (AD) methods typically assume clean training
data, yet real-world datasets often contain undetected or mislabeled anomalies,
leading to significant performance degradation. Existing solutions require access
to the training pipelines, data or prior knowledge of the proportions of anomalies
in the data, limiting their real-world applicability. To address this challenge, we
propose EPHAD, a simple yet effective test-time adaptation framework that updates
the outputs of AD models trained on contaminated datasets using evidence gath-
ered at test time. Our approach integrates the prior knowledge captured by the
AD model trained on contaminated datasets with evidence derived from multi-
modal foundation models like Contrastive Language-Image Pre-training (CLIP),
classical AD methods like the Local Outlier Factor or domain-specific knowledge.
We illustrate the intuition behind EPHAD using a synthetic toy example and vali-
date its effectiveness through comprehensive experiments across eight visual AD
datasets, twenty-six tabular AD datasets, and a real-world industrial AD dataset.
Additionally, we conduct an ablation study to analyse hyperparameter influence
and robustness to varying contamination levels, demonstrating the versatility and
robustness of EPHAD across diverse AD models and evidence pairs. To ensure
reproducibility, our code is publicly available2.
1
Introduction
Anomaly detection (AD) is the basis of many critical applications, including cybersecurity (Xiao
et al., 2024; Li et al., 2023a), healthcare (Bijlani et al., 2024; Huang et al., 2024), and i...
Conference: ICML 2025

Multiaccuracy and Multicalibration via Proxy Groups
Beepul Bharti 1 2 Mary Versa Clemens-Sewall 3 Paul Yi 4 Jeremias Sulam 1 2 5
Abstract
As the use of predictive machine learning algo-
rithms increases in high-stakes decision-making,
it is imperative that these algorithms are fair
across sensitive groups. However, measuring and
enforcing fairness in real-world applications can
be challenging due to missing or incomplete sensi-
tive group information. Proxy-sensitive attributes
have been proposed as a practical and effective so-
lution in these settings, but only for parity-based
fairness notions. Knowing how to evaluate and
control for fairness with missing sensitive group
data for newer, different, and more flexible frame-
works, such as multiaccuracy and multicalibra-
tion, remains unexplored. In this work, we ad-
dress this gap by demonstrating that in the ab-
sence of sensitive group data, proxy-sensitive at-
tributes can provably be used to derive actionable
upper bounds on the true multiaccuracy and mul-
ticalibration violations, providing insights into a
predictive model’s potential worst-case fairness
violations. Additionally, we show that adjusting
models to satisfy multiaccuracy and multicalibra-
tion across proxy-sensitive attributes can signifi-
cantly mitigate these violations for the true, but
unknown, sensitive groups. Through several ex-
periments on real-world datasets, we illustrate that
approximate multiaccuracy and multicalibration
can be achieved even when sensitive group data
is incomplete or unavailable.
1Department of Biomedical Engineering, Johns Hopkins Uni-
versity, Baltimore, USA 2Mathematical Institute of Data Science,
Johns Hopkins University, Baltimore, USA 3Department of Ap-
plied Mathematics & Statistics, Johns Hopkins University, Bal-
timore, USA 4St. Jude Children’s Research Hospital, Arlington,
USA 5Department of Computer Science, Johns Hopkins Univer-
sity, Baltimore, USA. Correspondence to: Beepul Bharti <bb-
harti1@jhu.edu>, J...
Conference: ICML 2025

LENSLLM: Unveiling Fine-Tuning Dynamics for LLM Selection
Xinyue Zeng 1 Haohui Wang 1 Junhong Lin 2 Jun Wu 3 Tyler Cody 4 Dawei Zhou 1
Abstract
The proliferation of open-sourced Large Lan-
guage Models (LLMs) and diverse downstream
tasks necessitates efficient model selection, given
the impracticality of fine-tuning all candidates due
to computational constraints. Despite the recent
advances in LLM selection, a fundamental re-
search question largely remains nascent: how can
we model the dynamic behaviors of LLMs during
fine-tuning, thereby enhancing our understand-
ing of their generalization performance across
diverse downstream tasks? In this work, we pro-
pose a novel theoretical framework that provides
a proper lens to assess the generalization capabili-
ties of LLMs, thereby enabling accurate and effi-
cient LLM selection for downstream applications.
In particular, we first derive a PAC-Bayesian Gen-
eralization Bound that unveils fine-tuning dynam-
ics of LLMs and then introduce LENSLLM, a
Neural Tangent Kernel (NTK)-based Rectified
Scaling Model that enables accurate performance
predictions across diverse tasks while maintain-
ing computational efficiency. Extensive empirical
results on 3 large-scale benchmarks demonstrate
that our model achieves up to 91.1% accuracy and
reduces up to 88.5% computational cost in LLM
selection, outperforming 5 state-of-the-art meth-
ods. We open-source our proposed LENSLLM
model and corresponding results at LensLLM.io.
1. Introduction
The rise of large language models (LLMs) has revolu-
tionized natural language processing, driving remarkable
1Department of Computer Science, Virginia Tech, Blacksburg,
VA, USA. 2Department of Electrical Engineering and Computer
Science, Massachusetts Institute of Technology, Cambridge, MA,
USA. 3Department of Computer Science and Engineering, Michi-
gan State University, East Lansing, MI, USA. 4Department of In-
telligent Systems Division, Virginia Tech, Blacksburg, VA, USA..
Correspondence to...
Conference: ICML 2025

Wyckoff Transformer: Generation of Symmetric Crystals
Nikita Kazeev 1 Wei Nong 2 Ignat Romanov 3 Ruiming Zhu 2 Andrey Ustyuzhanin 4 5 1 Shuya Yamazaki 2
Kedar Hippalgaonkar 1 2 6
Abstract
Crystal symmetry plays a fundamental role in
determining its physical, chemical, and elec-
tronic properties such as electrical and thermal
conductivity, optical and polarization behavior,
and mechanical strength.
Almost all known
crystalline materials have internal symmetry.
However, this is often inadequately addressed by
existing generative models, making the consistent
generation of stable and symmetrically valid
crystal structures a significant challenge.
We
introduce WyFormer, a generative model that
directly tackles this by formally conditioning
on space group symmetry. It achieves this by
using Wyckoff positions as the basis for an
elegant,
compressed,
and discrete structure
representation. To model the distribution, we
develop a permutation-invariant autoregressive
model based on the Transformer encoder and
an absence of positional encoding. Extensive
experimentation
demonstrates
WyFormer’s
compelling combination of attributes: it achieves
best-in-class symmetry-conditioned generation,
incorporates a physics-motivated inductive bias,
produces structures with competitive stability,
predicts material properties with competitive
accuracy even without atomic coordinates, and
exhibits unparalleled inference speed. https:
//github.com/SymmetryAdvantage/
WyckoffTransformer
1Institute for Functional Intelligent Materials University of Sin-
gapore, Block S9, Level 9, 4 Science Drive 2, Singapore 117544
2School of Materials Science and Engineering, Nanyang Techno-
logical University, Singapore 639798 3HSE University, Myasnit-
skaya Ulitsa, 20, Moscow, Russia, 101000 4Constructor University,
Bremen, Campus Ring 1, 28759, Germany 5Constructor Knowl-
edge Labs, Bremen, Campus Ring 1, 28759, Germany 6Institute of
Materials Research and Engineering, Agency for Science Technol-
ogy and Re...
Conference: ICML 2025

Multilayer Matrix Factorization via Dimension-Reducing
Diffusion Variational Inference
Junbin Liu 1 Farzan Farnia 2 Wing-Kin Ma 1
Abstract
Multilayer matrix factorization (MMF) has re-
cently emerged as a generalized model of, and
potentially a more expressive approach than, the
classic matrix factorization. This paper considers
MMF under a probabilistic formulation, and our
focus is on inference methods under variational
inference. The challenge in this context lies in
determining a variational process that leads to a
computationally efficient and accurate approxima-
tion of the maximum likelihood inference. One
well-known example is the variational autoen-
coder (VAE), which uses neural networks for the
variational process. In this work, we take insight
from variational diffusion models in the context
of generative models to develop variational infer-
ence for MMF. We propose a dimension-reducing
diffusion process that results in a new way to
interact with the layered structures of the MMF
model. Experimental results demonstrate that the
proposed diffusion variational inference method
leads to improved performance scores compared
to several existing methods, including the VAE.
1. Introduction
Over decades, matrix factorization (MF) methods have
played a crucial role in a wide variety of problems such
as dimensionality reduction, low-dimension representation
learning, blind source separation (Hyv¨arinen et al., 2023),
hyperspectral unmixing (Ma et al., 2013), topic model-
ing (Arora et al., 2012), community detection (Yang &
Leskovec, 2013), to name a few. The broad interest of
1Department of Electronic Engineering, the Chinese Univer-
sity of Hong Kong, Hong Kong SAR of China 2Department
of Computer Science and Engineering, the Chinese Univer-
sity of Hong Kong, Hong Kong SAR of China.
Cor-
respondence to:
Junbin Liu <liujunbin@link.cuhk.edu.hk>,
Farzan
Farnia
<farnia@cse.cuhk.edu.hk>,
Wing-Kin
Ma
<wkma@ee.cuhk.edu.hk>.
Proceedings of the 42 nd International Confe...