AIS-Bench/AIS-Bench huggingface dataset:
Revisiting Convergence: Shuffling Complexity Beyond Lipschitz Smoothness
Qi He 1 Peiran Yu 2 Ziyi Chen 1 Heng Huang 1
Abstract
Shuffling-type gradient methods are favored in
practice for their simplicity and rapid empirical
performance. Despite extensive development of
convergence guarantees under various assump-
tions in recent years, most require the Lipschitz
smoothness condition, which is often not met in
common machine learning models. We highlight
this issue with specific counterexamples. To ad-
dress this gap, we revisit the convergence rates of
shuffling-type gradient methods without assuming
Lipschitz smoothness. Using our stepsize strategy,
the shuffling-type gradient algorithm not only con-
verges under weaker assumptions but also match
the current best-known convergence rates, thereby
broadening its applicability. We prove the conver-
gence rates for nonconvex, strongly convex, and
non-strongly convex cases, each under both ran-
dom reshuffling and arbitrary shuffling schemes,
under a general bounded variance condition. Nu-
merical experiments further validate the perfor-
mance of our shuffling-type gradient algorithm,
underscoring its practical efficacy.
1. Introduction
Gradient-based optimization has always been a critical area
due to its extensive practical applications in machine learn-
ing, including reinforcement learning (Sutton and Barto,
2018), hyperparameter optimization (Feurer and Hutter,
2019), and large language models (Radford et al., 2018).
While numerous gradient-based algorithms have been devel-
oped for convex functions (Nemirovskij and Yudin, 1983;
Nesterov, 2013; d’Aspremont et al., 2021), research on non-
convex functions has become particularly active in recent
years, driven by advances in deep learning. Notably, with
unbiased stochastic gradients and bounded variance, SGD
1Department of Computer Science, University of Maryland,
College Park. 2Department of Computer Science and Engineer-
ing, University of Texas Arlington. Correspo...
Algorithms with Calibrated Machine Learning Predictions
Judy Hanwen Shen 1 Ellen Vitercik 1 2 Anders Wikum 2
Abstract
The field of algorithms with predictions incorpo-
rates machine learning advice in the design of
online algorithms to improve real-world perfor-
mance. A central consideration is the extent to
which predictions can be trusted—while existing
approaches often require users to specify an aggre-
gate trust level, modern machine learning models
can provide estimates of prediction-level uncer-
tainty. In this paper, we propose calibration as
a principled and practical tool to bridge this gap,
demonstrating the benefits of calibrated advice
through two case studies: the ski rental and on-
line job scheduling problems. For ski rental, we
design an algorithm that achieves near-optimal
prediction-dependent performance and prove that,
in high-variance settings, calibrated advice offers
more effective guidance than alternative methods
for uncertainty quantification. For job scheduling,
we demonstrate that using a calibrated predictor
leads to significant performance improvements
over existing methods. Evaluations on real-world
data validate our theoretical findings, highlighting
the practical impact of calibration for algorithms
with predictions.
1. Introduction
In recent years, advances in machine learning (ML) mod-
els have inspired researchers to revisit the design of classic
online algorithms, incorporating insights from ML-based
advice to improve decision-making in real-world environ-
ments. This research area, termed algorithms with predic-
tions, seeks to design algorithms that are both robust to
worst-case inputs and achieve performance that improves
with prediction accuracy (a desideratum termed consistency)
(Lykouris & Vassilvitskii, 2018). Many learning-augmented
1Department of Computer Science, Stanford University, Stan-
ford, CA, USA 2Department of Management Science & Engineer-
ing, Stanford University, Stanford, CA, USA. Correspondence to:
Anders ...
Published as a conference paper at ICLR 2025
EFFICIENT SPARSE PCA VIA BLOCK-
DIAGONALIZATION
Alberto Del Pia & Dekun Zhou
University of Wisconsin-Madison
{delpia,dzhou44}@wisc.edu
Yinglun Zhu
University of California, Riverside
yzhu@ucr.edu
ABSTRACT
Sparse Principal Component Analysis (Sparse PCA) is a pivotal tool in data analy-
sis and dimensionality reduction. However, Sparse PCA is a challenging problem
in both theory and practice: it is known to be NP-hard and current exact methods
generally require exponential runtime. In this paper, we propose a novel framework
to efficiently approximate Sparse PCA by (i) approximating the general input
covariance matrix with a re-sorted block-diagonal matrix, (ii) solving the Sparse
PCA sub-problem in each block, and (iii) reconstructing the solution to the original
problem. Our framework is simple and powerful: it can leverage any off-the-shelf
Sparse PCA algorithm and achieve significant computational speedups, with a
minor additive error that is linear in the approximation error of the block-diagonal
matrix. Suppose g(k, d) is the runtime of an algorithm (approximately) solving
Sparse PCA in dimension d and with sparsity constant k. Our framework, when in-
tegrated with this algorithm, reduces the runtime to O
 d
d⋆· g(k, d⋆) + d2
, where
d⋆≤d is the largest block size of the block-diagonal matrix. For instance, integrat-
ing our framework with the Branch-and-Bound algorithm reduces the complexity
from g(k, d) = O(k3 · dk) to O(k3 · d · (d⋆)k−1), demonstrating exponential
speedups if d⋆is small. We perform large-scale evaluations on many real-world
datasets: for exact Sparse PCA algorithm, our method achieves an average speedup
factor of 100.50, while maintaining an average approximation error of 0.61%;
for approximate Sparse PCA algorithm, our method achieves an average speedup
factor of 6.00 and an average approximation error of -0.91%, meaning that our
method oftentimes finds better solutions.
1
INTRODUCTION
In this ...
Hierarchical Masked Autoregressive Models with Low-Resolution Token Pivots
Guangting Zheng 1 † Yehao Li 2 Yingwei Pan 2 Jiajun Deng 3 Ting Yao 2 Yanyong Zhang 1 Tao Mei 2
Abstract
Autoregressive models have emerged as a pow-
erful generative paradigm for visual generation.
The current de-facto standard of next token pre-
diction commonly operates over a single-scale
sequence of dense image tokens, and is incapable
of utilizing global context especially for early
tokens prediction. In this paper, we introduce
a new autoregressive design to model a hierar-
chy from a few low-resolution image tokens to
the typical dense image tokens, and delve into a
thorough hierarchical dependency across multi-
scale image tokens. Technically, we present a
Hierarchical Masked Autoregressive models (Hi-
MAR) that pivot on low-resolution image tokens
to trigger hierarchical autoregressive modeling
in a multi-phase manner. Hi-MAR learns to pre-
dict a few image tokens in low resolution, func-
tioning as intermediary pivots to reflect global
structure, in the first phase. Such pivots act as
the additional guidance to strengthen the next au-
toregressive modeling phase by shaping global
structural awareness of typical dense image to-
kens. A new Diffusion Transformer head is fur-
ther devised to amplify the global context among
all tokens for mask token prediction. Extensive
evaluations on both class-conditional and text-to-
image generation tasks demonstrate that Hi-MAR
outperforms typical AR baselines, while requiring
fewer computational costs. Code is available at
https://github.com/HiDream-ai/himar.
† This work was performed when Guangting Zheng was visit-
ing HiDream.ai as a research intern. 1University of Science and
Technology of China, Anhui, China 2HiDream.ai Inc, Beijing,
China 3The University of Adelaide, Adelaide, Australia. Corre-
spondence to: Yingwei Pan <panyw.ustc@gmail.com>, Yanyong
Zhang <yanyongz@ustc.edu.cn>.
Proceedings of the 42 nd International Conference on Machin...
Closed-form Solutions: A New Perspective on Solving Differential Equations
Shu Wei 1 2 Yanjie Li 1 Lina Yu 1 2 Weijun Li 1 3 4 Min Wu 1 Linjun Sun 1 Jingyi Liu 1 Hong Qin 1 3
Yusong Deng 1 2 Jufeng Han 1 3 Yan Pang 1 4
Abstract
The quest for analytical solutions to differential
equations has traditionally been constrained by
the need for extensive mathematical expertise.
Machine learning methods like genetic algorithms
have shown promise in this domain, but are hin-
dered by significant computational time and the
complexity of their derived solutions. This paper
introduces SSDE (Symbolic Solver for Differen-
tial Equations), a novel reinforcement learning-
based approach that derives symbolic closed-form
solutions for various differential equations. Eval-
uations across a diverse set of ordinary and par-
tial differential equations demonstrate that SSDE
outperforms existing machine learning methods,
delivering superior accuracy and efficiency in ob-
taining analytical solutions.
1. Introduction
Differential equations (DEs) are foundational to mathemat-
ics, physics, and the natural sciences, providing abstract
models for diverse physical phenomena. For example, Pois-
son’s equation governs the distribution of electrostatic poten-
tial, while the heat equation describes temperature variations
within an object (Evans, 2022). Solving these DEs is critical
for understanding and predicting the dynamics of complex
physical systems (Alkhadhr & Almekkawy, 2021; Savovi´c
et al., 2023).
Deriving analytical solutions requires rigorous analysis of
their existence, uniqueness, stability, and behavioral proper-
ties, often leveraging advanced mathematical theories and
1AnnLab, Institute of Semiconductors, Chinese Academy of
Sciences, Beijing, China 2College of Materials Science and Opto-
Electronic Technology, University of Chinese Academy of Sci-
ences, Beijing, China 3School of Integrated Circuits, University
of Chinese Academy of Sciences, Beijing, China 4School of
Industry-ed...
Alpha-SQL: Zero-Shot Text-to-SQL using Monte Carlo Tree Search
Boyan Li 1 Jiayi Zhang 1 Ju Fan 2 Yanwei Xu 3 Chong Chen 3 Nan Tang 1 Yuyu Luo 1
Abstract
Text-to-SQL, which enables natural language
interaction with databases, serves as a pivotal
method across diverse industries. With new, more
powerful large language models (LLMs) emerg-
ing every few months, fine-tuning has become
incredibly costly, labor-intensive, and error-prone.
As an alternative, zero-shot Text-to-SQL, which
leverages the growing knowledge and reason-
ing capabilities encoded in LLMs without task-
specific fine-tuning, presents a promising and
more challenging direction. To address this chal-
lenge, we propose Alpha-SQL, a novel approach
that leverages a Monte Carlo Tree Search (MCTS)
framework to iteratively infer SQL construction
actions based on partial reasoning states.
To
enhance the framework’s reasoning capabilities,
we introduce LLM-as-Action-Model to dynami-
cally generate SQL construction actions during
the MCTS process, steering the search toward
more promising SQL queries. Moreover, Alpha-
SQL employs a self-supervised reward function
to evaluate the quality of candidate SQL queries,
ensuring more accurate and efficient query gener-
ation. Experimental results show that Alpha-SQL
achieves 69.7% execution accuracy on the BIRD
development set, using a 32B open-source LLM
without fine-tuning. Alpha-SQL outperforms the
best previous zero-shot approach based on GPT-
4o by 2.5% on the BIRD development set. The
code is available at https://github.com/
HKUSTDial/Alpha-SQL.
1. Introduction
Text-to-SQL (a.k.a. NL2SQL) converts natural language
queries into SQL, simplifying access to relational databases
1The Hong Kong University of Science and Technology
(Guangzhou) 2Renmin University of China 3Huawei Technologies
Ltd. Correspondence to: Yuyu Luo <yuyuluo@hkust-gz.edu.cn>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
...
Lookahead Routing for Large Language Models
Canbin Huang1, Tianyuan Shi1, Yuhua Zhu1, Ruijun Chen1, Xiaojun Quan1,2,∗
1School of Computer Science and Engineering, Sun Yat-sen University, China
2Shenzhen Loop Area Institute, China
{huangcb3, shity6, zhuyh53, chenrj8}@mail2.sysu.edu.cn
quanxj3@mail.sysu.edu.cn
Abstract
Large language model (LLM) routers improve the efficiency of multi-model sys-
tems by directing each query to the most appropriate model while leveraging the
diverse strengths of heterogeneous LLMs. Most existing approaches frame routing
as a classification problem based solely on the input query. While this reduces
overhead by avoiding inference across all models, it overlooks valuable information
that could be gleaned from potential outputs and fails to capture implicit intent
or contextual nuances that often emerge only during response generation. These
limitations can result in suboptimal routing decisions, particularly for complex
or ambiguous queries that require deeper semantic understanding. To address
this challenge, we propose Lookahead, a routing framework that “foresees” po-
tential model outputs by predicting their latent representations and uses these
predictions to guide model selection, thus enabling more informed routing without
full inference. Within this framework, we implement two approaches based on
causal and masked language models. Empirical evaluations across seven public
benchmarks—spanning instruction following, mathematical reasoning, and code
generation—show that Lookahead consistently outperforms existing routing base-
lines, achieving an average performance gain of 7.7% over the state-of-the-art. Our
code is available at https://github.com/huangcb01/lookahead-routing.
+13.4
+8.6
+20.8
+13.1
AlpacaEval-2
Arena-Hard
MT-Bench
GSM8K
MATH
HumanEval
MBPP
18
28
38
12
20
28
4
19
34
10
20
30
32
35
38
60
64
68
51
52
53
Methods
SMOOTHIE
RouterDC (SmolLM2)
RouterDC (ModernBERT)
Lookahead (SmolLM2)
Lookahead (ModernBERT)
Task categories...
Eigenspectrum Analysis of Neural Networks without Aspect Ratio Bias
Yuanzhe Hu 1 Kinshuk Goel 2 3 Vlad Killiakov 4 Yaoqing Yang 2
Abstract
Diagnosing deep neural networks (DNNs) by an-
alyzing the eigenspectrum of their weights has
been an active area of research in recent years.
One of the main approaches involves measuring
the heavytailness of the empirical spectral densi-
ties (ESDs) of weight matrices. This analysis has
been shown to provide insights to help diagnose
whether a model is well-trained or undertrained,
and has been used to guide training methods in-
volving layer-wise hyperparameter assignment.
In this paper, we address an often-overlooked
challenge in estimating the heavytailness of these
ESDs: the impact of the aspect ratio of weight
matrices. We demonstrate that matrices of vary-
ing sizes (and aspect ratios) introduce a non-
negligible bias in estimating the heavytailness of
ESDs, leading to inaccurate model diagnosis and
layer-wise hyperparameter assignment. To over-
come this challenge, we propose FARMS (Fixed-
Aspect-Ratio Matrix Subsampling), a method that
normalizes the weight matrices by subsampling
submatrices with a fixed aspect ratio. Instead of
measuring the heavytailness of the original ESD,
we measure the average ESD of these subsampled
submatrices. We show that this method effectively
mitigates the aspect ratio bias. We validate our
approach across various optimization techniques
and application domains that involve eigenspec-
trum analysis of weights, including image classifi-
cation in computer vision (CV) models, scientific
machine learning (SciML) model training, and
large language model (LLM) pruning. Our results
show that despite its simplicity, FARMS uniformly
improves the accuracy of eigenspectrum analysis
1Department of Computer Science and Engineering, Univer-
sity of California, San Diego, CA, USA 2Department of Com-
puter Science, Dartmouth College, NH, USA 3Department of
Computer Science & Engineering, SRM Institute of ...
Efficiently Vectorized MCMC on Modern Accelerators
Hugh Dance 1 Pierre Glaser 1 Peter Orbanz 1 Ryan Adams 2
Abstract
With the advent of automatic vectorization tools
(e.g., JAX’s vmap), writing multi-chain MCMC
algorithms is often now as simple as invoking
those tools on single-chain code. Whilst conve-
nient, for various MCMC algorithms this results
in a synchronization problem—loosely speaking,
at each iteration all chains running in parallel must
wait until the last chain has finished drawing its
sample. In this work, we show how to design
single-chain MCMC algorithms in a way that
avoids synchronization overheads when vectoriz-
ing with tools like vmap, by using the framework
of finite state machines (FSMs). Using a simpli-
fied model, we derive an exact theoretical form of
the obtainable speed-ups using our approach, and
use it to make principled recommendations for
optimal algorithm design. We implement several
popular MCMC algorithms as FSMs, including
Elliptical Slice Sampling, HMC-NUTS, and De-
layed Rejection, demonstrating speed-ups of up
to an order of magnitude in experiments.
1. Introduction
Automatic vectorization is the act of transforming a func-
tion to handle batches of inputs without user intervention.
Implementations of automatic vectorization algorithms are
now available in many mainstream scientific computing li-
braries, and have dramatically simplified the task of running
multiple instances of a single algorithm concurrently. They
are routinely used to train neural networks (Flax, 2023) and
in other scientific applications, e.g., Schoenholz and Cubuk
(2021); Oktay et al. (2023); Pfau et al. (2020).
This paper focuses on the use of automatic vectorization
for Markov chain Monte Carlo (MCMC) algorithms. Tools
like JAX’s vmap provide a convenient way to run multiple
1Gatsby Unit, University College London, UK 2Department of
Computer Science, Princeton University, USA. Correspondence
to: Hugh Dance <hugh.dance.15@ucl.ac.uk>.
Proceedings of the 4...
Structured Sparse Transition Matrices
to Enable State Tracking in State-Space Models
Aleksandar Terzi´c1,2∗
aleksandar.terzic1@ibm.com
Nicolas Menet1,2∗
nicolas.menet@ibm.com
Michael Hersche1
michael.hersche@ibm.com
Thomas Hofmann2
thomas.hofmann@inf.ethz.ch
Abbas Rahimi1
abr@zurich.ibm.com
1IBM Research – Zurich, 2Department of Computer Science, ETH Zürich
Abstract
Modern state-space models (SSMs) often utilize structured transition matrices
which enable efficient computation but pose restrictions on the model’s expressivity,
as measured in terms of the ability to emulate finite-state automata (FSA). While
unstructured transition matrices are optimal in terms of expressivity, they come
at a prohibitively high compute and memory cost, even for moderate state sizes.
We propose a structured sparse parametrization of transition matrices in SSMs
that enables FSA state tracking with provably optimal state size and depth, while
keeping the computational cost of the recurrence comparable to that of diagonal
SSMs. Our method, PD-SSM, parametrizes the transition matrix as the product
of a column one-hot matrix (P) and a complex-valued diagonal matrix (D). As
a result, the computational cost of parallel scans scales linearly with the state
size. Theoretically, the model is BIBO-stable and can emulate any N-state FSA
with one layer of dimension N and a linear readout of size N × N, significantly
improving on all current structured SSM guarantees. Experimentally, the model
significantly outperforms a wide collection of modern SSM variants on various FSA
state tracking tasks. On multivariate time-series classification, it outperforms neural
controlled differential equations, a paradigm explicitly built for time-series analysis.
Finally, we integrate PD-SSM into a hybrid Transformer-SSM architecture and
demonstrate that the model can effectively track the states of a complex FSA in
which transitions are encoded into sets of variable-length English sentences. The
code is available...
On the creation of narrow AI:
hierarchy and nonlocality of neural network skills
Eric J. Michaud1,3∗
Asher Parker-Sartori2
Max Tegmark1,3
1 Department of Physics, Massachusetts Institute of Technology
2 Department of EECS, Massachusetts Institute of Technology
3 The NSF AI Institute for Artificial Intelligence and Fundamental Interactions
Abstract
We study the problem of creating strong, yet narrow, AI systems. While recent
AI progress has been driven by the training of large general-purpose foundation
models, the creation of smaller models specialized for narrow domains could be
valuable for both efficiency and safety. In this work, we explore two challenges
involved in creating narrow AI systems, having to do with basic properties of how
neural networks learn and structure their representations. The first challenge regards
when it is possible to train narrow models from scratch. Through experiments
on a synthetic task, we find that it is sometimes necessary to train networks on
a wide distribution of data to learn certain narrow skills within that distribution.
This effect arises when skills depend on each other hierarchically, and training on a
broad distribution introduces a curriculum which substantially accelerates learning.
The second challenge regards how to transfer particular skills from large general
models into small specialized models. We find that model skills are often not
perfectly localized to a particular set of prunable components. However, we find
that methods based on pruning can still outperform distillation. We investigate the
use of a regularization objective to align desired skills with prunable components
while unlearning unnecessary skills.
1
Introduction
Today, the most competent AI systems in any particular domain are general systems that are relatively
competent in every domain. The best models at math and coding are also broadly knowledgeable about
a very diverse array of topics, from Roman history to home cooking recipes to medical di...
Fair Clustering via Alignment
Kunwoong Kim 1 Jihu Lee 1 Sangchul Park 2 Yongdai Kim 1
Abstract
Algorithmic fairness in clustering aims to bal-
ance the proportions of instances assigned to
each cluster with respect to a given sensitive at-
tribute. While recently developed fair clustering
algorithms optimize clustering objectives under
specific fairness constraints, their inherent com-
plexity or approximation often results in subopti-
mal clustering utility or numerical instability in
practice. To resolve these limitations, we propose
a new fair clustering algorithm based on a novel
decomposition of the fair K-means clustering ob-
jective function. The proposed algorithm, called
Fair Clustering via Alignment (FCA), operates
by alternately (i) finding a joint probability distri-
bution to align the data from different protected
groups, and (ii) optimizing cluster centers in the
aligned space. A key advantage of FCA is that
it theoretically guarantees approximately optimal
clustering utility for any given fairness level with-
out complex constraints, thereby enabling high-
utility fair clustering in practice. Experiments
show that FCA outperforms existing methods by
(i) attaining a superior trade-off between fairness
level and clustering utility, and (ii) achieving near-
perfect fairness without numerical instability.
1. Introduction
As artificial intelligence (AI) technology has advanced and
been successfully applied to diverse domains and tasks, the
requirement for AI systems to make fair decisions (i.e., al-
gorithmic fairness) has emerged as an important societal
issue. This requirement is particularly necessary when ob-
served data possess historical biases with respect to specific
sensitive attributes, leading to unfair outcomes of learned
models based on such biased data (Angwin et al., 2016;
1Department of Statistics, Seoul National University, Re-
public of Korea 2School of Law, Seoul National University,
Republic of Korea.
Correspondence to: Yongdai Kim <yd...