AIS-Bench/AIS-Bench huggingface dataset:
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
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
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...
MDNS: Masked Diffusion Neural Sampler via
Stochastic Optimal Control
Yuchen Zhu1,∗, Wei Guo1,∗, Jaemoo Choi1, Guan-Horng Liu2, Yongxin Chen1, Molei Tao1
1Georgia Institute of Technology
2FAIR at Meta
{yzhu738, wei.guo, jchoi843, yongchen, mtao}@gatech.edu, ghliu@meta.com
Abstract
We study the problem of learning a neural sampler to generate samples from
discrete state spaces where the target probability mass function π ∝e−U is
known up to a normalizing constant, which is an important task in fields such as
statistical physics, machine learning, combinatorial optimization, etc. To better
address this challenging task when the state space has a large cardinality and
the distribution is multi-modal, we propose Masked Diffusion Neural Sampler
(MDNS), a novel framework for training discrete neural samplers by aligning
two path measures through a family of learning objectives, theoretically grounded
in the stochastic optimal control of the continuous-time Markov chains. We
validate the efficiency and scalability of MDNS through extensive experiments
on various distributions with distinct statistical properties, where MDNS learns to
accurately sample from the target distributions despite the extremely high problem
dimensions and outperforms other learning-based baselines by a large margin. A
comprehensive study of ablations and extensions is also provided to demonstrate
the efficacy and potential of the proposed framework. Our code is available at
https://github.com/yuchen-zhu-zyc/MDNS.
1
Introduction
Drawing samples from an unnormalized target distribution π ∝e−U on some state space X0 is a
fundamental problem with wide-ranging applications across fields such as statistical physics [LB14],
Bayesian inference [Gel+13], machine learning [And+03; WJ08], etc. For decades, Markov chain
Monte Carlo (MCMC) methods, such as the Langevin Monte Carlo, the Metropolis-Hastings al-
gorithm, and the Glauber dynamics, have been a cornerstone. These methods simulate a Markov
chain whose ...
FairICP: Encouraging Equalized Odds via Inverse Conditional Permutation
Yuheng Lai 1 Leying Guan 2
Abstract
Equalized odds, an important notion of algorith-
mic fairness, aims to ensure that sensitive vari-
ables, such as race and gender, do not unfairly
influence the algorithm’s prediction when condi-
tioning on the true outcome. Despite rapid ad-
vancements, current research primarily focuses
on equalized odds violations caused by a sin-
gle sensitive attribute, leaving the challenge of
simultaneously accounting for multiple attributes
under-addressed. We bridge this gap by intro-
ducing an in-processing fairness-aware learning
approach, FairICP, which integrates adversarial
learning with a novel inverse conditional permu-
tation scheme. FairICP offers a flexible and effi-
cient scheme to promote equalized odds under
fairness conditions described by complex and
multi-dimensional sensitive attributes. The ef-
ficacy and adaptability of our method are demon-
strated through both simulation studies and em-
pirical analyses of real-world datasets.
1. Introduction
Machine learning models are increasingly important in aid-
ing decision-making across various applications. Ensuring
fairness in these models—preventing discrimination against
minorities or other protected groups—remains a significant
challenge (Mehrabi et al., 2021). To address different needs,
several fairness metrics have been developed in the litera-
ture (Mehrabi et al., 2021; Castelnovo et al., 2022). The
equalized odds metric defines fairness by requiring that the
predicted outcome ˆY provides the same level of informa-
tion about the true response Y across different sensitive
attribute(s) A (e.g. gender/race/age) (Hardt et al., 2016):
ˆY ⊥⊥A | Y.
(1)
1Department of Statistics, University of Wisconsin-Madison
2Department of Biostatistics, Yale University. Correspondence to:
Leying Guan <leying.guan@yale.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR...
Emergent Symbolic Mechanisms Support
Abstract Reasoning in Large Language Models
Yukang Yang 1 Declan Campbell 2 Kaixuan Huang 1 Mengdi Wang 1 Jonathan Cohen 2 Taylor Webb 3
Abstract
Many recent studies have found evidence for
emergent reasoning capabilities in large language
models (LLMs), but debate persists concerning
the robustness of these capabilities, and the ex-
tent to which they depend on structured reason-
ing mechanisms. To shed light on these issues,
we study the internal mechanisms that support
abstract reasoning in LLMs. We identify an emer-
gent symbolic architecture that implements ab-
stract reasoning via a series of three computations.
In early layers, symbol abstraction heads con-
vert input tokens to abstract variables based on
the relations between those tokens. In interme-
diate layers, symbolic induction heads perform
sequence induction over these abstract variables.
Finally, in later layers, retrieval heads predict
the next token by retrieving the value associated
with the predicted abstract variable. These re-
sults point toward a resolution of the longstanding
debate between symbolic and neural network ap-
proaches, suggesting that emergent reasoning in
neural networks depends on the emergence of
symbolic mechanisms.
1. Introduction
Large language models (LLMs) have become the domi-
nant paradigm in artificial intelligence, but there is ongoing
debate concerning the limits and reliability of their capabili-
ties. One major focus of this debate has been the question
of whether they can reason systematically in an abstract or
human-like manner. Many studies have documented im-
pressive performance on various reasoning tasks (Wei et al.,
1Department of Electrical and Computer Engineering, Prince-
ton University, Princeton, NJ 2Princeton Neuroscience Insti-
tute, Princeton University, Princeton, NJ 3Microsoft Research,
New York, NY. Correspondence to:
Taylor Webb <tay-
lor.w.webb@gmail.com>.
Proceedings of the 42 nd International Conference on...
Zero-shot Meta-learning for Tabular Prediction Tasks
with Adversarially Pre-trained Transformer
Yulun Wu 1 Doron L. Bergman 2
Abstract
We present an Adversarially Pre-trained Trans-
former (APT) that is able to perform zero-shot
meta-learning on tabular prediction tasks without
pre-training on any real-world dataset, extend-
ing on the recent development of Prior-Data Fit-
ted Networks (PFNs) and TabPFN. Specifically,
APT is pre-trained with adversarial synthetic data
agents, who continue to shift their underlying data
generating distribution and deliberately challenge
the model with different synthetic datasets. In
addition, we propose a mixture block architecture
that is able to handle classification tasks with ar-
bitrary number of classes, addressing the class
size limitation – a crucial weakness of prior deep
tabular zero-shot learners. In experiments, we
show that our framework matches state-of-the-art
performance on small classification tasks without
filtering on dataset characteristics such as number
of classes and number of missing values, while
maintaining an average runtime under one second.
On common benchmark dataset suites in both clas-
sification and regression, we show that adversarial
pre-training was able to enhance TabPFN’s perfor-
mance. In our analysis, we demonstrate that the
adversarial synthetic data agents were able to gen-
erate a more diverse collection of data compared
to the ordinary random generator in TabPFN. In
addition, we demonstrate that our mixture block
neural design has improved generalizability and
greatly accelerated pre-training.
1. Introduction
In standard deep learning workflows, models are either
trained per dataset, or employed on data in a form com-
1University of California, Berkeley 2Capital One. Correspon-
dence to: Yulun Wu <yulun wu@berkeley.edu>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
patible with, and drawn from, th...
A Black-Box Debiasing Framework for Conditional
Sampling
Han Cui
University of Illinois at Urbana-Champaign
Champaign, IL
hancui5@illinois.edu
Jingbo Liu
University of Illinois at Urbana-Champaign
Champaign, IL
jingbol@illinois.edu
Abstract
Conditional sampling is a fundamental task in Bayesian statistics and genera-
tive modeling. Consider the problem of sampling from the posterior distribution
PX|Y =y∗for some observation y∗, where the likelihood PY |X is known, and we
are given n i.i.d. samples D = {Xi}n
i=1 drawn from an unknown prior distribution
πX. Suppose that f(ˆπXn) is the distribution of a posterior sample generated by
an algorithm (e.g. a conditional generative model or the Bayes rule) when ˆπXn
is the empirical distribution of the training data. Although averaging over the
randomness of the training data D, we have ED (ˆπXn) = πX, we do not have
ED {f(ˆπXn)} = f(πX) due to the nonlinearity of f, leading to a bias. In this
paper we propose a black-box debiasing scheme that improves the accuracy of
such a naive plug-in approach. For any integer k and under boundedness of the
likelihood and smoothness of f, we generate samples ˆX(1), . . . , ˆX(k) and weights
w1, . . . , wk such that Pk
i=1 wiP ˆ
X(i) is a k-th order approximation of f(πX), where
the generation process treats f as a black-box. Our generation process achieves
higher accuracy when averaged over the randomness of the training data, with-
out degrading the variance, which can be interpreted as improving memorization
without compromising generalization in generative models.
1
Introduction
Conditional sampling is a major task in Bayesian statistics and generative modeling. Given an
observation y∗, the objective is to draw samples from the posterior distribution PX|Y =y∗, where the
likelihood PY |X is known but the prior distribution πX is unknown. Instead, we are provided with a
dataset D = {Xi}n
i=1 consisting of n i.i.d. samples drawn from πX.
The setting is common in a wide range of applicati...
Published as a conference paper at ICLR 2025
TOWARDS HIERARCHICAL RECTIFIED FLOW
Yichi Zhang1,
Yici Yan1,
Alex Schwing1,
Zhizhen Zhao1
1University of Illinois Urbana-Champaign
ABSTRACT
We formulate a hierarchical rectified flow to model data distributions. It hierarchi-
cally couples multiple ordinary differential equations (ODEs) and defines a time-
differentiable stochastic process that generates a data distribution from a known
source distribution. Each ODE resembles the ODE that is solved in a classic recti-
fied flow, but differs in its domain, i.e., location, velocity, acceleration, etc. Unlike
the classic rectified flow formulation, which formulates a single ODE in the loca-
tion domain and only captures the expected velocity field (sufficient to capture a
multi-modal data distribution), the hierarchical rectified flow formulation models
the multi-modal random velocity field, acceleration field, etc., in their entirety.
This more faithful modeling of the random velocity field enables integration paths
to intersect when the underlying ODE is solved during data generation. Inter-
secting paths in turn lead to integration trajectories that are more straight than
those obtained in the classic rectified flow formulation, where integration paths
cannot intersect. This leads to modeling of data distributions with fewer neural
function evaluations. We empirically verify this on synthetic 1D and 2D data
as well as MNIST, CIFAR-10, and ImageNet-32 data. Our code is available at:
https://riccizz.github.io/HRF/.
1
INTRODUCTION
Diffusion models (Ho et al., 2020; Song et al., 2021a;b) and particularly also flow matching (Liu
et al., 2023; Lipman et al., 2023; Albergo & Vanden-Eijnden, 2023; Albergo et al., 2023) have
gained significant attention recently. This is partly due to impressive results that have been reported
across domains from computer vision (Ho et al., 2020) and medical imaging (Song et al., 2022) to
robotics (Kapelyukh et al., 2023) and computational biology...
Tensor-Var: Efficient Four-Dimensional Variational Data Assimilation
Yiming Yang 1 Xiaoyuan Cheng 1 Daniel Giles 1 Sibo Cheng 2 Yi He 1 Xiao Xue 1 Boli Chen 1 Yukun Hu 1
Abstract
Variational data assimilation estimates the dynam-
ical system states by minimizing a cost function
that fits the numerical models with the obser-
vational data. Although four-dimensional vari-
ational assimilation (4D-Var) is widely used, it
faces high computational costs in complex non-
linear systems and depends on imperfect state-
observation mappings. Deep learning (DL) offers
more expressive approximators, while integrating
DL models into 4D-Var is challenging due to their
nonlinearities and lack of theoretical guarantees
in assimilation results. In this paper, we propose
Tensor-Var, a novel framework that integrates ker-
nel conditional mean embedding (CME) with 4D-
Var to linearize nonlinear dynamics, achieving
convex optimization in a learned feature space.
Moreover, our method provides a new perspec-
tive for solving 4D-Var in a linear way, offering
theoretical guarantees of consistent assimilation
results between the original and feature spaces.
To handle large-scale problems, we propose a
method to learn deep features (DFs) using neural
networks within the Tensor-Var framework. Ex-
periments on chaotic systems and global weather
prediction with real-time observations show that
Tensor-Var outperforms conventional and DL hy-
brid 4D-Var baselines in accuracy while achieving
a 10- to 20-fold speed improvement.
1. Introduction
Forecasting of dynamical systems is an initial value problem
of practical significance. Many real-world systems, such as
the ocean and atmosphere, are chaotic, which means that
minor errors in current estimations in computational models
can lead to rapid divergence and substantial forecast errors
1University College London, London, United Kingdom
2CEREA, ENPC and EDF R&D, Institut Polytechnique de
Paris, Paris, France.
Correspondence to: Yiming Yang <zc-
ahyy1...
Orthogonal Subspace Decomposition for Generalizable AI-Generated Image
Detection
Zhiyuan Yan * 1 2 Jiangming Wang * 2 Peng Jin 1 Ke-Yue Zhang 2 Chengchun Liu 1 Shen Chen 2 Taiping Yao 2
Shouhong Ding 2 Baoyuan Wu 3 Li Yuan 1
Abstract
Detecting AI-generated images (AIGIs), such as
natural images or face images, has become in-
creasingly important yet challenging. In this pa-
per, we start from a new perspective to exca-
vate the reason behind the failure generalization
in AIGI detection, named the asymmetry phe-
nomenon, where a naively trained detector tends
to favor overfitting to the limited and monotonous
fake patterns, causing the feature space to be-
come highly constrained and low-ranked, which
is proved seriously limiting the expressivity and
generalization. One potential remedy is incorpo-
rating the pre-trained knowledge within the vision
foundation models (higher-ranked) to expand the
feature space, alleviating the model’s overfitting
to fake. To this end, we employ Singular Value
Decomposition (SVD) to decompose the original
feature space into two orthogonal subspaces. By
freezing the principal components and adapting
only the remained components, we preserve the
pre-trained knowledge while learning fake pat-
terns. Compared to existing full-parameters and
LoRA-based tuning methods, we explicitly en-
sure orthogonality, enabling the higher rank of
the whole feature space, effectively minimizing
overfitting and enhancing generalization. We fi-
nally identify a crucial insight: our method im-
plicitly learns a vital prior that fakes are actually
derived from the real, indicating a hierarchical
relationship rather than independence. Modeling
this prior, we believe, is essential for achieving
superior generalization. Our codes are publicly
available at GitHub.
*Equal contribution
1Peking University Shenzhen Gradu-
ate School 2Tencent Youtu Lab 3The Chinese University of
Hong Kong, Shenzhen. Correspondence to: Taiping Yao <taip-
ingyao@tencent.cn>, Li Yuan <yu...
SparseLoRA: Accelerating LLM Fine-Tuning with Contextual Sparsity
Samir Khaki * 1 Xiuyu Li * † 2 Junxian Guo * 3 Ligeng Zhu 3 Konstantinos N. Plataniotis 1
Amir Yazdanbakhsh 4 Kurt Keutzer 2 Song Han 3 Zhijian Liu 3
https://z-lab.ai/projects/sparselora
Abstract
Fine-tuning LLMs is both computationally and
memory-intensive. While parameter-efficient fine-
tuning methods, such as QLoRA and DoRA,
reduce the number of trainable parameters and
lower memory usage, they do not decrease com-
putational cost. In some cases, they may even
slow down fine-tuning. In this paper, we introduce
SparseLoRA, a method that accelerates LLM fine-
tuning through contextual sparsity. We propose
a lightweight, training-free SVD sparsity estima-
tor that dynamically selects a sparse subset of
weights for loss and gradient computation. Also,
we systematically analyze and address sensitivity
across layers, tokens, and training steps. Our ex-
perimental results show that SparseLoRA reduces
computational cost by up to 2.2× and a measured
speedup of up to 1.6× while maintaining accu-
racy across various downstream tasks, including
commonsense and arithmetic reasoning, code gen-
eration, and instruction following.
1. Introduction
Large language models (LLMs) are trained on vast, general-
domain datasets. They are often fine-tuned to improve their
performance in specific domains (Saab et al., 2024) or to
align their predictions with user preferences (Zhang et al.,
2024a). However, fine-tuning very large models can be pro-
hibitively expensive, both in terms of memory requirements
and computational costs.
Extensive efforts have been made in parameter-efficient fine-
tuning (PEFT) to reduce memory consumption of LLM fine-
tuning. LoRA (Hu et al., 2022) represents weight updates
using low-rank approximations. Building upon this, many
*Equal contribution †Project lead 1University of Toronto 2UC
Berkeley 3MIT 4Google DeepMind. Correspondence to: Samir
Khaki <samir.khaki@mail.utoronto.ca>.
Proceedings o...