AIS-Bench/AIS-Bench huggingface dataset:
Published as a conference paper at ICLR 2025
THE BELIEF STATE TRANSFORMER
Edward S. Hu1,2
Kwangjun Ahn1
Qinghua Liu1
Haoran Xu1,3
Manan Tomar1,4
Ada Langford1
Dinesh Jayaraman2
Alex Lamb1
John Langford1
1Microsoft Research
2University of Pennsylvania
3UT Austin
4 University of Alberta
ABSTRACT
We introduce the “Belief State Transformer”, a next-token predictor that takes both
a prefix and suffix as inputs, with a novel objective of predicting both the next token
for the prefix and the previous token for the suffix. The Belief State Transformer
effectively learns to solve challenging problems that conventional forward-only
transformers struggle with, in a domain-independent fashion. Key to this success is
learning a compact belief state that captures all relevant information necessary for
accurate predictions. Empirical ablations show that each component of the model is
essential in difficult scenarios where standard Transformers fall short. For the task
of story writing with known prefixes and suffixes, our approach outperforms the
Fill-in-the-Middle method for reaching known goals and demonstrates improved
performance even when the goals are unknown. Altogether, the Belief State
Transformer enables more efficient goal-conditioned decoding, better test-time
inference, and high-quality text representations on small scale problems. Website:
https://edwhu.github.io/bst-website
1
INTRODUCTION
Transformer models (Vaswani et al., 2017) have created a revolution in language modeling (Achiam
et al., 2023) with the capability to generate language with many emergent properties at large scale.
Examining these models for flaws in the pursuit of further progress, it’s notable that they struggle
with planning-heavy problems (Bubeck et al., 2023; Momennejad et al., 2024). How can we modify
the architecture, objectives, and algorithms to create a model more capable of reaching goals?
To make progress, we propose the new Belief State Transformer (BST) architecture and objective
in S...
Diffuse Everything: Multimodal Diffusion Models on Arbitrary State Spaces
Kevin Rojas * 1 2 Yuchen Zhu * 1 2 Sichen Zhu 2 Felix X.-F. Ye 3 Molei Tao 1 2
Abstract
Diffusion models have demonstrated remarkable
performance in generating unimodal data across
various tasks, including image, video, and text
generation. On the contrary, the joint generation
of multimodal data through diffusion models is
still in the early stages of exploration. Existing
approaches heavily rely on external preprocess-
ing protocols, such as tokenizers and variational
autoencoders, to harmonize varied data represen-
tations into a unified, unimodal format.
This
process heavily demands the high accuracy of
encoders and decoders, which can be problem-
atic for applications with limited data. To lift
this restriction, we propose a novel framework
for building multimodal diffusion models on arbi-
trary state spaces, enabling native generation of
coupled data across different modalities. By in-
troducing an innovative decoupled noise schedule
for each modality, we enable both unconditional
and modality-conditioned generation within a sin-
gle model simultaneously. We empirically vali-
date our approach for text-image generation and
mixed-type tabular data synthesis, demonstrating
that it achieves competitive performance. Code is
available at Diffuse-Everything.
1. Introduction
Recent years have witnessed the tremendous success of
diffusion generative models in various applications. The
seminal works of continuous diffusion models on Euclidean
spaces (Sohl-Dickstein et al., 2015; Song et al., 2020; Ho
et al., 2020) have led to state-of-the-art methods for tasks
such as image generation (Dhariwal & Nichol, 2021; Bao
et al., 2023a; Karras et al., 2022; 2024b), video generation
*Equal contribution, random order by coin flip 1School of Math-
ematics, Georgia Institute of Technology, Atlanta, GA 2Machine
Learning Center, Georgia Institute of Technology, Atlanta, GA
3Department of Mathematics & Statisti...
Published as a conference paper at ICLR 2025
TPO: ALIGNING LARGE LANGUAGE MODELS WITH
MULTI-BRANCH & MULTI-STEP PREFERENCE TREES
Weibin Liao♠, Xu Chu♠♢♡∗, Yasha Wang♠♡
♠School of Computer Science, Peking University
♢Center on Frontiers of Computing Studies, Peking University
♡National Research and Engineering Center of Software Engineering, Peking University
liaoweibin@stu.pku.edu.cn, chu xu@pku.edu.cn
https://github.com/MrBlankness/TPO.git
ABSTRACT
In the domain of complex reasoning tasks, such as mathematical reasoning, recent
advancements have proposed the use of Direct Preference Optimization (DPO)
to suppress output of dispreferred responses, thereby enhancing the long-chain
reasoning capabilities of large language models (LLMs). To this end, these stud-
ies employed LLMs to generate preference trees via Tree-of-thoughts (ToT) and
sample the paired preference responses required by the DPO algorithm. However,
the DPO algorithm based on binary preference optimization was unable to learn
multiple responses with varying degrees of preference/dispreference that provided
by the preference trees, resulting in incomplete preference learning. In this work,
we introduce Tree Preference Optimization (TPO), which does not sample paired
preference responses from the preference tree; instead, it directly learns from the
entire preference tree during the fine-tuning. Specifically, TPO formulates the lan-
guage model alignment as a Preference List Ranking problem, where the policy
can potentially learn more effectively from a ranked preference list of responses
given the prompt. In addition, to further assist LLMs in identifying discriminative
steps within long-chain reasoning and increase the relative reward margin in the
preference list, TPO utilizes Adaptive Step Reward to adjust the reward values of
each step in the trajectory for performing fine-grained preference optimization.
We carry out extensive experiments on mathematical reasoning tasks to evaluate
TPO. The experim...
∆Flow: An Efficient Multi-frame Scene Flow
Estimation Method
Qingwen Zhang1,B
Xiaomeng Zhu1,3
Yushan Zhang2,B
Yixi Cai1
Olov Andersson1
Patric Jensfelt1
1KTH Royal Institute of Technology
2Linköping University
3Scania CV AB
qingwen@kth.se, yushan.zhang@liu.se
Abstract
Previous dominant methods for scene flow estimation focus mainly on input from
two consecutive frames, neglecting valuable information in the temporal domain.
While recent trends shift towards multi-frame reasoning, they suffer from rapidly
escalating computational costs as the number of frames grows. To leverage tem-
poral information more efficiently, we propose DeltaFlow (∆Flow), a lightweight
3D framework that captures motion cues via a ∆scheme, extracting temporal
features with minimal computational cost, regardless of the number of frames.
Additionally, scene flow estimation faces challenges such as imbalanced object
class distributions and motion inconsistency. To tackle these issues, we introduce a
Category-Balanced Loss to enhance learning across underrepresented classes and
an Instance Consistency Loss to enforce coherent object motion, improving flow
accuracy. Extensive evaluations on the Argoverse 2, Waymo and nuScenes datasets
show that ∆Flow achieves state-of-the-art performance with up to 22% lower error
and 2× faster inference compared to the next-best multi-frame supervised method,
while also demonstrating a strong cross-domain generalization ability. The code
is open-sourced at https://github.com/Kin-Zhang/DeltaFlow along
with trained model weights.
1
Introduction
Scene flow estimation determines the 3D motion of each point between consecutive point clouds,
making it an important task in computer vision and essential for autonomous driving [31, 3, 43, 41]
and motion compensation [8, 46]. While early approaches focused on per-point feature learning [49,
47, 25, 26, 48, 40] and achieve high accuracy on small-scale datasets, they become computationally
expensive when processing large-sca...
Cross-environment Cooperation Enables Zero-shot Multi-agent Coordination
Kunal Jha 1 Wilka Carvalho 2 Yancheng Liang 1 Simon S. Du 1 Max Kleiman-Weiner * 1 Natasha Jaques * 1
Abstract
Zero-shot coordination (ZSC), the ability to adapt
to a new partner in a cooperative task, is a criti-
cal component of human-compatible AI. While
prior work has focused on training agents to co-
operate on a single task, these specialized mod-
els do not generalize to new tasks, even if they
are highly similar.
Here, we study how rein-
forcement learning on a distribution of environ-
ments with a single partner enables learning
general cooperative skills that support ZSC with
many new partners on many new problems.
We introduce two Jax-based, procedural genera-
tors that create billions of solvable coordination
challenges. We develop a new paradigm called
Cross-Environment Cooperation (CEC), and
show that it outperforms competitive baselines
quantitatively and qualitatively when collaborat-
ing with real people. Our findings suggest that
learning to collaborate across many unique scenar-
ios encourages agents to develop general norms,
which prove effective for collaboration with dif-
ferent partners. Together, our results suggest a
new route toward designing generalist coopera-
tive agents capable of interacting with humans
without requiring human data. Code for environ-
ment, training, and testing scripts and more can
be found at https://kjha02.github.io/
publication/cross-env-coop.
1. Introduction
Humans excel at ad-hoc cooperation, readily adapting to
new partners and environments by attending jointly to rel-
evant objects, reasoning about shared intentions, and play-
ing their role within an implicit collective plan (Tomasello,
*Equal contribution
1Department of Computer Science, Uni-
versity of Washington, Seattle, WA 2Kempner Institute for the
Study of Natural and Artificial Intelligence, Harvard Univer-
sity, Cambridge, Massachusetts. Correspondence to: Kunal Jha
<kjha@uw.edu>....
CLAWS:Creativity detection for LLM-generated
solutions using Attention Window of Sections
Keuntae Kim1∗
Eunhye Jeong2∗
Sehyeon Lee3
Seohee Yoon1
Yong Suk Choi1†
1Department of Computer Science
2Department of Artificial Intelligence
3Department of Future Mobility
Hanyang University, Seoul, Korea
{ktkpv94, jeh0826, leesehyeon, seohee09}@hanyang.ac.kr
Abstract
Recent advances in enhancing the reasoning ability of Large Language Models
(LLMs) have been remarkably successful. LLMs trained with Reinforcement
Learning (RL) for reasoning demonstrate strong performance in challenging tasks
such as mathematics and coding, even with relatively small model sizes. How-
ever, despite these impressive improvements in task accuracy, the assessment of
creativity in LLM generations has been largely overlooked in reasoning tasks, in
contrast to writing tasks. The lack of research on creativity assessment in rea-
soning primarily stems from two challenges: (1) the difficulty of defining the
range of creativity, and (2) the necessity of human evaluation in the assessment
process. To address these challenges, we propose CLAWS, a novel method that
defines and classifies mathematical solutions into Typical, Creative, and Halluci-
nated categories without human evaluation, by leveraging attention weights across
prompt sections and output. CLAWS outperforms five existing white-box detec-
tion methods—Perplexity, Logit Entropy, Window Entropy, Hidden Score, and
Attention Score—on five 7–8B math RL models (DeepSeek, Qwen, Mathstral,
OpenMath2, and Oreal). We validate CLAWS on 4,545 math problems collected
from 181 math contests (A(J)HSME, AMC, AIME). Our code is available at
https://github.com/kkt94/CLAWS.
1
Introduction
In recent years, Large Language Models (LLMs) have achieved remarkable success across a wide
range of tasks. Among these, the most notable progress has been made in reasoning ability, particularly
in mathematical problem solving. Solving math problems requires cognitive proces...
Deep Neural Cellular Potts Models
Koen Minartz 1 Tim d’Hondt 1 Leon Hillmann 2 1 3 J¨orn Starruß 4 Lutz Brusch 4 Vlado Menkovski 1
Abstract
The cellular Potts model (CPM) is a power-
ful computational method for simulating collec-
tive spatiotemporal dynamics of biological cells.
To drive the dynamics, CPMs rely on physics-
inspired Hamiltonians. However, as first prin-
ciples remain elusive in biology, these Hamilto-
nians only approximate the full complexity of
real multicellular systems. To address this limi-
tation, we propose NeuralCPM, a more expres-
sive cellular Potts model that can be trained di-
rectly on observational data. At the core of Neu-
ralCPM lies the Neural Hamiltonian, a neural
network architecture that respects universal sym-
metries in collective cellular dynamics. More-
over, this approach enables seamless integration
of domain knowledge by combining known bi-
ological mechanisms and the expressive Neural
Hamiltonian into a hybrid model. Our evalua-
tion with synthetic and real-world multicellular
systems demonstrates that NeuralCPM is able to
model cellular dynamics that cannot be accounted
for by traditional analytical Hamiltonians.
1. Introduction
Cell migration and multicellular self-organization are cru-
cial biological processes that drive many phenomena of life,
such as embryo growth and the spread of cancer (Friedl
& Gilmour, 2009; Gottheil et al., 2023). Understanding
these cellular dynamics is not only a fundamental goal of
1Department of Mathematics and Computer Science, Eind-
hoven University of Technology, Eindhoven, the Netherlands
2Department of Applied Physics and Science Education, Eind-
hoven University of Technology, Eindhoven, the Netherlands
3Department of Biomedical Engineering, Eindhoven University of
Technology, Eindhoven, the Netherlands 4Center for Information
Services and High Performance Computing, TUD Dresden Univer-
sity of Technology, Dresden, Germany. Correspondence to: Koen
Minartz <k.minartz@tue.nl>, Lutz Bru...
Can NeRFs “See” without Cameras?
Chaitanya Amballa1
Sattwik Basu1∗
Yu-Lin Wei1∗
amballa2@illinois.edu
sattwik2@illinois.edu
yulinlw2@illinois.edu
Zhijian Yang2
Mehmet Ergezer2
Romit Roy Choudhury1,2
yzhijian@amazon.com
mergezer@amazon.com
croy@illinois.edu
1University of Illinois Urbana–Champaign
2Amazon
Abstract
Neural Radiance Fields (NeRFs) have been remarkably successful at synthesizing
novel views of 3D scenes by optimizing a volumetric scene function. This scene
function models how optical rays bring color information from a 3D object to the
camera pixels. Radio frequency (RF) or audio signals can also be viewed as a
vehicle for delivering information about the environment to a sensor. However,
unlike camera pixels, an RF/audio sensor receives a mixture of signals that contain
many environmental reflections (also called “multipath”). Is it still possible to infer
the environment using such multipath signals? We show that with redesign, NeRFs
can be taught to learn from multipath signals, and thereby “see” the environment.
As a grounding application, we aim to infer the indoor floorplan of a home from
sparse WiFi measurements made at multiple locations inside the home. Although
a difficult inverse problem, our implicitly learnt floorplans look promising, and
enables forward applications, such as indoor signal prediction and basic ray tracing.
1
Introduction
NeRFs [25, 43, 2, 45] have delivered impressive results in solving inverse problems, resulting in 3D
scene rendering. While NeRFs have mostly used pictures (from cameras or LIDARs) to infer a 3D
scene, we ask if the core ideas can generalize to the case of wireless signals (such as RF or audio).
Unlike camera pixels that receive line-of-sight (LoS) rays, a wireless receiver (e.g., a WiFi antenna
on a smartphone) would receive a mixture of LoS and many reflections, called multipath. If the
receiver moves, it receives a sequence of N measurements. Using these N wireless measurements, is
it possible to learn a ...
Skip the Equations: Learning Behavior of Personalized Dynamical Systems
Directly From Data
Krzysztof Kacprzyk 1 Julianna Piskorz 1 Mihaela van der Schaar 1
Abstract
While black-box approaches are commonly used
for data-driven modeling of dynamical systems,
they often obscure a system’s underlying behav-
ior and properties, limiting adoption in areas such
as medicine and pharmacology. A two-step pro-
cess of discovering ordinary differential equations
(ODEs) and their subsequent mathematical anal-
ysis can yield insights into the system’s dynam-
ics. However, this analysis may be infeasible for
complex equations, and refining the ODE to meet
certain behavioral requirements can be challeng-
ing. Direct semantic modeling has recently been
proposed to address these issues by predicting the
system’s behavior, such as the trajectory’s shape,
directly from data, bypassing post-hoc mathe-
matical analysis. In this work, we extend the
original instantiation, limited to one-dimensional
trajectories and inputs, to accommodate multi-
dimensional trajectories with additional personal-
ization, allowing evolution to depend on auxiliary
static features (e.g., patient covariates). In a se-
ries of experiments, we show how our approach
enables practitioners to integrate prior knowledge,
understand the dynamics, ensure desired behav-
iors, and revise the model when necessary.
1. Introduction
Personalized Dynamical Systems
In machine learning
(ML), modeling dynamical systems is a core activity with
extensive applications in fields such as physics (Rudy et al.,
2017), biology (Chen et al., 2019), engineering (Brunton
& Kutz, 2022), and medicine (Lee et al., 2020). We call
a dynamical system personalized when its evolution func-
tion depends on certain observable features. It is similar
to dynamical systems that depend on a parameter (as stud-
1University of Cambridge, Cambridge, UK. Correspondence
to: Krzysztof Kacprzyk <kk751@cam.ac.uk>.
Proceedings of the 42 nd International Conferen...
When Will It Fail?: Anomaly to Prompt
for Forecasting Future Anomalies in Time Series
Min-Yeong Park * 1 Won-Jeong Lee * 2 Seong Tae Kim 1 Gyeong-Moon Park 2
Abstract
Recently, forecasting future abnormal events has
emerged as an important scenario to tackle real-
world necessities. However, the solution of pre-
dicting specific future time points when anoma-
lies will occur, known as Anomaly Prediction
(AP), remains under-explored. Existing meth-
ods dealing with time series data fail in AP, fo-
cusing only on immediate anomalies or failing
to provide precise predictions for future anoma-
lies. To address the AP task, we propose a novel
framework called Anomaly to Prompt (A2P),
comprised of Anomaly-Aware Forecasting (AAF)
and Synthetic Anomaly Prompting (SAP). To en-
able the forecasting model to forecast abnormal
time points, we adopt a strategy to learn the re-
lationships of anomalies. For the robust detec-
tion of anomalies, our proposed SAP introduces
a learnable Anomaly Prompt Pool (APP) that
simulates diverse anomaly patterns using signal-
adaptive prompt. Comprehensive experiments on
multiple real-world datasets demonstrate the su-
periority of A2P over state-of-the-art methods,
showcasing its ability to predict future anoma-
lies. Our implementation code is available at
https://github.com/KU-VGI/AP.
1. Introduction
As deep learning methods have evolved rapidly, their appli-
cation to time series analysis has gained significant attention
due to their critical importance in real-world scenarios. As
part of this, recently, forecasting future abnormal events
has been newly proposed in time series analysis, such as in
*Equal contribution
1Department of Artificial Intelligence,
Kyung Hee University, Yongin, Republic of Korea 2Department
of Artificial Intelligence, Korea University, Seoul, Republic of
Korea. Correspondence to: Seong Tae Kim <st.kim@khu.ac.kr>,
Gyeong-Moon Park <gm-park@korea.ac.kr>.
Proceedings of the 42 nd International Conference on Machine
Lea...
Published as a conference paper at ICLR 2025
SCALING STICK-BREAKING ATTENTION:
AN EFFICIENT IMPLEMENTATION
AND IN-DEPTH STUDY
Shawn Tan
MIT-IBM Watson AI Lab
shawntan@ibm.com
Songlin Yang
MIT
yangsl66@mit.edu
Aaron Courville
Mila, Université de Montréal
courvila@mila.quebec
Rameswar Panda
MIT-IBM Watson AI Lab
rpanda@ibm.com
Yikang Shen
MIT-IBM Watson AI Lab
yikang.shen@ibm.com
ABSTRACT
The self-attention mechanism traditionally relies on the softmax operator, neces-
sitating positional embeddings like RoPE, or position biases to account for token
order. But current methods using still face length generalisation challenges. We
investigate an alternative attention mechanism based on the stick-breaking process
in larger scale settings. The method works as follows: For each token before the
current, we determine a break point, which represents the proportion of the stick,
the weight of the attention, to allocate to the current token. We repeat this on the
remaining stick, until all tokens are allocated a weight, resulting in a sequence
of attention weights. This process naturally incorporates recency bias, which has
linguistic motivations for grammar parsing (Shen et al., 2017). We study the im-
plications of replacing the conventional softmax-based attention mechanism with
stick-breaking attention. We then discuss implementation of numerically stable
stick-breaking attention and adapt Flash Attention to accommodate this mechanism.
When used as a drop-in replacement for current softmax+RoPE attention systems,
we find that stick-breaking attention performs competitively with current methods
on length generalisation and downstream tasks. Stick-breaking also performs well
at length generalisation, allowing a model trained with 211 context window to
perform well at 214 with perplexity improvements.
 https://github.com/shawntan/stickbreaking-attention
1
INTRODUCTION
The Transformer architecture (Vaswani et al., 2017) uses a self-attention mechanism based on the
softmax ope...
High-Fidelity Simultaneous Speech-To-Speech Translation
Tom Labiausse 1 Laurent Mazar´e 1 Edouard Grave 1 Alexandre D´efossez 1 Neil Zeghidour 1
Abstract
We introduce Hibiki, a decoder-only model for
simultaneous speech translation. Hibiki leverages
a multistream language model to synchronously
process source and target speech, and jointly
produces text and audio tokens to perform speech-
to-text and speech-to-speech translation.
We
furthermore address the fundamental challenge
of simultaneous interpretation, which unlike its
consecutive counterpart—where one waits for the
end of the source utterance to start translating—
adapts its flow to accumulate just enough context
to produce a correct translation in real-time,
chunk by chunk. To do so, we introduce a weakly-
supervised method that leverages the perplexity
of an off-the-shelf text translation system to
identify optimal delays on a per-word basis and
create aligned synthetic data. After supervised
training, Hibiki performs adaptive, simultaneous
speech translation with vanilla temperature
sampling.
On a French-English simultaneous
speech translation task, Hibiki demonstrates
state-of-the-art performance in translation quality,
speaker fidelity and naturalness.
Moreover,
the simplicity of its inference process makes it
compatible with batched translation and even
real-time on-device deployment.
We provide
examples1 as well as models and inference code.2
1. Introduction
We introduce Hibiki (“echo” in Japanese), a system for
streaming and expressive speech-to-speech (S2ST) and
speech-to-text (S2TT) translation. Most work in speech
translation has focused on the offline (or consecutive)
setting where the model has access to the full source
utterance before it translates, as it provides useful context
1Kyutai, Paris, France. Correspondence to:
Hibiki <hibiki@kyutai.org>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1https:/...