← Search

Prateek Jain

76 accepted papers

2025

Does Safety Training of LLMs Generalize to Semantically Related Natural Prompts?

ICLR 2025poster

Large Language Models (LLMs) are known to be susceptible to crafted adversarial attacks or jailbreaks that lead to the generation of objectionable content despite being aligned to human preferences using safety fine-tuning methods. While the large dimensionality of input token space makes it inevita…

Cited by 2SourcePDFScholar
2025

Masked Generative Nested Transformers with Decode Time Scaling

ICML 2025poster

Recent advances in visual generation have made significant strides in producing content of exceptional quality. However, most methods suffer from a fundamental problem - a bottleneck of inference computational efficiency. Most of these algorithms involve multiple passes over a transformer model to g…

Cited by 0SourcePDFScholar
2025

Spark Transformer: Reactivating Sparsity in Transformer FFN and Attention

NeurIPS 2025poster

The discovery of the *lazy neuron phenomenon* (Li et al., 2022), where fewer than 10% of the feedforward networks (FFN) parameters in trained Transformers are activated per token, has spurred significant interests in *activation sparsity* for enhancing large model efficiency. While notable progress…

Cited by 0SourceScholar
2024

Dual-Encoders for Extreme Multi-label Classification

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)…

2024

LLM Augmented LLMs: Expanding Capabilities through Composition

ICLR 2024poster

Foundational models with billions of parameters which have been trained on large corpus of data have demonstrated non-trivial skills in a variety of domains. However, due to their monolithic structure, it is challenging and expensive to augment them or impart new skills. On the other hand, due to th…

Cited by 44SourcePDFScholar
2024

LookupViT: Compressing visual information to a limited number of tokens

ECCV 2024poster

"Vision Transformers (ViT) have emerged as the de-facto choice for numerous industry grade vision solutions. But their inference cost can be prohibitive for many settings, as they compute self-attention in each layer which suffers from quadratic computational complexity in the number of tokens. On t…

Cited by 11SourcePDFScholar
2024

MatFormer: Nested Transformer for Elastic Inference

NeurIPS 2024poster

Foundation models are applied in a broad spectrum of settings with different inference constraints, from massive multi-accelerator clusters to resource-constrained standalone mobile devices. However, the substantial costs associated with training these models often limit the number of unique model s…

Cited by 11SourcePDFScholar
2024

Mixture of Nested Experts: Adaptive Processing of Visual Tokens

NeurIPS 2024poster

The visual medium (images and videos) naturally contains a large amount of information redundancy, thereby providing a great opportunity for leveraging efficiency in processing. While Vision Transformer (ViT) based models scale effectively to large data regimes, they fail to capitalize on this inher…

Cited by 8SourcePDFScholar
2024

Sample-Efficient Personalization: Modeling User Parameters as Low Rank Plus Sparse Components

AISTATS 2024poster

Personalization of machine learning (ML) predictions for individual users/domains/enterprises is critical for practical recommendation systems. Standard personalization approaches involve learning a user/domain specific \emph{embedding} that is fed into a fixed global model which can be limiting. On…

Cited by 1SourcePDFScholar
2024

Tandem Transformers for Inference Efficient LLMs

ICML 2024poster

The autoregressive nature of conventional large language models (LLMs) inherently limits inference speed, as tokens are generated sequentially. While speculative (Leviathan et al., 2023) and parallel (Stern et al., 2018) decoding techniques attempt to mitigate this, they face limitations: either rel…

Cited by 5SourcePDFScholar
2024

Time-Reversal Provides Unsupervised Feedback to LLMs

NeurIPS 2024spotlight

Large Language Models (LLMs) are typically trained to predict in the forward direction of time. However, recent works have shown that prompting these models to look back and critique their own generations can produce useful feedback. Motivated by this, we explore the question of whether LLMs can be…

Cited by 0SourcePDFScholar
2023

AdANNS: A Framework for Adaptive Semantic Search

NeurIPS 2023poster

Web-scale search systems learn an encoder to embed a given query which is then hooked into an approximate nearest neighbor search (ANNS) pipeline to retrieve similar data points. To accurately capture tail queries and data points, learned representations typically are _rigid, high-dimensional_ vecto…

2023

Blocked Collaborative Bandits: Online Collaborative Filtering with Per-Item Budget Constraints

NeurIPS 2023poster

We consider the problem of \emph{blocked} collaborative bandits where there are multiple users, each with an associated multi-armed bandit problem. These users are grouped into \emph{latent} clusters such that the mean reward vectors of users within the same cluster are identical. Our goal is to des…

Cited by 2SourcePDFScholar
2023

Feature Reconstruction From Outputs Can Mitigate Simplicity Bias in Neural Networks

ICLR 2023poster

Deep Neural Networks are known to be brittle to even minor distribution shifts compared to the training distribution. While one line of work has demonstrated that \emph{Simplicity Bias} (SB) of DNNs -- bias towards learning only the simplest features -- is a key reason for this brittleness, another…

Cited by 11SourcePDFScholar
2023

Label Robust and Differentially Private Linear Regression: Computational and Statistical Efficiency

NeurIPS 2023poster

We study the canonical problem of linear regression under $(\varepsilon,\delta)$-differential privacy when the datapoints are sampled i.i.d.~from a distribution and a fraction of response variables are adversarially corrupted. We provide the first provably efficient -- both computationally and stati…

Cited by 2SourcePDFScholar
2023

Multi-Task Differential Privacy Under Distribution Skew

ICML 2023poster

We study the problem of multi-task learning under user-level differential privacy, in which n users contribute data to m tasks, each involving a subset of users. One important aspect of the problem, that can significantly impact quality, is the distribution skew among tasks. Tasks that have much few…

Cited by 5SourcePDFScholar
2023

Multi-User Reinforcement Learning with Low Rank Rewards

ICML 2023poster

We consider collaborative multi-user reinforcement learning, where multiple users have the same state-action space and transition probabilities but different rewards. Under the assumption that the reward matrix of the $N$ users has a low-rank structure -- a standard and practically successful assump…

Cited by 1SourcePDFScholar
2023

Optimal Algorithms for Latent Bandits with Cluster Structure

AISTATS 2023poster

We consider the problem of latent bandits with cluster structure where there are multiple users, each with an associated multi-armed bandit problem. These users are grouped into latent clusters such that the mean reward vectors of users within the same cluster are identical. At each round, a user, s…

Cited by 14SourcePDFScholar
2023

Simplicity Bias in 1-Hidden Layer Neural Networks

NeurIPS 2023poster

Recent works have demonstrated that neural networks exhibit extreme *simplicity bias* (SB). That is, they learn *only the simplest* features to solve a task at hand, even in the presence of other, more robust but more complex features. Due to the lack of a general and rigorous definition of *featu…

Cited by 15SourcePDFScholar
2023

Treeformer: Dense Gradient Trees for Efficient Attention Computation

ICLR 2023poster

Standard inference and training with transformer based architectures scale quadratically with input sequence length. This is prohibitively large for a variety of applications especially in web-page translation, query-answering etc. Consequently, several approaches have been developed recently to spe…

Cited by 6SourcePDFScholar
2022

IGLU: Efficient GCN Training via Lazy Updates

ICLR 2022poster

Training multi-layer Graph Convolution Networks (GCN) using standard SGD techniques scales poorly as each descent step ends up updating node embeddings for a large portion of the graph. Recent attempts to remedy this sub-sample the graph that reduces compute but introduce additional variance and may…

2022

Matryoshka Representation Learning

NeurIPS 2022accept

Learned representations are a central component in modern ML systems, serving a multitude of downstream tasks. When training such representations, it is often the case that computational and statistical constraints for each downstream task are unknown. In this context rigid, fixed capacity represent…

2022

Online Target Q-learning with Reverse Experience Replay: Efficiently finding the Optimal Policy for Linear MDPs

ICLR 2022poster

Q-learning is a popular Reinforcement Learning (RL) algorithm which is widely used in practice with function approximation (Mnih et al., 2015). In contrast, existing theoretical results are pessimistic about Q-learning. For example, (Baird, 1995) shows that Q-learning does not converge even with lin…

Cited by 31SourcePDFScholar
2022

Reproducibility in Optimization: Theoretical Framework and Limits

NeurIPS 2022accept

We initiate a formal study of reproducibility in optimization. We define a quantitative measure of reproducibility of optimization procedures in the face of noisy or error-prone operations such as inexact or stochastic gradient computations or inexact initialization. We then analyze several convex o…

Cited by 25SourcePDFScholar
2022

Robust Training in High Dimensions via Block Coordinate Geometric Median Descent

AISTATS 2022poster

Geometric median (GM) is a classical method in statistics for achieving robust estimation of the uncorrupted data; under gross corruption, it achieves the optimal breakdown point of 1/2. However, its computational complexity makes it infeasible for robustifying stochastic gradient descent (SGD) in h…

2021

Differentially Private Model Personalization

NeurIPS 2021spotlight

We study personalization of supervised learning with user-level differential privacy. Consider a setting with many users, each of whom has a training data set drawn from their own distribution $P_i$. Assuming some shared structure among the problems $P_i$, can users collectively learn the shared str…

Cited by 41SourcePDFScholar
2021

Do Input Gradients Highlight Discriminative Features?

NeurIPS 2021poster

Post-hoc gradient-based interpretability methods [Simonyan et al., 2013, Smilkov et al., 2017] that provide instance-specific explanations of model predictions are often based on assumption (A): magnitude of input gradients—gradients of logits with respect to input—noisily highlight discriminative t…

2021

LLC: Accurate, Multi-purpose Learnt Low-dimensional Binary Codes

NeurIPS 2021poster

Learning binary representations of instances and classes is a classical problem with several high potential applications. In modern settings, the compression of high-dimensional neural representations to low-dimensional binary codes is a challenging task and often require large bit-codes to be accur…

2021

Near-optimal Offline and Streaming Algorithms for Learning Non-Linear Dynamical Systems

NeurIPS 2021spotlight

We consider the setting of vector valued non-linear dynamical systems $X_{t+1} = \phi(A^{*} X_t) + \eta_t$, where $\eta_t$ is unbiased noise and $\phi : \mathbb{R} \to \mathbb{R}$ is a known link function that satisfies certain {\em expansivity property}. The goal is to learn $A^{*}$ from a single t…

Cited by 39SourcePDFScholar
2021

Optimal regret algorithm for Pseudo-1d Bandit Convex Optimization

ICML 2021spotlight

We study online learning with bandit feedback (i.e. learner has access to only zeroth-order oracle) where cost/reward functions $\f_t$ admit a "pseudo-1d" structure, i.e. $\f_t(\w) = \loss_t(\pred_t(\w))$ where the output of $\pred_t$ is one-dimensional. At each round, the learner observes context $…

Cited by 7SourcePDFScholar
2021

Private Alternating Least Squares: Practical Private Matrix Completion with Tighter Rates

ICML 2021oral

We study the problem of differentially private (DP) matrix completion under user-level privacy. We design a joint differentially private variant of the popular Alternating-Least-Squares (ALS) method that achieves: i) (nearly) optimal sample complexity for matrix completion (in terms of number of ite…

Cited by 23SourcePDFScholar
2021

Statistically and Computationally Efficient Linear Meta-representation Learning

NeurIPS 2021poster

In typical few-shot learning, each task is not equipped with enough data to be learned in isolation. To cope with such data scarcity, meta-representation learning methods train across many related tasks to find a shared (lower-dimensional) representation of the data where all tasks can be solved acc…

Cited by 16SourcePDFScholar
2021

Streaming Linear System Identification with Reverse Experience Replay

NeurIPS 2021poster

We consider the problem of estimating a linear time-invariant (LTI) dynamical system from a single trajectory via streaming algorithms, which is encountered in several applications including reinforcement learning (RL) and time-series analysis. While the LTI system estimation problem is well-studie…

Cited by 21SourcePDFScholar
2020

DROCC: Deep Robust One-Class Classification

ICML 2020poster

Classical approaches for one-class problems such as one-class SVM and isolation forest require careful feature engineering when applied to structured domains like images. State-of-the-art methods aim to leverage deep learning to learn appropriate features via two main approaches. The first approach…

2020

Least Squares Regression with Markovian Data: Fundamental Limits and Algorithms

NeurIPS 2020spotlight

We study the problem of least squares linear regression where the datapoints are dependent and are sampled from a Markov chain. We establish sharp information theoretic minimax lower bounds for this problem in terms of $\tmix$, the mixing time of the underlying Markov chain, under different noise se…

Cited by 86SourcePDFScholar
2020

Optimization and Analysis of the pAp@k Metric for Recommender Systems

ICML 2020poster

Modern recommendation and notification systems must be robust to data imbalance, limitations on the number of recommendations/notifications, and heterogeneous engagement profiles across users. The pAp@k metric, which combines the partial-AUC and the precision@k metrics, was recently proposed to eval…

2020

Projection Efficient Subgradient Method and Optimal Nonsmooth Frank-Wolfe Method

NeurIPS 2020spotlight

We consider the classical setting of optimizing a nonsmooth Lipschitz continuous convex function over a convex constraint set, when having access to a (stochastic) first-order oracle (FO) for the function and a projection oracle (PO) for the constraint set. It is well known that to achieve $\epsilon…

2020

RNNPool: Efficient Non-linear Pooling for RAM Constrained Inference

NeurIPS 2020spotlight

Standard Convolutional Neural Networks (CNNs) designed for computer vision tasks tend to have large intermediate activation maps. These require large working memory and are thus unsuitable for deployment on resource-constrained devices typically used for inference on the edge. Aggressively downsampl…

2020

Soft Threshold Weight Reparameterization for Learnable Sparsity

ICML 2020poster

Sparsity in Deep Neural Networks (DNNs) is studied extensively with the focus of maximizing prediction accuracy given an overall parameter budget. Existing methods rely on uniform or heuristic non-uniform sparsity budgets which have sub-optimal layer-wise parameter allocation resulting in a) lower p…

2020

The Pitfalls of Simplicity Bias in Neural Networks

NeurIPS 2020poster

Several works have proposed Simplicity Bias (SB)---the tendency of standard training procedures such as Stochastic Gradient Descent (SGD) to find simple models---to justify why neural networks generalize well [Arpit et al. 2017, Nakkiran et al. 2019, Valle-Perez et al. 2019]. However, the precise no…

2019

Cost aware Inference for IoT Devices

AISTATS 2019poster

Networked embedded devices (IoTs) of limited CPU, memory and power resources are revolutionizing data gathering, remote monitoring and planning in many consumer and business applications. Nevertheless, resource limitations place a significant burden on their service life and operation, warranting co…

Cited by 15SourcePDFScholar
2019

Efficient Algorithms for Smooth Minimax Optimization

NeurIPS 2019poster

This paper studies first order methods for solving smooth minimax optimization problems $\min_x \max_y g(x,y)$ where $g(\cdot,\cdot)$ is smooth and $g(x,\cdot)$ is concave for each $x$. In terms of $g(\cdot,y)$, we consider two settings -- strongly convex and nonconvex -- and improve upon the best k…

2019

Globally-convergent Iteratively Reweighted Least Squares for Robust Regression Problems

AISTATS 2019poster

We provide the first global model recovery results for the IRLS (iteratively reweighted least squares) heuristic for robust regression problems. IRLS is known to offer excellent performance, despite bad initializations and data corruption, for several parameter estimation problems. Existing analyses…

Cited by 32SourcePDFScholar
2019

Learning Natural Programs from a Few Examples in Real-Time

AISTATS 2019poster

Programming by examples (PBE) is a rapidly growing subfield of AI, that aims to synthesize user-intended programs using input-output examples from the task. As users can provide only a few I/O examples, capturing user-intent accurately and ranking user-intended programs over other programs is challe…

Cited by 11SourcePDFScholar
2019

Shallow RNN: Accurate Time-series Classification on Resource Constrained Devices

NeurIPS 2019poster

Recurrent Neural Networks (RNNs) capture long dependencies and context, and 2 hence are the key component of typical sequential data based tasks. However, the sequential nature of RNNs dictates a large inference cost for long sequences even if the hardware supports parallelization. To induce long-te…

2018

FastGRNN: A Fast, Accurate, Stable and Tiny Kilobyte Sized Gated Recurrent Neural Network

NeurIPS 2018poster

This paper develops the FastRNN and FastGRNN algorithms to address the twin RNN limitations of inaccurate training and inefficient prediction. Previous approaches have improved accuracy at the expense of prediction costs making them infeasible for resource-constrained and real-time applications. Uni…

2018

Multiple Instance Learning for Efficient Sequential Data Classification on Resource-constrained Devices

NeurIPS 2018poster

We study the problem of fast and efficient classification of sequential data (such as time-series) on tiny devices, which is critical for various IoT related applications like audio keyword detection or gesture detection. Such tasks are cast as a standard classification task by sliding windows over…

2018

Neural-Guided Deductive Search for Real-Time Program Synthesis from Examples

ICLR 2018poster

Synthesizing user-intended programs from a small number of input-output exam- ples is a challenging problem with several important applications like spreadsheet manipulation, data wrangling and code refactoring. Existing synthesis systems either completely rely on deductive logic techniques that are…

Cited by 190SourcePDFScholar
2018

On the insufficiency of existing momentum schemes for Stochastic Optimization

ICLR 2018oral

Momentum based stochastic gradient methods such as heavy ball (HB) and Nesterov's accelerated gradient descent (NAG) method are widely used in practice for training deep networks and other supervised learning models, as they often provide significant improvements over stochastic gradient descent (SG…

2018

Support Recovery for Orthogonal Matching Pursuit: Upper and Lower bounds

NeurIPS 2018spotlight

This paper studies the problem of sparse regression where the goal is to learn a sparse vector that best optimizes a given objective function. Under the assumption that the objective function satisfies restricted strong convexity (RSC), we analyze orthogonal matching pursuit (OMP), a greedy algorith…

Cited by 8SourcePDFScholar
2017

Global Convergence of Non-Convex Gradient Descent for Computing Matrix Squareroot

AISTATS 2017poster

While there has been a significant amount of work studying gradient descent techniques for non-convex optimization problems over the last few years, all existing results establish either local convergence with good rates or global convergence with highly suboptimal rates, for many problems of intere…

Cited by 33SourcePDFScholar
2017

ProtoNN: Compressed and Accurate kNN for Resource-scarce Devices

ICML 2017poster

Several real-world applications require real-time prediction on resource-scarce devices such as an Internet of Things (IoT) sensor. Such applications demand prediction models with small storage and computational complexity that do not compromise significantly on accuracy. In this work, we propose Pr…

2016

Tensor vs. Matrix Methods: Robust Tensor Decomposition under Block Sparse Perturbations

AISTATS 2016poster

Robust tensor CP decomposition involves decomposing a tensor into low rank and sparse components. We propose a novel non-convex iterative algorithm with guaranteed recovery. It alternates between low-rank CP decomposition through gradient ascent (a variant of the tensor power method), and hard…

Cited by 106SourcePDFScholar
2015

Optimizing Non-decomposable Performance Measures: A Tale of Two Classes

ICML 2015poster

Modern classification problems frequently present mild to severe label imbalance as well as specific requirements on classification characteristics, and require optimizing performance measures that are non-decomposable over the dataset, such as F-measure. Such measures have spurred much interest and…

Cited by 65SourcePDFScholar
2015

Predtron: A Family of Online Algorithms for General Prediction Problems

NeurIPS 2015poster

Modern prediction problems arising in multilabel learning and learning to rank pose unique challenges to the classical theory of supervised learning. These problems have large prediction and label spaces of a combinatorial nature and involve sophisticated loss functions. We offer a general framework…

Cited by 3SourcePDFScholar
2015

Sparse Local Embeddings for Extreme Multi-label Classification

NeurIPS 2015poster

The objective in extreme multi-label learning is to train a classifier that can automatically tag a novel data point with the most relevant subset of labels from an extremely large label set. Embedding based approaches make training and prediction tractable by assuming that the training label matrix…

Cited by 562SourcePDFScholar