dataset-review/dataset-review huggingface dataset:
Integrating Intermediate Layer Optimization and Projected Gradient Descent for Solving Inverse Problems with Diffusion Models
ICML 2025

Integrating Intermediate Layer Optimization and Projected Gradient Descent
for Solving Inverse Problems with Diffusion Models
Yang Zheng 1 Wen Li 1 Zhaoqiang Liu 1
Abstract
Inverse problems (IPs) involve reconstructing sig-
nals from noisy observations. Recently, diffu-
sion models (DMs) have emerged as a powerful
framework for solving IPs, achieving remarkable
reconstruction performance. However, existing
DM-based methods frequently encounter issues
such as heavy computational demands and subop-
timal convergence. In this work, building upon
the idea of the recent work DMPlug (Wang et al.,
2024), we propose two novel methods, DMILO
and DMILO-PGD, to address these challenges.
Our first method, DMILO, employs intermediate
layer optimization (ILO) to alleviate the mem-
ory burden inherent in DMPlug. Additionally,
by introducing sparse deviations, we expand the
range of DMs, enabling the exploration of under-
lying signals that may lie outside the range of
the diffusion model. We further propose DMILO-
PGD, which integrates ILO with projected gra-
dient descent (PGD), thereby reducing the risk
of suboptimal convergence. We provide an intu-
itive theoretical analysis of our approaches under
appropriate conditions and validate their superi-
ority through extensive experiments on diverse
image datasets, encompassing both linear and
nonlinear IPs. Our results demonstrate signifi-
cant performance gains over state-of-the-art meth-
ods, highlighting the effectiveness of DMILO and
DMILO-PGD in addressing common challenges
in DM-based IP solvers.
1. Introduction
Inverse problems (IPs) represent a broad class of challenges
focused on reconstructing degraded signals, particularly
images, from noisy observations. These problems are of
1University of Electronic Science and Technology of China.
Correspondence to: Zhaoqiang Liu <zqliu12@gmail.com>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s...
ELMO : Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
ICML 2025

ELMO : Efficiency via Low-precision and Peak Memory Optimization
in Large Output Spaces
Jinbin Zhang * 1 Nasib Ullah * 1 Erik Schultheis 1 2 Rohit Babbar 3
Abstract
Large output spaces, also referred to as Extreme
multilabel classification (XMC), is a setting that
arises, e.g., in large-scale tagging and product-
to-product recommendation, and is characterized
by the number of labels ranging from hundreds
of thousands to millions. This means that the
linear classification head, usually only a tiny
fraction of the overall model, turns into the main
driver for compute and memory demand. Current
state-of-the-art XMC methods predominantly
rely on FP16-FP32 mixed-precision training,
which we show can be unstable, and inefficient
in terms of memory usage and computational
overhead.
Meanwhile, existing low-precision
methods typically retain higher precision for
the classification layer.
In this work, we
propose ELMO, a pure low-precision training
framework for XMC models using BFloat16
and Float8 data types.
By leveraging Ka-
han summation and stochastic rounding, we
demonstrate that XMC models can be effectively
trained entirely in Float8, without relying
on single-precision master weights or tensor
scaling. Low-precision training, combined with
our proposed memory optimizations—gradient
fusion
and
chunking—enables
significant
reductions in GPU memory usage. For example,
we train a 3-million-label XMC model with
only 6.6 GiB of GPU memory,
compared
to the 39.7 GiB required by the optimized
SOTA method, Renee (Jain et al., 2023) with-
out compromising accuracy. Code available at
https://github.com/xmc-aalto/elmo.
*Equal contribution 1Department of Computer Science, Aalto
University, Espoo, Finland 2IST Austria 3Department of Computer
Science, University of Bath, Bath, UK. Correspondence to: Jinbin
Zhang <jinbin.zhang@aalto.fi>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1. Intro...
Data Selection Matters: Towards Robust Instruction Tuning of Large Multimodal Models
NeurIPS 2025

Data Selection Matters: Towards Robust Instruction
Tuning of Large Multimodal Models
Xu Yang1
Chen Liu† 1
Ying Wei† 2
1 City University of Hong Kong
2 Zhejiang University
xyang337-c@my.cityu.edu.hk
chen.liu@cityu.edu.hk
ying.wei@zju.edu.cn
Abstract
Selecting a compact subset of visual instruction–following data has emerged as
an effective way to align large multimodal models with human intentions while
avoiding the high cost of full-dataset training. Yet we observe that both full-data
training and existing state-of-the-art data selection methods tend to inherit underly-
ing dataset biases such as position bias and spurious correlations, leading to biased
model behaviors. To address this issue, we introduce ARDS, a robustness-aware
targeted visual instruction-selection framework that explicitly mitigates these weak-
nesses, sidestepping the need for access to downstream data or time-consuming
gradient computation. Specifically, we first identify the worst-case evaluation sub-
groups through visual and textual task-specific perturbations. The robust training
mixture is then constructed by prioritizing samples that are semantically closer to
these subgroups in a rich multimodal embedding space. Extensive experiments
demonstrate that ARDS substantially boosts both robustness and data efficiency for
visual instruction tuning. We also showcase that the robust mixtures produced with
a smaller model transfer effectively to larger architectures. Our code and selected
datasets that have been demonstrated transferable across models are available at
https://github.com/xyang583/ARDS.
1
Introduction
Large multimodal models (LMMs) [23, 9, 30, 123, 64, 115, 75] have garnered significant attention
due to their strong zero-shot capabilities, enabling a wide range of applications in real-world scenarios.
Visual instruction tuning [64] has proven effective in enhancing these models’ abilities to follow user
instructions and reason deeply based on visual cues. Recent efforts have demons...
UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
ICML 2025

UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules
Ziyang Yu 1 Wenbing Huang 2 3 Yang Liu 1 4
Abstract
Molecular Dynamics (MD) simulations are essen-
tial for understanding the atomic-level behavior of
molecular systems, giving insights into their tran-
sitions and interactions. However, classical MD
techniques are limited by the trade-off between ac-
curacy and efficiency, while recent deep learning-
based improvements have mostly focused on
single-domain molecules, lacking transferability
to unfamiliar molecular systems. Therefore, we
propose Unified Simulator (UniSim), which lever-
ages cross-domain knowledge to enhance the un-
derstanding of atomic interactions. First, we em-
ploy a multi-head pretraining approach to learn a
unified atomic representation model from a large
and diverse set of molecular data. Then, based on
the stochastic interpolant framework, we learn the
state transition patterns over long timesteps from
MD trajectories, and introduce a force guidance
module for rapidly adapting to different chemical
environments. Our experiments demonstrate that
UniSim achieves highly competitive performance
across small molecules, peptides, and proteins.
1. Introduction
Molecular Dynamics (MD) simulations, an in silico method
used to comprehend the time evolution of molecular sys-
tems in given environments, serve as a fundamental and
essential tool in various fields like computational chem-
istry, pharmacology, material design, and condensed mat-
ter physics (Van Gunsteren & Berendsen, 1990; Lindorff-
Larsen et al., 2011; Hollingsworth & Dror, 2018; Lau et al.,
*Equal contribution
1Department of Computer Science and
Technology, Tsinghua University, Beijing, China 2Gaoling School
of Artificial intelligence, Renmin University of China, Beijing,
China 3Engineering Research Center of Next-Generation Intelli-
gent Search and Recommendation, MOE 4Institute for AI Industry
Research (AIR), Tsinghua University, Beijing, China. Correspon-
dence to...
Let LLM Tell What to Prune and How Much to Prune
ICML 2025

Let LLM Tell What to Prune and How Much to Prune
Mingzhe Yang * 1 Sihao Lin * 2 Changlin Li 3 Xiaojun Chang 1
Abstract
Large language models (LLMs) have revolution-
ized various AI applications. However, their bil-
lions of parameters pose significant challenges
for practical deployment. Structured pruning is
a hardware-friendly compression technique and
receives widespread attention. Nonetheless, exist-
ing literature typically targets a single structure
of LLMs. We observe that the structure units of
LLMs differ in terms of inference cost and func-
tionality. Therefore, pruning a single structure
unit in isolation often results in an imbalance be-
tween performance and efficiency. In addition,
previous works mainly employ a prescribed prun-
ing ratio. Since the significance of LLM modules
may vary, it is ideal to distribute the pruning load
to a specific structure unit according to its role
within LLMs. To address the two issues, we pro-
pose a pruning method that targets multiple LLM
modules with dynamic pruning ratios. Specifi-
cally, we find the intrinsic properties of LLMs
can guide us to determine the importance of each
module and thus distribute the pruning load on de-
mand, i.e., what to prune and how much to prune.
This is achieved by quantifying the complex inter-
actions within LLMs. Extensive experiments on
multiple benchmarks and LLM variants demon-
strate that our method effectively balances the
trade-off between efficiency and performance.
1. Introduction
Large language models (LLMs) have demonstrated excep-
tional performance in the field of natural language process-
ing (Mann et al., 2020; Thoppilan et al., 2022; Zhu et al.,
2024). Nonetheless, the billions of parameters (Zhao et al.,
2023) characteristic of LLMs substantially increase storage
requirements and computational overhead. Therefore, de-
*Equal contribution 1University of Science and Technology of
China 2RMIT University 3Stanford University. Correspondence
to: Xiaojun Chang <xjchang@ustc...
CALM: Consensus-Aware Localized Merging for Multi-Task Learning
ICML 2025

CALM: Consensus-Aware Localized Merging for Multi-Task Learning
Kunda Yan * 1 Min Zhang * 2 Sen Cui † 1 Zikun Qu 3 Bo Jiang 2 Feng Liu 4 Changshui Zhang † 1
Abstract
Model merging aims to integrate the strengths of
multiple fine-tuned models into a unified model
while preserving task-specific capabilities. Ex-
isting methods, represented by task arithmetic,
are typically classified into global- and local-
aware methods. However, global-aware meth-
ods inevitably cause parameter interference, while
local-aware methods struggle to maintain the ef-
fectiveness of task-specific details in the merged
model. To address these limitations, we propose
a Consensus-Aware Localized Merging (CALM)
method which incorporates localized informa-
tion aligned with global task consensus, ensur-
ing its effectiveness post-merging. CALM con-
sists of three key components: (1) class-balanced
entropy minimization sampling, providing a
more flexible and reliable way to leverage un-
supervised data; (2) an efficient-aware frame-
work, selecting a small set of tasks for sequential
merging with high scalability; (3) a consensus-
aware mask optimization, aligning localized bi-
nary masks with global task consensus and merg-
ing them conflict-free. Experiments demonstrate
the superiority and robustness of our CALM, sig-
nificantly outperforming existing methods and
achieving performance close to traditional MTL.
1. Introduction
Multi-task learning (MTL) facilitates knowledge transfer
across different tasks through a shared backbone, improving
*Equal contribution † Corresponding author 1Institute for Artifi-
cial Intelligence, Tsinghua University (THUAI), Beijing National
Research Center for Information Science and Technology (BN-
Rist), Department of Automation, Tsinghua University, Beijing,
P.R.China 2East China Normal University (ECNU) 3The Chinese
University of Hong Kong, Shenzhen 4School of Computing and In-
formation Systems, The University of Melbourne. Correspondence
to: Sen Cui <cuis@ma...
Online Optimization for Offline Safe Reinforcement Learning
NeurIPS 2025

Online Optimization for Offline Safe
Reinforcement Learning
Yassine Chemingui
Washington State University
yassine.chemingui@wsu.edu
Aryan Deshwal
University of Minnesota
adeshwal@umn.edu
Alan Fern
Oregon State University
alan.fern@oregonstate.edu
Thanh Nguyen-Tang
New Jersey Institute of Technology
thanh.nguyen@njit.edu
Janardhan Rao Doppa
Washington State University
jana.doppa@wsu.edu
Abstract
We study the problem of Offline Safe Reinforcement Learning (OSRL), where the
goal is to learn a reward-maximizing policy from fixed data under a cumulative
cost constraint. We propose a novel OSRL approach that frames the problem as a
minimax objective and solves it by combining offline RL with online optimization
algorithms. We prove the approximate optimality of this approach when integrated
with an approximate offline RL oracle and no-regret online optimization. We
also present a practical approximation that can be combined with any offline RL
algorithm, eliminating the need for offline policy evaluation. Empirical results
on the DSRL benchmark demonstrate that our method reliably enforces safety
constraints under stringent cost budgets, while achieving high rewards. The code is
available at https://github.com/yassineCh/O3SRL.
1
Introduction
Offline reinforcement learning (RL) [Levine et al., 2020] is a powerful paradigm to learn decision-
making policies from logged datasets without the need for additional interaction with the environment.
Offline RL has shown good success in domains including autonomous driving, robotics, control
systems, and black-box optimization [Fang et al., 2022, Lin et al., 2024, Li et al., 2024, Zhan et al.,
2022, Chemingui et al., 2024] where executing exploratory actions is not practical. However, in
safety-critical domains such as healthcare and smart grid, the decision-making agent needs to also
satisfy some cost/safety constraints. Such problems are studied under the sub-area referred to as
offline safe RL (OSRL) where the goal is to learn r...
DyCodeEval: Dynamic Benchmarking of Reasoning Capabilities in Code Large Language Models Under Data Contamination
ICML 2025

DyCodeEval: Dynamic Benchmarking of Reasoning Capabilities
in Code Large Language Models Under Data Contamination
Simin Chen 1 Pranav Pusarla 1 Baishakhi Ray 1
1Columbia University
Abstract
The rapid advancement of code large language
models (Code LLMs) underscores the critical
need for effective and transparent benchmarking
methods. However, current benchmarking pre-
dominantly relies on publicly available, human-
created datasets. The widespread use of these
static benchmark datasets makes the evaluation
process particularly susceptible to data contam-
ination—an unavoidable consequence of the ex-
tensive data collection processes employed dur-
ing LLM training.
Existing methods for ad-
dressing data contamination typically face sig-
nificant limitations, including reliance on sub-
stantial human effort and difficulty in manag-
ing class imbalances. To overcome these chal-
lenges, we propose DyCodeEval, a novel bench-
marking suite specifically designed to evaluate
Code LLMs under realistic contamination scenar-
ios. Given an initial seed programming problem,
DyCodeEval utilizes multiple agents to system-
atically extract and modify contextual informa-
tion without changing the core logic, generating
semantically equivalent variations. We introduce
a dynamic data generation method and conduct
extensive empirical studies on two seed datasets
involving 18 Code LLMs. The results demon-
strate that DyCodeEval effectively assesses the
reasoning capabilities of Code LLMs under con-
tamination conditions while producing diverse
problem variants, thereby ensuring robust and
consistent benchmarking outcomes. Our project
webpage can be found at this link1.
1Department of Computer Science, Columbia University. Cor-
respondence to: Simin Chen <sc5687@columbia.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1https://codekaleidoscope.github.io/
dycodeeval.html
1. Introduction
Large ...
COLA: Towards Efficient Multi-Objective Reinforcement Learning with Conflict Objective Regularization in Latent Space
NeurIPS 2025

COLA: Towards Efficient Multi-Objective
Reinforcement Learning with Conflict Objective
Regularization in Latent Space
Pengyi Li*, Hongyao Tang†, Yifu Yuan, Jianye Hao†, Zibin Dong, Yan Zheng
College of Intelligence and Computing, Tianjin University
Abstract
Many real-world control problems require continual policy adjustments to bal-
ance multiple objectives, which requires the acquisition of high-quality policies
to cover diverse preferences. Multi-Objective Reinforcement Learning (MORL)
provides a general framework to solve such problems. However, current MORL
methods suffer from high sample complexity, primarily due to the neglect of effi-
cient knowledge sharing and conflicts in optimization with different preferences.
To this end, this paper introduces a novel framework, Conflict Objective Regular-
ization in Latent Space (COLA). To enable efficient knowledge sharing, COLA
establishes a shared latent representation space for common knowledge, which
can avoid redundant learning under different preferences. Besides, COLA in-
troduces a regularization term for the value function to mitigate the negative ef-
fects of conflicting preferences on the value function approximation, thereby im-
proving the accuracy of value estimation. The experimental results across vari-
ous multi-objective continuous control tasks demonstrate the significant superi-
ority of COLA over the state-of-the-art MORL baselines. Code is available at
https://github.com/yeshenpy/COLA.
1
Introduction
Real-world problems often involve multiple performance metrics, where conflicts among these met-
rics make it infeasible to maximize all metrics simultaneously with a single policy. This necessitates
a trade-off among different metrics, requiring the learning of multiple high-quality policies to meet
diverse preferences [1–3]. For instance, in robot control problems, performance metrics may include
both forward speed and energy consumption. In such cases, two reward signals are returned, and
multipl...
Tortoise and Hare Guidance: Accelerating Diffusion Model Inference with Multirate Integration
NeurIPS 2025

Tortoise and Hare Guidance: Accelerating Diffusion
Model Inference with Multirate Integration
Yunghee Lee
Byeonghyun Pak
Junwha Hong
Hoseong Kim†
Agency for Defense Development
{yhl, byeonghyun_pak, qbit, hoseongkim}@add.re.kr
Abstract
In this paper, we propose Tortoise and Hare Guidance (THG), a training-free
strategy that accelerates diffusion sampling while maintaining high-fidelity gen-
eration. We demonstrate that the noise estimate and the additional guidance
term exhibit markedly different sensitivity to numerical error by reformulating
the classifier-free guidance (CFG) ODE as a multirate system of ODEs. Our
error-bound analysis shows that the additional guidance branch is more robust to
approximation, revealing substantial redundancy that conventional solvers fail to
exploit. Building on this insight, THG significantly reduces the computation of the
additional guidance: the noise estimate is integrated with the tortoise equation on
the original, fine-grained timestep grid, while the additional guidance is integrated
with the hare equation only on a coarse grid. We also introduce (i) an error-bound-
aware timestep sampler that adaptively selects step sizes and (ii) a guidance-scale
scheduler that stabilizes large extrapolation spans. THG reduces the number of
function evaluations (NFE) by up to 30% with virtually no loss in generation
fidelity (∆ImageReward ≤0.032) and outperforms state-of-the-art CFG-based
training-free accelerators under identical computation budgets. Our findings high-
light the potential of multirate formulations for diffusion solvers, paving the way for
real-time high-quality image synthesis without any model retraining. The source
code is available at https://github.com/yhlee-add/THG.
1
Introduction
Diffusion models (DMs) have become the state-of-the-art generative model for images [10, 40, 47]
and, more recently, for video [20, 1, 52, 21] and audio-visual content [5, 41]. Despite their impressive
quality, sampling is costly: each outp...
Inference-time Alignment in Continuous Space
NeurIPS 2025

Inference-time Alignment in Continuous Space
Yige Yuan1,2∗,
Teng Xiao3,4∗,
Yunfan Li1,2,
Bingbing Xu1†,
Shuchang Tao5,
Yunqi Qiu1,2,
Huawei Shen1,2,
Xueqi Cheng1,2
1Institute of Computing Technology, Chinese Academy of Sciences
2University of Chinese Academy of Sciences
3University of Washington, 4Allen Institute for AI, 5Alibaba Group
yuanyige20z@ict.ac.cn, tengxiao01@gmail.com, liyunfan24s@ict.ac.cn
Abstract
Aligning large language models with human feedback at inference time has received
increasing attention due to its flexibility. Existing methods rely on generating
multiple responses from the base policy for search using a reward model, which
can be considered as searching in a discrete response space. However, these
methods struggle to explore informative candidates when the base policy is weak
or the candidate set is small, resulting in limited effectiveness. In this paper, to
address this problem, we propose Simple Energy Adaptation (SEA), a simple yet
effective algorithm for inference-time alignment. In contrast to expensive search
over the discrete space, SEA directly adapts original responses from the base policy
toward the optimal one via gradient-based sampling in continuous latent space.
Specifically, SEA formulates inference as an iterative optimization procedure on an
energy function over actions in the continuous space defined by the optimal policy,
enabling simple and effective alignment. For instance, despite its simplicity, SEA
outperforms the second-best baseline with a relative improvement of up to 77.51%
on AdvBench and 16.36% on MATH. Our code is publicly available at this link.
1
Introduction
The alignment of large language models (LLMs) plays a crucial role to ensure the model outputs
meet human expectations and reflect human values [1, 2, 3, 4, 5]. Reinforcement Learning from
Human Feedback (RLHF) [2, 6] has emerged as a widely adopted method for LLM alignment. RLHF
typically involves training a reward model based on human feedback and sub...