AIS-Bench/AIS-Bench huggingface dataset:
Online Laplacian-Based Representation Learning in Reinforcement Learning
Maheed H. Ahmed 1 Jayanth Bhargav 1 Mahsa Ghasemi 1
Abstract
Representation learning plays a crucial role in re-
inforcement learning, especially in complex envi-
ronments with high-dimensional and unstructured
states. Effective representations can enhance the
efficiency of learning algorithms by improving
sample efficiency and generalization across tasks.
This paper considers the Laplacian-based frame-
work for representation learning, where the eigen-
vectors of the Laplacian matrix of the underlying
transition graph are leveraged to encode mean-
ingful features from raw sensory observations of
the states. Despite the promising algorithmic ad-
vances in this framework, it remains an open ques-
tion whether the Laplacian-based representations
can be learned online and with theoretical guaran-
tees along with policy learning. We address this
by formulating an online optimization approach
using the Asymmetric Graph Drawing Objective
(AGDO) and analyzing its convergence via online
projected gradient descent under mild assump-
tions. Our extensive simulation studies empiri-
cally validate the convergence guarantees to the
true Laplacian representation. Furthermore, we
provide insights into the compatibility of differ-
ent reinforcement learning algorithms with online
representation learning.
1. Introduction
Representation learning is an important part of machine
learning that involves learning compact and useful represen-
tations of data. The quality of these representations signifi-
cantly impacts the performance and efficiency of machine
learning algorithms (Bengio et al., 2013). In reinforcement
learning (RL), agents often deal with complex environments
characterized by high-dimensional and unstructured states.
1Electrical and Computer Engineering, Purdue University, West
Lafayette, IN 47907, USA. Correspondence to: Maheed H. Ahmed
<ahmed237@purdue.edu>.
Proceedings of the 42 nd International C...
Harnessing Heterogeneous Statistical Strength for Personalized Federated
Learning via Hierarchical Bayesian Inference
Mahendra Singh Thapa 1 Rui Li 1
Abstract
Personalized federated learning (PFL) based on
Bayesian approach tackle the challenges from sta-
tistical heterogeneity of client data by comput-
ing a personalized posterior distribution over the
parameters of each client’s local model and con-
structing a global distribution by aggregating the
parameters of these personalized posteriors. How-
ever, the heuristic aggregation methods introduce
strong biases and result in global models with
poor generalization. We thus propose a novel hi-
erarchical Bayesian inference framework for PFL
by specifying a conjugate hyper-prior over the
parameters of the personalized posteriors. This
allows us to jointly compute a global posterior
distribution for aggregation and the personalized
ones at local level. This hierarchical Bayesian
inference framework achieves elegant balance be-
tween local personalization and global model ro-
bustness. Extensive empirical study shows that by
effectively sharing the heterogeneous statistical
strength across the local models while retaining
their distinctive characteristics, our framework
yields state-of-the-art performance. We also show
that existing Bayesian PFLs are special cases of
our framework.
1. Introduction
Federated learning (FL) emerges as a promising framework
that enables collaborative model training across decentral-
ized client devices while preserving data privacy (McMahan
et al., 2017). This approach demonstrates significant poten-
tial in various real-world applications, such as healthcare
(Yang et al., 2021; Dayan et al., 2021; Yang et al., 2024),
recommendation systems (Muhammad et al., 2020; Yuan
1College of Computing and Information Sciences, Rochester
Institute of Technology, Rochester, New York, USA. Correspon-
dence to: Rui Li <rxlics@rit.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning...
Finding Wasserstein Ball Center: Efficient Algorithm and The Applications in
Fairness
Yuntao Wang * 1 Yuxuan li * 2 Qingyuan Yang 1 Hu Ding 1
Abstract
Wasserstein Barycenter (WB) is a fundamental
geometric optimization problem in machine learn-
ing, whose objective is to find a representative
probability measure that minimizes the sum of
Wasserstein distances to given distributions. WB
has a number of applications in various areas.
However, WB may lead to unfair outcome to-
wards underrepresented groups in some applica-
tions (e.g., a “minority” distribution may be far
away from the obtained WB under Wasserstein
distance). To address this issue, we propose an
alternative objective called “Wasserstein Ball Cen-
ter (WBC)”. Specifically, WBC is a distribution
that encompasses all input distributions within
the minimum Wasserstein distance, which can
be formulated as a “minmax” optimization prob-
lem. We show that the WBC problem with fixed
support is equivalent to solving a large-scale lin-
ear programming (LP) instance, which is quite
different from the previously studied LP model
for WB. By incorporating some novel observa-
tions on the induced normal equation, we propose
an efficient algorithm that accelerates the inte-
rior point method by O(min{N 2m, Nm2, m4})
times (“N” is the number of distributions and “m”
is the support size). Finally, we conduct a set
of experiments on both synthetic and real-world
datasets, demonstrating the computational effi-
ciency of our algorithm, and showing its ability
to provide more fairness for input distributions.
*Equal contribution 1School of Computer Science and Technol-
ogy, University of Science and Technology of China, Hefei, China
2School of Artificial Intelligence and Data Science, University of
Science and Technology of China, Hefei, China. Correspondence
to: Hu Ding <huding@ustc.edu.cn>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author...
EcoMapper: Generative Modeling for Climate-Aware Satellite Imagery
Muhammed Goktepe * 1 Amir Hossein Shamseddin * 1 Erencan Uysal 1 Javier Muinelo Monteagudo 1
Lukas Drees 2 Aysim Toker 3 Senthold Asseng 4 Malte von Bloh 4
Abstract
Satellite imagery is essential for Earth observa-
tion, enabling applications like crop yield pre-
diction, environmental monitoring, and climate
change assessment. However, integrating satellite
imagery with climate data remains a challenge,
limiting its utility for forecasting and scenario
analysis. We introduce a novel dataset of 2.9 mil-
lion Sentinel-2 images spanning 15 land cover
types with corresponding climate records, form-
ing the foundation for two satellite image genera-
tion approaches using fine-tuned Stable Diffusion
3 models. The first is a text-to-image generation
model that uses textual prompts with climate and
land cover details to produce realistic synthetic im-
agery for specific regions. The second leverages
ControlNet for multi-conditional image genera-
tion, preserving spatial structures while mapping
climate data or generating time-series to simulate
landscape evolution. By combining synthetic im-
age generation with climate and land cover data,
our work advances generative modeling in remote
sensing, offering realistic inputs for environmen-
tal forecasting and new possibilities for climate
adaptation and geospatial analysis.
1. Introduction
Satellite imagery is a critical resource for a variety of re-
search and industrial applications, ranging from environ-
mental monitoring and climate research (de Ara´ujo et al.,
*Equal contribution 1Technical University of Munich, School of
Computation, Information and Technology, Germany 2University
of Zurich, Department of Mathematical Modeling and Machine
Learning, EcoVision Lab, Switzerland 3Technical University of
Munich, School of Computation, Information and Technology, Dy-
namic Vision and Learning Group, Germany 4Technical University
of Munich, School of Life Scien...
Published as a conference paper at ICLR 2025
PROGRESSIVE COMPRESSION WITH
UNIVERSALLY QUANTIZED DIFFUSION MODELS
Yibo Yang∗
Justus C. Will∗
Stephan Mandt
Department of Computer Science
University of California, Irvine
{yibo.yang, jcwill, mandt}@uci.edu
ABSTRACT
Diffusion probabilistic models have achieved mainstream success in many gen-
erative modeling tasks, from image generation to inverse problem solving. A
distinct feature of these models is that they correspond to deep hierarchical latent
variable models optimizing a variational evidence lower bound (ELBO) on the
data likelihood. Drawing on a basic connection between likelihood modeling and
compression, we explore the potential of diffusion models for progressive coding,
resulting in a sequence of bits that can be incrementally transmitted and decoded
with progressively improving reconstruction quality. Unlike prior work based on
Gaussian diffusion or conditional diffusion models, we propose a new form of
diffusion model with uniform noise in the forward process, whose negative ELBO
corresponds to the end-to-end compression cost using universal quantization. We
obtain promising first results on image compression, achieving competitive rate-
distortion and rate-realism results on a wide range of bit-rates with a single model,
bringing neural codecs a step closer to practical deployment. Our code can be
found at https://github.com/mandt-lab/uqdm.
1
INTRODUCTION
A diffusion probabilistic model can be equivalently viewed as a deep latent-variable model (Sohl-
Dickstein et al., 2015; Ho et al., 2020; Kingma et al., 2021), a cascade of denoising autoencoders that
perform score matching at different noise levels (Vincent, 2011; Song & Ermon, 2019), or a neural
SDE (Song et al., 2021b). Here we take the latent-variable model view and explore the potential of
diffusion models for communicating information. Given the strong performance of these models on
likelihood estimation (Kingma et al., 2021; Nichol & Dhariwal, 202...
Incentivizing LLMs to Self-Verify Their Answers
Fuxiang Zhang1,2
Jiacheng Xu1,2
Chaojie Wang2
Ce Cui2
Yang Liu2
Bo An1,2∗
1 Nanyang Technological University, Singapore
2 Skywork AI
Abstract
Large Language Models (LLMs) have demonstrated remarkable progress in com-
plex reasoning tasks through both post-training and test-time scaling laws. While
prevalent test-time scaling approaches are often realized by using external reward
models to guide the model generation process, we find that only marginal gains
can be acquired when scaling a model post-trained on specific reasoning tasks.
We identify that the limited improvement stems from distribution discrepancies
between the specific post-trained generator and the general reward model. To
address this, we propose a framework that incentivizes LLMs to self-verify their
own answers. By unifying answer generation and verification within a single
reinforcement learning (RL) process, we train models that can effectively assess
the correctness of their own solutions. The trained model can further scale its
performance at inference time by verifying its generations, without the need for
external verifiers. We train our self-verification models based on Qwen2.5-Math-
7B and DeepSeek-R1-Distill-Qwen-1.5B, demonstrating their capabilities across
varying reasoning context lengths. Experiments on multiple mathematical rea-
soning benchmarks show that our models can not only improve post-training
performance but also enable effective test-time scaling. Our code is available at
https://github.com/mansicer/self-verification.
MATH500 Score (avg@1)
Base Model: Qwen2.5-Math-7B
AIME24 Score (avg@10)
Base Model: DeepSeek-R1-Distill-Qwen-1.5B
Figure 1: Average performance of post-trained models (dotted lines) and test-time scaling methods
(solid lines) on the MATH500 and AIME24 benchmarks. Our self-verification framework not only
enhances post-training performance with RL on both problem-solving and verification, but also
enables effective t...
Attributes Shape the Embedding Space of Face Recognition Models
Pierrick Leroy* 1 Antonio Mastropietro* 2 Marco Nurisso* 1 3 Francesco Vaccarino 1
Abstract
Face Recognition (FR) tasks have made signif-
icant progress with the advent of Deep Neu-
ral Networks, particularly through margin-based
triplet losses that embed facial images into high-
dimensional feature spaces. During training, these
contrastive losses focus exclusively on identity
information as labels. However, we observe a
multiscale geometric structure emerging in the
embedding space, influenced by interpretable fa-
cial (e.g., hair color) and image attributes (e.g.,
contrast). We propose a geometric approach to
describe the dependence or invariance of FR mod-
els to these attributes and introduce a physics-
inspired alignment metric. We evaluate the pro-
posed metric on controlled, simplified models and
widely used FR models fine-tuned with synthetic
data for targeted attribute augmentation. Our find-
ings reveal that the models exhibit varying degrees
of invariance across different attributes, provid-
ing insight into their strengths and weaknesses
and enabling deeper interpretability. Code avail-
able here: https://github.com/mantonios107/attrs-
fr-embs.
1. Introduction
Deep Learning (DL) models have achieved state-of-the-art
performance for Face Recognition (FR) tasks, due to in-
creased availability of curated datasets (Deng et al., 2019b)
and improved recognition architectures and losses (Deng
et al., 2022). In the challenging open-set scenario, novel
face identities can appear at testing time, hence the problem
is framed as a metric learning task (Liu et al., 2017). There
are three main innovation to reach outstanding FR results in
this scenario. The first consists of learning an embedding
*Equal contribution
1Department of Mathematical Sciences,
Politecnico di Torino, Turin, Italy 2Department of Computer
Science, University of Pisa, Pisa, Italy
3CENTAI Institute,
Turin, Italy. Correspondence to:...
DeFoG: Discrete Flow Matching for Graph Generation
Yiming Qin * 1 Manuel Madeira * 1 Dorina Thanou 1 Pascal Frossard 1
Abstract
Graph generative models are essential across
diverse scientific domains by capturing complex
distributions over relational data. Among them,
graph diffusion models achieve superior perfor-
mance but face inefficient sampling and limited
flexibility due to the tight coupling between
training and sampling stages.
We introduce
DeFoG, a novel graph generative framework that
disentangles sampling from training, enabling
a broader design space for more effective and
efficient model optimization. DeFoG employs a
discrete flow-matching formulation that respects
the inherent symmetries of graphs.
We theo-
retically ground this disentangled formulation
by explicitly relating the training loss to the
sampling algorithm and showing that DeFoG
faithfully replicates the ground truth graph
distribution. Building on these foundations, we
thoroughly investigate DeFoG’s design space
and propose novel sampling methods that sig-
nificantly enhance performance and reduce the
required number of refinement steps. Extensive
experiments demonstrate state-of-the-art perfor-
mance across synthetic, molecular, and digital
pathology datasets, covering both unconditional
and conditional generation settings.
It also
outperforms most diffusion-based models with
just 5–10% of their sampling steps.
1. Introduction
Graph generation has become a fundamental task across
diverse fields, from molecular chemistry to social net-
work analysis, due to graphs’ capacity to represent com-
plex relationships and generate realistic structured data.
Diffusion-based graph generative models (Niu et al., 2020;
Jo et al., 2022), particularly those tailored for discrete data
(Vignac et al., 2022), have emerged as compelling ap-
*Equal contribution
1EPFL, Lausanne, Switzerland.
Cor-
respondence to: Yiming Qin <yiming.qin@epfl.ch>, Manuel
Madeira <manuel.madeira@epfl.ch>.
Proceedings of the 42 ...
Published as a conference paper at ICLR 2025
SCALABLE DECISION-MAKING IN STOCHASTIC EN-
VIRONMENTS THROUGH LEARNED TEMPORAL AB-
STRACTION
Baiting Luo1, Ava Pettet3, Aron Laszka2, Abhishek Dubey1, Ayan Mukhopadhyay1
1Vanderbilt University, 2Pennsylvania State University, 3Nissan Advanced Technology Center
{baiting.luo, abhishek.dubey, ayan.mukhopadhyay}@vanderbilt.edu
alaszka@psu.edu, ava.pettet@nissan-usa.com
ABSTRACT
Sequential decision-making in high-dimensional continuous action spaces, partic-
ularly in stochastic environments, faces significant computational challenges. We
explore this challenge in the traditional offline RL setting, where an agent must
learn how to make decisions based on data collected through a stochastic behav-
ior policy. We present Latent Macro Action Planner (L-MAP), which addresses
this challenge by learning a set of temporally extended macro-actions through
a state-conditional Vector Quantized Variational Autoencoder (VQ-VAE), effec-
tively reducing action dimensionality. L-MAP employs a (separate) learned prior
model that acts as a latent transition model and allows efficient sampling of plau-
sible actions. During planning, our approach accounts for stochasticity in both the
environment and the behavior policy by using Monte Carlo tree search (MCTS).
In offline RL settings, including stochastic continuous control tasks, L-MAP ef-
ficiently searches over discrete latent actions to yield high expected returns. Em-
pirical results demonstrate that L-MAP maintains low decision latency despite
increased action dimensionality. Notably, across tasks ranging from continuous
control with inherently stochastic dynamics to high-dimensional robotic hand ma-
nipulation, L-MAP significantly outperforms existing model-based methods and
performs on-par with strong model-free actor-critic baselines, highlighting the ef-
fectiveness of the proposed approach in planning in complex and stochastic envi-
ronments with high-dimensional action spaces.
1
INT...
Connecting Neural Models Latent Geometries with
Relative Geodesic Representations
Hanlin Yu1
University of Helsinki
Berfin Inal
University of Amsterdam
Georgios Arvanitidis
DTU
Søren Hauberg
DTU
Francesco Locatello
IST Austria
Marco Fumero1
IST Austria
Abstract
Neural models learn representations of high-dimensional data on low-dimensional
manifolds. Multiple factors, including stochasticities in the training process, model
architectures, and additional inductive biases, may induce different representations,
even when learning the same task on the same data. However, it has recently been
shown that when a latent structure is shared between distinct latent spaces, relative
distances between representations can be preserved, up to distortions. Building
on this idea, we demonstrate that exploiting the differential-geometric structure of
latent spaces of neural models, it is possible to capture precisely the transformations
between representational spaces trained on similar data distributions. Specifically,
we assume that distinct neural models parametrize approximately the same un-
derlying manifold, and introduce a representation based on the pullback metric
that captures the intrinsic structure of the latent space, while scaling efficiently
to large models. We validate experimentally our method on model stitching and
retrieval tasks, covering autoencoders and vision foundation discriminative models,
across diverse architectures, datasets, pretraining schemes and modalities. Code is
available at https://github.com/marc0git/RelativeGeodesics.
1
Introduction
Figure 1:
Neural models trained on similar
data learn parametrizations of the same mani-
fold. NNs learn parametrizations (D1, D2) of the
same underlying manifold Y up to isometries T.
Pulling back the metric from Y makes relative
geodesic representations invariant to transforma-
tions T between latent spaces Z1 and Z2.
Neural models learn meaningful representations
of high-dimensional data generalizing to many
task...
Potemkin Understanding in Large Language Models
Marina Mancoridis 1 Bec Weeks 2 Keyon Vafa 3 Sendhil Mullainathan 4
Abstract
Large language models (LLMs) are regularly eval-
uated using benchmark datasets. But what justi-
fies making inferences about an LLM’s capabili-
ties based on its answers to a curated set of ques-
tions? This paper first introduces a formal frame-
work to address this question. The key is to note
that the benchmarks used to test LLMs—such
as AP exams—are also those used to test people.
However, this raises an implication: such bench-
marks are only valid tests if LLMs misunderstand
concepts in ways that mirror human misunder-
standings. Otherwise, success on benchmarks
only demonstrates potemkin understanding: the
illusion of understanding driven by answers irrec-
oncilable with how any human would interpret
a concept. We present two procedures for quan-
tifying the existence of potemkins: one using a
specially designed benchmark in three domains,
the other using a general procedure that provides
a lower-bound on their prevalence. We find that
potemkins are ubiquitous across models, tasks,
and domains. We also find that these failures re-
flect not just incorrect understanding, but deeper
internal incoherence in concept representations.
1. Introduction
There has been a marked change in how we interpret ma-
chine learning evaluations. Today, large language models
(LLMs) are evaluated on benchmark datasets: curated ques-
tions with rubrics for grading. Success on these questions
is treated as evidence of broader conceptual understanding.
Previously, while benchmarks were also used to evaluate
supervised learning models, success was interpreted more
narrowly. A pathology classifier that performs well on X-
ray classification is not credited with an understanding of
1Massachusetts Institute of Technology 2University of Chicago
3Harvard University 4Massachusetts Institute of Technology. Cor-
respondence to: Marina Mancoridis <marinam@mit.edu>.
Proc...
Latent Variable Causal Discovery under Selection Bias
Haoyue Dai 1 2 Yiwen Qiu 1 Ignavier Ng 1 Xinshuai Dong 1 Peter Spirtes 1 Kun Zhang 1 2
Abstract
Addressing selection bias in latent variable causal
discovery is important yet underexplored, largely
due to a lack of suitable statistical tools: While
various tools beyond basic conditional indepen-
dencies have been developed to handle latent vari-
ables, none have been adapted for selection bias.
We make an attempt by studying rank constraints,
which, as a generalization to conditional inde-
pendence constraints, exploits the ranks of co-
variance submatrices in linear Gaussian models.
We show that although selection can significantly
complicate the joint distribution, interestingly, the
ranks in the biased covariance matrices still pre-
serve meaningful information about both causal
structures and selection mechanisms. We provide
a graph-theoretic characterization of such rank
constraints. Using this tool, we demonstrate that
the one-factor model, a classical latent variable
model, can be identified under selection bias. Sim-
ulations and real-world experiments confirm the
effectiveness of using our rank constraints.
1
Introduction
At the core of understanding complex systems lies causal
discovery, the identification of causal relations from obser-
vational data (Spirtes et al., 2000; Pearl, 2009). In many
real-world scenarios, the variables of interest are latent
constructs that cannot be directly observed or quantized,
while the observed variables serve merely as indirect
measurements. For instance, in psychological or political-
economic surveys, measured responses serve as proxies for
latent personality traits or political orientations. Recovering
the causal structure among these latent variables—referred
to as latent variable causal discovery—is essential for un-
derstanding and reasoning, yet remains a challenging task.
Furthermore, a typical assumption in causal discovery,
1Carnegie Mellon University 2Moham...