dataset-review/dataset-review huggingface dataset:
Hierarchical Planning for Complex Tasks with Knowledge Graph-RAG and Symbolic Verification
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...
Variational Learning of Fractional Posteriors
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...
ROOT: Rethinking Offline Optimization as Distributional Translation via Probabilistic Bridge
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...
When can in-context learning generalize out of task distribution?
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...
Montessori-Instruct: Generate Influential Training Data Tailored for Student Learning
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...
Fine-grained List-wise Alignment for Generative Medication Recommendation
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...
A Generalist Hanabi Agent
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...
Finite-Time Bounds for Average-Reward Fitted Q-Iteration
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...
Continuous Semi-Implicit Models
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 ...
Causal Effect Identification in lvLiNGAM from Higher-Order Cumulants
ICML 2025

Causal Effect Identification in lvLiNGAM from Higher-Order Cumulants
Daniele Tramontano 1 Yaroslav Kivva 2 Saber Salehkaleybar 3 Negar Kiyavash 2 Mathias Drton 1 4
Abstract
This paper investigates causal effect identification
in latent variable Linear Non-Gaussian Acyclic
Models (lvLiNGAM) using higher-order cumu-
lants, addressing two prominent setups that are
challenging in the presence of latent confound-
ing: (1) a single proxy variable that may causally
influence the treatment and (2) underspecified in-
strumental variable cases where fewer instruments
exist than treatments. We prove that causal effects
are identifiable with a single proxy or instrument
and provide corresponding estimation methods.
Experimental results demonstrate the accuracy
and robustness of our approaches compared to
existing methods, advancing the theoretical and
practical understanding of causal inference in lin-
ear systems with latent confounders.
1. Introduction
Predicting the impact of an unseen intervention in a sys-
tem is a crucial challenge in many fields, such as medicine
(Sanchez et al., 2022; Michoel & Zhang, 2023), policy evalu-
ation (Athey & Imbens, 2017), fair decision-making (Kilber-
tus et al., 2017), and finance (de Prado, 2023). Randomized
experiments/interventional studies are the gold standard for
addressing this challenge but are often infeasible due to a
variety of reasons, such as ethical concerns or prohibitively
high costs. Thus, when merely observational data is avail-
able, additional assumptions on the underlying causal sys-
tem are needed to compensate for the lack of interventional
data. The field of causal inference seeks to formalize such
assumptions. One notable approach in causal inference is
modeling causal relationships through structural causal mod-
els (SCM) (Pearl, 2009). In this framework, a random vector
1Technical University of Munich, Munich, Germany 2Ecole
Polytechnique F´ed´erale de Lausanne, Lausanne, Switzerland
3Leiden Institute of Advanced...
BARNN: A Bayesian Autoregressive and Recurrent Neural Network
ICML 2025

BARNN: A Bayesian Autoregressive and Recurrent Neural Network
Dario Coscia 1 2 Max Welling 2 Nicola Demo 3 Gianluigi Rozza 1
Abstract
Autoregressive and recurrent networks have
achieved remarkable progress across various
fields, from weather forecasting to molecular gen-
eration and Large Language Models. Despite
their strong predictive capabilities, these mod-
els lack a rigorous framework for addressing un-
certainty, which is key in scientific applications
such as PDE solving, molecular generation and
Machine Learning Force Fields. To address this
shortcoming we present BARNN: a variational
Bayesian Autoregressive and Recurrent Neural
Network. BARNNs aim to provide a principled
way to turn any autoregressive or recurrent model
into its Bayesian version. BARNN is based on
the variational dropout method, allowing to apply
it to large recurrent neural networks as well. We
also introduce a temporal version of the “Vari-
ational Mixtures of Posteriors” prior (tVAMP-
prior) to make Bayesian inference efficient and
well-calibrated. Extensive experiments on PDE
modelling and molecular generation demonstrate
that BARNN not only achieves comparable or
superior accuracy compared to existing methods,
but also excels in uncertainty quantification and
modelling long-range dependencies.
1. Introduction
Autoregressive and recurrent models have demonstrated
impressive advancements in different disciplines, from
weather prediction (Bodnar et al., 2024; Lam et al., 2022),
to molecules generation (Shi et al., 2020; Simm et al., 2020)
and Large Language Models (LLMs) (Bengio et al., 2000;
Vaswani, 2017; Radford et al., 2019). However, while au-
toregressive models show strong predictive abilities, they are
also prone to overfitting to the specific tasks they are trained
1Mathematics Area, International School of Advanced Studies,
Italy 2Informatics Institute, University of Amsterdam, The Nether-
lands 3FAST Computing Srl, Italy. Correspondence to: Dario
Coscia <dario.coscia@sissa.it>...