AIConfPaper
← Search

Ankit Singh Rawat

41 accepted papers

2026

Continuous Chain of Thought Enables Parallel Exploration and Reasoning

Halil Alperen Gozeten, Muhammed Emrullah Ildiz, Xuechen Zhang, Hrayr Harutyunyan, Ankit Singh Rawat, Samet Oymak
ICLR 2026poster

Modern language models generate chain-of-thought traces by autoregressively sampling tokens from a finite vocabulary. While this discrete sampling has achieved remarkable success, conducting chain-of-thought with continuously-valued tokens (CoT2) offers a richer and more expressive alternative. Our…

Cited by 0Source ↗code ↗Scholar ↗
2026

Universal Model Routing for Efficient LLM Inference

Wittawat Jitkrittum, Harikrishna Narasimhan, Ankit Singh Rawat, Jeevesh Juneja, Congchao Wang, Zifeng Wang +6
ICLR 2026poster

Model routing is a simple technique for reducing the inference cost of large language models (LLMs), wherein one maintains a pool of candidate LLMs, and learns to route each prompt to the smallest feasible LLM. Existing works focus on learning a router for a fixed pool of LLMs. In this paper, we con…

Cited by 0Source ↗Scholar ↗
2025

Faster Cascades via Speculative Decoding

Harikrishna Narasimhan, Wittawat Jitkrittum, Ankit Singh Rawat, Seungyeon Kim, Neha Gupta, Aditya Krishna Menon +1
ICLR 2025oral

Cascades and speculative decoding are two common approaches to improving language models' inference efficiency. Both approaches interleave two models, but via fundamentally distinct mechanisms: deferral rule that invokes the larger model only for “hard” inputs, while speculative decoding uses spec…

Cited by 4Source ↗PDF ↗Scholar ↗
2024

A Statistical Framework for Data-dependent Retrieval-Augmented Models

Soumya Basu, Ankit Singh Rawat, Manzil Zaheer
ICML 2024poster

Modern ML systems increasingly augment input instances with additional relevant information to enhance final prediction. Despite growing interest in such retrieval-augmented models, their fundamental properties and training are not well understood. We propose a statistical framework to study such mo…

Cited by 0Source ↗PDF ↗Scholar ↗
2024

Analysis of Plan-based Retrieval for Grounded Text Generation

Ameya Godbole, Nicholas Monath, Seungyeon Kim, Ankit Singh Rawat, Andrew McCallum, Manzil Zaheer
EMNLP 2024main

In text generation, hallucinations refer to the generation of seemingly coherent text that contradicts established knowledge. One compelling hypothesis is that hallucinations occur when a language model is given a generation task outside its parametric knowledge (due to rarity, recency, domain, etc.…

Cited by 1Source ↗PDF ↗Scholar ↗
2024

DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar +2
ICLR 2024poster

Speculative decoding~(SD) accelerates large language model inference by employing a faster {\em draft} model for generating multiple tokens, which are then verified in parallel by the larger {\em target} model, resulting in the text generated according to the target model distribution. However, iden…

Cited by 86Source ↗PDF ↗Scholar ↗
2024

Dual-Encoders for Extreme Multi-label Classification

Nilesh Gupta, Fnu Devvrit, Ankit Singh Rawat, Srinadh Bhojanapalli, Prateek Jain, Inderjit S Dhillon
ICLR 2024poster

Dual-encoder (DE) models are widely used in retrieval tasks, most commonly studied on open QA benchmarks that are often characterized by multi-class and limited training data. In contrast, their performance in multi-label and data-rich retrieval settings like extreme multi-label classification (XMC)…

Cited by 15Source ↗PDF ↗code ↗Scholar ↗
2024

Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond

Yingcong Li, Ankit Singh Rawat, Samet Oymak
NeurIPS 2024poster

Recent research has shown that Transformers with linear attention are capable of in-context learning (ICL) by implementing a linear estimator through gradient descent steps. However, the existing results on the optimization landscape apply under stylized settings where task and feature vectors are a…

Cited by 4Source ↗PDF ↗Scholar ↗
2024

From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers

Muhammed Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, Samet Oymak
ICML 2024poster

Modern language models rely on the transformer architecture and attention mechanism to perform language understanding and text generation. In this work, we study learning a 1-layer self-attention model from a set of prompts and the associated outputs sampled from the model. We first establish a form…

Cited by 19Source ↗PDF ↗Scholar ↗
2024

Language Model Cascades: Token-Level Uncertainty And Beyond

Neha Gupta, Harikrishna Narasimhan, Wittawat Jitkrittum, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar
ICLR 2024poster

Recent advances in language models (LMs) have led to significant improvements in quality on complex NLP tasks, but at the expense of increased inference costs. A simple strategy to achieve more favorable cost-quality tradeoffs is cascading: here, a small model is invoked for most “easy” instances, w…

Cited by 36Source ↗PDF ↗Scholar ↗
2024

Mechanics of Next Token Prediction with Self-Attention

Yingcong Li, Yixiao Huang, Muhammed E. Ildiz, Ankit Singh Rawat, Samet Oymak
AISTATS 2024poster

Transformer-based language models are trained on large datasets to predict the next token given an input sequence. Despite this simple training objective, they have led to revolutionary advances in natural language processing. Underlying this success is the self-attention mechanism. In this work, we…

Cited by 34Source ↗PDF ↗Scholar ↗
2024

Think before you speak: Training Language Models With Pause Tokens

Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, Vaishnavh Nagarajan
ICLR 2024poster

Language models generate responses by producing a series of tokens in immediate succession: the $(K+1)^{\rm th}$ token is an outcome of manipulating $K$ hidden vectors per layer, one vector per preceding token. What if instead we were to let the model manipulate say, $K+10$ hidden vectors, before it…

Cited by 76Source ↗PDF ↗Scholar ↗
2024

USTAD: Unified Single-model Training Achieving Diverse Scores for Information Retrieval

Seungyeon Kim, Ankit Singh Rawat, Manzil Zaheer, Wittawat Jitkrittum, Veeranjaneyulu Sadhanala, Sadeep Jayasumana +3
ICML 2024poster

Modern information retrieval (IR) systems consists of multiple stages like retrieval and ranking, with Transformer-based models achieving state-of-the-art performance at each stage. In this paper, we challenge the tradition of using separate models for different stages and ask if a single Transforme…

Cited by 0Source ↗PDF ↗Scholar ↗
2023

Large Language Models with Controllable Working Memory

Daliang Li, Ankit Singh Rawat, Manzil Zaheer, Xin Wang, Michal Lukasik, Andreas Veit +2
ACL 2023findings

Large language models (LLMs) have led to a series of breakthroughs in natural language processing (NLP), partly owing to the massive amounts of world knowledge they memorize during pretraining. While many downstream applications provide the model with an informational context to aid its underlying t…

Cited by 148Source ↗PDF ↗Scholar ↗
2023

On the Role of Attention in Prompt-tuning

Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, Christos Thrampoulidis
ICML 2023poster

Prompt-tuning is an emerging strategy to adapt large language models (LLM) to downstream tasks by learning a (soft-)prompt parameter from data. Despite its success in LLMs, there is limited theoretical understanding of the power of prompt-tuning and the role of the attention mechanism in prompting.…

Cited by 60Source ↗PDF ↗Scholar ↗
2023

ResMem: Learn what you can and memorize the rest

Zitong Yang, Michal Lukasik, Vaishnavh Nagarajan, Zonglin Li, Ankit Singh Rawat, Manzil Zaheer +2
NeurIPS 2023poster

The impressive generalization performance of modern neural networks is attributed in part to their ability to implicitly memorize complex training patterns. Inspired by this, we explore a novel mechanism to improve model generalization via explicit memorization. Specifically, we propose the residual…

Cited by 12Source ↗PDF ↗Scholar ↗
2023

Serving Graph Compression for Graph Neural Networks

Si Si, Felix Yu, Ankit Singh Rawat, Cho-Jui Hsieh, Sanjiv Kumar
ICLR 2023top-25%

Serving a GNN model online is challenging --- in many applications when testing nodes are connected to training nodes, one has to propagate information from training nodes to testing nodes to achieve the best performance, and storing the whole training set (including training graph and node features…

Cited by 9Source ↗PDF ↗Scholar ↗
2023

Supervision Complexity and its Role in Knowledge Distillation

Hrayr Harutyunyan, Ankit Singh Rawat, Aditya Krishna Menon, Seungyeon Kim, Sanjiv Kumar
ICLR 2023poster

Despite the popularity and efficacy of knowledge distillation, there is limited understanding of why it helps. In order to study the generalization behavior of a distilled student, we propose a new theoretical framework that leverages supervision complexity: a measure of alignment between teacher-pr…

Cited by 14Source ↗PDF ↗Scholar ↗
2023

Teacher Guided Training: An Efficient Framework for Knowledge Transfer

Manzil Zaheer, Ankit Singh Rawat, Seungyeon Kim, Chong You, Himanshu Jain, Andreas Veit +2
ICLR 2023poster

The remarkable performance gains realized by large pretrained models, e.g., GPT-3, hinge on the massive amounts of data they are exposed to during training. Analogously, distilling such large models to compact models for efficient deployment also necessitates a large amount of (labeled or unlabeled)…

Cited by 2Source ↗PDF ↗Scholar ↗
2023

The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers

Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J. Reddi +5
ICLR 2023poster

This paper studies a curious phenomenon that machine learning model with Transformer architectures have sparse activation maps. By activation map we refer to the intermediate output of the multi-layer perceptrons (MLPs) after a ReLU activation function, and by "sparse" we mean that on average very f…

Cited by 92Source ↗PDF ↗Scholar ↗
2023

When Does Confidence-Based Cascade Deferral Suffice?

Wittawat Jitkrittum, Neha Gupta, Aditya Krishna Menon, Harikrishna Narasimhan, Ankit Singh Rawat, Sanjiv Kumar
NeurIPS 2023poster

Cascades are a classical strategy to enable inference cost to vary adaptively across samples, wherein a sequence of classifiers are invoked in turn. A deferral rule determines whether to invoke the next classifier in the sequence, or to terminate prediction. One simple deferral rule employs the co…

Cited by 17Source ↗PDF ↗Scholar ↗
2022

A Fourier Approach to Mixture Learning

Mingda Qiao, Guru Guruganesh, Ankit Singh Rawat, Kumar Avinava Dubey, Manzil Zaheer
NeurIPS 2022accept

We revisit the problem of learning mixtures of spherical Gaussians. Given samples from a mixture $\frac{1}{k}\sum_{j=1}^{k}\mathcal{N}(\mu_j, I_d)$, the goal is to estimate the means $\mu_1, \mu_2, \ldots, \mu_k \in \mathbb{R}^d$ up to a small error. The hardness of this learning problem can be meas…

Cited by 9Source ↗PDF ↗Scholar ↗
2022

In defense of dual-encoders for neural ranking

Aditya Menon, Sadeep Jayasumana, Ankit Singh Rawat, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar
ICML 2022spotlight

Transformer-based models such as BERT have proven successful in information retrieval problem, which seek to identify relevant documents for a given query. There are two broad flavours of such models: cross-attention (CA) models, which learn a joint embedding for the query and document, and dual-enc…

Cited by 33Source ↗PDF ↗Scholar ↗
2022

Post-hoc estimators for learning to defer to an expert

Harikrishna Narasimhan, Wittawat Jitkrittum, Aditya Krishna Menon, Ankit Singh Rawat, Sanjiv Kumar
NeurIPS 2022accept

Many practical settings allow a learner to defer predictions to one or more costly experts. For example, the learning to defer paradigm allows a learner to defer to a human expert, at some monetary cost. Similarly, the adaptive inference paradigm allows a base model to defer to one or more large mod…

Cited by 33Source ↗PDF ↗Scholar ↗
2021

A statistical perspective on distillation

Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, Seungyeon Kim, Sanjiv Kumar
ICML 2021spotlight

Knowledge distillation is a technique for improving a “student” model by replacing its one-hot training labels with a label distribution obtained from a “teacher” model. Despite its broad success, several basic questions — e.g., Why does distillation help? Why do more accurate teachers not necessari…

Cited by 107Source ↗PDF ↗Scholar ↗
2021

Disentangling Sampling and Labeling Bias for Learning in Large-output Spaces

Ankit Singh Rawat, Aditya K Menon, Wittawat Jitkrittum, Sadeep Jayasumana, Felix Yu, Sashank Reddi +1
ICML 2021spotlight

Negative sampling schemes enable efficient training given a large number of classes, by offering a means to approximate a computationally expensive loss function that takes all labels into account. In this paper, we present a new connection between these schemes and loss modification techniques for…

Cited by 11Source ↗PDF ↗Scholar ↗
2021

Long-tail learning via logit adjustment

Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, Sanjiv Kumar
ICLR 2021spotlight

Real-world classification problems typically exhibit an imbalanced or long-tailed label distribution, wherein many labels have only a few associated samples. This poses a challenge for generalisation on such labels, and also makes naive learning biased towards dominant labels. In this paper, we pr…

Cited by 907Source ↗PDF ↗code ↗Scholar ↗
2021

Overparameterisation and worst-case generalisation: friend or foe?

Aditya Krishna Menon, Ankit Singh Rawat, Sanjiv Kumar
ICLR 2021poster

Overparameterised neural networks have demonstrated the remarkable ability to perfectly fit training samples, while still generalising to unseen test samples. However, several recent works have revealed that such models' good average performance does not always translate to good worst-case performan…

Cited by 42Source ↗PDF ↗Scholar ↗
2021

RankDistil: Knowledge Distillation for Ranking

Sashank Reddi, Rama Kumar Pasumarthi, Aditya Menon, Ankit Singh Rawat, Felix Yu, Seungyeon Kim +2
AISTATS 2021poster

Knowledge distillation is an approach to improve the performance of a student model by using the knowledge of a complex teacher. Despite its success in several deep learning applications, the study of distillation is mostly confined to classification settings. In particular, the use of distillation…

Cited by 38Source ↗PDF ↗Scholar ↗
2020

Adversarial robustness via robust low rank representations

Pranjal Awasthi, Himanshu Jain, Ankit Singh Rawat, Aravindan Vijayaraghavan
NeurIPS 2020poster

Adversarial robustness measures the susceptibility of a classifier to imperceptible perturbations made to the inputs at test time. In this work we highlight the benefits of natural low rank representations that often exist for real data such as images, for training neural networks with certified rob…

Cited by 27Source ↗PDF ↗Scholar ↗
2020

Are Transformers universal approximators of sequence-to-sequence functions?

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, Sanjiv Kumar
ICLR 2020poster

Despite the widespread adoption of Transformer models for NLP tasks, the expressive power of these models is not well-understood. In this paper, we establish that Transformer models are universal approximators of continuous permutation equivariant sequence-to-sequence functions with compact support,…

Cited by 432Source ↗Scholar ↗
2020

Low-Rank Bottleneck in Multi-head Attention Models

Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, Sanjiv Kumar
ICML 2020poster

Attention based Transformer architecture has enabled significant advances in the field of natural language processing. In addition to new pre-training techniques, recent improvements crucially rely on working with a relatively larger embedding dimension for tokens. Unfortunately, this leads to model…

Cited by 112Source ↗PDF ↗code ↗Scholar ↗
2020

O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformers

Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, Sanjiv Kumar
NeurIPS 2020poster

Recently, Transformer networks have redefined the state of the art in many NLP tasks. However, these models suffer from quadratic computational cost in the input sequence length $n$ to compute pairwise attention in each layer. This has prompted recent research into sparse Transformers that sparsify…

Cited by 94Source ↗PDF ↗Scholar ↗
2020

Robust large-margin learning in hyperbolic space

Melanie Weber, Manzil Zaheer, Ankit Singh Rawat, Aditya K Menon, Sanjiv Kumar
NeurIPS 2020poster

Recently, there has been a surge of interest in representation learning in hyperbolic spaces, driven by their ability to represent hierarchical data with significantly fewer dimensions than standard Euclidean spaces. However, the viability and benefits of hyperbolic spaces for downstream machine lea…

Cited by 40Source ↗PDF ↗Scholar ↗
2019

Lifting high-dimensional non-linear models with Gaussian regressors

Christos Thrampoulidis, Ankit Singh Rawat
AISTATS 2019poster

We study the problem of recovering a structured signal $\mathbf{x}_0$ from high-dimensional data $\mathbf{y}_i=f(\mathbf{a}_i^T\mathbf{x}_0)$ for some nonlinear (and potentially unknown) link function $f$, when the regressors $\mathbf{a}_i$ are iid Gaussian. Brillinger (1982) showed that ordinary le…

Cited by 13Source ↗PDF ↗Scholar ↗
2019

Multilabel reductions: what is my loss optimising?

Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, Sanjiv Kumar
NeurIPS 2019spotlight

Multilabel classification is a challenging problem arising in applications ranging from information retrieval to image tagging. A popular approach to this problem is to employ a reduction to a suitable series of binary or multiclass problems (e.g., computing a softmax based cross-entropy over the re…

Cited by 68Source ↗PDF ↗Scholar ↗
2019

Sampled Softmax with Random Fourier Features

Ankit Singh Rawat, Jiecao Chen, Felix Xinnan X Yu, Ananda Theertha Suresh, Sanjiv Kumar
NeurIPS 2019poster

The computational cost of training with softmax cross entropy loss grows linearly with the number of classes. For the settings where a large number of classes are involved, a common method to speed up training is to sample a subset of classes and utilize an estimate of the loss gradient based on the…

Cited by 69Source ↗PDF ↗Scholar ↗