dataset-review/dataset-review huggingface dataset:
Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
ICML 2025

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Christian Pedersen 1 2 Laure Zanna 1 2 Joan Bruna 1 2
Abstract
Autoregressive surrogate models (or emulators)
of spatiotemporal systems provide an avenue for
fast, approximate predictions, with broad applica-
tions across science and engineering. At inference
time, however, these models are generally unable
to provide predictions over long time rollouts due
to accumulation of errors leading to diverging tra-
jectories. In essence, emulators operate out of
distribution, and controlling the online distribu-
tion quickly becomes intractable in large-scale
settings. To address this fundamental issue, and
focusing on time-stationary systems admitting an
invariant measure, we leverage diffusion models
to obtain an implicit estimator of the score of this
invariant measure. We show that this model of
the score function can be used to stabilize autore-
gressive emulator rollouts by applying on-the-fly
denoising during inference, a process we call ther-
malization. Thermalizing an emulator rollout is
shown to extend the time horizon of stable pre-
dictions by two orders of magnitude in complex
systems exhibiting turbulent and chaotic behav-
ior, opening up a novel application of diffusion
models in the context of neural emulation.
1. Introduction
The modeling of dynamical systems is a cornerstone task
in the physical sciences and engineering, with applications
across weather (Pathak et al., 2022; Watt-Meyer et al., 2023;
Lam et al., 2023; Lang et al., 2024) and climate modeling
(Kochkov et al., 2024; Subel & Zanna, 2024; Lupin-Jimenez
et al., 2025). The standard approach to modeling these sys-
tems is to solve the underlying partial differential equation
(PDE) describing the system using numerical methods. In
the case of high-dimensional systems, the computational
1Courant Institute of Mathematical Sciences, New York Univer-
sity, USA 2Center for Data Science, New York University, USA.
Correspondence ...
Spatial Reasoning with Denoising Models
ICML 2025

Spatial Reasoning with Denoising Models
Christopher Wewer 1 Bart Pogodzinski 1 Bernt Schiele 1 Jan Eric Lenssen 1
SRMs (Ours)
Standard DDPM
Playing MNIST Sudoku with Image Generators….
Figure 1: We introduce Spatial Reasoning Models (SRMs), a framework to systematically investigate diffusion/flow-based
generative models with respect to their reasoning capabilities over multiple variables. We introduce benchmarks that allow
quantification of higher-level reasoning capabilities and show that manual and automatic schemes for sequentialization can
heavily reduce hallucination. Here, we show the solving process of one of the benchmarks, a Sudoku game consisting of
MNIST images, which is solved correctly by our SRMs, while standard diffusion models fail.
Abstract
We introduce Spatial Reasoning Models (SRMs),
a framework to perform reasoning over sets of
continuous variables via denoising generative
models. SRMs infer continuous representations
on a set of unobserved variables, given observa-
tions on observed variables. Current generative
models on spatial domains, such as diffusion and
flow matching models, often collapse to hallucina-
tion in case of complex distributions. To measure
this, we introduce a set of benchmark tasks that
test the quality of complex reasoning in generative
models and can quantify hallucination. The SRM
framework allows to report key findings about
importance of sequentialization in generation, the
associated order, as well as the sampling strate-
gies during training. It demonstrates, for the first
time, that order of generation can successfully be
predicted by the denoising network itself. Using
these findings, we can increase the accuracy of
specific reasoning tasks from < 1% to > 50%.
Our project website provides additional videos,
code, and the benchmark datasets.
1Max Planck Institute for Informatics, Saarland Informat-
ics Campus, Germany. Correspondence to: Christopher Wewer
<cwewer@mpi-inf.mpg.de>, Bart Pogodzinski <bpogodzi@mpi-
inf.m...
In-context Time Series Predictor
ICLR 2025

Published as a conference paper at ICLR 2025
IN-CONTEXT TIME SERIES PREDICTOR
Jiecheng Lu, Yan Sun, Shihao Yang
Georgia Institute of Technology
{jlu414,yansun}@gatech.edu, shihao.yang@isye.gatech.edu
ABSTRACT
Recent Transformer-based large language models (LLMs) demonstrate in-context
learning ability to perform various functions based solely on the provided context,
without updating model parameters. To fully utilize the in-context capabilities
in time series forecasting (TSF) problems, unlike previous Transformer-based or
LLM-based time series forecasting methods, we reformulate "time series fore-
casting tasks" as input tokens by constructing a series of (lookback, future) pairs
within the tokens. This method aligns more closely with the inherent in-context
mechanisms, and is more parameter-efficient without the need of using pre-trained
LLM parameters. Furthermore, it addresses issues such as overfitting in existing
Transformer-based TSF models, consistently achieving better performance across
full-data, few-shot, and zero-shot settings compared to previous architectures 1.
1
INTRODUCTION
Transformer-based large language models (LLMs) have significantly impacted various research
and application areas (Brown et al., 2020). Their inherent in-context learning (ICL) capabilities
highlighted by previous studies (Müller et al., 2022; Min et al., 2022; Wei et al., 2023; Xie et al.,
2022) allow them to adapt and generalize from context examples provided in input prompts without
any parameter updates. This enables LLMs to effectively handle few-shot and zero-shot tasks. Recent
research on ICL (Zhang et al., 2023; Garg et al., 2022; Akyürek et al., 2023; Li et al., 2023; Dai et al.,
2023; Bai et al., 2023) has shown that Transformers can adaptively learn to perform various functions
including linear predictors, shallow MLPs, gradient descent, algorithm selection, etc., based on a
series of (input, label) pairs as input tokens, as shown in Figure 1 (a).
Timestep
Input
Labe...
KVTuner: Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache Quantization for Efficient and Nearly Lossless LLM Inference
ICML 2025

KVTuner: Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache
Quantization for Efficient and Nearly Lossless LLM Inference
Xing Li * 1 Zeyu Xing * 2 Yiming Li 1 Linping Qu 1 Hui-Ling Zhen 1 Yiwu Yao 3 Wulong Liu 1
Sinno Jialin Pan 2 Mingxuan Yuan 1
Abstract
KV cache quantization can improve Large Lan-
guage Models (LLMs) inference throughput and
latency in long contexts and large batch-size
scenarios while preserving LLMs effectiveness.
However, current methods have three unsolved
issues: overlooking layer-wise sensitivity to KV
cache quantization, high overhead of online fine-
grained decision-making, and low flexibility to
different LLMs and constraints. Therefore, we
theoretically analyze the inherent correlation of
layer-wise transformer attention patterns to KV
cache quantization errors and study why key
cache is generally more important than value
cache for quantization error reduction. We fur-
ther propose a simple yet effective framework
KVTuner to adaptively search for the optimal
hardware-friendly layer-wise KV quantization
precision pairs for coarse-grained KV cache with
multi-objective optimization and directly utilize
the offline searched configurations during on-
line inference. To reduce the computational cost
of offline calibration, we utilize the intra-layer
KV precision pair pruning and inter-layer clus-
tering to reduce the search space. Experimen-
tal results show that we can achieve nearly loss-
less 3.25-bit mixed precision KV cache quantiza-
tion for LLMs like Llama-3.1-8B-Instruct and
4.0-bit for sensitive models like Qwen2.5-7B-
Instruct on mathematical reasoning tasks. The
maximum inference throughput can be improved
by 21.25% compared with KIVI-KV8 quantiza-
tion over various context lengths. Our code and
searched configurations are available at https:
//github.com/cmd2001/KVTuner.
*Equal contribution
1Huawei Noah’s Ark Lab 2The Chi-
nese University of Hong Kong 3Huawei Computing Product
Line.
Correspondence to: Xing Li <li.xing2@huawei.c...
Optimizing Test-Time Compute via Meta Reinforcement Finetuning
ICML 2025

Optimizing Test-Time Compute via Meta Reinforcement Finetuning
Yuxiao Qu * 1 Matthew Y. R. Yang * 1 Amrith Setlur 1 Lewis Tunstall 2 Edward Emanuel Beeching 2
Ruslan Salakhutdinov 1 Aviral Kumar 1
Abstract
Training models to effectively use test-time com-
pute is crucial for improving the reasoning per-
formance of LLMs. Current methods mostly do
so via fine-tuning on search traces or running RL
with 0/1 outcome reward, but do these approaches
efficiently utilize test-time compute? Would these
approaches continue to scale as the budget im-
proves? To answer these questions, in this paper,
we formalize the problem of optimizing test-time
compute as a meta-reinforcement learning (RL)
problem, which provides a principled perspective
on spending test-time compute. This perspective
enables us to view the long output stream from the
LLM as consisting of several episodes run at test
time and leads us to use a notion akin to cumula-
tive regret over output tokens as a way to measure
the efficacy of test-time compute. Akin to how RL
algorithms can best tradeoff exploration and ex-
ploitation over training, minimizing regret should
also provide the best balance between exploration
and exploitation in the token stream. While we
show that state-of-the-art models do not minimize
regret, one can do so by maximizing a dense re-
ward bonus in conjunction with the outcome 0/1
reward RL. This bonus is the “progress” made
by each subsequent block in the output stream,
quantified by the change in the likelihood of even-
tual success. Using these insights, we develop
Meta Reinforcment Fine-Tuning, or MRT, a new
class of fine-tuning methods for optimizing test-
time compute. MRT leads to a 2-3x relative gain
in performance and roughly a 1.5x gain in token
efficiency for math reasoning.
1. Introduction
Recent results in LLM reasoning (Snell et al., 2024) il-
lustrate the potential to improve reasoning capabilities by
*Equal contribution 1CMU 2Hugging Face. Correspondence to:
Yuxiao Qu <yu...
STAMP Your Content: Proving Dataset Membership via Watermarked Rephrasings
ICML 2025

STAMP Your Content:
Proving Dataset Membership via Watermarked Rephrasings
Saksham Rastogi 1 Pratyush Maini 2 3 Danish Pruthi 1
Abstract
Given how large parts of publicly available text
are crawled to pretrain large language mod-
els (LLMs), data creators increasingly worry
about the inclusion of their proprietary data for
model training without attribution or licensing.
Their concerns are also shared by benchmark cu-
rators whose test-sets might be compromised. In
this paper, we present STAMP, a framework for de-
tecting dataset membership—i.e., determining the
inclusion of a dataset in the pretraining corpora of
LLMs. Given an original piece of content, our pro-
posal involves first generating multiple rephrases,
each embedding a watermark with a unique se-
cret key. One version is to be released publicly,
while others are to be kept private. Subsequently,
creators can compare model likelihoods between
public and private versions using paired statis-
tical tests to prove membership. We show that
our framework can successfully detect contami-
nation across four benchmarks which appear only
once in the training data and constitute less than
0.001% of the total tokens, outperforming sev-
eral contamination detection and dataset inference
baselines. We verify that STAMP preserves both
the semantic meaning and utility of the original
data. We apply STAMP to two real-world scenar-
ios to confirm the inclusion of paper abstracts and
blog articles in the pretraining corpora.
1. Introduction
To train large language models, much of the available text
from the internet is crawled, allegedly including copyrighted
material such as news articles and blogs (Grynbaum & Mac,
1Indian
Institute
of
Science
2Carnegie
Mellon
Uni-
versity
3DatologyAI.
Correspondence
to:
Saksham
Ras-
togi <iitdsaksham@gmail.com>, Pratyush Maini <pratyush-
maini.@cmu.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author...
Organize the Web: Constructing Domains Enhances Pre-Training Data Curation
ICML 2025

Organize the Web: Constructing Domains Enhances
Pre-Training Data Curation
Alexander Wettig 1 2 Kyle Lo 2 Sewon Min 3 2 Hannaneh Hajishirzi 2 4 Danqi Chen 1 Luca Soldaini 2
Abstract
Modern language models are trained on large,
unstructured datasets consisting of trillions of to-
kens and obtained by crawling the web.
The
unstructured nature makes it difficult to reason
about their contents and develop systematic ap-
proaches to data curation. In this paper, we un-
pack monolithic web corpora by developing tax-
onomies of their contents and organizing them
into domains. We introduce WebOrganizer, a
framework for organizing web pages in terms of
both their topic and format.
Using these two
complementary notions of domains, we automat-
ically annotate pre-training data by distilling an-
notations from a large language model into effi-
cient classifiers. This allows us to study how data
from different domains should be mixed to im-
prove models on downstream tasks, and we show
that we can combine insights about effective top-
ics and formats to further boost performance.
We demonstrate that our domain mixing also im-
proves existing methods that select data based on
quality. Furthermore, we study and compare how
quality-based methods will implicitly change the
domain mixture. Overall, our work demonstrates
that constructing and mixing domains provides a
valuable complement to quality-based data cura-
tion methods, opening new avenues for effective
and insightful pre-training data curation.
Website
weborganizer.allen.ai
Artifacts
hf.co/WebOrganizer
Code
CodeCreator/WebOrganizer
1Princeton Language and Intelligence, Princeton University
2Allen Institute for Artificial Intelligence 3University of Califor-
nia, Berkeley 4Paul G. Allen School of Computer Science & Engi-
neering, University of Washington. Correspondence to: Alexan-
der Wettig <awettig@cs.princeton.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. C...
Sequential Attention-based Sampling for Histopathological Analysis
NeurIPS 2025

Sequential Attention-based Sampling for
Histopathological Analysis
Tarun G
Indian Institute of Science, Bangalore, India
tarung@iisc.ac.in
Naman Malpani
Indian Institute of Science, Bangalore, India
namanmalpani@iisc.ac.in
Gugan Thoppe
Indian Institute of Science, Bangalore, India
gthoppe@iisc.ac.in
Devarajan Sridharan
Indian Institute of Science, Bangalore, India
sridhar@iisc.ac.in
Abstract
Deep neural networks are increasingly applied in automated histopathology. Yet,
whole-slide images (WSIs) are often acquired at gigapixel sizes, rendering them
computationally infeasible to analyze entirely at high resolution. Diagnostic labels
are largely available only at the slide-level, because expert annotation of images
at a finer (patch) level is both laborious and expensive. Moreover, regions with
diagnostic information typically occupy only a small fraction of the WSI, mak-
ing it inefficient to examine the entire slide at full resolution. Here, we propose
SASHA – Sequential Attention-based Sampling for Histopathological Analysis – a
deep reinforcement learning approach for efficient analysis of histopathological
images. First, SASHA learns informative features with a lightweight hierarchical,
attention-based multiple instance learning (MIL) model. Second, SASHA samples
intelligently and zooms selectively into a small fraction (10-20%) of high-resolution
patches to achieve reliable diagnoses. We show that SASHA matches state-of-the-
art methods that analyze the WSI fully at high resolution, albeit at a fraction of
their computational and memory costs. In addition, it significantly outperforms
competing, sparse sampling methods. We propose SASHA as an intelligent sam-
pling model for medical imaging challenges that involve automated diagnosis with
exceptionally large images containing sparsely informative features1.
1
Introduction
Deep learning models have emerged as a major frontier in automated medical imaging diagnosis. For
example, such models accurately identify tum...
Geometry-Informed Neural Networks
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...
Learning Dynamics under Environmental Constraints via Measurement-Induced Bundle Structures
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...
Constrained Best Arm Identification
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...