AIS-Bench/AIS-Bench huggingface dataset:
Universal Neural Optimal Transport
Jonathan Geuter 1 2 * Gregor Kornhardt 3 * Ingimar Tomasson 3 * Vaios Laschos 4
Abstract
Optimal Transport (OT) problems are a corner-
stone of many applications, but solving them
is computationally expensive. To address this
problem, we propose UNOT (Universal Neural
Optimal Transport), a novel framework capable
of accurately predicting (entropic) OT distances
and plans between discrete measures for a given
cost function. UNOT builds on Fourier Neural
Operators, a universal class of neural networks
that map between function spaces and that are
discretization-invariant, which enables our net-
work to process measures of variable resolutions.
The network is trained adversarially using a sec-
ond, generating network and a self-supervised
bootstrapping loss. We ground UNOT in an ex-
tensive theoretical framework. Through experi-
ments on Euclidean and non-Euclidean domains,
we show that our network not only accurately
predicts OT distances and plans across a wide
range of datasets, but also captures the geometry
of the Wasserstein space correctly. Furthermore,
we show that our network can be used as a state-
of-the-art initialization for the Sinkhorn algorithm
with speedups of up to 7.4×, significantly outper-
forming existing approaches.
1. Introduction
Optimal Transport (Villani, 2009; Peyr´e & Cuturi, 2019)
plays an increasing role in various areas in machine learn-
ing, such as domain adaptation (Courty et al., 2017), single-
cell genomics (Schiebinger et al., 2019), imitation learn-
ing (Dadashi et al., 2020), imaging (Schmitz et al., 2018),
dataset adaptation (Alvarez-Melis & Fusi, 2021), and signal
processing (Kolouri et al., 2017). Oftentimes, an entropic
*Equal contribution
1Harvard John A. Paulson School of En-
gineering and Applied Sciences 2Kempner Institute at Harvard
University 3Department of Mathematics, Technische Universit¨at
Berlin, Germany 4Weierstrass Institute, Berlin, Germany. Corre-
spondence to: Jonathan Geuter...
Benign Overfitting in Single-Head Attention
Roey Magen∗
Weizmann Institute of Science
roey.magen@weizmann.ac.il
Shuning Shang∗†
Princeton University
shuning@princeton.edu
Zhiwei Xu
University of Michigan
zhiweixu@umich.edu
Spencer Frei‡
UC Davis
sfrei@ucdavis.edu
Wei Hu§
University of Michigan
vvh@umich.edu
Gal Vardi§
Weizmann Institute of Science
gal.vardi@weizmann.ac.il
Abstract
The phenomenon of benign overfitting, where a trained neural network perfectly fits
noisy training data but still achieves near-optimal test performance, has been exten-
sively studied in recent years for linear models and fully-connected/convolutional
networks. In this work, we study benign overfitting in a single-head softmax
attention model, which is the fundamental building block of Transformers. We
prove that under appropriate conditions, the model exhibits benign overfitting in
a classification setting already after two steps of gradient descent. Moreover, we
show conditions where a minimum-norm/maximum-margin interpolator exhibits
benign overfitting. We study how the overfitting behavior depends on the signal-
to-noise ratio (SNR) of the data distribution, namely, the ratio between norms of
signal and noise tokens, and prove that a sufficiently large SNR is both necessary
and sufficient for benign overfitting.
1
Introduction
Neural networks often exhibit a remarkable phenomenon, known as benign overfitting, where they
achieve a perfect fit to noisy training examples and still generalize well to unseen data [1, 2]. This
phenomenon contradicts classical wisdom in machine learning, and has become a central research
question in the theory of deep learning. Existing works on benign overfitting study under what
conditions the phenomenon occurs in different architectures. These works focus on linear models,
and on shallow fully-connected and convolutional neural networks.
In recent years, Transformers [3] have emerged as a leading neural network architecture, with
impactful applications ac...
Adjusting Model Size in Continual Gaussian Processes:
How Big is Big Enough?
Guiomar Pescador-Barrios 1 Sarah Filippi 1 Mark van der Wilk 2
Abstract
Many machine learning models require setting
a parameter that controls their size before train-
ing, e.g. number of neurons in DNNs, or inducing
points in GPs. Increasing capacity typically im-
proves performance until all the information from
the dataset is captured. After this point, compu-
tational cost keeps increasing, without improved
performance. This leads to the question “How big
is big enough?” We investigate this problem for
Gaussian processes (single-layer neural networks)
in continual learning. Here, data becomes avail-
able incrementally, and the final dataset size will
therefore not be known before training, prevent-
ing the use of heuristics for setting a fixed model
size. We develop a method to automatically ad-
just model size while maintaining near-optimal
performance. Our experimental procedure fol-
lows the constraint that any hyperparameters must
be set without seeing dataset properties, and we
show that our method performs well across di-
verse datasets without the need to adjust its hy-
perparameter, showing it requires less tuning than
others.
1. Introduction
Continual learning aims to train models when the data ar-
rives in a stream of batches, without storing data after it has
been processed, and while obtaining predictive performance
that is as high as possible at each point in time (Ring, 1997).
Selecting the size of the model is challenging in this set-
ting, since typical non-continual training procedures do this
by trial-and-error (cross-validation) using repeated training
runs, which is not possible under our requirement of not
storing any data. This is a crucial parameter to set well,
1Department of Mathematics, Imperial College London, UK
2Department of Computer Science, University of Oxford, UK.
Correspondence to: Guiomar Pescador-Barrios <glp22@ic.ac.uk>.
Proceedings of the 42 nd Int...
Large Language Models are Demonstration Pre-Selectors for Themselves
Jiarui Jin 1 2 * † Yuwei Wu 3 * Haoxuan Li 4 Xiaoting He 5 Weinan Zhang 1 Yiming Yang 3 Yong Yu 1
Jun Wang 6 Mengyue YangB 7
Abstract
In-context learning (ICL) with large language
models (LLMs) delivers strong few-shot per-
formance by choosing few-shot demonstrations
from the entire training data. However, exist-
ing ICL methods, which rely on similarity or di-
versity scores to choose demonstrations, incur
high computational costs due to repeatedly re-
trieval from large-scale datasets for each query.
To this end, we propose FEEDER (FEw yet Es-
sential Demonstration prE-selectoR), a novel pre-
selection framework that identifies a core sub-
set of demonstrations containing the most repre-
sentative examples in the training data, tailored
to specific LLMs. To construct this subset, we
introduce the “sufficiency” and “necessity” met-
rics in the pre-selection stage and design a tree-
based algorithm to identify representative exam-
ples efficiently. Once pre-selected, this represen-
tative subset can effectively replace the full train-
ing data, improving efficiency while maintain-
ing comparable performance in ICL. Additionally,
our pre-selected subset also benefits fine-tuning
LLMs, where we introduce a bi-level optimization
method that enhances training efficiency without
sacrificing performance. Experiments with LLMs
ranging from 300M to 8B parameters show that
FEEDER can reduce training data size by over 20%
while maintaining performance and seamlessly in-
tegrating with various downstream demonstration
selection strategies in ICL.
* Equal contributions: Jiarui Jin and Yuwei Wu. † Work done
during Jiarui Jin’s visit at University College London. 1Shanghai
Jiao Tong University 2Xiaohongshu Inc. 3Carnegie Mellon Uni-
versity 4Peking University 5No Affiliation 6University College
London 7University of Bristol. Correspondence to: Mengyue Yang
<mengyue.yang@bristol.ac.uk>.
Proceedings of the 42 nd ...
Protein Structure Tokenization: Benchmarking and New Recipe
Xinyu Yuan * 1 2 Zichen Wang † 3 Marcus Collins 3 Huzefa Rangwala † ‡ 3
Abstract
Recent years have witnessed a surge in the devel-
opment of protein structural tokenization meth-
ods, which chunk protein 3D structures into dis-
crete or continuous representations. Structure to-
kenization enables the direct application of pow-
erful techniques like language modeling for pro-
tein structures, and large multimodal models to
integrate structures with protein sequences and
functional texts. Despite the progress, the capa-
bilities and limitations of these methods remain
poorly understood due to the lack of a unified
evaluation framework. We first introduce Struct-
TokenBench, a framework that comprehensively
evaluates the quality and efficiency of structure
tokenizers, focusing on fine-grained local sub-
structures rather than global structures, as typical
in existing benchmarks. Our evaluations reveal
that no single model dominates all benchmarking
perspectives. Observations of codebook under-
utilization led us to develop AminoAseed, a sim-
ple yet effective strategy that enhances codebook
gradient updates and optimally balances codebook
size and dimension for improved tokenizer utiliza-
tion and quality. Compared to the leading VQ-
VAE model ESM3, our method achieves an aver-
age of 6.31% performance improvement across
24 supervised tasks, with sensitivity and utiliza-
tion rates increased by 12.83% and 124.03%,
respectively. Source code and model weights
are
available
at
https://github.com/
KatarinaYuan/StructTokenBench.
*This work was completed while the author was an intern
at Amazon.
†Corresponding author.
‡Huzefa Rangwala is on
LOA as a Professor of Computer Science at George Mason
University.
This paper describes work performed at Ama-
zon.
1Mila - Quebec AI Institute
2University of Montreal
3Amazon.
Emails: Xinyu Yuan <xinyu.yuan@mila.quebec>,
Zichen
Wang
<wangzc921@gmail.com>,
Marcus
Collins
<collmr@...
AutoAL: Automated Active Learning with Differentiable Query Strategy Search
Yifeng Wang 1 Xueying Zhan 2 Siyu Huang 3
Abstract
As deep learning continues to evolve, the need
for data efficiency becomes increasingly impor-
tant. Considering labeling large datasets is both
time-consuming and expensive, active learning
(AL) provides a promising solution to this chal-
lenge by iteratively selecting the most informa-
tive subsets of examples to train deep neural net-
works, thereby reducing the labeling cost. How-
ever, the effectiveness of different AL algorithms
can vary significantly across data scenarios, and
determining which AL algorithm best fits a given
task remains a challenging problem. This work
presents the first differentiable AL strategy search
method, named AutoAL, which is designed on
top of existing AL sampling strategies. AutoAL
consists of two neural nets, named SearchNet
and FitNet, which are optimized concurrently un-
der a differentiable bi-level optimization frame-
work. For any given task, SearchNet and Fit-
Net are iteratively co-optimized using the la-
beled data, learning how well a set of candi-
date AL algorithms perform on that task. With
the optimal AL strategies identified, SearchNet
selects a small subset from the unlabeled pool
for querying their annotations, enabling efficient
training of the task model. Experimental results
demonstrate that AutoAL consistently achieves
superior accuracy compared to all candidate AL
algorithms and other selective AL approaches,
showcasing its potential for adapting and inte-
grating multiple existing AL methods across di-
verse tasks and domains. Code is available at:
https://github.com/haizailache999/AutoAL.
1Department of Electrical and Computer Engineering, Carnegie
Mellon University, Pittsburgh, USA 2Ray and Stephanie Lane
Computational Biology Department, Carnegie Mellon University,
Pittsburgh, USA 3Visual Computing Division, School of Com-
puting, Clemson University, Clemson, USA. Correspondence t...
QoS-Efficient Serving of Multiple Mixture-of-Expert LLMs Using Partial
Runtime Reconfiguration
HamidReza Imani 1 Jiaxin Peng 1 Peiman Mohseni 2 Abdolah Amirany 1 Tarek El-Ghazawi 1
https://github.com/hamid-74/Multi-MoE
Abstract
The deployment of mixture-of-experts (MoE)
large language models (LLMs) presents signif-
icant challenges due to their high memory de-
mands.
These challenges become even more
pronounced in multi-tenant environments, where
shared resources must accommodate multiple
models, limiting the effectiveness of conventional
virtualization techniques. This paper addresses
the problem of efficiently serving multiple fine-
tuned MoE-LLMs on a single GPU. We propose a
serving system that employs similarity-based ex-
pert consolidation to reduce the overall memory
footprint by sharing similar experts across mod-
els. To ensure output quality, we introduce run-
time partial reconfiguration, dynamically replac-
ing non-expert layers when processing requests
from different models. As a result, our approach
achieves competitive output quality while main-
taining throughput comparable to serving a sin-
gle model, and incurs only a negligible increase
in time-to-first-token (TTFT). Experiments on a
server with a single NVIDIA A100 GPU (80GB)
using Mixtral-8x7B models demonstrate an 85%
average reduction in turnaround time compared to
NVIDIA’s multi-instance GPU (MIG). Further-
more, experiments on Google’s Switch Trans-
former Base-8 model with up to four variants
demonstrate the scalability and resilience of our
approach in maintaining output quality compared
to other model merging baselines, highlighting its
effectiveness.
1Department of Electrical and Computer Engineering, The
George Washington University 2Computer Science and Engineer-
ing Department, Texas A&M University. Correspondence to:
HamidReza Imani <hamidreza@gwu.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the auth...
Resolving Lexical Bias in Model Editing
Hammad Rizwan 1 Domenic Rosati 1 Ga Wu 1 Hassan Sajjad 1
Abstract
Model editing aims to modify the knowledge
of a pre-trained language model. Previous ap-
proaches have often involved direct alterations to
model weights, which can result in model degrada-
tion. Recent weight-preserving techniques avoid
making modifications to the model’s weights
by employing an adapter that implements edits
through auxiliary components. These rely heavily
on scoping mechanisms based on distance func-
tions on the model’s representation space to de-
termine when to trigger edits. We demonstrate
that current adapter methods are critically vul-
nerable to strong lexical biases, leading to is-
sues such as applying edits to irrelevant prompts
with overlapping words. This paper presents a
principled approach to learning a disentangled
representation space that facilitates precise lo-
calization of edits by maintaining distance be-
tween irrelevant prompts while preserving prox-
imity among paraphrases. In our empirical study,
we show that our method, Projector Editor Net-
works for Model Editing - PENME, achieves
state-of-the-art model editing results while be-
ing computationally efficient during inference
compared to previous methods and adaptable
across different architectures. We provide the
codebase of PENME here: https://github.
com/hammadrizwan/PENME.git
1. Introduction
Large Language Models (LLMs) are successful in solving a
diverse range of natural language processing tasks (Devlin
et al., 2019; Liu et al., 2019; Touvron et al., 2023b; Radford
et al., 2019). Despite their successes, LLMs are fallible in
large part due to the noisy and imperfect nature of the data
used for training (Zhu et al., 2020). Moreover, as the world
1Department of Computer Science, Dalhousie University, Hal-
ifax, Canada.
Correspondence to: Hammad Rizwan <ham-
mad.rizwan@dal.ca>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canad...
Leveraging Predictive Equivalence in Decision Trees
Hayden McTavish * 1 Zachery Boner * 1 Jon Donnelly * 1 Margo Seltzer 2 Cynthia Rudin 1
Abstract
Decision trees are widely used for interpretable
machine learning due to their clearly structured
reasoning process. However, this structure belies
a challenge we refer to as predictive equivalence:
a given tree’s decision boundary can be repre-
sented by many different decision trees. The pres-
ence of models with identical decision boundaries
but different evaluation processes makes model
selection challenging. The models will have dif-
ferent variable importance and behave differently
in the presence of missing values, but most opti-
mization procedures will arbitrarily choose one
such model to return. We present a boolean logi-
cal representation of decision trees that does not
exhibit predictive equivalence and is faithful to the
underlying decision boundary. We apply our rep-
resentation to several downstream machine learn-
ing tasks. Using our representation, we show that
decision trees are surprisingly robust to test-time
missingness of feature values; we address predic-
tive equivalence’s impact on quantifying variable
importance; and we present an algorithm to opti-
mize the cost of reaching predictions.
1. Introduction
Decision trees are widely used for interpretable machine
learning (Rudin et al., 2022). Their structure of discrete
decisions has long been leveraged for difficult tasks such as
handling missing data (Therneau et al., 1997) and measur-
ing variable importance (Breiman, 1984). Recent advances
in decision tree optimization (Lin et al., 2020; Demirovi´c
*Equal contribution
1Department of Computer Science,
Duke University, Durham, North Carolina, USA 2Department
of Computer Science, University of British Columbia, Van-
couver,
BC, Canada.
Correspondence to:
Hayden Mc-
Tavish <hayden.mctavish@duke.edu>, Zachery Boner <zach-
ery.boner@duke.edu>.
Proceedings of the 42 nd International Conference on Mac...
PokéChamp: an Expert-level Minimax Language Agent
Seth Karten * 1 Andy Luu Nguyen 1 Chi Jin 1
Abstract
We introduce PokéChamp, a minimax agent
powered by Large Language Models (LLMs) for
Pokémon battles. Built on a general framework
for two-player competitive games, PokéChamp
leverages the generalist capabilities of LLMs to
enhance minimax tree search. Specifically, LLMs
replace three key modules: (1) player action sam-
pling, (2) opponent modeling, and (3) value func-
tion estimation, enabling the agent to effectively
utilize gameplay history and human knowledge to
reduce the search space and address partial observ-
ability. Notably, our framework requires no addi-
tional LLM training. We evaluate PokéChamp
in the popular Gen 9 OU format. When powered
by GPT-4o, it achieves a win rate of 76% against
the best existing LLM-based bot and 84% against
the strongest rule-based bot, demonstrating its su-
perior performance. Even with an open-source 8-
billion-parameter Llama 3.1 model, PokéChamp
consistently outperforms the previous best LLM-
based bot, Pokéllmon powered by GPT-4o, with
a 64% win rate. PokéChamp attains a projected
Elo of 1300-1500 on the Pokémon Showdown
online ladder, placing it among the top 30%-10%
of human players. In addition, this work compiles
the largest real-player Pokémon battle dataset, fea-
turing over 3 million games, including more than
500k high-Elo matches. Based on this dataset,
we establish a series of battle benchmarks and
puzzles to evaluate specific battling skills. We fur-
ther provide key updates to the local game engine.
This work establishes Pokémon as a benchmark to
integrate LLM technologies with game-theoretic
algorithms addressing general multi-agent prob-
lems. Videos, code, and dataset are available on-
line.
*Equal contribution 1Princeton University. Correspondence to:
Seth Karten <sethkarten@princeton.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright ...
What Limits Bidirectional Model’s Generative Capabilities?
A Uni-Bi-Directional Mixture-of-Expert Method For Bidirectional Fine-tuning
Zuchao Li 1 Yonghua Hei 2 Qiwei Li 3 Lefei Zhang 3 Ping Wang 4 Hai Zhao 5 Baoyuan Qi 6 Guoming Liu 6
Abstract
Large Language Models (LLMs) excel in genera-
tion tasks, yet their causal attention mechanisms
limit performance in embedding tasks. While
bidirectional modeling may enhance embeddings,
naively fine-tuning unidirectional models bidi-
rectionally severely degrades generative perfor-
mance. To investigate this trade-off, we analyze
attention weights as dependence indicators and
find that bidirectional fine-tuning increases subse-
quent dependence, impairing unidirectional gen-
eration. Through systematic Transformer module
evaluations, we discover the FFN layer is least
affected by such dependence. Leveraging this dis-
covery, we propose UBMoE-LLM, a novel Uni-
Bi-directional Mixture-of-Experts LLM, which
integrates the original unidirectional FFN with a
bidirectionally fine-tuned FFN via unsupervised
contrastive learning. This MoE-based approach
enhances embedding performance while preserv-
ing robust generation.
Extensive experiments
across diverse datasets and model scales vali-
date our attention dependence metric and demon-
strate UBMoE-LLM’s superior generative qual-
ity and reduced hallucination. Code is available
at: https://github.com/heiyonghua/
ubmoe_llm.
1. Introduction
Creating a model capable of both generation and text em-
bedding has been a long-term goal in the field of natural
1School of Artificial intelligence, Wuhan University 2This work
was conducted during Yonghua Hei’s internship at the School of
Artificial Intelligence, Wuhan University. 3School of Computer
Science, Wuhan University 4School of Information Management,
Wuhan University 5Department of Computer Science and Engineer-
ing, Shanghai Jiao Tong University 6Xiaomi, Beijing, China. Cor-
respondence to: Lefei Zhang <zhanglefei@whu.edu.cn>, Ping
Wan...
Adversarial Reasoning at Jailbreaking Time
Mahdi Sabbaghi 1 Paul Kassianik 2 George Pappas 1 Amin Karbasi 1 Hamed Hassani 1
Abstract
As large language models (LLMs) are becoming
more capable and widespread, the study of their
failure cases is becoming increasingly important.
Recent advances in standardizing, measuring, and
scaling test-time compute suggest new method-
ologies for optimizing models to achieve high per-
formance on hard tasks. In this paper, we apply
these advances to the task of “model jailbreaking”:
eliciting harmful responses from aligned LLMs.
We develop an adversarial reasoning approach to
automatic jailbreaking that leverages a loss signal
to guide the test-time compute, achieving SOTA
attack success rates against many aligned LLMs,
even those that aim to trade inference-time com-
pute for adversarial robustness. Our approach in-
troduces a new paradigm in understanding LLM
vulnerabilities, laying the foundation for the de-
velopment of more robust and trustworthy AI sys-
tems. Code is available at Github.
1. Introduction
Large language models (LLMs) are increasingly deployed
with various safety techniques to ensure alignment with hu-
man values. Common strategies include RLHF (Christiano
et al., 2023; Ouyang et al., 2022), DPO (Rafailov et al.,
2024), and the usage of dedicated guardrail models (Inan
et al., 2023; Rebedea et al., 2023). In nominal use cases,
alignment methods typically refuse to generate objection-
able content, but adversarially designed prompts can bypass
these guardrails. A challenge known as jailbreaking con-
sists of finding prompts that circumvent safety measures
and elicit undesirable behaviors.
Current jailbreaking methods fall into two categories: token-
space and semantic-space attacks.
Token-space attacks
(Shin et al., 2020; Wen et al., 2023; Zou et al., 2023; Hayase
et al., 2024; Andriushchenko et al., 2024a) focus on token-
1University of Pennsylvania 2Robust Intelligence @ Cisco.
Correspondence to: Mahdi Sabagghi ...