dataset-review/dataset-review huggingface dataset:
Lookahead Routing for Large Language Models
NeurIPS 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
ICML 2025

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
ICML 2025

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
NeurIPS 2025

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
NeurIPS 2025

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
ICML 2025

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...
Efficient Logit-based Knowledge Distillation of Deep Spiking Neural Networks for Full-Range Timestep Deployment
ICML 2025

Efficient Logit-based Knowledge Distillation of Deep Spiking Neural Networks
for Full-Range Timestep Deployment
Chengting Yu 1 2 † Xiaochen Zhao 2 † Lei Liu 2 Shu Yang 2 Gaoang Wang 2 Erping Li 1 2 Aili Wang 1 2 Abstract
Spiking Neural Networks (SNNs) are emerging
as a brain-inspired alternative to traditional Arti-
ficial Neural Networks (ANNs), prized for their
potential energy efficiency on neuromorphic hard-
ware. Despite this, SNNs often suffer from accu-
racy degradation compared to ANNs and face
deployment challenges due to fixed inference
timesteps, which require retraining for adjust-
ments, limiting operational flexibility.
To ad-
dress these issues, our work considers the spatio-
temporal property inherent in SNNs, and pro-
poses a novel distillation framework for deep
SNNs that optimizes performance across full-
range timesteps without specific retraining, en-
hancing both efficacy and deployment adaptabil-
ity. We provide both theoretical analysis and em-
pirical validations to illustrate that training guaran-
tees the convergence of all implicit models across
full-range timesteps.
Experimental results on
CIFAR-10, CIFAR-100, CIFAR10-DVS, and Im-
ageNet demonstrate state-of-the-art performance
among distillation-based SNNs training methods.
Our code is available at https://github.com/Intelli-
Chip-Lab/snn temporal decoupling distillation.
1. Introduction
Spiking Neural Networks (SNNs) are modeled after biologi-
cal neural systems and feature spiking neurons that replicate
the dynamics of biological neurons (Maass, 1997; Roy et al.,
2019). In contrast to Artificial Neural Networks (ANNs),
which utilize continuous data forms, SNNs employ a spike-
coding approach, using discrete binary spike trains for data
transmission (Panzeri & Schultz, 2001). This binary sig-
naling significantly reduces the multiply-accumulate op-
† Equal contribution.
1College of Information Science and
Electronic Engineering, Zhejiang University 2ZJU-UIUC Insti-
tute, Zhejiang Univ...
Hyper-Transforming Latent Diffusion Models
ICML 2025

Hyper-Transforming Latent Diffusion Models
Ignacio Peis 1 2 Batuhan Koyuncu 3 4 Isabel Valera 3 Jes Frellsen 1 2
Abstract
We introduce a novel generative framework for
functions by integrating Implicit Neural Represen-
tations (INRs) and Transformer-based hypernet-
works into latent variable models. Unlike prior ap-
proaches that rely on MLP-based hypernetworks
with scalability limitations, our method employs
a Transformer-based decoder to generate INR pa-
rameters from latent variables, addressing both
representation capacity and computational effi-
ciency. Our framework extends latent diffusion
models (LDMs) to INR generation by replacing
standard decoders with a Transformer-based hy-
pernetwork, which can be trained either from
scratch or via hyper-transforming—a strategy that
fine-tunes only the decoder while freezing the pre-
trained latent space. This enables efficient adapta-
tion of existing generative models to INR-based
representations without requiring full retraining.
We validate our approach across multiple modali-
ties, demonstrating improved scalability, expres-
siveness, and generalization over existing INR-
based generative models. Our findings establish a
unified and flexible framework for learning struc-
tured function representations.
1. Introduction
Generative modelling has seen remarkable advances in re-
cent years, with diffusion models achieving state-of-the-art
performance across multiple domains, including images,
videos, and 3D synthesis (Ho et al., 2020; Dhariwal &
Nichol, 2021; Rombach et al., 2022). A key limitation of ex-
isting generative frameworks, however, is their reliance on
structured output representations, such as pixel grids, which
constrain resolution and generalisation across data modal-
1Department of Applied Mathematics and Computer Science,
Technical University of Denmark, Copenhagen, Denmark 2Pioneer
Centre for Artificial Intelligence, Copenhagen, Denmark 3Saarland
University, Saarbr¨ucken, Germany 4Zuse School ELIZA. C...
Model-Based Exploration in Monitored Markov Decision Processes
ICML 2025

Model-Based Exploration in Monitored Markov Decision Processes
Alireza Kazemipour 1 2 Matthew E. Taylor 1 2 3 Michael Bowling 1 2 3
Abstract
A tenet of reinforcement learning is that the agent
always observes rewards. However, this is not
true in many realistic settings, e.g., a human ob-
server may not always be available to provide
rewards, sensors may be limited or malfunction-
ing, or rewards may be inaccessible during de-
ployment. Monitored Markov decision processes
(Mon-MDPs) have recently been proposed to
model such settings. However, existing Mon-
MDP algorithms have several limitations: they
do not fully exploit the problem structure, can-
not leverage a known monitor, lack worst-case
guarantees for “unsolvable” Mon-MDPs without
specific initialization, and offer only asymptotic
convergence proofs. This paper makes three con-
tributions. First, we introduce a model-based algo-
rithm for Mon-MDPs that addresses these short-
comings. The algorithm employs two instances
of model-based interval estimation: one to ensure
that observable rewards are reliably captured, and
another to learn a minimax-optimal policy. Sec-
ond, we empirically demonstrate the advantages.
We show faster convergence than prior algorithms
in more than four dozen benchmarks, and even
more dramatic improvements when the monitor-
ing process is known. Third, we present the first
finite sample-bound on performance. We show
convergence to a minimax-optimal policy even
when some rewards are never observable.
1. Introduction
Reinforcement learning (RL) is based on trial-and-error: in-
stead of being directly shown what to do, an agent receives
consistent numerical feedback in the form of rewards for
its decisions. However, this assumption is not always re-
1Department of Computing Science, University of Alberta
2Alberta Machine Intelligence Institute (Amii) 3Canada CIFAR
AI Chair, Amii.
Correspondence to:
Alireza Kazemipour
<kazemipo@ualberta.ca>.
Proceedings of the 42 nd International Confere...
TULIP: Token-length Upgraded CLIP
ICLR 2025

Published as a conference paper at ICLR 2025
TULIP: TOKEN-LENGTH UPGRADED CLIP
Ivona Najdenkoska∗Mohammad Mahdi Derakhshani∗Yuki M. Asano
Nanne van Noord
Marcel Worring† Cees G. M. Snoek†
University of Amsterdam
Amsterdam, the Netherlands
ABSTRACT
We address the challenge of representing long captions in vision-language mod-
els, such as CLIP. By design these models are limited by fixed, absolute posi-
tional encodings, restricting inputs to a maximum of 77 tokens and hindering
performance on tasks requiring longer descriptions. Although recent work has
attempted to overcome this limit, their proposed approaches struggle to model
token relationships over longer distances and simply extend to a fixed new to-
ken length. Instead, we propose a generalizable method, named TULIP, able
to upgrade the token length to any length for CLIP-like models.
We do so
by improving the architecture with relative position encodings, followed by a
training procedure that (i) distills the original CLIP text encoder into an en-
coder with relative position encodings and (ii) enhances the model for aligning
longer captions with images. By effectively encoding captions longer than the
default 77 tokens, our model outperforms baselines on cross-modal tasks such
as retrieval and text-to-image generation.
The code repository is available at
https://github.com/ivonajdenkoska/tulip.
1
INTRODUCTION
Obtaining text representations for long captions in vision-language models is an open research chal-
lenge. Within the context of text-only large language models, this problem has been studied ex-
tensively (Chung et al., 2024; Touvron et al., 2023; Bai et al., 2023; Gu & Dao, 2023; Su et al.,
2024; Dubey et al., 2024; Pawar et al., 2024; Jiang et al., 2024; Lieber et al., 2024), yet, methods
for context length expansion have only scarcely made their way into the vision-language domain.
For instance, contrastive vision-language models like CLIP (Radford et al., 2021; Jia et al., 2021; Li
et al., 2022)...
Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges
ICML 2025

Self-Improving Transformers Overcome Easy-to-Hard and Length
Generalization Challenges
Nayoung Lee * 1 Ziyang Cai * 1 Avi Schwarzschild 2 Kangwook Lee 1 Dimitris Papailiopoulos 1 3
Abstract
Large language models often struggle with length
generalization and solving complex problem in-
stances beyond their training distribution. We
present a self-improvement approach where mod-
els iteratively generate and learn from their own
solutions, progressively tackling harder problems
while maintaining a standard transformer archi-
tecture.
Across diverse tasks including arith-
metic, string manipulation, and maze solving,
our method enables models to solve problems far
beyond their initial training distribution—for in-
stance, generalizing from 10-digit to 100-digit
addition without apparent saturation.
We ob-
serve that filtering for correct self-generated exam-
ples leads to exponential improvements in out-of-
distribution performance across training rounds.
Additionally, starting from pretrained models sig-
nificantly accelerates this self-improvement pro-
cess for several tasks. Our results demonstrate
how controlled weak-to-strong curricula can sys-
tematically expand model capabilities while pre-
serving architectural simplicity.
1. Introduction
Despite the remarkable success of transformer-based lan-
guage models (Vaswani et al., 2017) across a wide range
of tasks, these models exhibit significant limitations in
length generalization—the ability to extrapolate to longer
sequences than those seen during training. Even in simple
algorithmic tasks such as arithmetic, standard transformer
models trained with autoregressive objectives struggle to
generalize to longer problem instances (Anil et al., 2022).
To address this, prior work has explored various modifica-
*Equal contribution
1University of Wisconsin-Madison
2Carnegie Mellon University
3Microsoft.
Correspondence
to:
Nayoung Lee <nayoung.lee@wisc.edu>,
Ziyang Cai
<zcai75@wisc.edu>.
Proceedings of the 42 nd Internatio...