← Search

Sham Kakade

38 accepted papers

2026

Exploring Relational Reasoning Capabilities in LLMs with REL

ICML 2026poster

Relational reasoning is the ability to infer relations that jointly bind multiple entities, attributes, or variables. While this capability is essential for scientific reasoning, most existing evaluations of relational reasoning in large language models focus on structured inputs such as tables, gra…

Cited by 0SourceScholar
2026

Fine-Tuning Masked Diffusion for Provable Self-Correction

ICML 2026poster

A natural desideratum for generative models is \emph{self-correction}--detecting and revising low-quality tokens at inference. While Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces, their capacity for self-correction remains poorly under…

Cited by 0SourceScholar
2026

Prescriptive Scaling Reveals the Evolution of Language Model Capabilities

ICML 2026oral

For deploying foundation models, practitioners increasingly need prescriptive scaling laws: given a pre-training compute budget, what downstream accuracy is attainable with contemporary post-training practice, and how stable is that mapping as the field evolves? Using large-scale observational evalu…

Cited by 0SourceScholar
2026

Stop Training for the Worst: Progressive Unmasking Accelerates Masked Diffusion Training

ICML 2026poster

Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces. By generating sequences in any order and allowing for parallel decoding, they enable fast inference and strong performance on non-causal tasks. However, this flexibility comes with a *trai…

Cited by 0SourceScholar
2025

LoRA Soups: Merging LoRAs for Practical Skill Composition Tasks

COLING 2025industry

Low-Rank Adaptation (LoRA) is a popular technique for parameter-efficient fine-tuning of Large Language Models (LLMs). We study how different LoRA modules can be merged to achieve skill composition—testing the performance of the merged model on a target task that involves combining multiple skills,…

2024

A Study on the Calibration of In-context Learning

NAACL 2024long

Accurate uncertainty quantification is crucial for the safe deployment of machine learning models, and prior research has demonstrated improvements in the calibration of modern language models (LMs). We study in-context learning (ICL), a prevalent method for adapting static LMs through tailored prom…

2022

Inductive Biases and Variable Creation in Self-Attention Mechanisms

ICML 2022spotlight

Self-attention, an architectural motif designed to model long-range interactions in sequential data, has driven numerous recent breakthroughs in natural language processing and beyond. This work provides a theoretical analysis of the inductive biases of self-attention modules. Our focus is to rigoro…

Cited by 159SourcePDFScholar
2022

Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regression

ICML 2022oral

Stochastic gradient descent (SGD) has been shown to generalize well in many deep learning applications. In practice, one often runs SGD with a geometrically decaying stepsize, i.e., a constant initial stepsize followed by multiple geometric stepsize decay, and uses the last iterate as the output. Th…

Cited by 32SourcePDFScholar
2022

Sparsity in Partially Controllable Linear Systems

ICML 2022spotlight

A fundamental concept in control theory is that of controllability, where any system state can be reached through an appropriate choice of control inputs. Indeed, a large body of classical and modern approaches are designed for controllable linear dynamical systems. However, in practice, we often en…

Cited by 14SourcePDFScholar
2022

Understanding Contrastive Learning Requires Incorporating Inductive Biases

ICML 2022spotlight

Contrastive learning is a popular form of self-supervised learning that encourages augmentations (views) of the same input to have more similar representations compared to augmentations of different inputs. Recent attempts to theoretically explain the success of contrastive learning on downstream cl…

Cited by 141SourcePDFScholar
2021

Bilinear Classes: A Structural Framework for Provable Generalization in RL

ICML 2021oral

This work introduces Bilinear Classes, a new structural framework, which permit generalization in reinforcement learning in a wide variety of settings through the use of function approximation. The framework incorporates nearly all existing models in which a polynomial sample complexity is achievabl…

Cited by 256SourcePDFScholar
2021

How Important is the Train-Validation Split in Meta-Learning?

ICML 2021spotlight

Meta-learning aims to perform fast adaptation on a new task through learning a “prior” from multiple existing tasks. A common practice in meta-learning is to perform a train-validation split (\emph{train-val method}) where the prior adapts to the task on one split of the data, and the resulting pred…

Cited by 92SourcePDFScholar
2021

Instabilities of Offline RL with Pre-Trained Neural Representation

ICML 2021spotlight

In offline reinforcement learning (RL), we seek to utilize offline data to evaluate (or learn) policies in scenarios where the data are collected from a distribution that substantially differs from that of the target policy to be evaluated. Recent theoretical advances have shown that such sample-eff…

Cited by 56SourcePDFScholar
2020

Calibration, Entropy Rates, and Memory in Language Models

ICML 2020poster

Building accurate language models that capture meaningful long-term dependencies is a core challenge in natural language processing. Towards this end, we present a calibration-based approach to measure long-term discrepancies between a generative sequence model and the true distribution, and use the…

Cited by 46SourcePDFScholar
2020

FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs

NeurIPS 2020oral

In order to deal with the curse of dimensionality in reinforcement learning (RL), it is common practice to make parametric assumptions where values or policies are functions of some low dimensional feature space. This work focuses on the representation learning question: how can we learn such featur…

Cited by 308SourcePDFScholar
2020

Information Theoretic Regret Bounds for Online Nonlinear Control

NeurIPS 2020poster

This work studies the problem of sequential control in an unknown, nonlinear dynamical system, where we model the underlying system dynamics as an unknown function in a known Reproducing Kernel Hilbert Space. This framework yields a general setting that permits discrete and continuous control input…

Cited by 154SourcePDFScholar
2020

Model-Based Multi-Agent RL in Zero-Sum Markov Games with Near-Optimal Sample Complexity

NeurIPS 2020spotlight

Model-based reinforcement learning (RL), which finds an optimal policy using an empirical model, has long been recognized as one of the cornerstones of RL. It is especially suitable for multi-agent RL (MARL), as it naturally decouples the learning and the planning phases, and avoids the non-stationa…

Cited by 169SourcePDFScholar
2020

PC-PG: Policy Cover Directed Exploration for Provable Policy Gradient Learning

NeurIPS 2020poster

Direct policy gradient methods for reinforcement learning are a successful approach for a variety of reasons: they are model free, they directly optimize the performance metric of interest, and they allow for richly parameterized policies. Their primary drawback is that, by being local in nature, th…

Cited by 153SourcePDFScholar
2020

Provable Representation Learning for Imitation Learning via Bi-level Optimization

ICML 2020poster

A common strategy in modern learning systems is to learn a representation that is useful for many tasks, a.k.a. representation learning. We study this strategy in the imitation learning setting for Markov decision processes (MDPs) where multiple experts’ trajectories are available. We formulate repr…

Cited by 78SourcePDFScholar
2020

Robust Meta-learning for Mixed Linear Regression with Small Batches

NeurIPS 2020poster

A common challenge faced in practical supervised learning, such as medical image processing and robotic interactions, is that there are plenty of tasks but each task cannot afford to collect enough labeled examples to be learned in isolation. However, by exploiting the similarities across those task…

Cited by 43SourcePDFScholar
2020

Sample-Efficient Reinforcement Learning of Undercomplete POMDPs

NeurIPS 2020spotlight

Partial observability is a common challenge in many reinforcement learning applications, which requires an agent to maintain memory, infer latent states, and integrate this past information into exploration. This challenge leads to a number of computational and statistical hardness results for learn…

Cited by 96SourcePDFScholar
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…

2019

Maximum Likelihood Estimation for Learning Populations of Parameters

ICML 2019oral

Consider a setting with $N$ independent individuals, each with an unknown parameter, $p_i \in [0, 1]$ drawn from some unknown distribution $P^\star$. After observing the outcomes of $t$ independent Bernoulli trials, i.e., $X_i \sim \text{Binomial}(t, p_i)$ per individual, our objective is to accurat…

Cited by 51SourcePDFScholar
2019

Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Control

ICLR 2019poster

We propose a "plan online and learn offline" framework for the setting where an agent, with an internal model, needs to continually act and learn in the world. Our work builds on the synergistic relationship between local model-based control, global value function learning, and exploration. We study…

Cited by 284SourcePDFScholar
2018

Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulator

ICML 2018oral

Direct policy gradient methods for reinforcement learning and continuous control problems are a popular approach for a variety of reasons: 1) they are easy to implement without explicit knowledge of the underlying model, 2) they are an “end-to-end” approach, directly optimizing the performance metri…

Cited by 765SourcePDFScholar
2018

Variance Reduction for Policy Gradient with Action-Dependent Factorized Baselines

ICLR 2018oral

Policy gradient methods have enjoyed great success in deep reinforcement learning but suffer from high variance of gradient estimates. The high variance problem is particularly exasperated in problems with long horizons or high-dimensional action spaces. To mitigate this issue, we derive a bias-free…

Cited by 187SourcePDFScholar
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
2015

Un-regularizing: approximate proximal point and faster stochastic algorithms for empirical risk minimization

ICML 2015poster

We develop a family of accelerated stochastic algorithms that optimize sums of convex functions. Our algorithms improve upon the fastest running time for empirical risk minimization (ERM), and in particular linear least-squares regression, across a wide range of problem settings. To achieve this, we…

Cited by 177SourcePDFScholar