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

Geometry-Informed Neural Networks
Arturs Berzins * 1 2 Andreas Radler * 1 Eric Volkmann * 1
Sebastian Sanokowski 1 Sepp Hochreiter 1 Johannes Brandstetter 1 3
Abstract
Geometry is a ubiquitous tool in computer graph-
ics, design, and engineering. However, the lack
of large shape datasets limits the application of
state-of-the-art supervised learning methods and
motivates the exploration of alternative learning
strategies. To this end, we introduce geometry-
informed neural networks (GINNs) – a frame-
work for training shape-generative neural fields
without data by leveraging user-specified design
requirements in the form of objectives and con-
straints. By adding diversity as an explicit con-
straint, GINNs avoid mode-collapse and can gen-
erate multiple diverse solutions, often required in
geometry tasks. Experimentally, we apply GINNs
to several problems spanning physics, geometry,
and engineering design, showing control over ge-
ometrical and topological properties, such as sur-
face smoothness or the number of holes. These
results demonstrate the potential of training shape-
generative models without data, paving the way
for new generative design approaches without
large datasets.
1. Introduction
Recent advances in deep learning have revolutionized fields
with abundant data, such as computer vision and natural
language processing. However, the scarcity of large datasets
in many other domains, including 3D computer graphics,
design, engineering, and physics, restricts the use of ad-
vanced supervised learning techniques, necessitating the
exploration of alternative learning strategies.
Fortunately, these disciplines are often equipped with formal
*Equal contribution 1LIT AI Lab, Institute for Machine Learn-
ing, JKU Linz, Austria 2Part of this work was done while at SIN-
TEF and Department of Mathematics, University of Oslo, Nor-
way 3Emmi AI GmbH, Linz, Austria. Correspondence to: Arturs
Berzins <berzins@ml.jku.at>.
Proceedings of the 42 nd International Confere...
Conference: ICML 2025

Learning Dynamics under Environmental Constraints via Measurement-Induced
Bundle Structures
Dongzhe Zheng 1 Wenjie Mei 2
Abstract
Learning unknown dynamics under environmen-
tal (or external) constraints is fundamental to
many fields (e.g., modern robotics), particularly
challenging when constraint information is only
locally available and uncertain.
Existing ap-
proaches requiring global constraints or using
probabilistic filtering fail to fully exploit the geo-
metric structure inherent in local measurements
(by using, e.g., sensors) and constraints. This
paper presents a geometric framework unifying
measurements, constraints, and dynamics learn-
ing through a fiber bundle structure over the state
space. This naturally induced geometric struc-
ture enables measurement-aware Control Barrier
Functions that adapt to local sensing (or measure-
ment) conditions. By integrating Neural ODEs,
our framework learns continuous-time dynam-
ics while preserving geometric constraints, with
theoretical guarantees of learning convergence
and constraint satisfaction dependent on sensing
quality. The geometric framework not only en-
ables efficient dynamics learning but also sug-
gests promising directions for integration with re-
inforcement learning approaches. Extensive sim-
ulations demonstrate significant improvements in
both learning efficiency and constraint satisfac-
tion over traditional methods, especially under
limited and uncertain sensing conditions.
1. Introduction
Learning unknown dynamics under measurement con-
straints is fundamental to many applications, from manipu-
lators with force sensing to autonomous vehicles with range
1Department of Computer Science and Engineering, Shanghai
Jiao Tong University, Shanghai, China 2School of Automation
and Key Laboratory of MCCSE of Ministry of Education, South-
east University, Nanjing, China. Correspondence to: Wenjie Mei
<wenjie.mei@seu.edu.cn; vmus1130@gmail.com>.
Proceedings of the 42 nd International Conference on Mach...
Conference: NeurIPS 2025

Constrained Best Arm Identification
Tyron Lardy
CWI and Leiden University
Christina Katsimerou
Booking.com
christina.katsimerou@booking.com
Wouter M. Koolen
CWI and University of Twente
wmkoolen@cwi.nl
Abstract
In real-world decision-making problems, one needs to pick among multiple policies
the one that performs best while respecting economic constraints. This motivates
the problem of constrained best-arm identification for bandit problems where every
arm is a joint distribution of reward and cost. We investigate the general case where
reward and cost are dependent. The goal is to accurately identify the arm with the
highest mean reward among all arms whose mean cost is below a given threshold.
We prove information-theoretic lower bounds on the sample complexity for three
models: Gaussian with fixed covariance, Gaussian with unknown covariance, and
non-parametric distributions of rectangular support. We propose a combination of
a sampling and a stopping rule that correctly identifies the constrained best arm and
matches the optimal sample complexities for each of the three models. Simulations
demonstrate the performance of our algorithms.
1
Introduction
In real-world decision-making systems, identifying the best policy is rarely a matter of optimizing
a single performance metric in isolation. Effective policies often involve inherent trade-offs: they
achieve desirable outcomes but incur costs. For instance, online platforms routinely deploy pro-
motions to drive customer engagement [Zhao and Harinen, 2019, Zhang et al., 2024]. To ensure
sustainability, such incentives must not only be effective but also economically viable—typically
captured through feasibility constraints on budget or return on investment (ROI) [Goldenberg et al.,
2020]. Similarly, bidding policies in online advertising need to drive traffic or purchases at an
acceptable incremental ROI [Chen and Au, 2022]. In healthcare, a treatment needs to improve health
outcomes under safety constraints. In al...
Conference: ICML 2025

Hierarchical Planning for Complex Tasks with
Knowledge Graph-RAG and Symbolic Verification
Flavio Petruzzellis 1 Cristina Cornelio 2 Pietro Li`o 3
Abstract
Large Language Models (LLMs) have shown
promise as robotic planners but often struggle
with long-horizon and complex tasks, especially
in specialized environments requiring external
knowledge.
While hierarchical planning and
Retrieval-Augmented Generation (RAG) address
some of these challenges, they remain insufficient
on their own and a deeper integration is required
for achieving more reliable systems. To this end,
we propose a neuro-symbolic approach that en-
hances LLMs-based planners with Knowledge
Graph-based RAG for hierarchical plan genera-
tion. This method decomposes complex tasks
into manageable subtasks, further expanded into
executable atomic action sequences. To ensure
formal correctness and proper decomposition, we
integrate a Symbolic Validator, which also func-
tions as a failure detector by aligning expected
and observed world states. Our evaluation against
baseline methods demonstrates the consistent sig-
nificant advantages of integrating hierarchical
planning, symbolic verification, and RAG across
tasks of varying complexity and different LLMs.
Additionally, our experimental setup and novel
metrics not only validate our approach for com-
plex planning but also serve as a tool for assess-
ing LLMs’ reasoning and compositional capa-
bilities. Code available at https://github.
com/corneliocristina/HVR .
1. Introduction
Humans naturally tackle complex tasks, such as boiling
water, by first conceptualizing high-level steps (e.g., filling
1Department of Mathematics, University of Padova, Padova,
Italy 2Samsung AI, Cambridge, UK 3Computer Science Depart-
ment, University of Cambridge, Cambridge, UK. Correspondence
to: Cristina Cornelio <c.cornelio@samsung.com>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
a...
Conference: ICML 2025

Variational Learning of Fractional Posteriors
Kian Ming A. Chai 1 Edwin V. Bonilla 2
Abstract
We introduce a novel one-parameter variational
objective that lower bounds the data evidence and
enables the estimation of approximate fractional
posteriors. We extend this framework to hierar-
chical construction and Bayes posteriors, offering
a versatile tool for probabilistic modelling. We
demonstrate two cases where gradients can be ob-
tained analytically and a simulation study on mix-
ture models showing that our fractional posteriors
can be used to achieve better calibration compared
to posteriors from the conventional variational
bound. When applied to variational autoencoders
(VAEs), our approach attains higher evidence
bounds and enables learning of high-performing
approximate Bayes posteriors jointly with frac-
tional posteriors. We show that VAEs trained with
fractional posteriors produce decoders that are
better aligned for generation from the prior.
1. Introduction
Exact Bayesian inference is intractable for most models of
interest in machine learning. Variational methods (Jordan
et al., 1999; Minka, 2001; Opper & Winther, 2005; Blei
et al., 2017) address this by casting the required integration
as optimisation. These methods have two objectives: to
estimate the marginal likelihood or data evidence (MacKay,
2003) for model comparison or model optimisation; and to
obtain an approximate Bayes posterior for prediction.
The widely used evidence lower bound (ELBO) (Jordan
et al., 1999) often leads to underestimated uncertainty
and suboptimal posterior calibration (Wang & Tittering-
ton, 2005; Bishop, 2006; Yao et al., 2018). These deficien-
cies can be compounded by challenges inherent to general
Bayesian modelling, such as misspecification. Fractional
1DSO
National
Laboratories,
Singapore
2CSIRO’s
Data61,
Australia.
Correspondence to:
Kian Ming A.
Chai
<ckianmin@dso.org.sg>,
Edwin
V.
Bonilla
<ed-
win.bonilla@data61.csiro.au>.
Proceedings of the 42 nd International...
Conference: NeurIPS 2025

ROOT: Rethinking Offline Optimization as
Distributional Translation via Probabilistic Bridge
Manh Cuong Dao∗
National University of Singapore
cuongdao@nus.edu.sg
The Hung Tran∗
Washington State University
hung.t.tran@wsu.edu
Phi Le Nguyen
Hanoi University of Science and Technology
lenp@soict.hust.edu.vn
Thao Nguyen Truong
National Institute of Advanced Industrial Science and Technology
nguyen.truong@aist.go.jp
Trong Nghia Hoang†
Washington State University
trongnghia.hoang@wsu.edu
Abstract
This paper studies the black-box optimization task which aims to find the maxima
of a black-box function using a static set of its observed input-output pairs. This is
often achieved via learning and optimizing a surrogate function with that offline
data. Alternatively, it can also be framed as an inverse modeling task that maps a
desired performance to potential input candidates that achieve it. Both approaches
are constrained by the limited amount of offline data. To mitigate this limitation,
we introduce a new perspective that casts offline optimization as a distributional
translation task. This is formulated as learning a probabilistic bridge transforming
an implicit distribution of low-value inputs (i.e., offline data) into another distribu-
tion of high-value inputs (i.e., solution candidates). Such probabilistic bridge can
be learned using low- and high-value inputs sampled from synthetic functions that
resemble the target function. These synthetic functions are constructed as the mean
posterior of multiple Gaussian processes fitted with different parameterizations on
the offline data, alleviating the data bottleneck. The proposed approach is evalu-
ated on an extensive benchmark comprising most recent methods, demonstrating
significant improvement and establishing a new state-of-the-art performance. Our
code is publicly available at https://github.com/cuong-dm/ROOT.
1
Introduction
Black-box optimization arises in scientific and engineering domains where evaluating each can...
Conference: ICML 2025

When can in-context learning generalize out of task distribution?
Chase Goddard 1 Lindsay M. Smith 1 Vudtiwat Ngampruetikorn * 2 David J. Schwab * 3
Abstract
In-context learning (ICL) is a remarkable capabil-
ity of pretrained transformers that allows models
to generalize to unseen tasks after seeing only
a few examples. We investigate empirically the
conditions necessary on the pretraining distribu-
tion for ICL to emerge and generalize out-of-
distribution. Previous work has focused on the
number of distinct tasks necessary in the pretrain-
ing dataset. Here, we use a different notion of task
diversity to study the emergence of ICL in trans-
formers trained on linear functions. We find that
as task diversity increases, transformers undergo
a transition from a specialized solution, which
exhibits ICL only within the pretraining task dis-
tribution, to a solution which generalizes out of
distribution to the entire task space. We also in-
vestigate the nature of the solutions learned by
the transformer on both sides of the transition,
and observe similar transitions in nonlinear re-
gression problems. We construct a phase diagram
to characterize how our concept of task diversity
interacts with the number of pretraining tasks. In
addition, we explore how factors such as the depth
of the model and the dimensionality of the regres-
sion problem influence the transition.
1. Introduction
The ability of transformers (Vaswani et al., 2017) to do
few-shot learning from examples seen in their context is a
striking phenomenon exhibited by modern large language
models called in-context learning (ICL) (Brown et al., 2020).
ICL has been extensively studied (Raventos et al., 2023; Lu
et al., 2024; Garg et al., 2023; Chan et al., 2022; Singh
et al., 2023) and enables models to solve certain new tasks
*Equal contribution
1Joseph Henry Laboratories of Physics,
Princeton University, Princeton, NJ, USA 2School of Physics, Uni-
versity of Sydney, Sydney, Australia 3Initiative for the Th...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
MONTESSORI-INSTRUCT:
GENERATE INFLUENTIAL
TRAINING DATA TAILORED FOR STUDENT LEARNING
Xiaochuan Li♦∗, Zichun Yu✦, Chenyan Xiong✦
♦School of Software, Tsinghua University
✦Language Technologies Institute, Carnegie Mellon University
li-xc20@mails.tsinghua.edu.cn
zichunyu@andrew.cmu.edu
cx@cs.cmu.edu
ABSTRACT
Synthetic data has been widely used to train large language models, but their gener-
ative nature inevitably introduces noisy, non-informative, and misleading learning
signals. In this paper, we propose MONTESSORI-INSTRUCT, a novel data synthesis
framework that tailors the data synthesis ability of the teacher language model
toward the student language model’s learning process. Specifically, we utilize local
data influence of synthetic training data points on students to characterize students’
learning preferences. Then, we train the teacher model with Direct Preference
Optimization (DPO) to generate synthetic data tailored toward student learning pref-
erences. Experiments with Llama3-8B-Instruct (teacher) and Llama3-8B (student)
on Alpaca Eval and MT-Bench demonstrate that Montessori-Instruct significantly
outperforms standard synthesis methods by 18.35% and 46.24% relatively. Our
method also beats data synthesized by a stronger teacher model, GPT-4o. Further
analysis confirms the benefits of teacher’s learning to generate more influential train-
ing data in the student’s improved learning, the advantages of local data influence
in accurately measuring student preferences, and the robustness of Montessori-
Instruct across different student models. Our code and data are open-sourced at
https://github.com/cxcscmu/Montessori-Instruct.
1
INTRODUCTION
Synthetic training data is highly effective in various applications of large language models (LLMs) (Lu
et al., 2023), spanning from general pretraining (Allal et al., 2024; Zhou et al., 2024), instruction-
tuning (Tong et al., 2024) to domain-specific scenarios such as mathe...
Conference: NeurIPS 2025

Fine-grained List-wise Alignment for Generative
Medication Recommendation
Chenxiao Fan1 Chongming Gao1∗Wentao Shi1
Yaxin Gong1 Zihao Zhao1 Fuli Feng1∗
1University of Science and Technology of China
{simonfan, shiwentao123, gyx2022, zzh1998}@mail.ustc.edu.cn
chongminggao@ustc.edu.cn, fulifeng93@gmail.com
Abstract
Accurate and safe medication recommendations are critical for effective clinical
decision-making, especially in multimorbidity cases. However, existing systems
rely on point-wise prediction paradigms that overlook synergistic drug effects
and potential adverse drug-drug interactions (DDIs). We propose FLAME, a fine-
grained list-wise alignment framework for large language models (LLMs), enabling
drug-by-drug generation of drug lists. FLAME formulates recommendation as
a sequential decision process, where each step adds or removes a single drug.
To provide fine-grained learning signals, we devise step-wise Group Relative
Policy Optimization (GRPO) with potential-based reward shaping, which explicitly
models DDIs and optimizes the contribution of each drug to the overall prescription.
Furthermore, FLAME enhances patient modeling by integrating structured clinical
knowledge and collaborative information into the representation space of LLMs.
Experiments on benchmark datasets demonstrate that FLAME achieves state-of-
the-art performance, delivering superior accuracy, controllable safety–accuracy
trade-offs, and strong generalization across diverse clinical scenarios. Our code is
available at https://github.com/cxfann/Flame.
1
Introduction
Accurate and safe medication recommendation is essential for clinical decision-making, especially in
complex cases involving multimorbidity [1, 2]. In these scenarios, clinicians must consider not only
the therapeutic effects of individual drugs, but also their potential interactions and cumulative safety
risks. Recent advances in AI have led to the development of automated medication recommendation
systems [3, 4, 5], yet their...
Conference: ICLR 2025

Published as a conference paper at ICLR 2025
A GENERALIST HANABI AGENT
Arjun Vaithilingam Sudhakar∗1,2,3 Hadi Nekoei∗1,2,4 Mathieu Reymond1,2,3
Janarthanan Rajendran6 Miao Liu5 Sarath Chandar1,2,3,7
1Chandar Research Lab
2Mila – Quebec AI Institute
3Polytechnique Montr´eal
4Universit´e de Montr´eal
5IBM Research
6Dalhousie University
7Canada CIFAR AI Chair
ABSTRACT
Traditional multi-agent reinforcement learning (MARL) systems can develop co-
operative strategies through repeated interactions. However, these systems are
unable to perform well on any other setting than the one they have been trained
on, and struggle to successfully cooperate with unfamiliar collaborators. This is
particularly visible in the Hanabi benchmark, a popular 2-to-5 player coopera-
tive card-game which requires complex reasoning and precise assistance to other
agents. Current MARL agents for Hanabi can only learn one specific game-setting
(e.g., 2-player games), and play with the same algorithmic agents. This is in
stark contrast to humans, who can quickly adjust their strategies to work with
unfamiliar partners or situations. In this paper, we introduce Recurrent Replay
Relevance Distributed DQN (R3D2), a generalist agent for Hanabi, designed to
overcome these limitations. We reformulate the task using text, as language has
been shown to improve transfer. We then propose a distributed MARL algorithm
that copes with the resulting dynamic observation- and action-space. In doing
so, our agent is the first that can play all game settings concurrently, and extend
strategies learned from one setting to other ones. As a consequence, our agent
also demonstrates the ability to collaborate with different algorithmic agents —
agents that are themselves unable to do so. The implementation code is available
at: R3D2-A-Generalist-Hanabi-Agent
1
INTRODUCTION
Humans were able to thrive as a society through their ability to cooperate. Interactions among multi-
ple people or agents are essential components of...
Conference: NeurIPS 2025

Finite-Time Bounds for
Average-Reward Fitted Q-Iteration
Jongmin Lee
Seoul National University
Department of Mathematical Sciences
dlwhd2000@snu.ac.kr
Ernest K. Ryu
UCLA
Department of Mathematics
eryu@math.ucla.edu
Abstract
Although there is an extensive body of work characterizing the sample complexity
of discounted-return offline RL with function approximations, prior work on the
average-reward setting has received significantly less attention, and existing ap-
proaches rely on restrictive assumptions, such as ergodicity or linearity of the MDP.
In this work, we establish the first sample complexity results for average-reward
offline RL with function approximation for weakly communicating MDPs, a much
milder assumption. To this end, we introduce Anchored Fitted Q-Iteration, which
combines the standard Fitted Q-Iteration with an anchor mechanism. We show
that the anchor, which can be interpreted as a form of weight decay, is crucial
for enabling finite-time analysis in the average-reward setting. We also extend
our finite-time analysis to the setup where the dataset is generated from a single-
trajectory rather than IID transitions, again leveraging the anchor mechanism.
1
Introduction
The goal of offline Reinforcement Learning (RL) is to find a near-optimal policy using a precollected
dataset without any direct interaction with the environment. Characterizing the sample complexity
for finding an ϵ-optimal policy using function approximation under assumptions that the offline
data has sufficient coverage over the whole state-action space has been an active area of theoretical
RL research. However, more prior work focuses on the discounted cumulative reward setup, and
research on obtaining sample complexity in the average reward has been limited due to the absence
of the discount factor and the complexity of the Bellman equation. Specifically, all prior works with
function approximation rely on restrictive assumptions such as ergodicity or linearity of the MDP.
Alth...
Conference: ICML 2025

Continuous Semi-Implicit Models
Longlin Yu∗1 Jiajun Zha∗2 Tong Yang 3 4 Tianyu Xie 1 Xiangyu Zhang 4 S.-H. Gary Chan 2 Cheng Zhang 1 5
Abstract
Semi-implicit distributions have shown great
promise in variational inference and generative
modeling.
Hierarchical semi-implicit models,
which stack multiple semi-implicit layers, en-
hance the expressiveness of semi-implicit distri-
butions and can be used to accelerate diffusion
models given pretrained score networks. How-
ever, their sequential training often suffers from
slow convergence. In this paper, we introduce
CoSIM, a continuous semi-implicit model that
extends hierarchical semi-implicit models into a
continuous framework. By incorporating a contin-
uous transition kernel, CoSIM enables efficient,
simulation-free training. Furthermore, we show
that CoSIM achieves consistency with a care-
fully designed transition kernel, offering a novel
approach for multistep distillation of generative
models at the distributional level. Extensive ex-
periments on image generation demonstrate that
CoSIM performs on par or better than existing
diffusion model acceleration methods, achieving
superior performance on FD-DINOv2.
1. Introduction
Semi-implicit distributions, which are constructed through
the convolution of explicit conditional distributions and im-
plicit mixing distributions, have gained significant traction
in variational inference and generative modeling. Unlike
traditional approximating distributions with explicit density
forms, semi-implicit distributions enable a more expressive
family of variational posteriors, leading to improved approx-
*Equal contribution. This work was done during an internship
at StepFun.
1School of Mathematical Sciences, Peking University,
Beijing, China 2Department of Computer Science and Engineering,
Hong Kong University of Science and Technology 3School of
Computer Science, Fudan University, Shanghai, China 4Megvii
Technology Inc., Beijing, China 5Center for Statistical Science,
Peking ...