Lee, Jason D.

86 publications

NeurIPS 2025 Accelerating RL for LLM Reasoning with Optimal Advantage Regression Kianté Brantley, Mingyu Chen, Zhaolin Gao, Jason D. Lee, Wen Sun, Wenhao Zhan, Xuezhou Zhang

ICLR 2025 Correcting the Mythos of KL-Regularization: Direct Alignment Without Overoptimization via Chi-Squared Preference Optimization Audrey Huang, Wenhao Zhan, Tengyang Xie, Jason D. Lee, Wen Sun, Akshay Krishnamurthy, Dylan J Foster

NeurIPS 2025 Deployment Efficient Reward-Free Exploration with Linear Function Approximation Zihan Zhang, Yuxin Chen, Jason D. Lee, Simon Shaolei Du, Lin Yang, Ruosong Wang

ICML 2025 Discrepancies Are Virtue: Weak-to-Strong Generalization Through Lens of Intrinsic Dimension Yijun Dong, Yicheng Li, Yunai Li, Jason D. Lee, Qi Lei

NeurIPS 2025 Emergence and Scaling Laws in SGD Learning of Shallow Neural Networks Yunwei Ren, Eshaan Nichani, Denny Wu, Jason D. Lee

ICLR 2025 Exploiting Structure in Offline Multi-Agent RL: The Benefits of Low Interaction Rank Wenhao Zhan, Scott Fujimoto, Zheqing Zhu, Jason D. Lee, Daniel Jiang, Yonathan Efroni

AISTATS 2025 How Well Can Transformers Emulate In-Context Newton’s Method? Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, Jason D. Lee

COLT 2025 Learning Compositional Functions with Transformers from Easy-to-Hard Data Zixuan Wang, Eshaan Nichani, Alberto Bietti, Alex Damian, Daniel Hsu, Jason D Lee, Denny Wu

ICLR 2025 Learning Hierarchical Polynomials of Multiple Nonlinear Features Hengyu Fu, Zihao Wang, Eshaan Nichani, Jason D. Lee

NeurIPS 2025 Learning Orthogonal Multi-Index Models: A Fine-Grained Information Exponent Analysis Yunwei Ren, Jason D. Lee

ICML 2025 Metastable Dynamics of Chain-of-Thought Reasoning: Provable Benefits of Search, RL and Distillation Juno Kim, Denny Wu, Jason D. Lee, Taiji Suzuki

ICML 2025 Minimax Optimal Regret Bound for Reinforcement Learning with Trajectory Feedback Zihan Zhang, Yuxin Chen, Jason D. Lee, Simon Shaolei Du, Ruosong Wang

ICLR 2025 Regressing the Relative Future: Efficient Policy Optimization for Multi-Turn RLHF Zhaolin Gao, Wenhao Zhan, Jonathan Daniel Chang, Gokul Swamy, Kianté Brantley, Jason D. Lee, Wen Sun

ICML 2025 Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding Jiajun Zhu, Peihao Wang, Ruisi Cai, Jason D. Lee, Pan Li, Zhangyang Wang

TMLR 2025 Task Diversity Shortens the In-Context Learning Plateau Jaeyeon Kim, Sehyun Kwon, Joo Young Choi, Jongho Park, Jaewoong Cho, Jason D. Lee, Ernest K. Ryu

NeurIPS 2025 The Generative Leap: Tight Sample Complexity for Efficiently Learning Gaussian Multi-Index Models Alex Damian, Jason D. Lee, Joan Bruna

ICLR 2025 Transformers Learn to Implement Multi-Step Gradient Descent with Chain of Thought Jianhao Huang, Zixuan Wang, Jason D. Lee

ICLR 2025 Transformers Provably Learn Two-Mixture of Linear Classification via Gradient Flow Hongru Yang, Zhangyang Wang, Jason D. Lee, Yingbin Liang

ICLR 2025 Understanding Factual Recall in Transformers via Associative Memories Eshaan Nichani, Jason D. Lee, Alberto Bietti

ICLR 2025 Understanding Optimization in Deep Learning with Central Flows Jeremy Cohen, Alex Damian, Ameet Talwalkar, J Zico Kolter, Jason D. Lee

NeurIPS 2025 What Makes a Reward Model a Good Teacher? an Optimization Perspective Noam Razin, Zixuan Wang, Hubert Strauss, Stanley Wei, Jason D. Lee, Sanjeev Arora

NeurIPS 2025 What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains Chanakya Ekbote, Ashok Vardhan Makkuva, Marco Bondaschi, Nived Rajaraman, Michael Gastpar, Jason D. Lee, Paul Pu Liang

ICML 2024 An Information-Theoretic Analysis of In-Context Learning Hong Jun Jeon, Jason D. Lee, Qi Lei, Benjamin Van Roy

NeurIPS 2024 BitDelta: Your Fine-Tune May Only Be Worth One Bit James Liu, Guangxuan Xiao, Kai Li, Jason D. Lee, Song Han, Tri Dao, Tianle Cai

ICLR 2024 Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D. Lee, Wei Hu

ICLR 2024 Horizon-Free Regret for Linear Markov Decision Processes Zihan Zhang, Jason D. Lee, Yuxin Chen, Simon Shaolei Du

ICML 2024 How Transformers Learn Causal Structure with Gradient Descent Eshaan Nichani, Alex Damian, Jason D. Lee

ICLR 2024 Learning Hierarchical Polynomials with Three-Layer Neural Networks Zihao Wang, Eshaan Nichani, Jason D. Lee

NeurIPS 2024 Learning and Transferring Sparse Contextual Bigrams with Linear Transformers Yunwei Ren, Zixuan Wang, Jason D. Lee

ICML 2024 LoRA Training in the NTK Regime Has No Spurious Local Minima Uijeong Jang, Jason D. Lee, Ernest K. Ryu

ICML 2024 Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao

ICMLW 2024 Neural Network Learns Low-Dimensional Polynomials with SGD near the Information-Theoretic Limit Jason D. Lee, Kazusato Oko, Taiji Suzuki, Denny Wu

NeurIPS 2024 Neural Network Learns Low-Dimensional Polynomials with SGD near the Information-Theoretic Limit Jason D. Lee, Kazusato Oko, Taiji Suzuki, Denny Wu

COLT 2024 Optimal Multi-Distribution Learning Zihan Zhang, Wenhao Zhan, Yuxin Chen, Simon S Du, Jason D Lee

ICLR 2024 Provable Offline Preference-Based Reinforcement Learning Wenhao Zhan, Masatoshi Uehara, Nathan Kallus, Jason D. Lee, Wen Sun

ICLR 2024 Provable Reward-Agnostic Preference-Based Reinforcement Learning Wenhao Zhan, Masatoshi Uehara, Wen Sun, Jason D. Lee

ICLR 2024 Provably Efficient CVaR RL in Low-Rank MDPs Yulai Zhao, Wenhao Zhan, Xiaoyan Hu, Ho-fung Leung, Farzan Farnia, Wen Sun, Jason D. Lee

NeurIPS 2024 REBEL: Reinforcement Learning via Regressing Relative Rewards Zhaolin Gao, Jonathan D. Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, J. Andrew Bagnell, Jason D. Lee, Wen Sun

ICMLW 2024 REBEL: Reinforcement Learning via Regressing Relative Rewards Zhaolin Gao, Jonathan Daniel Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, J. Andrew Bagnell, Jason D. Lee, Wen Sun

ICMLW 2024 REBEL: Reinforcement Learning via Regressing Relative Rewards Zhaolin Gao, Jonathan Daniel Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, J. Andrew Bagnell, Jason D. Lee, Wen Sun

ICML 2024 Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark Yihua Zhang, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin-Yu Chen, Jason D. Lee, Wotao Yin, Mingyi Hong, Zhangyang Wang, Sijia Liu, Tianlong Chen

NeurIPS 2024 Scaling Laws in Linear Regression: Compute, Parameters, and Data Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter L. Bartlett, Jason D. Lee

COLT 2024 Settling the Sample Complexity of Online Reinforcement Learning Zihan Zhang, Yuxin Chen, Jason D Lee, Simon S Du

NeurIPS 2024 Stochastic Zeroth-Order Optimization Under Strongly Convexity and Lipschitz Hessian: Minimax Sample Complexity Qian Yu, Yining Wang, Baihe Huang, Qi Lei, Jason D. Lee

ICLR 2024 Teaching Arithmetic to Small Transformers Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos

ICML 2024 Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot Zixuan Wang, Stanley Wei, Daniel Hsu, Jason D. Lee

NeurIPSW 2024 Understanding Factual Recall in Transformers via Associative Memories Eshaan Nichani, Jason D. Lee, Alberto Bietti

ICLR 2023 Can We Find Nash Equilibria at a Linear Rate in Markov Games? Zhuoqing Song, Jason D. Lee, Zhuoran Yang

ICML 2023 Computationally Efficient PAC RL in POMDPs with Latent Determinism and Conditional Embeddings Masatoshi Uehara, Ayush Sekhari, Jason D. Lee, Nathan Kallus, Wen Sun

ICLR 2023 Decentralized Optimistic Hyperpolicy Mirror Descent: Provably No-Regret Learning in Markov Games Wenhao Zhan, Jason D. Lee, Zhuoran Yang

ICML 2023 Efficient Displacement Convex Optimization with Particle Gradient Descent Hadi Daneshmand, Jason D. Lee, Chi Jin

ICMLW 2023 Fine-Tuning Language Models with Just Forward Passes Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Jason D. Lee, Danqi Chen, Sanjeev Arora

ICMLW 2023 Fine-Tuning Language Models with Just Forward Passes Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, Sanjeev Arora

ICMLW 2023 How to Query Human Feedback Efficiently in RL? Wenhao Zhan, Masatoshi Uehara, Wen Sun, Jason D. Lee

ICMLW 2023 How to Query Human Feedback Efficiently in RL? Wenhao Zhan, Masatoshi Uehara, Wen Sun, Jason D. Lee

ICML 2023 Local Optimization Achieves Global Optimality in Multi-Agent Reinforcement Learning Yulai Zhao, Zhuoran Yang, Zhaoran Wang, Jason D. Lee

ICML 2023 Looped Transformers as Programmable Computers Angeliki Giannou, Shashank Rajput, Jy-Yong Sohn, Kangwook Lee, Jason D. Lee, Dimitris Papailiopoulos

ICLRW 2023 Looped Transformers as Programmable Computers Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, Dimitris Papailiopoulos

AISTATS 2023 Optimal Sample Complexity Bounds for Non-Convex Optimization Under Kurdyka-Lojasiewicz Condition Qian Yu, Yining Wang, Baihe Huang, Qi Lei, Jason D. Lee

ICLR 2023 PAC Reinforcement Learning for Predictive State Representations Wenhao Zhan, Masatoshi Uehara, Wen Sun, Jason D. Lee

ICMLW 2023 Provable Offline Reinforcement Learning with Human Feedback Wenhao Zhan, Masatoshi Uehara, Nathan Kallus, Jason D. Lee, Wen Sun

ICMLW 2023 Provable Offline Reinforcement Learning with Human Feedback Wenhao Zhan, Masatoshi Uehara, Nathan Kallus, Jason D. Lee, Wen Sun

L4DC 2023 Regret Guarantees for Online Deep Control Xinyi Chen, Edgar Minasyan, Jason D. Lee, Elad Hazan

ICMLW 2023 Reward Collapse in Aligning Large Language Models: A Prompt-Aware Approach to Preference Rankings Ziang Song, Tianle Cai, Jason D. Lee, Weijie J Su

ICMLW 2023 Scaling In-Context Demonstrations with Structured Attention Tianle Cai, Kaixuan Huang, Jason D. Lee, Mengdi Wang

ICLR 2023 Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability Alex Damian, Eshaan Nichani, Jason D. Lee

ICML 2023 Understanding Incremental Learning of Gradient Descent: A Fine-Grained Analysis of Matrix Sensing Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, Jason D. Lee

NeurIPSW 2022 Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability Alex Damian, Eshaan Nichani, Jason D. Lee

ICLR 2022 Towards General Function Approximation in Zero-Sum Markov Games Baihe Huang, Jason D. Lee, Zhaoran Wang, Zhuoran Yang

ICLR 2021 Few-Shot Learning via Learning the Representation, Provably Simon Shaolei Du, Wei Hu, Sham M. Kakade, Jason D. Lee, Qi Lei

ICLR 2021 Impact of Representation Learning in Linear Bandits Jiaqi Yang, Wei Hu, Jason D. Lee, Simon Shaolei Du

JMLR 2021 On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift Alekh Agarwal, Sham M. Kakade, Jason D. Lee, Gaurav Mahajan

ICLR 2020 Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks Yu Bai, Jason D. Lee

COLT 2020 Kernel and Rich Regimes in Overparametrized Models Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, Nathan Srebro

COLT 2020 Optimality and Approximation with Policy Gradient Methods in Markov Decision Processes Alekh Agarwal, Sham M Kakade, Jason D Lee, Gaurav Mahajan

ICLR 2018 Learning One-Hidden-Layer Neural Networks with Landscape Design Rong Ge, Jason D. Lee, Tengyu Ma

ICLR 2018 When Is a Convolutional Filter Easy to Learn? Simon S. Du, Jason D. Lee, Yuandong Tian

AISTATS 2017 Black-Box Importance Sampling Qiang Liu, Jason D. Lee

JMLR 2017 Communication-Efficient Sparse Regression Jason D. Lee, Qiang Liu, Yuekai Sun, Jonathan E. Taylor

JMLR 2017 Distributed Stochastic Variance Reduced Gradient Methods by Sampling Extra Data with Replacement Jason D. Lee, Qihang Lin, Tengyu Ma, Tianbao Yang

AISTATS 2017 On the Learnability of Fully-Connected Neural Networks Yuchen Zhang, Jason D. Lee, Martin J. Wainwright, Michael I. Jordan

AISTATS 2017 Sketching Meets Random Projection in the Dual: A Provable Recovery Algorithm for Big and High-Dimensional Data Jialei Wang, Jason D. Lee, Mehrdad Mahdavi, Mladen Kolar, Nati Srebro

COLT 2016 Gradient Descent Only Converges to Minimizers Jason D. Lee, Max Simchowitz, Michael I. Jordan, Benjamin Recht

ICML 2016 L1-Regularized Neural Networks Are Improperly Learnable in Polynomial Time Yuchen Zhang, Jason D. Lee, Michael I. Jordan

JMLR 2015 Matrix Completion and Low-Rank SVD via Fast Alternating Least Squares Trevor Hastie, Rahul Mazumder, Jason D. Lee, Reza Zadeh

AISTATS 2013 Structure Learning of Mixed Graphical Models Jason D. Lee, Trevor Hastie