dataset-review/dataset-review huggingface dataset:
How to build a consistency model: Learning flow maps via self-distillation
NeurIPS 2025

How to build a consistency model:
Learning flow maps via self-distillation
Nicholas M. Boffi
Carnegie Mellon University
Michael S. Albergo
Harvard University
Eric Vanden-Eijnden
Courant Institute of Mathematical Sciences
Abstract
Flow-based generative models achieve state-of-the-art sample quality, but require
the expensive solution of a differential equation at inference time. Flow map mod-
els, commonly known as consistency models, encompass many recent efforts to
improve inference-time efficiency by learning the solution operator of this differ-
ential equation. Yet despite their promise, these models lack a unified description
that clearly explains how to learn them efficiently in practice. Here, building on
the methodology proposed in Boffi et al. (2024), we present a systematic algo-
rithmic framework for directly learning the flow map associated with a flow or
diffusion model. By exploiting a relationship between the velocity field underlying
a continuous-time flow and the instantaneous rate of change of the flow map, we
show how to convert any distillation scheme into a direct training algorithm via
self-distillation, eliminating the need for pre-trained teachers. We introduce three
algorithmic families based on different mathematical characterizations of the flow
map: Eulerian, Lagrangian, and Progressive methods, which we show encompass
and extend all known distillation and direct training schemes for consistency mod-
els. We find that the novel class of Lagrangian methods, which avoid both spatial
derivatives and bootstrapping from small steps by design, achieve significantly
more stable training and higher performance than more standard Eulerian and Pro-
gressive schemes. Our methodology unifies existing training schemes under a single
common framework and reveals new design principles for accelerated generative
modeling. Associated code is available at https://github.com/nmboffi/flow-maps.
1
Introduction
Generative models based on dynamical systems, suc...
Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
NeurIPS 2025

Consistent Sampling and Simulation: Molecular
Dynamics with Energy-Based Diffusion Models
Michael Plainer† 1,2,3,4
Hao Wu† 5
Leon Klein 1
Stephan Günnemann 6
Frank Noé† 1,7,8
1Freie Universität Berlin
2Zuse School ELIZA
3Technische Universität Berlin
4Berlin Institute for the Foundations of Learning and Data 5Shanghai Jiao Tong University
6Technische Universität München
7Rice University
8Microsoft Research AI4Science
Abstract
In recent years, diffusion models trained on equilibrium molecular distributions
have proven effective for sampling biomolecules. Beyond direct sampling, the
score of such a model can also be used to derive the forces that act on molecular
systems. However, while classical diffusion sampling usually recovers the training
distribution, the corresponding energy-based interpretation of the learned score is
often inconsistent with this distribution, even for low-dimensional toy systems. We
trace this inconsistency to inaccuracies of the learned score at very small diffusion
timesteps, where the model must capture the correct evolution of the data distribu-
tion. In this regime, diffusion models fail to satisfy the Fokker–Planck equation,
which governs the evolution of the score. We interpret this deviation as one source
of the observed inconsistencies and propose an energy-based diffusion model with a
Fokker–Planck-derived regularization term to enforce consistency. We demonstrate
our approach by sampling and simulating multiple biomolecular systems, including
fast-folding proteins, and by introducing a state-of-the-art transferable Boltzmann
emulator for dipeptides that supports simulation and achieves improved consistency
and efficient sampling. Our code, model weights, and self-contained JAX and
PyTorch notebooks are available at https://github.com/noegroup/ScoreMD.
1
Introduction
Understanding biochemical systems requires modeling not only static molecular structures but also
their temporal evolution and interactions. Molecular dynamics (MD) si...
DRAG: Data Reconstruction Attack using Guided Diffusion
ICML 2025

DRAG: Data Reconstruction Attack using Guided Diffusion
Wa-Kin Lei 1 Jun-Cheng Chen 2 Shang-Tse Chen 1
Abstract
With the rise of large foundation models, split
inference (SI) has emerged as a popular compu-
tational paradigm for deploying models across
lightweight edge devices and cloud servers, ad-
dressing data privacy and computational cost con-
cerns. However, most existing data reconstruction
attacks have focused on smaller CNN classifica-
tion models, leaving the privacy risks of founda-
tion models in SI settings largely unexplored. To
address this gap, we propose a novel data recon-
struction attack based on guided diffusion, which
leverages the rich prior knowledge embedded in
a latent diffusion model (LDM) pre-trained on a
large-scale dataset. Our method performs iterative
reconstruction on the LDM’s learned image prior,
effectively generating high-fidelity images resem-
bling the original data from their intermediate rep-
resentations (IR). Extensive experiments demon-
strate that our approach significantly outperforms
state-of-the-art methods, both qualitatively and
quantitatively, in reconstructing data from deep-
layer IRs of the vision foundation model. The
results highlight the urgent need for more robust
privacy protection mechanisms for large models
in SI scenarios. Code is available at: https:
//github.com/ntuaislab/DRAG
1. Introduction
The rapid development of deep learning has revolution-
ized various aspects of daily life—from AI assistants to
autonomous vehicles. However, the substantial computa-
tional resources required by these emerging models often
hinder their deployment on edge devices. Therefore, offload-
ing intensive computation to cloud servers has become a
popular alternative. Following this paradigm, split inference
(SI) (Kang et al., 2017) has emerged as one of the most
1National Taiwan University 2Research Center for Information
Technology Innovation, Academia Sinica. Correspondence to:
Shang-Tse Chen <stchen@csie.ntu.edu.tw>.
Pr...
Expected Variational Inequalities
ICML 2025

Expected Variational Inequalities
Brian Hu Zhang * 1 Ioannis Anagnostides * 1 Emanuel Tewolde 1 2 Ratip Emin Berker 1 2 Gabriele Farina 3
Vincent Conitzer 1 2 4 Tuomas Sandholm 1 5 6 7
Abstract
Variational inequalities (VIs) encompass many
fundamental problems in diverse areas ranging
from engineering to economics and machine learn-
ing.
However, their considerable expressivity
comes at the cost of computational intractabil-
ity. In this paper, we introduce and analyze a
natural relaxation—which we refer to as expected
variational inequalities (EVIs)—where the goal
is to find a distribution that satisfies the VI con-
straint in expectation. By adapting recent tech-
niques from game theory, we show that, unlike
VIs, EVIs can be solved in polynomial time under
general (nonmonotone) operators. EVIs capture
the seminal notion of correlated equilibria, but
enjoy a greater reach beyond games. We also
employ our framework to capture and generalize
several existing disparate results, including from
settings such as smooth games, and games with
coupled constraints or nonconcave utilities.
1. Introduction
Variational inequalities (VIs) provide a unifying framework
for analyzing a wide range of optimization and equilib-
rium problems. They have a host of important applica-
tions in engineering and economics (Facchinei & Pang,
2003), including identifying stationary points in constrained
optimization; computing Nash equilibria in multi-player
games (Nash, 1951), such as Cournot’s classical model of
oligopoly (Cournot, 1838); predicting economic activity—
commodity prices and consumer consumption—in a closed,
competitive economy (Arrow & Debreu, 1954), which is at
the heart of general equilibrium theory; traffic equilibrium
*Equal contribution 1Carnegie Mellon University 2Foundations
of Cooperative AI Lab (FOCAL) 3Massachusetts Institute of Tech-
nology 4University of Oxford 5Strategy Robot, Inc. 6Strategic Ma-
chine, Inc. 7Optimized Markets, Inc. Correspondence to: Brian Hu
Zha...
RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding
ICML 2025

RAPID: Long-Context Inference with Retrieval-Augmented
Speculative Decoding
Guanzheng Chen * 1 2 3 Qilong Feng * 1 Jinjie Ni 1 Xin Li 2 3 Michael Qizhe Shieh 1
Code:
https://github.com/NUS-TRAIL/RAPID
Abstract
The emergence of long-context large language
models (LLMs) offers a promising alternative to
traditional retrieval-augmented generation (RAG)
for processing extensive documents. However, the
computational overhead of long-context inference
presents significant efficiency challenges. While
Speculative Decoding (SD) traditionally acceler-
ates inference using smaller draft models, its effec-
tiveness diminishes substantially in long-context
scenarios due to memory-bound KV cache op-
erations.
We introduce Retrieval-Augmented
SPeculatIve Decoding (RAPID), which leverages
RAG for both accelerating and enhancing gener-
ation quality in long-context inference. RAPID
introduces the RAG drafter—a draft LLM oper-
ating on shortened retrieval contexts—to spec-
ulate on the generation of long-context target
LLMs. Our approach enables a new paradigm
where same-scale or even larger LLMs can serve
as RAG drafters while maintaining computational
efficiency. To fully leverage the potentially supe-
rior capabilities from stronger RAG drafters, we
develop an inference-time knowledge transfer that
enriches the target distribution by RAG. Exten-
sive experiments on the LLaMA-3.1 and Qwen2.5
backbones demonstrate that RAPID effectively
integrates the strengths of both RAG and long-
context LLMs, achieving significant performance
improvements (e.g., from 39.33 to 42.83 on In-
finiteBench for LLaMA-3.1-8B) with more than
2× speedups for long-context inference. Our anal-
yses also reveal the robustness of RAPID across
various context lengths and retrieval quality.
*Equal contribution 1National University of Singapore 2DAMO
Academy, Alibaba Group 3Hupan Lab, 310023, Hangzhou, China.
Correspondence to: Guanzheng Chen <gc.chen@u.nus.edu>,
Michael Qizhe Shieh <michaelshieh@comp.nus.edu.s...
Policy Gradient with Tree Expansion
ICML 2025

Policy Gradient with Tree Expansion
Gal Dalal * 1 Assaf Hallak * 1 Gugan Thoppe 2 Shie Mannor 1 3 Gal Chechik 1 4
Abstract
Policy gradient methods are notorious for hav-
ing a large variance and high sample complexity.
To mitigate this, we introduce SoftTreeMax—a
generalization of softmax that employs planning.
In SoftTreeMax, we extend the traditional logits
with the multi-step discounted cumulative reward,
topped with the logits of future states. We an-
alyze SoftTreeMax and explain how tree expan-
sion helps to reduce its gradient variance. We
prove that the variance depends on the chosen
tree-expansion policy. Specifically, we show that
the closer the induced transitions are to being
state-independent, the stronger the variance decay.
With approximate forward models, we prove that
the resulting gradient bias diminishes with the
approximation error while retaining the same vari-
ance reduction. Ours is the first result to bound
the gradient bias for an approximate model. In
a practical implementation of SoftTreeMax, we
utilize a parallel GPU-based simulator for fast and
efficient tree expansion. Using this implementa-
tion in Atari, we show that SoftTreeMax reduces
the gradient variance by three orders of magni-
tude. This leads to better sample complexity and
improved performance compared to distributed
PPO.
1. Introduction
Policy Gradient (PG) methods (Sutton et al., 1999) for Re-
inforcement Learning (RL) are often the first choice for
environments that allow numerous interactions at a fast pace
(Schulman et al., 2017). Their success is attributed to sev-
eral factors: they are easy to distribute to multiple workers,
require no assumptions on the underlying value function,
and have both on-policy and off-policy variants.
*Equal contribution
1NVIDIA Research 2Indian Institute of
Science 3Technion University 4Bar-Ilan University. Correspon-
dence to: Gal dalal <gdalal@nvidia.com>, Assaf Hallak <ahal-
lak@nvidia.com>.
Proceedings of the 42 nd International Confer...
Lightweight Protocols for Distributed Private Quantile Estimation
ICML 2025

Lightweight Protocols for Distributed Private Quantile Estimation
Anders Aamand 1 Fabrizio Boninsegna 2 Abigail Gentle 3 Jacob Imola 1 Rasmus Pagh 1
Abstract
Distributed data analysis is a large and growing
field driven by a massive proliferation of user
devices, and by privacy concerns surrounding
the centralised storage of data. We consider two
adaptive algorithms for estimating one quantile
(e.g. the median) when each user holds a sin-
gle data point lying in a domain [B] that can
be queried once through a private mechanism;
one under local differential privacy (LDP) and
another for shuffle differential privacy (shuffle-
DP). In the adaptive setting we present an ε-
LDP algorithm which can estimate any quan-
tile within error α only requiring O( log B
ε2α2 ) users,
and an (ε, δ)-shuffle DP algorithm requiring only
eO(( 1
ε2 +
1
α2 ) log B) users. Prior (nonadaptive)
algorithms require more users by several logarith-
mic factors in B. We further provide a match-
ing lower bound for adaptive protocols, showing
that our LDP algorithm is optimal in the low-ε
regime. Additionally, we establish lower bounds
against non-adaptive protocols which paired with
our understanding of the adaptive case, proves a
fundamental separation between these models.
1. Introduction
A strong trend in recent years has been towards federated
computations (Kairouz et al., 2021) in which algorithms are
run on a distributed dataset rather than by collecting data and
performing the computation in a centralized manner. This
trend is motivated by the wish to protect individuals’ data,
as well as organizations’ wish steer clear of liability issues
stemming from collecting and handling private data. The
1BARC and Department of Computer Science, Univer-
sity of Copenhagen, Denmark 2Department of Information En-
gineering, University of Padova, Italy 3School of Computer
Science, University of Sydney, Australia.
Correspondence
to:
Anders Aamand <aa@di.ku.dk>, Abigail Gentle <abi-
gail.gentle@sydney.ed...
ADIOS: Antibody Development via Opponent Shaping
ICML 2025

ADIOS: Antibody Development via Opponent Shaping
Sebastian Towers * 1 2 Aleksandra Kalisz * 1 2 Philippe A. Robert 3 Alicia Higueruelo 4 Francesca Vianello 5
Ming-Han Chloe Tsai 6 Harrison Steel 2 Jakob Foerster 1 2
Abstract
Anti-viral therapies are typically designed to tar-
get only the current strains of a virus, a my-
opic response. However, therapy-induced selec-
tive pressures drive the emergence of new viral
strains, against which the original myopic ther-
apies are no longer effective.
This evolution-
ary response presents an opportunity: our ther-
apies could both defend against and actively influ-
ence viral evolution. This motivates our method
ADIOS: Antibody Development vIa Opponent
Shaping. ADIOS is a meta-learning framework
where the process of antibody therapy design, the
outer loop, accounts for the virus’s adaptive re-
sponse, the inner loop. With ADIOS, antibodies
are not only robust against potential future vari-
ants, they also influence, i.e., shape, which fu-
ture variants emerge. In line with the opponent
shaping literature, we refer to our optimised an-
tibodies as shapers. To demonstrate the value of
ADIOS, we build a viral evolution simulator us-
ing the Absolut! framework, in which shapers
successfully target both current and future viral
variants, outperforming myopic antibodies. Fur-
thermore, we show that shapers modify the dis-
tribution over viral evolutionary trajectories to
result in weaker variants. We believe that our
ADIOS paradigm will facilitate the discovery of
long-lived vaccines and antibody therapies while
also generalising to other domains. Specifically,
domains such as antimicrobial resistance, cancer
treatment, and others with evolutionarily adap-
tive opponents. Our code is available at https:
//github.com/olakalisz/adios.
*Equal contribution
1FLAIR, Foerster Lab for AI Research
2Department of Engineering, University of Oxford, Oxford, UK
3Department of Biomedicine, University of Basel, Basel, Switzer-
land 4Isomorphic ...
The Diffusion Duality
ICML 2025

The Diffusion Duality
Subham Sekhar Sahoo 1 Justin Deschenaux 2 Aaron Gokaslan 1 Guanghan Wang 1 Justin Chiu 3
Volodymyr Kuleshov 1
Abstract
Uniform-state discrete diffusion models hold
the promise of fast text generation due to their
inherent ability to self-correct. However, they are
typically outperformed by autoregressive models
and masked diffusion models. In this work, we
narrow this performance gap by leveraging a
key insight: Uniform-state diffusion processes
naturally emerge from an underlying Gaussian
diffusion. Our method, Duo, transfers powerful
techniques from Gaussian diffusion to improve
both training and sampling. First, we introduce
a curriculum learning strategy guided by the
Gaussian process, doubling training speed by
reducing variance. Models trained with curricu-
lum learning surpass autoregressive models in
zero-shot perplexity on 3 of 7 benchmarks. Sec-
ond, we present Discrete Consistency Distillation,
which adapts consistency distillation from the
continuous to the discrete setting. This algorithm
unlocks few-step generation in diffusion
language models by accelerating sampling by
two orders of magnitude. We provide the code
and model checkpoints on the project page:
https://s-sahoo.com/duo
1. Introduction
An eternal theme in mathematics is that discreteness
emerges from underlying continuity. From quantum me-
chanics, where the quantized energy states of electrons arise
as solutions to continuous wave equations, to the binary
logic of digital circuits, fundamentally driven by smooth
analog currents, discreteness has repeatedly and naturally
emerged from an underlying continuum. Our work contin-
1Computer and Information Science, Cornell Tech, NY, USA.
2School of Computer and Communication Sciences, EPFL Lau-
sanne, Switzerland 3Cohere, NY, USA. Correspondence to: Sub-
ham Sekhar Sahoo <ssahoo@cs.cornell.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the auth...
Optimistic Algorithms for Adaptive Estimation of the Average Treatment Effect
ICML 2025

Optimistic Algorithms for
Adaptive Estimation of the Average Treatment Effect
Ojash Neopane 1 Aaditya Ramdas 1 2 Aarti Singh 1
Abstract
Estimation and inference for the Average Treat-
ment Effect (ATE) is a cornerstone of causal infer-
ence and often serves as the foundation for devel-
oping procedures for more complicated settings.
Although traditionally analyzed in a batch setting,
recent advances in martingale theory have paved
the way for adaptive methods that can enhance
the power of downstream inference. Despite these
advances, progress in understanding and develop-
ing adaptive algorithms remains in its early stages.
Existing work either focus on asymptotic analy-
ses that overlook exploration-exploitation trade-
offs relevant in finite-sample regimes or rely on
simpler but suboptimal estimators. In this work,
we address these limitations by studying adaptive
sampling procedures that take advantage of the
asymptotically optimal Augmented Inverse Prob-
ability Weighting (AIPW) estimator. Our analy-
sis uncovers challenges obscured by asymptotic
approaches and introduces a novel algorithmic
design principle reminiscent of optimism in multi-
armed bandits. This principled approach enables
our algorithm to achieve significant theoretical
and empirical gains compared to previous meth-
ods. Our findings mark a step forward in the ad-
vancement of adaptive causal inference methods
in theory and practice.
1. Introduction
The problem of estimating the average treatment effect
(ATE) is central to causal inference and has been extensively
studied. We have a precise understanding of the difficulty of
this problem in both asymptotic and nonasymptotic regimes.
1Machine Learning Department, Carnegie Mellon Uni-
versity
2Department of Statistics & Data Science, Carnegie
Mellon University.
Correspondence to:
Ojash Neopane
<oneopane@andrew.cmu.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the a...
Partially Observable Reinforcement Learning with Memory Traces
ICML 2025

Partially Observable Reinforcement Learning with Memory Traces
Onno Eberhard 1 2 Michael Muehlebach 1 Claire Vernade 2
Abstract
Partially observable environments present a con-
siderable computational challenge in reinforce-
ment learning due to the need to consider long his-
tories. Learning with a finite window of observa-
tions quickly becomes intractable as the window
length grows. In this work, we introduce memory
traces. Inspired by eligibility traces, these are
compact representations of the history of observa-
tions in the form of exponential moving averages.
We prove sample complexity bounds for the prob-
lem of offline on-policy evaluation that quantify
the return errors achieved with memory traces
for the class of Lipschitz continuous value esti-
mates. We establish a close connection to the
window approach, and demonstrate that, in cer-
tain environments, learning with memory traces
is significantly more sample efficient. Finally,
we underline the effectiveness of memory traces
empirically in online reinforcement learning ex-
periments for both value prediction and control.
1. Introduction
Learning and acting in partially observable environments
requires memory. It is generally not enough to act simply
according to the most recent observation, and an agent must
therefore keep track of the past, as every piece of informa-
tion is likely to be relevant for future actions. In the absence
of additional assumptions, optimal behavior requires the
agent to reason over its entire history of observations. This
space quickly becomes intractable, and just like with contin-
uous state spaces in fully observed systems, it is therefore
necessary to resort to approximation. This article inves-
tigates the fundamental limitations and trade-offs arising
from these approximations. We focus on extracting features
from the history of observations and analyze the complexity
of learning functions of these features. In contrast to meth-
1Max Planck Institute for Intelligent Sy...