dataset-review/dataset-review huggingface dataset:
DEALing with Image Reconstruction: Deep Attentive Least Squares
ICML 2025

DEALing with Image Reconstruction: Deep Attentive Least Squares
Mehrsa Pourya 1 Erich Kobler 2 Michael Unser 1 Sebastian Neumayer 3
Abstract
State-of-the-art image reconstruction often relies
on complex, abundantly parameterized deep ar-
chitectures. We propose an alternative: a data-
driven reconstruction method inspired by the clas-
sic Tikhonov regularization. Our approach itera-
tively refines intermediate reconstructions by solv-
ing a sequence of quadratic problems. These up-
dates have two key components: (i) learned filters
to extract salient image features; and (ii) an atten-
tion mechanism that locally adjusts the penalty of
the filter responses. Our method matches leading
plug-and-play and learned regularizer approaches
in performance while offering interpretability, ro-
bustness, and convergent behavior. In effect, we
bridge traditional regularization and deep learning
with a principled reconstruction approach.
Code: https://github.com/mehrsapo/DEAL.
1. Introduction
Image reconstruction plays a fundamental role in compu-
tational imaging and computer vision (McCann & Unser,
2019; Zeng, 2001). The task is to recover an unknown im-
age of interest x ∈Rd from noisy measurements y ∈RM.
Their relation is modeled as y = Hx, where the forward op-
erator H ∈RM×d encodes the acquisition process. When
H is ill-conditioned, one can resort to the regularized recon-
struction
ˆx ∈arg min
x∈Rd

∥Hx −y∥2
2 + λR(x)

.
(1)
The data fidelity ∥Hx −y∥2
2 controls the consistency of
the reconstruction with the measurements. The regularizer
R: Rd →R≥0 encodes prior information about the solu-
tion and is also intended to make the problem well-posed.
1Biomedical Imaging group, EPFL Lausanne, Switzerland
2Institute for Machine Learning, JKU Linz, Austria 3Faculty
of Mathematics, TU Chemnitz, Germany. Correspondence to:
Mehrsa Pourya <mehrsa.pourya@epfl.ch>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025...
Private Federated Learning using Preference-Optimized Synthetic Data
ICML 2025

Private Federated Learning using Preference-Optimized Synthetic Data
Charlie Hou * 1 Mei-Yu Wang * 2 Yige Zhu * Daniel Lazar 3 Giulia Fanti 1
Abstract
In practical settings, differentially private feder-
ated learning (DP-FL) is the dominant method
for training models from private, on-device client
data. Recent work has suggested that DP-FL may
be enhanced or outperformed by methods that use
DP synthetic data (Wu et al., 2024; Hou et al.,
2024). The primary algorithms for generating DP
synthetic data for FL applications require careful
prompt engineering based on public information
and/or iterative private client feedback. Our key
insight is that the private client feedback collected
by prior DP synthetic data methods (Hou et al.,
2024; Xie et al., 2024) can be viewed as an RL
reward. Our algorithm, Policy Optimization for
Private Data (POPri) harnesses client feedback
using policy optimization algorithms such as Di-
rect Preference Optimization (DPO) to fine-tune
LLMs to generate high-quality DP synthetic data.
To evaluate POPri, we release LargeFedBench, a
new federated text benchmark for uncontaminated
LLM evaluations on federated client data. POPri
closes the gap in performance between the fully-
private and non-private settings by up to 58%,
compared to 28% for prior synthetic data methods,
and 3% for state-of-the-art DP federated learn-
ing methods. The code and data are available at
https://github.com/meiyuw/POPri.
1. Introduction
Many important machine learning (ML) applications feature
sensitive datasets that are distributed across client devices
(e.g. mobile devices). Such ML models are often hosted
on client devices. These on-device models offer privacy, la-
tency, and storage benefits relative to centrally-hosted mod-
*Equal contribution
1Department of ECE, Carnegie Mellon
University, Pittsburgh, PA 2Pittsburgh Supercomputing Center,
Pittsburgh, USA 3Coldrays, Tucson, AZ. Correspondence to: Char-
lie Hou <hou.charlie2@gmail.com>.
Proceedings of the 42 nd...
Lossy Compression with Pretrained Diffusion Models
ICLR 2025

Published as a conference paper at ICLR 2025
LOSSY COMPRESSION WITH
PRETRAINED DIFFUSION MODELS
Jeremy Vonderfecht & Feng Liu
Department of Computer Science
Portland State University
{vonder2,fliu}@pdx.edu
ABSTRACT
We apply the DiffC algorithm (Theis et al., 2022) to Stable Diffusion 1.5, 2.1,
XL, and Flux-dev, and demonstrate that these pretrained models are remarkably
capable lossy image compressors. A principled algorithm for lossy compression
using pretrained diffusion models has been understood since at least (Ho et al.,
2020), but challenges in reverse-channel coding have prevented such algorithms
from ever being fully implemented. We introduce simple workarounds that lead
to the first complete implementation of DiffC, which is capable of compressing
and decompressing images using Stable Diffusion in under 10 seconds. Despite
requiring no additional training, our method is competitive with other state-of-the-
art generative compression methods at low ultra-low bitrates.
1
INTRODUCTION
Today’s AI labs are spending millions of dollars training state-of-the-art diffusion models. The
original Stable Diffusion cost around $600,000 to train, while Stable Diffusion 3 is rumored to have
cost around $10 million (Wikipedia, 2024). The sophistication and expense of these models are only
likely to increase as generative modelling expands into the video domain.
While diffusion models are used primarily for prompt-based image generation, they are also highly
powerful and flexible image priors. They provide many other affordances, such as easy multiplica-
tion with other distributions (i.e., posterior sampling), and log-likelihood estimation (Sohl-Dickstein
et al., 2015). Researchers have discovered many remarkable applications of pretrained diffusion
models, including image restoration (Chung et al., 2022), optical illusions (Geng et al., 2024b;a),
and generative 3D modeling (Poole et al., 2022).
Ho et al. (2020) proposed an algorithm by which pretrained diffusion models ca...
Diffusion on Language Model Encodings for Protein Sequence Generation
ICML 2025

Diffusion on Language Model Encodings for Protein Sequence Generation
Viacheslav Meshchaninov * 1 Pavel Strashnov * 2 Andrey Shevtsov * 2 Fedor Nikolaev 2 Nikita Ivanisenko 2
Olga Kardymon 2 Dmitry Vetrov 1
Abstract
Protein sequence design has seen significant ad-
vances through discrete diffusion and autore-
gressive approaches, yet the potential of con-
tinuous diffusion remains underexplored. Here,
we present DiMA, a latent diffusion framework
that operates on protein language model rep-
resentations.
Through systematic exploration
of architectural choices and diffusion compo-
nents, we develop a robust methodology that gen-
eralizes across multiple protein encoders rang-
ing from 8M to 3B parameters.
We demon-
strate that our framework achieves consistently
high performance across sequence-only (ESM-
2, ESMc), dual-decodable (CHEAP), and mul-
timodal (SaProt) representations using the same
architecture and training approach. We exten-
sively evaluate existing methods alongside DiMA
using multiple metrics across two protein modali-
ties, covering quality, diversity, novelty, and dis-
tribution matching of generated proteins. DiMA
consistently produces novel, high-quality and di-
verse protein sequences and achieves strong re-
sults compared to baselines such as autoregres-
sive, discrete diffusion and flow matching lan-
guage models. The model demonstrates versa-
tile functionality, supporting conditional gener-
ation tasks including protein family-generation,
motif scaffolding and infilling, and fold-specific
sequence design. This work provides a univer-
sal continuous diffusion framework for protein
sequence generation, offering both architectural
insights and practical applicability across vari-
ous protein design scenarios. Code is released at
GitHub.
*Equal contribution
1Constructor University, Bremen, Ger-
many 2AIRI, Moscow, Russia.
Correspondence to:
Viach-
eslav Meshchaninov <meshchaninov.viacheslav@gmail.com>,
Pavel
Strashnov
<strashnov@airi.net>,
Andre...
Robot-Gated Interactive Imitation Learning with Adaptive Intervention Mechanism
ICML 2025

Robot-Gated Interactive Imitation Learning with Adaptive Intervention
Mechanism
Haoyuan Cai 1 Zhenghao Peng 1 Bolei Zhou 1
Abstract
Interactive Imitation Learning (IIL) allows agents
to acquire desired behaviors through human in-
terventions, but current methods impose high
cognitive demands on human supervisors. We
propose the Adaptive Intervention Mechanism
(AIM), a novel robot-gated IIL algorithm that
learns an adaptive criterion for requesting hu-
man demonstrations. AIM utilizes a proxy Q-
function to mimic the human intervention rule
and adjusts intervention requests based on the
alignment between agent and human actions. By
assigning high Q-values when the agent deviates
from the expert and decreasing these values as the
agent becomes proficient, the proxy Q-function
enables the agent to assess the real-time align-
ment with the expert and request assistance when
needed. Our expert-in-the-loop experiments re-
veal that AIM significantly reduces expert moni-
toring efforts in both continuous and discrete con-
trol tasks. Compared to the uncertainty-based
baseline Thrifty-DAgger, our method achieves a
40% improvement in terms of human take-over
cost and learning efficiency. Furthermore, AIM
effectively identifies safety-critical states for ex-
pert assistance, thereby collecting higher-quality
expert demonstrations and reducing overall ex-
pert data and environment interactions needed.
Code and demo video are available at https:
//github.com/metadriverse/AIM.
1. Introduction
Reinforcement learning (RL) has been successful in a wide
variety of tasks, from playing Atari games (Mnih et al.,
2013) to autonomous driving (Li et al., 2022a) and robotic
1Computer Science Department, University of Califor-
nia, Los Angeles (UCLA). Correspondence to:
Bolei Zhou
<bolei@cs.ucla.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
I can have a rest.
I can handle it.
Can you help me?
...
Diversity By Design: Leveraging Distribution Matching for Offline Model-Based Optimization
ICML 2025

Diversity By Design: Leveraging Distribution Matching
for Offline Model-Based Optimization
Michael S. Yao 1 James C. Gee 1 Osbert Bastani 1
Abstract
The goal of offline model-based optimization
(MBO) is to propose new designs that maximize
a reward function given only an offline dataset.
However, an important desiderata is to also pro-
pose a diverse set of final candidates that capture
many optimal and near-optimal design config-
urations. We propose Diversity in Adversarial
Model-based Optimization (DynAMO) as a
novel method to introduce design diversity as an
explicit objective into any MBO problem. Our
key insight is to formulate diversity as a distribu-
tion matching problem where the distribution of
generated designs captures the inherent diversity
contained within the offline dataset. Extensive
experiments spanning multiple scientific domains
show that DynAMO can be used with common
optimization methods to significantly improve the
diversity of proposed designs while still discover-
ing high-quality candidates.
1. Introduction
Discovering designs that optimize certain desirable proper-
ties is a ubiquitous task that spans a wide range of scientific
and engineering domains. For example, we might seek
to design a drug with the most potent therapeutic efficacy
(Brown et al., 2019; Kong et al., 2023; Du et al., 2024); build
a robot that is most capable of navigating complex environ-
ments (Ahn et al., 2020; Trabucco et al., 2021; Wang et al.,
2023); or engineer a material with a certain desirable prop-
erty (Stanev et al., 2018; Pogue et al., 2023; Gashmard et al.,
2024; Ma et al., 2024). However, experimentally validating
every proposed design can be expensive, time-intensive, or
even impossible in many applications. These limitations
can preclude the use of conventional ‘online’ optimization
methods for such generative design tasks.
1University of Pennsylvania, Philadelphia, PA, USA. Corre-
spondence to: Michael Yao <myao2199@seas.upenn.edu>.
Proceedings of the...
BEST-Route: Adaptive LLM Routing with Test-Time Optimal Compute
ICML 2025

BEST-Route: Adaptive LLM Routing with Test-Time Optimal Compute
Dujian Ding 1 Ankur Mallick 2 Shaokun Zhang 3 Chi Wang 4 Daniel Madrigal 2
Mirian Del Carmen Hipolito Garcia 2 Menglin Xia 2 Laks V. S. Lakshmanan 5 Qingyun Wu 3 6 Victor R¨uhle 2
Abstract
Large language models (LLMs) are powerful
tools but are often expensive to deploy at scale.
LLM query routing mitigates this by dynamically
assigning queries to models of varying cost and
quality to obtain a desired tradeoff. Prior query
routing approaches generate only one response
from the selected model and a single response
from a small (inexpensive) model was often not
good enough to beat a response from a large (ex-
pensive) model due to which they end up overus-
ing the large model and missing out on poten-
tial cost savings. However, it is well known that
for small models, generating multiple responses
and selecting the best can enhance quality while
remaining cheaper than a single large-model re-
sponse. We leverage this idea to propose BEST-
Route, a novel routing framework that chooses
a model and the number of responses to sample
from it based on query difficulty and the quality
thresholds. Experiments on real-world datasets
demonstrate that our method reduces costs by up
to 60% with less than 1% performance drop.
1. Introduction
Large language models (LLMs) have revolutionized natural
language processing (NLP) by delivering state-of-the-art
performance across a wide range of tasks, from language
understanding to creative writing, code generation, and be-
yond (Zhao et al., 2023). Their widespread deployment in
applications like ChatGPT (OpenAI, a) and other conversa-
tional agents (Zheng et al., 2023; Zhang et al., 2024a;b) has
made them a cornerstone of modern NLP systems. However,
the superior performance of these models often comes with
1University of British Columbia (work performed during in-
ternship at Microsoft) 2Microsoft 3Pennsylvania State Univer-
sity 4Google (work performed while at Microsoft...
Oscillatory State-Space Models
ICLR 2025

Published as a conference paper at ICLR 2025
OSCILLATORY STATE-SPACE MODELS
T. Konstantin Rusch
MIT
tkrusch@mit.edu
Daniela Rus
MIT
ABSTRACT
We propose Linear Oscillatory State-Space models (LinOSS) for efficiently learn-
ing on long sequences. Inspired by cortical dynamics of biological neural net-
works, we base our proposed LinOSS model on a system of forced harmonic os-
cillators. A stable discretization, integrated over time using fast associative paral-
lel scans, yields the proposed state-space model. We prove that LinOSS produces
stable dynamics only requiring nonnegative diagonal state matrix. This is in stark
contrast to many previous state-space models relying heavily on restrictive param-
eterizations. Moreover, we rigorously show that LinOSS is universal, i.e., it can
approximate any continuous and causal operator mapping between time-varying
functions, to desired accuracy. In addition, we show that an implicit-explicit dis-
cretization of LinOSS perfectly conserves the symmetry of time reversibility of
the underlying dynamics. Together, these properties enable efficient modeling of
long-range interactions, while ensuring stable and accurate long-horizon forecast-
ing. Finally, our empirical results, spanning a wide range of time-series tasks from
mid-range to very long-range classification and regression, as well as long-horizon
forecasting, demonstrate that our proposed LinOSS model consistently outper-
forms state-of-the-art sequence models. Notably, LinOSS outperforms Mamba
and LRU by nearly 2x on a sequence modeling task with sequences of length 50k.
Code: https://github.com/tk-rusch/linoss.
1
INTRODUCTION
State-space models (Gu et al., 2021; Hasani et al., 2022; Smith et al., 2023; Orvieto et al., 2023) have
recently emerged as a powerful tool for learning on long sequences. These models posses the state-
fullness and fast inference capabilities of Recurrent Neural Networks (RNNs) together with many of
the benefits of Transformers (Vaswani, 2017; ...
AttentionPredictor: Temporal Patterns Matter for KV Cache Compression
NeurIPS 2025

AttentionPredictor: Temporal Patterns Matter
for KV Cache Compression
Qingyue Yang1*, Jie Wang1†, Xing Li2‡, Zhihai Wang1, Chen Chen2, Lei Chen2,
Xianzhi Yu2, Wulong Liu2, Jianye Hao2, 3, Mingxuan Yuan2, Bin Li1
1MoE Key Laboratory of Brain-inspired Intelligent Perception and Cognition,
University of Science and Technology of China
2Huawei Noah’s Ark Lab
3College of Intelligence and Computing, Tianjin University
{yangqingyue, zhwangx}@mail.ustc.edu.cn,
jiewangx@ustc.edu.cn, li.xing2@huawei.com
Abstract
With the development of large language models (LLMs), efficient inference through
Key-Value (KV) cache compression has attracted considerable attention, especially
for long-context generation. To compress the KV cache, recent methods identify
critical KV tokens through static modeling of attention scores. However, these
methods often struggle to accurately determine critical tokens as they neglect the
temporal patterns in attention scores, resulting in a noticeable degradation in LLM
performance. To address this challenge, we propose AttentionPredictor, which
is the first learning-based method to directly predict attention patterns for
KV cache compression and critical token identification. Specifically, Attention-
Predictor learns a lightweight, unified convolution model to dynamically capture
spatiotemporal patterns and predict the next-token attention scores. An appealing
feature of AttentionPredictor is that it accurately predicts the attention score and
shares the unified prediction model, which consumes negligible memory, among
all transformer layers. Moreover, we propose a cross-token critical cache prefetch-
ing framework that hides the token estimation time overhead to accelerate the
decoding stage. By retaining most of the attention information, AttentionPredictor
achieves 13× KV cache compression and 5.6× speedup in a cache offloading
scenario with comparable LLM performance, significantly outperforming the state-
of-the-arts. The code is available at https...
The Disparate Benefits of Deep Ensembles
ICML 2025

The Disparate Benefits of Deep Ensembles
Kajetan Schweighofer 1 Adrian Arnaiz-Rodriguez 2 Sepp Hochreiter 1 3 Nuria Oliver 2
Abstract
Ensembles of Deep Neural Networks, Deep En-
sembles, are widely used as a simple way to boost
predictive performance. However, their impact
on algorithmic fairness is not well understood yet.
Algorithmic fairness examines how a model’s per-
formance varies across socially relevant groups
defined by protected attributes such as age, gen-
der, or race. In this work, we explore the interplay
between the performance gains from Deep Ensem-
bles and fairness. Our analysis reveals that they
unevenly favor different groups, a phenomenon
that we term the disparate benefits effect. We
empirically investigate this effect using popular
facial analysis and medical imaging datasets with
protected group attributes and find that it affects
multiple established group fairness metrics, in-
cluding statistical parity and equal opportunity.
Furthermore, we identify that the per-group differ-
ences in predictive diversity of ensemble members
can explain this effect. Finally, we demonstrate
that the classical Hardt post-processing method is
particularly effective at mitigating the disparate
benefits effect of Deep Ensembles by leveraging
their better-calibrated predictive distributions.
1. Introduction
Deep Ensembles (Lakshminarayanan et al., 2017) have
demonstrated their efficacy as a simple and robust method
to improve the performance of individual Deep Neural Net-
works (DNNs). Their superior performance has made them
a popular choice for real-world applications (Bhusal et al.,
2021; Dolezal et al., 2022), including high-stakes scenar-
ios where the impact on people’s lives of Machine Learning
(ML) supported decisions can be profound, such as in health-
care, education, finance or the law. In such applications,
1ELLIS Unit, LIT AI Lab, Institute for Machine Learning,
JKU Linz, Austria 2ELLIS Alicante, Alicante, Spain 3NXAI
GmbH, Linz, Austria. Correspon...
pLSTM: parallelizable Linear Source Transition Mark networks
NeurIPS 2025

pLSTM: parallelizable Linear Source Transition Mark
networks
Korbinian Pöppel1,2
Richard Freinschlag1
Thomas Schmied1
Wei Lin1
Sepp Hochreiter1,3
1 ELLIS Unit Linz, LIT AI Lab, Institute for Machine Learning
Johannes Kepler University Linz, Austria
2 Zuse School ELIZA
3 NXAI GmbH
poeppel@ml.jku.at
Abstract
Modern recurrent architectures, such as xLSTM and Mamba, have recently chal-
lenged the Transformer in language modeling. However, their structure constrains
their applicability to sequences only or requires processing multi-dimensional
data structures, such as images or molecular graphs, in a pre-defined sequential
order. In contrast, Multi-Dimensional RNNs (MDRNNs) are well suited for data
with a higher level structure, like 2D grids, trees, and directed acyclic graphs
(DAGs). In this work, we extend the notion of multi-dimensionality to linear RNNs.
We introduce parallelizable Linear Source Transition Mark networks (pLSTMs)
using Source, Transition, and Mark gates that act on the linegraph of a general
DAG. This enables parallelization in analogy to parallel associative scans and the
chunkwise-recurrent form of sequential linear RNNs, but for DAGs. For regular
grids (1D and 2D), like images, this scheme can be efficiently implemented using
einsum operations, concatenations, and padding in logarithmic time. pLSTMs
tackle the vanishing/exploding activation/gradient problem for long distances in
DAGs via two distinct modes: a directed propagation mode (P-mode) and a dif-
fusive distribution mode (D-mode). To showcase the long-range capabilities of
pLSTM, we introduce arrow-pointing extrapolation as a synthetic computer vi-
sion task that contains long-distance directional information. We demonstrate
that pLSTMs generalize well to larger image sizes, whereas Transformers strug-
gle to extrapolate. On established molecular graph and computer vision bench-
marks, pLSTMs also show strong performance. The complete code is available at
https://github.com/ml-jku/plstm_ex...