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

Feature-aware Modulation for Learning from
Temporal Tabular Data
Hao-Run Cai
Han-Jia Ye
School of Artificial Intelligence, Nanjing University, China
National Key Laboratory for Novel Software Technology, Nanjing University, China
{caihr, yehj}@lamda.nju.edu.cn
Abstract
While tabular machine learning has achieved remarkable success, temporal distribu-
tion shifts pose significant challenges in real-world deployment, as the relationships
between features and labels continuously evolve. Static models assume fixed map-
pings to ensure generalization, whereas adaptive models may overfit to transient
patterns, creating a dilemma between robustness and adaptability. In this paper,
we analyze key factors essential for constructing an effective dynamic mapping for
temporal tabular data. We discover that evolving feature semantics—particularly
objective and subjective meanings—introduce concept drift over time. Crucially,
we identify that feature transformation strategies are able to mitigate discrepancies
in feature representations across temporal stages. Motivated by these insights, we
propose a feature-aware temporal modulation mechanism that conditions feature
representations on temporal context, modulating statistical properties such as scale
and skewness. By aligning feature semantics across time, our approach achieves
a lightweight yet powerful adaptation, effectively balancing generalizability and
adaptability. Benchmark evaluations validate the effectiveness of our method in
handling temporal shifts in tabular data.
1
Introduction
Tabular data, structured in rows and columns, forms the foundation of critical decision-making across
diverse domains such as healthcare [39], finance [42], and e-commerce [33]. Traditional machine
learning models for tabular data—including tree-based methods [5, 7, 22, 36] and modern deep
architectures [2, 10–13, 16–18, 52]—commonly assume independent and identically distributed
(i.i.d.) data [3, 30]. However, in real-world scenarios [54],...
Conference: ICML 2025

Laplace Transform Based Low-Complexity Learning
of Continuous Markov Semigroups
Vladimir R. Kostic * 1 2 Karim Lounici * 3 H´el`ene Halconruy * 4 5
Timoth´ee Devergne 1 Pietro Novelli 1 Massimiliano Pontil 1 6
Abstract
Markov processes serve as universal models for
many real-world random processes. This paper
presents a data-driven approach to learning these
models through the spectral decomposition of the
infinitesimal generator (IG) of the Markov semi-
group. Its unbounded nature complicates tradi-
tional methods such as vector-valued regression
and Hilbert-Schmidt operator analysis. Existing
techniques, including physics-informed kernel re-
gression, are computationally expensive and lim-
ited in scope, with no recovery guarantees for
transfer operator methods when the time-lag is
small. We propose a novel method leveraging the
IG’s resolvent, characterized by the Laplace trans-
form of transfer operators. This approach is robust
to time-lag variations, ensuring accurate eigen-
value learning even for small time-lags. Our statis-
tical analysis applies to a broader class of Markov
processes than current methods while reducing
computational complexity from quadratic to linear
in the state dimension. Finally, we demonstrate
our theoretical findings in several experiments.
1. Introduction
Markov semigroups play a critical role in modeling dynam-
ics of complex systems across various fields, including op-
tion pricing in finance (Karatzas & Shreve, 1991), molecular
dynamics (Sch¨utte & Huisinga, 2003) and climate modeling
(Majda et al., 2009), where understanding long-term behav-
ior is essential for accurate forecasting and interpretation.
*Equal contribution
1CSML, Istituto Italiano di Tecnologia,
Genova, Italy 2Faculty of Science, University of Novi Sad, Ser-
bia 3CMAP, ´Ecole Polytechnique, Paris, France 4SAMOVAR,
T´el´ecom Sud-Paris, Institut Polytechnique de Paris, France
5MODAL’X, Universit´e Paris Nanterre, France 6AI Center, Univer-
sity College London, UK. ...
Conference: ICML 2025

Schwarz–Schur Involution: Lightspeed Differentiable Sparse Linear Solvers
Yu Wang 1 2 3 S. Mazdak Abulnaga 1 2 Ya¨el Balbastre 1 4 Bruce Fischl 1 2
Abstract
Sparse linear solvers or generalized deconvolution
are fundamental to science and engineering, ap-
plied in partial differential equations (PDEs), sci-
entific computing, computer vision, and beyond.
Indirect solvers possess characteristics that make
them undesirable as stable differentiable modules;
existing direct solvers, though reliable, are too
expensive to be adopted in neural architectures.
We substantially accelerate direct sparse solvers
by up to three orders of magnitude, violating com-
mon assumptions that direct solvers are too slow.
We “condense” a sparse Laplacian matrix into a
dense tensor, a compact data structure that batch-
wise stores the Dirichlet-to-Neumann matrices,
reducing the sparse solving to recursively merg-
ing pairs of dense matrices that are much smaller.
The batched small dense systems are sliced and
inverted in parallel to take advantage of dense
GPU BLAS kernels, highly optimized in the era
of deep learning. Our method is efficient, quali-
fied as a strong zero-shot baseline for AI-based
PDE solving, and a reliable differentiable mod-
ule integrable into machine learning pipelines.
Our code is available at https://github.com/
wangyu9/Schwarz_Schur_Involution.
1. Introduction
In this paper, we propose a conceptually simple approach
that reduces solving certain sparse linear system (A, b) to
“involuting” tensors (α, β), leading to significant improve-
ments. Global linear systems in computer vision and scien-
tific computing, such as the Laplacian or Hessian systems on
image domains, are usually very large yet sparse, as induced
by pixel affinity. For a sparse A, direct solvers find the
exact solution: x = A−1b, for which we present the first
1Athinoula Martinos Center, Massachusetts General Hospital,
Harvard Medical School 2Massachusetts Institute of Technology,
MIT CSAIL, USA 3S...
Conference: ICML 2025

Wasserstein Flow Matching: Generative Modeling Over Families of Distributions
Doron Haviv * 1 2 Aram-Alexandre Pooladian * 3 Dana Pe’er 2 4 Brandon Amos 5
Abstract
Generative modeling typically concerns transport-
ing a single source distribution to a target distri-
bution via simple probability flows. However, in
fields like computer graphics and single-cell ge-
nomics, samples themselves can be viewed as dis-
tributions, where standard flow matching ignores
their inherent geometry. We propose Wasserstein
flow matching (WFM), which lifts flow matching
onto families of distributions using the Wasser-
stein geometry. Notably, WFM is the first algo-
rithm capable of generating distributions in high
dimensions, whether represented analytically (as
Gaussians) or empirically (as point-clouds). Our
theoretical analysis establishes that Wasserstein
geodesics constitute proper conditional flows over
the space of distributions, making for a valid FM
objective. Our algorithm leverages optimal trans-
port theory and the attention mechanism, demon-
strating versatility across computational regimes:
exploiting closed-form optimal transport paths for
Gaussian families, while using entropic estimates
on point-clouds for general distributions. WFM
successfully generates both 2D & 3D shapes
and high-dimensional cellular microenvironments
from spatial transcriptomics data. Code is avail-
able at Wasserstein Flow Matching.
1. Introduction
Today’s abundance of data and scalability of training mas-
sive neural networks has made it possible to generate hyper-
realistic images on the basis of training examples (OpenAI,
2022), as well as video and audio clips (Vyas et al., 2023;
Xing et al., 2023), and, of course, text (Bubeck et al., 2023).
All of these are instances of generative modeling: given
*Equal contribution 1Weill Cornell 2Memorial Sloan–Kettering
Cancer Center 3Center for Data Science, New York University
4Howard Hughes Medical Institute 5Meta AI. Correspondence to:
Doron Haviv ...
Conference: NeurIPS 2025

ALMGuard: Safety Shortcuts and Where to Find
Them as Guardrails for Audio–Language Models
Weifei Jin1, Yuxin Cao2, Junjie Su1, Minhui Xue3,4, Jie Hao1∗,
Ke Xu5, Jin Song Dong2, Derui Wang3
1Beijing University of Posts and Telecommunications
2National University of Singapore
3CSIRO’s Data61
4 Responsible AI Research (RAIR) Centre, The University of Adelaide
5Tsinghua University
Abstract
Recent advances in Audio-Language Models (ALMs) have significantly improved
multimodal understanding capabilities. However, the introduction of the audio
modality also brings new and unique vulnerability vectors. Previous studies have
proposed jailbreak attacks that specifically target ALMs, revealing that defenses di-
rectly transferred from traditional audio adversarial attacks or text-based Large Lan-
guage Model (LLM) jailbreaks are largely ineffective against these ALM-specific
threats. To address this issue, we propose ALMGuard, the first defense framework
tailored to ALMs. Based on the assumption that safety-aligned shortcuts naturally
exist in ALMs, we design a method to identify universal Shortcut Activation Pertur-
bations (SAPs) that serve as triggers that activate the safety shortcuts to safeguard
ALMs at inference time. To better sift out effective triggers while preserving the
model’s utility on benign tasks, we further propose Mel-Gradient Sparse Mask
(M-GSM), which restricts perturbations to Mel-frequency bins that are sensitive to
jailbreaks but insensitive to speech understanding. Both theoretical analyses and
empirical results demonstrate the robustness of our method against both seen and
unseen attacks. Overall, ALMGuard reduces the average success rate of advanced
ALM-specific jailbreak attacks to 4.6% across four models, while maintaining com-
parable utility on benign benchmarks, establishing it as the new state of the art. Our
code and data are available at https://github.com/WeifeiJin/ALMGuard.
1
Introduction
Audio-Language Models (ALMs) [18, 52] are revolutio...
Conference: NeurIPS 2025

FedMGP: Personalized Federated Learning with
Multi-Group Text-Visual Prompts
Weihao Bo1,Yanpeng Sun2∗, Yu Wang3, Xinyu Zhang4, Zechao Li1
1Nanjing University of Science and Technology
2National University of Singapore
3Baidu VIS
4University of Auckland
Abstract
In this paper, we introduce FedMGP, a new paradigm for personalized federated
prompt learning in vision-language models (VLMs). Existing federated prompt
learning (FPL) methods often rely on a single, text-only prompt representation,
which leads to client-speciﬁc overﬁtting and unstable aggregation under heteroge-
neous data distributions. Toward this end, FedMGP equips each client with mul-
tiple groups of paired textual and visual prompts, enabling the model to capture
diverse, ﬁne-grained semantic and instance-level cues. A diversity loss is intro-
duced to drive each prompt group to specialize in distinct and complementary se-
mantic aspects, ensuring that the groups collectively cover a broader range of local
characteristics. During communication, FedMGP employs a dynamic prompt ag-
gregation strategy based on similarity-guided probabilistic sampling: each client
computes the cosine similarity between its prompt groups and the global prompts
from the previous round, then samples s groups via a softmax-weighted distribu-
tion. This soft selection mechanism preferentially aggregates semantically aligned
knowledge while still enabling exploration of underrepresented patternseffectively
balancing the preservation of common knowledge with client-speciﬁc features.
Notably, FedMGP maintains parameter efﬁciency by redistributing a ﬁxed prompt
capacity across multiple groups, achieving state-of-the-art performance with the
lowest communication parameters (5.1k) among all federated prompt learning
methods. Theoretical analysis shows that our dynamic aggregation strategy pro-
motes robust global representation learning by reinforcing shared semantics while
suppressing client-speciﬁc noise. Extensive experiments dem...
Conference: ICML 2025

Modulated Diffusion:
Accelerating Generative Modeling with Modulated Quantization
Weizhi Gao 1 Zhichao Hou 1 Junqi Yin 2 Feiyi Wang 2 Linyu Peng 3 Xiaorui Liu 1
Abstract
Diffusion models have emerged as powerful gen-
erative models, but their high computational cost
in iterative sampling remains a significant bot-
tleneck. In this work, we present an in-depth
and insightful study of state-of-the-art accelera-
tion techniques for diffusion models, including
caching and quantization, and reveal their limita-
tions in computation error and generation quality.
To break these limits, this work introduces Modu-
lated Diffusion (MoDiff), an innovative, rigorous,
and principled framework that accelerates gener-
ative modeling through modulated quantization
and error compensation. MoDiff not only inherits
the advantages of existing caching and quantiza-
tion methods but also serves as a general frame-
work to accelerate all diffusion models. The ad-
vantages of MoDiff are supported by solid theoret-
ical insight and analysis. In addition, extensive ex-
periments on CIFAR-10 and LSUN demonstrate
that MoDiff significantly reduces activation quan-
tization from 8 bits to 3 bits without performance
degradation in post-training quantization (PTQ).
Our code implementation is available at https:
//github.com/WeizhiGao/MoDiff.
1. Introduction
Diffusion models have emerged as powerful generative mod-
els for producing high-quality data samples, ranging from
images to audio and beyond (Ho et al., 2020; Song et al.,
2021a;b). These models work by iteratively transforming
a simple noise distribution into complex, structured out-
puts, guided by a learned reverse diffusion process. Despite
their effectiveness, diffusion models come with significant
1Department of Computer Science, North Carolina State Uni-
versity 2National Center for Computational Science, Oak Ridge
National Lab 3Department of Mechanical Engineering, Keio Uni-
versity. Correspondence to: Xiaorui Liu <xliu96@ncsu.edu>.
P...
Conference: ICML 2025

Joint Localization and Activation Editing for Low-Resource Fine-Tuning
Wen Lai 1 2 Alexander Fraser 1 2 Ivan Titov 3 4
Abstract
Parameter-efficient fine-tuning (PEFT) methods,
such as LoRA, are commonly used to adapt LLMs.
However, the effectiveness of standard PEFT
methods is limited in low-resource scenarios with
only a few hundred examples. Recent advances
in interpretability research have inspired the emer-
gence of activation editing (or steering) tech-
niques, which modify the activations of specific
model components. Due to their extremely small
parameter counts, these methods show promise
for small datasets. However, their performance is
highly dependent on identifying the correct mod-
ules to edit and often lacks stability across differ-
ent datasets. In this paper, we propose Joint Local-
ization and Activation Editing (JOLA), a method
that jointly learns (1) which heads in the Trans-
former to edit (2) whether the intervention should
be additive, multiplicative, or both and (3) the
intervention parameters themselves - the vectors
applied as additive offsets or multiplicative scal-
ings to the head output. Through evaluations on
three benchmarks spanning commonsense reason-
ing, natural language understanding, and natural
language generation, we demonstrate that JOLA
consistently outperforms existing methods.1
1. Introduction
Parameter-efficient fine-tuning (PEFT; Han et al., 2024)
methods are widely used to adapt large language models
(LLMs). However, popular PEFT approaches like LoRA
(Hu et al., 2021) often struggle in low-resource settings with
*Work done during Wen Lai’s visit to The University of Edin-
burgh. 1Technical University of Munich 2Munich Center for Ma-
chine Learning 3ILLC, University of Edinburgh 4ILLC, University
of Amsterdam. Correspondence to: Wen Lai <wen.lai@tum.de>,
Ivan Titov <ititov@inf.ed.ac.uk>, Alexander Fraser <alexan-
der.fraser@tum.de>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. ...
Conference: NeurIPS 2025

Ditch the Denoiser: Emergence of Noise Robustness in
Self-Supervised Learning from Data Curriculum
Wenquan Lu1
Jiaqi Zhang1
Hugues Van Assel2
Randall Balestriero1
1Brown University
2Genentech
{wenquan_lu, jiaqi_zhang6, randall_balestriero}@brown.edu
van_assel.hugues@gene.com
Abstract
Self-Supervised Learning (SSL) has become a powerful solution to extract rich
representations from unlabeled data. Yet, SSL research is mostly focused on
clean, curated and high-quality datasets. As a result, applying SSL on noisy data
remains a challenge, despite being crucial to applications such as astrophysics,
medical imaging, geophysics or finance. In this work, we present a fully self-
supervised framework that enables noise-robust representation learning without
requiring a denoiser at inference or downstream fine-tuning. Our method first
trains an SSL denoiser on noisy data, then uses it to construct a denoised-to-
noisy data curriculum (i.e., training first on denoised, then noisy samples) for
pretraining a SSL backbone (e.g., DINOv2), combined with a teacher-guided
regularization that anchors noisy embeddings to their denoised counterparts. This
process encourages the model to internalize noise robustness. Notably, the denoiser
can be discarded after pretraining, simplifying deployment. On ImageNet-1k with
ViT-B under extreme Gaussian noise (σ = 255, SNR = 0.72 dB), our method
improves linear probing accuracy by 4.8% over DINOv2, demonstrating that
denoiser-free robustness can emerge from noise-aware pretraining. The code is
available at https://github.com/wenquanlu/noisy_dinov2.
1
Introduction
Self-supervised methods like DINOv2 [31, 7, 16, 9] have demonstrated remarkable success by
leveraging unlabeled data to learn visual representations able to solve many tasks in zero or few
shot manners [3]. However, the performances of those methods heavily rely on the availability
of clean and high-quality datasets. As a result, some prior works have focused on building data
curation ...
Conference: NeurIPS 2025

Towards General Continuous Memory for
Vision-Language Models
Wenyi Wu∗, Zixuan Song∗, Kun Zhou†, Yifei Shao, Zhiting Hu, Biwei Huang
University of California, San Diego.
kuzhou@ucsd.edu
Abstract
Language models (LMs) and their extension, vision-language models (VLMs),
have achieved remarkable performance across various tasks. However, they still
struggle with complex reasoning tasks that require multimodal or multilingual
real-world knowledge. To support such capabilities, an external memory system
that can efficiently provide relevant multimodal information is essential. Existing
approaches generally concatenate image and text tokens into a long sequence as
memory, which, however, may drastically increase context length and even de-
grade performance. In contrast, we propose using continuous memory-a compact
set of dense embeddings-to more effectively and efficiently represent multimodal
and multilingual knowledge. Our key insight is that a VLM can serve as its own
continuous memory encoder. We empirically show that this design improves per-
formance on complex multimodal reasoning tasks. Building on this, we introduce a
data-efficient and parameter-efficient method to fine-tune the VLM into a memory
encoder, requiring only 1.2% of the model’s parameters and a small corpus of
15.6K self-synthesized samples. Our approach CoMEM utilizes VLM’s origi-
nal capabilities to encode arbitrary multimodal and multilingual knowledge into
just 8 continuous embeddings. Since the inference-time VLM remains frozen,
our memory module is plug-and-play and can be flexibly integrated as needed.
Extensive experiments across eight multimodal reasoning benchmarks demon-
strate the effectiveness of our approach. Code and data is publicly released here
https://github.com/WenyiWU0111/CoMEM.
1
Introduction
Figure 1: CoMEM architecture in comparison
to the traditional RAG method.
Through
large-scale
training,
language
mod-
els (LMs) [1, 2] have demonstrated remarkable
performance across diver...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
BOOTSTRAPPED MODEL PREDICTIVE CONTROL
Yuhang Wang, Hanwei Guo, Sizhe Wang, Long Qian, Xuguang Lan∗
National Key Laboratory of Human-Machine Hybrid Augmented Intelligence
National Engineering Research Center for Visual Information and Application
Institute of Artificial Intelligence and Robotics, Xi’an Jiaotong University, Xi’an, China
{wertyuilife,clannnagisa,wangsizhe,qianlongym}@stu.xjtu.edu.cn
xglan@mail.xjtu.edu.cn
ABSTRACT
Model Predictive Control (MPC) has been demonstrated to be effective in con-
tinuous control tasks. When a world model and a value function are available,
planning a sequence of actions ahead of time leads to a better policy. Existing
methods typically obtain the value function and the corresponding policy in a
model-free manner. However, we find that such an approach struggles with com-
plex tasks, resulting in poor policy learning and inaccurate value estimation. To
address this problem, we leverage the strengths of MPC itself. In this work, we
introduce Bootstrapped Model Predictive Control (BMPC), a novel algorithm that
performs policy learning in a bootstrapped manner. BMPC learns a network pol-
icy by imitating an MPC expert, and in turn, uses this policy to guide the MPC
process. Combined with model-based TD-learning, our policy learning yields bet-
ter value estimation and further boosts the efficiency of MPC. We also introduce
a lazy reanalyze mechanism, which enables computationally efficient imitation
learning. Our method achieves superior performance over prior works on diverse
continuous control tasks. In particular, on challenging high-dimensional loco-
motion tasks, BMPC significantly improves data efficiency while also enhancing
asymptotic performance and training stability, with comparable training time and
smaller network sizes. Code is available at https://github.com/wertyuilife2/bmpc.
1
INTRODUCTION
Model-based reinforcement learning (MBRL) algorithms that incorporate online pla...
Conference: ICML 2025

Reducing Tool Hallucination via Reliability Alignment
Hongshen Xu 1 2 3 Zichen Zhu 1 2 3 Lei Pan 4 Zihan Wang 1 2 3 Su Zhu 4
Da Ma 1 2 3 Ruisheng Cao 1 2 3 Lu Chen 1 2 3 5 Kai Yu 1 2 3 5
Abstract
Large Language Models (LLMs) have expanded
their capabilities beyond language generation to
interact with external tools, enabling automation
and real-world applications. However, tool hallu-
cinations—where models either select inappropri-
ate tools or misuse them—pose significant chal-
lenges, leading to erroneous task execution, in-
creased computational costs, and reduced system
reliability. To systematically address this issue,
we define and categorize tool hallucinations into
two main types: tool selection hallucination and
tool usage hallucination. To evaluate and mitigate
these issues, we introduce RelyToolBench, which
integrates specialized test cases and novel metrics
to assess hallucination-aware task success and ef-
ficiency. Finally, we propose Relign, a reliability
alignment framework that expands the tool-use
action space to include indecisive actions, allow-
ing LLMs to defer tool use, seek clarification, or
adjust tool selection dynamically. Through exten-
sive experiments, we demonstrate that Relign sig-
nificantly reduces tool hallucinations, improves
task reliability, and enhances the efficiency of
LLM tool interactions. The code and data are
publicly available at https://github.com/
X-LANCE/ToolHallucination.
1. Introduction
The primary goal of tool learning is to enable Large Lan-
guage Models (LLMs; Gemini Team, 2023; Achiam et al.,
2023; Dubey et al., 2024) to understand and effectively use
1X-LANCE Lab, School of Computer Science, Shanghai Jiao
Tong University, Shanghai, China. 2MoE Key Lab of Artificial
Intelligence, Shanghai, China.
3Jiangsu Key Lab of Language
Computing, Suzhou, China. 4AISpeech Co., Ltd., Suzhou, China.
5Suzhou Laboratory, Suzhou, China. xuhongshen@sjtu.edu.cn.
Correspondence to: Lu Chen <chenlusz@sjtu.edu.cn>, Kai Yu
<kai.yu@sj...