AIS-Bench/AIS-Bench huggingface dataset:
Programming Every Example:
Lifting Pre-training Data Quality Like Experts at Scale
Fan Zhou * 1 2 Zengzhi Wang * 1 2 Qian Liu 3 Junlong Li 1 Pengfei Liu ‡ 1 2 4
Abstract
Large language model pre-training has tradition-
ally relied on human experts to craft heuristics for
improving the corpora quality, resulting in numer-
ous rules developed to date. However, these fixed
rules lack the flexibility to address the unique
characteristics of individual examples, yet craft-
ing sample-wise rules is impractical for human
experts. In this paper, we show that even small
language models, with only 0.3B parameters, can
exhibit substantial data refining capabilities. We
propose Programming Every Example (PROX),
a novel framework that treats data refinement as
a programming task, and enables the model to
refine corpora by generating and executing fine-
grained operations, such as string normalization,
for each individual example at scale. Experiments
show that models trained on PROX-refined data
consistently outperform other baselines across 10
benchmarks, demonstrating effectiveness across
model sizes (up to 1.7B) and pre-training cor-
pora (C4, RedPajama-V2, FineWeb, FineWeb-
Edu, and DCLM). PROX also shows great po-
tential in continual pre-training: on math domain,
PROX boosts 7B models by up to 20% within
10B tokens—results typically achieved with much
larger scale training (e.g., 200B tokens). We be-
lieve PROX offers a way to curate high-quality
pre-training data, and finally contributes to effi-
cient LLM development.
1. Introduction
Large Language Models (LLMs) have made significant
strides in capabilities (Meta, 2024; Achiam et al., 2023;
Anthropic, 2024; Reid et al., 2024), excelling in tasks such
*Equal
contribution
1Shanghai
Jiao
Tong
University
2Generative AI Research Lab (GAIR) 3Sea AI Lab 4Shanghai
Artificial Intelligence Laboratory. Correspondence to: Pengfei Liu
<pengfei@sjtu.edu.cn>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancou...
Optimal Transport-based Conformal Prediction
Gauthier Thurin 1 Kimia Nadjahi 1 Claire Boyer 2
Abstract
Conformal Prediction (CP) is a principled frame-
work for quantifying uncertainty in black-box
learning models, by constructing prediction sets
with finite-sample coverage guarantees. Tradi-
tional approaches rely on scalar nonconformity
scores, which fail to fully exploit the geometric
structure of multivariate outputs, such as in multi-
output regression or multiclass classification. Re-
cent methods addressing this limitation impose
predefined convex shapes for the prediction sets,
potentially misaligning with the intrinsic data ge-
ometry. We introduce a novel CP procedure hand-
ling multivariate score functions through the lens
of optimal transport. Specifically, we leverage
Monge-Kantorovich vector ranks and quantiles to
construct prediction region with flexible, poten-
tially non-convex shapes, better suited to the com-
plex uncertainty patterns encountered in multivari-
ate learning tasks. We prove that our approach
ensures finite-sample, distribution-free coverage
properties, similar to typical CP methods. We
then adapt our method for multi-output regression
and multiclass classification, and also propose
simple adjustments to generate adaptive predic-
tion regions with asymptotic conditional coverage
guarantees. Finally, we evaluate our method on
practical regression and classification problems,
illustrating its advantages in terms of (conditional)
coverage and efficiency.
1. Introduction
In various domains, including high-stakes applications,
state-of-the-art performances are often achieved by black-
box machine learning models. As a result, accurately quan-
tifying the uncertainty of their predictions has become a crit-
1CNRS, Ecole Normale Sup´erieure, Paris, France 2Laboratoire
de Math´ematiques d’Orsay (LMO), Universit´e Paris Saclay,
France, and Institut universitaire de France. Correspondence to:
Gauthier Thurin <gthurin@mail.di.ens.fr>.
Proceeding...
Symmetry-Aware GFlowNets
Hohyun Kim 1 Seunggeun Lee 1 Min-hwan Oh 1
Abstract
Generative Flow Networks (GFlowNets) offer a
powerful framework for sampling graphs in pro-
portion to their rewards. However, existing ap-
proaches suffer from systematic biases due to
inaccuracies in state transition probability com-
putations. These biases, rooted in the inherent
symmetries of graphs, impact both atom-based
and fragment-based generation schemes. To ad-
dress this challenge, we introduce Symmetry-
Aware GFlowNets (SA-GFN), a method that in-
corporates symmetry corrections into the learning
process through reward scaling. By integrating
bias correction directly into the reward structure,
SA-GFN eliminates the need for explicit state
transition computations. Empirical results show
that SA-GFN enables unbiased sampling while
enhancing diversity and consistently generating
high-reward graphs that closely match the target
distribution.
1. Introduction
GFlowNets have emerged as a powerful framework for learn-
ing generative models capable of sampling complex, compo-
sitional objects with probabilities proportional to a given re-
ward. Inspired by reinforcement learning (RL), GFlowNets
generate these objects through a sequence of actions that it-
eratively modify the structure of the object being built. This
approach is particularly well-suited for generating compo-
sitional objects, such as graphs, where there are multiple
paths for constructing an object. A prominent application
of GFlowNets is molecule generation, where molecules are
sequentially constructed as graphs (Bengio et al., 2021; Jain
et al., 2023a).
However, GFlowNet training objectives rely on the accurate
computation of the transition probability of a policy, which
1Graduate School of Data Science, Seoul National University,
Seoul, Republic of Korea. Correspondence to: Seunggeun Lee
<lee7801@snu.ac.kr>, Min-hwan Oh <minoh@snu.ac.kr>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver...
Score-based Pullback Riemannian Geometry:
Extracting the Data Manifold Geometry using Anisotropic Flows
Willem Diepeveen * 1 Georgios Batzolis * 2 Zakhar Shumaylov 3 Carola-Bibiane Sch¨onlieb 3
Abstract
Data-driven Riemannian geometry has emerged
as a powerful tool for interpretable representation
learning, offering improved efficiency in down-
stream tasks. Moving forward, it is crucial to
balance cheap manifold mappings with efficient
training algorithms. In this work, we integrate
concepts from pullback Riemannian geometry
and generative models to propose a framework
for data-driven Riemannian geometry that is scal-
able in both geometry and learning: score-based
pullback Riemannian geometry. Focusing on uni-
modal distributions as a first step, we propose
a score-based Riemannian structure with closed-
form geodesics that pass through the data prob-
ability density. With this structure, we construct
a Riemannian autoencoder (RAE) with error
bounds for discovering the correct data manifold
dimension. This framework can naturally be used
with anisotropic normalizing flows by adopting
isometry regularization during training. Through
numerical experiments on diverse datasets, includ-
ing image data, we demonstrate that the proposed
framework produces high-quality geodesics pass-
ing through the data support, reliably estimates
the intrinsic dimension of the data manifold, and
provides a global chart of the manifold. To the
best of our knowledge, this is the first scalable
framework for extracting the complete geometry
of the data manifold.
*Equal contribution
1Department of Mathematics, Univer-
sity of California, Los Angeles, USA 2Department of Engi-
neering, University of Cambridge, UK 3Faculty of Mathemat-
ics, University of Cambridge, UK. Correspondence to: Willem
Diepeveen <wdiepeveen@math.ucla.edu>, Georgios Batzolis
<gb511@cam.ac.uk>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by th...
Published as a conference paper at ICLR 2025
SCALING LAWS FOR PRECISION
Tanishq Kumar∗1
Zachary Ankner* 3,4
Benjamin F. Spector2
Blake Bordelon1
Niklas Muennighoff2
Mansheej Paul4
Cengiz Pehlevan1
Christopher R´e2
Aditi Raghunathan5
1Harvard University
2Stanford University
3MIT
4Databricks
5Carnegie Mellon University
ABSTRACT
Low precision training and inference affect both the quality and cost of language
models, but current scaling laws do not account for this. In this work, we devise
“precision-aware” scaling laws for both training and inference. We propose that
training in lower precision reduces the model’s effective parameter count, allow-
ing us to predict the additional loss incurred from training in low precision and
post-train quantization. For inference, we find that the degradation introduced by
post-training quantization increases as models are trained on more data, eventu-
ally making additional pretraining data actively harmful. For training, our scaling
laws allow us to predict the loss of a model with different parts in different preci-
sions, and suggest that training larger models in lower precision may be compute
optimal. We unify the scaling laws for post and pretraining quantization to arrive
at a single functional form that predicts degradation from training and inference in
varied precisions. We fit on over 465 pretraining runs and validate our predictions
on model sizes up to 1.7B parameters trained on up to 26B tokens.
1
INTRODUCTION
Scale has emerged as a central driver of progress in deep learning (Brown et al., 2020). Key work
on scaling (Kaplan et al., 2020; Hoffmann et al., 2022) studied tradeoffs between model/dataset
size to balance performance and compute. However, the precision in which models are trained and
served is an important third factor that contributes to both cost and performance. Deep learning
is trending towards lower precision: current frontier models like the Llama-3 series are trained in
BF16 (Dubey et al., 2024), an...
Published as a conference paper at ICLR 2025
VIDEO ACTION DIFFERENCING
James Burgess1, Xiaohan Wang1, Yuhui Zhang1, Anita Rau1, Alejandro Lozano1,
Lisa Dunlap2, Trevor Darrell2, Serena Yeung-Levy1
1Stanford, 2 UC Berkeley
ABSTRACT
How do two individuals differ when performing the same action? In this work, we
introduce Video Action Differencing (VidDiff), the novel task of identifying subtle
differences between videos of the same action, which has many applications, such
as coaching and skill learning. To enable development on this new task, we first
create VidDiffBench, a benchmark dataset containing 549 video pairs, with human
annotations of 4,469 fine-grained action differences and 2,075 localization times-
tamps indicating where these differences occur. Our experiments demonstrate that
VidDiffBench poses a significant challenge for state-of-the-art large multimodal
models (LMMs), such as GPT-4o and Qwen2-VL. By analyzing failure cases of
LMMs on VidDiffBench, we highlight two key challenges for this task: localiz-
ing relevant sub-actions over two videos and fine-grained frame comparison. To
overcome these, we propose the VidDiff method, an agentic workflow that breaks
the task into three stages: action difference proposal, keyframe localization, and
frame differencing, each stage utilizing specialized foundation models. To en-
courage future research in this new task, we release the benchmark1 and code2.
1
INTRODUCTION
The ability to compare two videos of the same action and discern their detailed differences plays a
critical role in a wide variety of applications. For instance, in fitness coaching, a novice learning
to perform a barbell squat typically watches instructional videos and then compares their actions
in a recorded video to identify discrepancies between their movements and those of an expert. In
medical training, junior surgeons compare videos of themselves performing surgical procedures with
reference videos from experts to identify errors and imp...
Go With the Flow: Fast Diffusion for Gaussian
Mixture Models
George Rapakoulias 1 ∗
Ali Reza Pedram 1,2
Fengjiao Liu 3
Lingjiong Zhu 4
Panagiotis Tsiotras 1
1 Department of Aerospace Engineering, Georgia Institute of Technology, Atlanta, GA
2 School of Computer Science, University of Oklahoma, Norman, OK
3 Department of ECE, FAMU-FSU College of Engineering, Tallahassee, FL
4 Department of Mathematics, Florida State University, Tallahassee, FL
Abstract
Schr¨odinger Bridges (SBs) are diffusion processes that steer, in finite time, a given
initial distribution to another final one while minimizing a suitable cost functional.
Although various methods for computing SBs have recently been proposed in the
literature, most of these approaches require computationally expensive training
schemes, even for solving low-dimensional problems. In this work, we propose
an analytic parametrization of a set of feasible policies for steering the distribution
of a dynamical system from one Gaussian Mixture Model (GMM) to another. In-
stead of relying on standard non-convex optimization techniques, the optimal pol-
icy within the set can be approximated as the solution of a low-dimensional linear
program whose dimension scales linearly with the number of components in each
mixture. The proposed method generalizes naturally to more general classes of
dynamical systems, such as controllable linear time-varying systems, enabling ef-
ficient solutions to multi-marginal momentum SBs between GMMs, a challenging
distribution interpolation problem. We showcase the potential of this approach in
low-to-moderate dimensional problems such as image-to-image translation in the
latent space of an autoencoder, learning of cellular dynamics using multi-marginal
momentum SBs, and various other examples. The implementation is publicly
available at https://github.com/georgeRapa/GMMflow.
1
Introduction and Background
The problem of finding mappings between distributions of data, originally known as the Optim...
FlyLoRA: Boosting Task Decoupling and Parameter
Efficiency via Implicit Rank-Wise Mixture-of-Experts
Heming Zou1∗
Yunliang Zang2∗
Wutong Xu1
Yao Zhu1
Xiangyang Ji1†
1Department of Automation, Tsinghua University
2Academy of Medical Engineering and Translational Medicine, Tianjin University
{zouhm24, xwt22}@mails.tsinghua.edu.cn
yunliangzang@tju.edu.cn, ee_zhuy@zju.edu.cn
xyji@tsinghua.edu.cn
Abstract
Low-Rank Adaptation (LoRA) is a widely used parameter-efficient fine-tuning
method for foundation models, but it suffers from parameter interference, resulting
in suboptimal performance. Although Mixture-of-Experts (MoE)-based LoRA vari-
ants show promise in mitigating intra-task correlations in single-task instruction tun-
ing, they introduce additional router parameters and remain ineffective in multi-task
model merging where inter-task interference arises. Inspired by the fly olfactory
circuit, we propose FlyLoRA, an implicit MoE-based LoRA variant that introduces:
(1) rank-wise expert activation in the up-projection matrix, and (2) an implicit
router that unifies expert routing and down-projection, where a frozen sparse ran-
dom projection matrix replaces the traditional dense trainable version. This design
resolves the trade-off between intra-task decorrelation and computational efficiency
by eliminating the need for an explicit router, while inherently mitigating inter-
task interference due to the orthogonality property of random matrices. Extensive
experiments across four domains—general knowledge understanding, scientific
question answering, mathematical reasoning, and code generation—demonstrate
consistent performance improvements over existing methods. Beyond empirical
gains, FlyLoRA highlights how biological structures can inspire innovations in AI
technologies. Code is available at https://github.com/gfyddha/FlyLoRA.
1
Introduction
Foundation models have demonstrated remarkable cross-domain capabilities with the scaling of
model parameters [1, 4, 15, 43, 57...
Published as a conference paper at ICLR 2025
MRT5: DYNAMIC TOKEN MERGING FOR EFFICIENT
BYTE-LEVEL LANGUAGE MODELS
Julie Kallini, Shikhar Murty, Christopher D. Manning, Christopher Potts, R´obert Csord´as
Stanford University
{kallini, jsmurty, manning, cgpotts, rcsordas}@stanford.edu
ABSTRACT
Models that rely on subword tokenization have significant drawbacks, such as sen-
sitivity to character-level noise like spelling errors and inconsistent compression
rates across different languages and scripts. While character- or byte-level models
like ByT5 attempt to address these concerns, they have not gained widespread
adoption—processing raw byte streams without tokenization results in signifi-
cantly longer sequence lengths, making training and inference inefficient. This
work introduces MrT5 (MergeT5), a more efficient variant of ByT5 that inte-
grates a token deletion mechanism in its encoder to dynamically shorten the input
sequence length. After processing through a fixed number of encoder layers, a
learned delete gate determines which tokens are to be removed and which are to
be retained for subsequent layers. MrT5 effectively “merges” critical information
from deleted tokens into a more compact sequence, leveraging contextual infor-
mation from the remaining tokens. In continued pre-training experiments, we
find that MrT5 can achieve significant gains in inference runtime with minimal
effect on performance, as measured by bits-per-byte. Additionally, with multilin-
gual training, MrT5 adapts to the orthographic characteristics of each language,
learning language-specific compression rates. Furthermore, MrT5 shows compa-
rable accuracy to ByT5 on downstream evaluations such as XNLI, TyDi QA, and
character-level tasks while reducing sequence lengths by up to 75%. Our approach
presents a solution to the practical limitations of existing byte-level models.
https://github.com/jkallini/mrt5
1
INTRODUCTION
Subword tokenization, typically via algorithms such as byte-pair e...
Calibrated Physics-Informed Uncertainty Quantification
Vignesh Gopakumar 1 2 Ander Gray 3 Lorenzo Zanisi 2 Timothy Nunn 2 Daniel Giles 1 Matt J. Kusner 4 5
Stanislas Pamela 2 Marc Peter Deisenroth 1
Abstract
Simulating complex physical systems is cru-
cial for understanding and predicting phenomena
across diverse fields, such as fluid dynamics and
heat transfer, as well as plasma physics and struc-
tural mechanics. Traditional approaches rely on
solving partial differential equations (PDEs) using
numerical methods, which are computationally ex-
pensive and often prohibitively slow for real-time
applications or large-scale simulations. Neural
PDEs have emerged as efficient alternatives to
these costly numerical solvers, offering significant
computational speed-ups. However, their lack of
robust uncertainty quantification (UQ) limits de-
ployment in critical applications. We introduce a
model-agnostic, physics-informed conformal pre-
diction (CP) framework that provides guaranteed
uncertainty estimates without requiring labelled
data. By utilising a physics-based approach, we
can quantify and calibrate the model’s inconsisten-
cies with the physics rather than the uncertainty
arising from the data. Our approach utilises con-
volutional layers as finite-difference stencils and
leverages physics residual errors as nonconfor-
mity scores, enabling data-free UQ with marginal
and joint coverage guarantees across prediction
domains for a range of complex PDEs. We further
validate the efficacy of our method on neural PDE
models for plasma modelling and shot design in
fusion reactors.
1. Introduction
Numerical PDE solvers are essential tools in scientific and
engineering simulations (Danabasoglu et al., 2020; Giudi-
1Centre for Artificial Intelligence, University College Lon-
don
2Computing Division,
UK Atomic Energy Authority
3Heudiasyc Laboratory 4Polytechnique Montr´eal 5Mila - Que-
bec AI Institute.
Correspondence to:
Vignesh Gopakumar
<v.gopakumar@ucl.ac.uk>.
Proceedings...
Selective Preference Aggregation
Shreyas Kadekodi 1 * Hayden McTavish 2 * Berk Ustun 1
Abstract
Many applications in machine learning and de-
cision making rely on procedures to aggregate
human preferences. In such tasks, individuals ex-
press ordinal preferences over a set of items by
voting, ratings, or comparing them. We then ag-
gregate these data into a ranking that reveals their
collective preferences. Standard methods for pref-
erence aggregation are designed to return rankings
that arbitrate conflicting preferences between in-
dividuals. In this work, we introduce a paradigm
for selective aggregation where we abstain from
comparison rather than arbitrate dissent. We sum-
marize collective preferences as a selective rank-
ing – i.e., a partial order that reflects all collective
preferences where at least 100 · (1 −τ)% of in-
dividuals agree. We develop algorithms to build
selective rankings that achieve all possible trade-
offs between comparability and disagreement, and
derive formal guarantees on their recovery and
robustness. We conduct an extensive set of ex-
periments on real-world datasets to benchmark
our approach and demonstrate its functionality.
Our results show how selective rankings can pro-
mote transparency and robustness by revealing
disagreement and abstaining from arbitration.
1
Introduction
Many of our most important systems rely on procedures
where we elicit and aggregate human preferences.In such
systems, we ask a group of individuals to express their
preferences over a set of items through votes, ratings, or
pairwise comparisons. We then use these data to order items
in a way that represents their collective preferences as a
group. Over the past century, we have applied this pattern to
reap transformative benefits from collective intelligence in
elections [21], online search [26], and model alignment [22].
*Equal contribution 1UCSD 2Duke University. Correspondence
to: Berk Ustun <berk@ucsd.edu>.
Proceedings of the 42 nd International Conf...
GMAIL: Generative Modality Alignment for generated Image Learning
Shentong Mo 1 2 Sukmin Yun 3
Abstract
Generative models have made it possible to syn-
thesize highly realistic images, potentially provid-
ing an abundant data source for training machine
learning models. Despite the advantages of these
synthesizable data sources, the indiscriminate use
of generated images as real images for training
can even cause mode collapse due to modality
discrepancies between real and synthetic domains.
In this paper, we propose a novel framework for
discriminative use of generated images, coined
GMAIL, that explicitly treats generated images as
a separate modality from real images. Instead of
indiscriminately replacing real images with gener-
ated ones in the pixel space, our approach bridges
the two distinct modalities in the same latent space
through a multi-modal learning approach. To be
specific, we first fine-tune a model exclusively on
generated images using a cross-modality align-
ment loss and then employ this aligned model to
further train various vision-language models with
generated images. By aligning the two modalities,
our approach effectively leverages the benefits
of recent advances in generative models, thereby
boosting the effectiveness of generated image
learning across a range of vision-language tasks.
Our framework can be easily incorporated with
various vision-language models, and we demon-
strate its efficacy throughout extensive experi-
ments. For example, our framework significantly
improves performance on image captioning, zero-
shot image retrieval, zero-shot image classifica-
tion, and long caption retrieval tasks. It also shows
positive generated data scaling trends and notable
enhancements in the captioning performance of
the large multimodal model, LLaVA.
1Department of Machine Learning, CMU, USA 2Department
of Machine Learning, MBZUAI, UAE 3Department of Artificial
Intelligence, Hanyang University ERICA, South Korea. Corre-
spondence to: Sukmin ...