Publications (76)
Spectrum Approximation Beyond Fast Matrix Multiplication: Algorithms and Hardness
Cameron Musco, Praneeth Netrapalli, Aaron Sidford +2
Understanding the singular value spectrum of a matrix is a fundamental task in countless applications. In matrix multiplication time, it is possible…
A Faster Generalized Two-Stage Approximate Top-K
Yashas Samaga, Varun Yerram, Spandana Raj Babbula +2
We consider the Top- selection problem, which aims to identify the largest elements in an array. Top- selection arises in many machine learning algorithms and often becom…
Smoothed analysis for low-rank solutions to semidefinite programs in quadratic penalty form
Srinadh Bhojanapalli, Nicolas Boumal, Prateek Jain +1
Semidefinite programs (SDP) are important in learning and combinatorial optimization with numerous applications. In pursuit of low-rank solutions and low complexity algorithms, we…
A Clustering Approach to Learn Sparsely-Used Overcomplete Dictionaries
Alekh Agarwal, Animashree Anandkumar, Praneeth Netrapalli
We consider the problem of learning overcomplete dictionaries in the context of sparse coding, where each sample selects a sparse subset of dictionary elements. Our main result is…
Non-Gaussianity of Stochastic Gradient Noise
Abhishek Panigrahi, Raghav Somani, Navin Goyal +1
What enables Stochastic Gradient Descent (SGD) to achieve better generalization than Gradient Descent (GD) in Neural Network training? This question has attracted much attention. I…
Finding the Graph of Epidemic Cascades
Praneeth Netrapalli, Sujay Sanghavi
We consider the problem of finding the graph on which an epidemic cascade spreads, given only the times when each node gets infected. While this is a problem of importance in sever…
A Short Note on Concentration Inequalities for Random Vectors with SubGaussian Norm
Chi Jin, Praneeth Netrapalli, Rong Ge +2
In this note, we derive concentration inequalities for random vectors with subGaussian norm (a generalization of both subGaussian random vectors and norm bounded random vectors), w…
Phase Retrieval using Alternating Minimization
Praneeth Netrapalli, Prateek Jain, Sujay Sanghavi
Phase retrieval problems involve solving linear equations, but with missing sign (or phase, for complex numbers) information. More than four decades after it was first proposed, th…
Simplicity Bias in 1-Hidden Layer Neural Networks
Depen Morwani, Jatin Batra, Prateek Jain +1
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 presenc…
Provable Efficient Online Matrix Completion via Non-convex Stochastic Gradient Descent
Chi Jin, Sham M. Kakade, Praneeth Netrapalli
Matrix completion, where we wish to recover a low rank matrix by observing a few entries from it, is a widely studied problem in both theory and practice with wide applications. Mo…
Learning Planar Ising Models
Jason K. Johnson, Praneeth Netrapalli, Michael Chertkov
Inference and learning of graphical models are both well-studied problems in statistics and machine learning that have found many applications in science and engineering. However,…
Efficient Algorithms for Smooth Minimax Optimization
Kiran Koshy Thekumparampil, Prateek Jain, Praneeth Netrapalli +1
This paper studies first order methods for solving smooth minimax optimization problems where is smooth and is concave for each…
Parallelizing Stochastic Gradient Descent for Least Squares Regression: mini-batching, averaging, and model misspecification
Prateek Jain, Sham M. Kakade, Rahul Kidambi +2
This work characterizes the benefits of averaging schemes widely used in conjunction with stochastic gradient descent (SGD). In particular, this work provides a sharp analysis of:…
Near-optimal Offline and Streaming Algorithms for Learning Non-Linear Dynamical Systems
Prateek Jain, Suhas S Kowshik, Dheeraj Nagaraj +1
We consider the setting of vector valued non-linear dynamical systems , where is unbiased noise and is a known…
Optimal Regret Algorithm for Pseudo-1d Bandit Convex Optimization
Aadirupa Saha, Nagarajan Natarajan, Praneeth Netrapalli +1
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)…
Projection Efficient Subgradient Method and Optimal Nonsmooth Frank-Wolfe Method
Kiran Koshy Thekumparampil, Prateek Jain, Praneeth Netrapalli +1
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 orac…
Near-Optimal Lower Bounds For Convex Optimization For All Orders of Smoothness
Ankit Garg, Robin Kothari, Praneeth Netrapalli +1
We study the complexity of optimizing highly smooth convex functions. For a positive integer , we want to find an -approximate minimum of a convex function , given oracle…
Leverage Score Sampling for Faster Accelerated Regression and ERM
Naman Agarwal, Sham Kakade, Rahul Kidambi +3
Given a matrix and a vector , we show how to compute an -approximate solution to the regression problem $ \min_{x\in\…
Do Input Gradients Highlight Discriminative Features?
Harshay Shah, Prateek Jain, Praneeth Netrapalli
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 a…
On the insufficiency of existing momentum schemes for Stochastic Optimization
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain +1
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 a…
Fast Exact Matrix Completion with Finite Samples
Prateek Jain, Praneeth Netrapalli
Matrix completion is the problem of recovering a low rank matrix by observing a small fraction of its entries. A series of recent works [KOM12,JNS13,HW14] have proposed fast non-co…
Non-convex Robust PCA
Praneeth Netrapalli, U N Niranjan, Sujay Sanghavi +2
We propose a new method for robust PCA -- the task of recovering a low-rank matrix from sparse corruptions that are of unknown value and support. Our method involves alternating be…
MET: Masked Encoding for Tabular Data
Kushal Majmundar, Sachin Goyal, Praneeth Netrapalli +1
We consider the task of self-supervised representation learning (SSL) for tabular data: tabular-SSL. Typical contrastive learning based SSL methods require instance-wise data augme…
Follow the Perturbed Leader: Optimism and Fast Parallel Algorithms for Smooth Minimax Games
Arun Sai Suggala, Praneeth Netrapalli
We consider the problem of online learning and its application to solving minimax games. For the online learning problem, Follow the Perturbed Leader (FTPL) is a widely studied alg…
Online Target Q-learning with Reverse Experience Replay: Efficiently finding the Optimal Policy for Linear MDPs
Naman Agarwal, Syomantak Chaudhuri, Prateek Jain +2
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 re…
SGD without Replacement: Sharper Rates for General Smooth Convex Functions
Prateek Jain, Dheeraj Nagaraj, Praneeth Netrapalli
We study stochastic gradient descent {\em without replacement} (\sgdwor) for smooth convex functions. \sgdwor is widely observed to converge faster than true \sgd where each sample…
P-SIF: Document Embeddings Using Partition Averaging
Vivek Gupta, Ankit Saw, Pegah Nokhiz +3
Simple weighted averaging of word vectors often yields effective representations for sentences which outperform sophisticated seq2seq neural models in many tasks. While it is desir…
The Pitfalls of Simplicity Bias in Neural Networks
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan +2
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 neu…
The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares
Rong Ge, Sham M. Kakade, Rahul Kidambi +1
Minimax optimal convergence rates for classes of stochastic convex optimization problems are well characterized, where the majority of results utilize iterate averaged stochastic g…
DAFT: Distilling Adversarially Fine-tuned Models for Better OOD Generalization
Anshul Nasery, Sravanti Addepalli, Praneeth Netrapalli +1
We consider the problem of OOD generalization, where the goal is to train a model that performs well on test distributions that are different from the training distribution. Deep l…
Learning Sparsely Used Overcomplete Dictionaries via Alternating Minimization
Alekh Agarwal, Animashree Anandkumar, Prateek Jain +1
We consider the problem of sparse coding, where each sample consists of a sparse linear combination of a set of dictionary atoms, and the task is to learn both the dictionary eleme…
Reproducibility in Optimization: Theoretical Framework and Limits
Kwangjun Ahn, Prateek Jain, Ziwei Ji +3
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 op…
Robust Shift-and-Invert Preconditioning: Faster and More Sample Efficient Algorithms for Eigenvector Computation
Chi Jin, Sham M. Kakade, Cameron Musco +2
We provide faster algorithms and improved sample complexities for approximating the top eigenvector of a matrix. Offline Setting: Given an matrix , we show how to c…
Online Non-Convex Learning: Following the Perturbed Leader is Optimal
Arun Sai Suggala, Praneeth Netrapalli
We study the problem of online learning with non-convex losses, where the learner has access to an offline optimization oracle. We show that the classical Follow the Perturbed Lead…
Accelerating Stochastic Gradient Descent For Least Squares Regression
Prateek Jain, Sham M. Kakade, Rahul Kidambi +2
There is widespread sentiment that it is not possible to effectively utilize fast gradient methods (e.g. Nesterov's acceleration, conjugate gradient, heavy ball) for the purposes o…
Greedy Learning of Markov Network Structure
Praneeth Netrapalli, Siddhartha Banerjee, Sujay Sanghavi +1
We propose a new yet natural algorithm for learning the graph structure of general discrete graphical models (a.k.a. Markov random fields) from samples. Our algorithm finds the nei…
Consistent Multiclass Algorithms for Complex Metrics and Constraints
Harikrishna Narasimhan, Harish G. Ramaswamy, Shiv Kumar Tavker +3
We present consistent algorithms for multiclass learning with complex performance metrics and constraints, where the objective and constraints are defined by arbitrary functions of…
Learning Planar Ising Models
Jason K. Johnson, Diane Oyen, Michael Chertkov +1
Inference and learning of graphical models are both well-studied problems in statistics and machine learning that have found many applications in science and engineering. However,…
Non-Reconstructability in the Stochastic Block Model
Joe Neeman, Praneeth Netrapalli
We consider the problem of clustering (or reconstruction) in the stochastic block model, in the regime where the average degree is constant. For the case of two clusters with equal…
Streaming PCA: Matching Matrix Bernstein and Near-Optimal Finite Sample Guarantees for Oja's Algorithm
Prateek Jain, Chi Jin, Sham M. Kakade +2
This work provides improved guarantees for streaming principle component analysis (PCA). Given sampled independently from distributions…
The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks
Lénaïc Chizat, Praneeth Netrapalli
Deep learning succeeds by doing hierarchical feature learning, yet tuning hyper-parameters (HP) such as initialization scales, learning rates etc., only give indirect control over…
HiRE: High Recall Approximate Top- Estimation for Efficient LLM Inference
Yashas Samaga B L, Varun Yerram, Chong You +4
Autoregressive decoding with generative Large Language Models (LLMs) on accelerators (GPUs/TPUs) is often memory-bound where most of the time is spent on transferring model paramet…
Global Convergence of Non-Convex Gradient Descent for Computing Matrix Squareroot
Prateek Jain, Chi Jin, Sham M. Kakade +1
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…
Making the Last Iterate of SGD Information Theoretically Optimal
Prateek Jain, Dheeraj Nagaraj, Praneeth Netrapalli
Stochastic gradient descent (SGD) is one of the most widely used algorithms for large scale optimization problems. While classical theoretical analysis of SGD for convex problems s…
Thresholding based Efficient Outlier Robust PCA
Yeshwanth Cherapanamjeri, Prateek Jain, Praneeth Netrapalli
We consider the problem of outlier robust PCA (OR-PCA) where the goal is to recover principal directions despite the presence of outlier data points. That is, given a data matrix $…
Near Optimal Heteroscedastic Regression with Symbiotic Learning
Dheeraj Baby, Aniket Das, Dheeraj Nagaraj +1
We consider the problem of heteroscedastic linear regression, where, given samples from $y_i = \langle \mathbf{w}^{*}, \mathbf{x}_i \rangle + ε_i \cdot \…
Learning Minimax Estimators via Online Learning
Kartik Gupta, Arun Sai Suggala, Adarsh Prasad +2
We consider the problem of designing minimax estimators for estimating the parameters of a probability distribution. Unlike classical approaches such as the MLE and minimum distanc…
Multi-User Reinforcement Learning with Low Rank Rewards
Naman Agarwal, Prateek Jain, Suhas Kowshik +2
In this work, we consider the problem of collaborative multi-user reinforcement learning. In this setting there are multiple users with the same state-action space and transition p…
Faster Eigenvector Computation via Shift-and-Invert Preconditioning
Dan Garber, Elad Hazan, Chi Jin +4
We give faster algorithms and improved sample complexities for estimating the top eigenvector of a matrix -- i.e. computing a unit vector such that $x^T Σx \ge (1-ε)λ_1…
Learning an Invertible Output Mapping Can Mitigate Simplicity Bias in Neural Networks
Sravanti Addepalli, Anshul Nasery, R. Venkatesh Babu +2
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 Simplicity Bias (…
Streaming Linear System Identification with Reverse Experience Replay
Prateek Jain, Suhas S Kowshik, Dheeraj Nagaraj +1
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…
Efficient Domain Generalization via Common-Specific Low-Rank Decomposition
Vihari Piratla, Praneeth Netrapalli, Sunita Sarawagi
Domain generalization refers to the task of training a model which generalizes to new domains that are not seen during training. We present CSD (Common Specific Decomposition), for…
On Nonconvex Optimization for Machine Learning: Gradients, Stochasticity, and Saddle Points
Chi Jin, Praneeth Netrapalli, Rong Ge +2
Gradient descent (GD) and stochastic gradient descent (SGD) are the workhorses of large-scale machine learning. While classical theory focused on analyzing the performance of these…
Information-theoretic thresholds for community detection in sparse networks
Jess Banks, Cristopher Moore, Joe Neeman +1
We give upper and lower bounds on the information-theoretic threshold for community detection in the stochastic block model. Specifically, consider the symmetric stochastic block m…
Focus on the Common Good: Group Distributional Robustness Follows
Vihari Piratla, Praneeth Netrapalli, Sunita Sarawagi
We consider the problem of training a classification model with group annotated training data. Recent work has established that, if there is distribution shift across different gro…
Second Order Methods for Bandit Optimization and Control
Arun Suggala, Y. Jennifer Sun, Praneeth Netrapalli +1
Bandit convex optimization (BCO) is a general framework for online decision making under uncertainty. While tight regret bounds for general convex losses have been established, exi…
Efficient Algorithms for Large-scale Generalized Eigenvector Computation and Canonical Correlation Analysis
Rong Ge, Chi Jin, Sham M. Kakade +2
This paper considers the problem of canonical-correlation analysis (CCA) (Hotelling, 1936) and, more broadly, the generalized eigenvector problem for a pair of symmetric matrices.…
Inferring Asteroseismic Parameters from Short Observations Using Deep Learning: Application to TESS and K2 Red Giants
Nipun Ghanghas, Siddharth Dhanpal, Shravan Hanasoge +2
Asteroseismology is the study of resonant oscillations of stars to infer their internal structure and dynamics. It is also a powerful tool for precisely determining stellar paramet…
Sample Efficient Linear Meta-Learning by Alternating Minimization
Kiran Koshy Thekumparampil, Prateek Jain, Praneeth Netrapalli +1
Meta-learning synthesizes and leverages the knowledge from a given set of tasks to rapidly learn new tasks using very little data. Meta-learning of linear regression tasks, where t…
Low-rank Matrix Completion using Alternating Minimization
Prateek Jain, Praneeth Netrapalli, Sujay Sanghavi
Alternating minimization represents a widely applicable and empirically successful approach for finding low-rank matrices that best fit the given data. For example, for the problem…
Compressing Many-Shots in In-Context Learning
Devvrit Khatri, Pranamya Kulkarni, Nilesh Gupta +9
Large Language Models (LLMs) have been shown to be able to learn different tasks without explicit finetuning when given many input-output examples / demonstrations through In-Conte…
Optimistic MLE -- A Generic Model-based Algorithm for Partially Observable Sequential Decision Making
Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvári +1
This paper introduces a simple efficient learning algorithms for general sequential decision making. The algorithm combines Optimism for exploration with Maximum Likelihood Estimat…
MOReL : Model-Based Offline Reinforcement Learning
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli +1
In offline reinforcement learning (RL), the goal is to learn a highly rewarding policy based solely on a dataset of historical interactions with the environment. The ability to tra…
Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities
Gheorghe Comanici, Eric Bieber, Mike Schaekermann +3431
In this report, we introduce the Gemini 2.X model family: Gemini 2.5 Pro and Gemini 2.5 Flash, as well as our earlier Gemini 2.0 Flash and Flash-Lite models. Gemini 2.5 Pro is our…
Accelerated Gradient Descent Escapes Saddle Points Faster than Gradient Descent
Chi Jin, Praneeth Netrapalli, Michael I. Jordan
Nesterov's accelerated gradient descent (AGD), an instance of the general family of "momentum methods", provably achieves faster convergence rate than gradient descent (GD) in the…
Spark Transformer: Reactivating Sparsity in FFN and Attention
Chong You, Kan Wu, Zhipeng Jia +16
The discovery of the lazy neuron phenomenon in trained Transformers, where the vast majority of neurons in their feed-forward networks (FFN) are inactive for each token, has spurre…
Tandem Transformers for Inference Efficient LLMs
Aishwarya P S, Pranav Ajit Nair, Yashas Samaga +4
The autoregressive nature of conventional large language models (LLMs) inherently limits inference speed, as tokens are generated sequentially. While speculative and parallel decod…
Minimax Optimization with Smooth Algorithmic Adversaries
Tanner Fiez, Chi Jin, Praneeth Netrapalli +1
This paper considers minimax optimization in the challenging setting where can be both nonconvex in and nonconcave in . Though such optimization…
Least Squares Regression with Markovian Data: Fundamental Limits and Algorithms
Guy Bresler, Prateek Jain, Dheeraj Nagaraj +2
We study the problem of least squares linear regression where the data-points are dependent and are sampled from a Markov chain. We establish sharp information theoretic minimax lo…
A model of errors in transformers
Suvrat Raju, Praneeth Netrapalli
We study the error rate of LLMs on tasks like arithmetic that require a deterministic output, and repetitive processing of tokens drawn from a small set of alternatives. We argue t…
All Mistakes Are Not Equal: Comprehensive Hierarchy Aware Multi-label Predictions (CHAMP)
Ashwin Vaswani, Gaurav Aggarwal, Praneeth Netrapalli +1
This paper considers the problem of Hierarchical Multi-Label Classification (HMC), where (i) several labels can be present for each example, and (ii) labels are related via a domai…
No quantum speedup over gradient descent for non-smooth convex optimization
Ankit Garg, Robin Kothari, Praneeth Netrapalli +1
We study the first-order convex optimization problem, where we have black-box access to a (not necessarily smooth) function and its (sub)gradient. O…
Convergence Rates of Active Learning for Maximum Likelihood Estimation
Kamalika Chaudhuri, Sham Kakade, Praneeth Netrapalli +1
An active learner is given a class of models, a large set of unlabeled examples, and the ability to interactively query labels of a subset of these examples; the goal of the learne…
A Markov Chain Theory Approach to Characterizing the Minimax Optimality of Stochastic Gradient Descent (for Least Squares)
Prateek Jain, Sham M. Kakade, Rahul Kidambi +3
This work provides a simplified proof of the statistical minimax optimality of (iterate averaged) stochastic gradient descent (SGD), for the special case of least squares. This res…
What is Local Optimality in Nonconvex-Nonconcave Minimax Optimization?
Chi Jin, Praneeth Netrapalli, Michael I. Jordan
Minimax optimization has found extensive applications in modern machine learning, in settings such as generative adversarial networks (GANs), adversarial training and multi-agent r…
How to Escape Saddle Points Efficiently
Chi Jin, Rong Ge, Praneeth Netrapalli +2
This paper shows that a perturbed form of gradient descent converges to a second-order stationary point in a number iterations which depends only poly-logarithmically on dimension…