Lee, Jason

76 publications

TMLR 2026 LZ Penalty: An Information-Theoretic Repetition Penalty for Autoregressive Language Models. Tony A Ginart, Naveen Kodali, Jason Lee, Caiming Xiong, Silvio Savarese, John Emmons

COLT 2025 Anytime Acceleration of Gradient Descent Zihan Zhang, Jason Lee, Simon Du, Yuxin Chen

CoRL 2025 BranchOut: Capturing Realistic Multimodality in Autonomous Driving Decisions Hee Jae Kim, Zekai Yin, Lei Lai, Jason Lee, Eshed Ohn-Bar

COLT 2024 Computational-Statistical Gaps in Gaussian Single-Index Models (Extended Abstract) Alex Damian, Loucas Pillaud-Vivien, Jason Lee, Joan Bruna

NeurIPS 2024 MUVERA: Multi-Vector Retrieval via Fixed Dimensional Encoding Laxman Dhulipala, Majid Hadian, Rajesh Jayaram, Jason Lee, Vahab Mirrokni

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

NeurIPS 2023 Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability Jingfeng Wu, Vladimir Braverman, Jason Lee

NeurIPS 2023 Offline Minimax Soft-Q-Learning Under Realizability and Partial Coverage Masatoshi Uehara, Nathan Kallus, Jason Lee, Wen Sun

NeurIPS 2023 Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks Eshaan Nichani, Alex Damian, Jason Lee

AISTATS 2023 Provable Hierarchy-Based Meta-Reinforcement Learning Kurtland Chua, Qi Lei, Jason Lee

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

AISTATS 2023 Provably Efficient Reinforcement Learning via Surprise Bound Hanlin Zhu, Ruosong Wang, Jason Lee

AISTATS 2023 Reconstructing Training Data from Model Gradient, Provably Zihan Wang, Jason Lee, Qi Lei

NeurIPS 2023 Reward-Agnostic Fine-Tuning: Provable Statistical Benefits of Hybrid Reinforcement Learning Gen Li, Wenhao Zhan, Jason Lee, Yuejie Chi, Yuxin Chen

NeurIPS 2023 Sample Complexity for Quadratic Bandits: Hessian Dependent Bounds and Optimal Algorithms Qian Yu, Yining Wang, Baihe Huang, Qi Lei, Jason Lee

NeurIPS 2023 Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models Alex Damian, Eshaan Nichani, Rong Ge, Jason Lee

NeurIPSW 2023 Teaching Arithmetic to Small Transformers Nayoung Lee, Kartik Sreenivasan, Jason Lee, Kangwook Lee, Dimitris Papailiopoulos

NeurIPSW 2023 Towards Optimal Statistical Watermarking Baihe Huang, Banghua Zhu, Hanlin Zhu, Jason Lee, Jiantao Jiao, Michael Jordan

AISTATS 2022 Provably Efficient Policy Optimization for Two-Player Zero-Sum Markov Games Yulai Zhao, Yuandong Tian, Jason Lee, Simon Du

NeurIPS 2022 From Gradient Flow on Population Loss to Learning with Stochastic Gradient Descent Christopher M De Sa, Satyen Kale, Jason Lee, Ayush Sekhari, Karthik Sridharan

NeurIPS 2022 Identifying Good Directions to Escape the NTK Regime and Efficiently Learn Low-Degree Plus Sparse Polynomials Eshaan Nichani, Yu Bai, Jason Lee

NeurIPS 2022 Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent Zhiyuan Li, Tianhao Wang, Jason Lee, Sanjeev Arora

COLT 2022 Neural Networks Can Learn Representations with Gradient Descent Alexandru Damian, Jason Lee, Mahdi Soltanolkotabi

COLT 2022 Offline Reinforcement Learning with Realizability and Single-Policy Concentrability Wenhao Zhan, Baihe Huang, Audrey Huang, Nan Jiang, Jason Lee

NeurIPS 2022 On the Effective Number of Linear Regions in Shallow Univariate ReLU Networks: Convergence Guarantees and Implicit Bias Itay Safran, Gal Vardi, Jason Lee

COLT 2022 Optimization-Based Separations for Neural Networks Itay Safran, Jason Lee

NeurIPS 2022 Provably Efficient Reinforcement Learning in Partially Observable Dynamical Systems Masatoshi Uehara, Ayush Sekhari, Jason Lee, Nathan Kallus, Wen Sun

ICML 2021 A Theory of Label Propagation for Subpopulation Shift Tianle Cai, Ruiqi Gao, Jason Lee, Qi Lei

ICML 2021 Bilinear Classes: A Structural Framework for Provable Generalization in RL Simon Du, Sham Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, Ruosong Wang

NeurIPS 2021 Going Beyond Linear RL: Sample Efficient Neural Function Approximation Baihe Huang, Kaixuan Huang, Sham Kakade, Jason Lee, Qi Lei, Runzhe Wang, Jiaqi Yang

NeurIPS 2021 How Fine-Tuning Allows for Effective Meta-Learning Kurtland Chua, Qi Lei, Jason Lee

ICML 2021 How Important Is the Train-Validation Split in Meta-Learning? Yu Bai, Minshuo Chen, Pan Zhou, Tuo Zhao, Jason Lee, Sham Kakade, Huan Wang, Caiming Xiong

NeurIPS 2021 Label Noise SGD Provably Prefers Flat Global Minimizers Alex Damian, Tengyu Ma, Jason Lee

COLT 2021 Modeling from Features: A Mean-Field Framework for Over-Parameterized Deep Neural Networks Cong Fang, Jason Lee, Pengkun Yang, Tong Zhang

ICML 2021 Near-Optimal Linear Regression Under Distribution Shift Qi Lei, Wei Hu, Jason Lee

NeurIPS 2021 Optimal Gradient-Based Algorithms for Non-Concave Bandit Optimization Baihe Huang, Kaixuan Huang, Sham Kakade, Jason Lee, Qi Lei, Runzhe Wang, Jiaqi Yang

NeurIPS 2021 Predicting What You Already Know Helps: Provable Self-Supervised Learning Jason Lee, Qi Lei, Nikunj Saunshi, Jiacheng Zhuo

COLT 2021 Shape Matters: Understanding the Implicit Bias of the Noise Covariance Jeff Z. HaoChen, Colin Wei, Jason Lee, Tengyu Ma

NeurIPS 2020 Agnostic $q$-Learning with Function Approximation in Deterministic Systems: Near-Optimal Bounds on Approximation Error and Sample Complexity Simon S Du, Jason Lee, Gaurav Mahajan, Ruosong Wang

NeurIPS 2020 Beyond Lazy Training for Over-Parameterized Tensor Decomposition Xiang Wang, Chenwei Wu, Jason Lee, Tengyu Ma, Rong Ge

NeurIPS 2020 Convergence of Meta-Learning with Task-Specific Adaptation over Partial Parameters Kaiyi Ji, Jason Lee, Yingbin Liang, H. Vincent Poor

NeurIPS 2020 Generalized Leverage Score Sampling for Neural Networks Jason Lee, Ruoqi Shen, Zhao Song, Mengdi Wang, Zheng Yu

NeurIPS 2020 How to Characterize the Landscape of Overparameterized Convolutional Neural Networks Yihong Gu, Weizhong Zhang, Cong Fang, Jason Lee, Tong Zhang

NeurIPS 2020 Implicit Bias in Deep Linear Classification: Initialization Scale vs Training Accuracy Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason Lee, Nati Srebro, Daniel Soudry

AAAI 2020 Latent-Variable Non-Autoregressive Neural Machine Translation with Deterministic Inference Using a Delta Posterior Raphael Shu, Jason Lee, Hideki Nakayama, Kyunghyun Cho

ICML 2020 Optimal Transport Mapping via Input Convex Neural Networks Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, Jason Lee

ICML 2020 SGD Learns One-Layer Networks in WGANs Qi Lei, Jason Lee, Alex Dimakis, Constantinos Daskalakis

NeurIPS 2020 Sanity-Checking Pruning Methods: Random Tickets Can Win the Jackpot Jingtong Su, Yihang Chen, Tianle Cai, Tianhao Wu, Ruiqi Gao, Liwei Wang, Jason Lee

NeurIPS 2020 Towards Understanding Hierarchical Learning: Benefits of Neural Representations Minshuo Chen, Yu Bai, Jason Lee, Tuo Zhao, Huan Wang, Caiming Xiong, Richard Socher

NeurIPS 2019 Convergence of Adversarial Training in Overparametrized Neural Networks Ruiqi Gao, Tianle Cai, Haochuan Li, Cho-Jui Hsieh, Liwei Wang, Jason Lee

AISTATS 2019 Convergence of Gradient Descent on Separable Data Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, Daniel Soudry

ICML 2019 Gradient Descent Finds Global Minima of Deep Neural Networks Simon Du, Jason Lee, Haochuan Li, Liwei Wang, Xiyu Zhai

ICML 2019 Lexicographic and Depth-Sensitive Margins in Homogeneous and Non-Homogeneous Deep Models Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, Daniel Soudry

NeurIPS 2019 Neural Temporal-Difference Learning Converges to Global Optima Qi Cai, Zhuoran Yang, Jason Lee, Zhaoran Wang

NeurIPS 2019 Regularization Matters: Generalization and Optimization of Neural Nets V.s. Their Induced Kernel Colin Wei, Jason Lee, Qiang Liu, Tengyu Ma

NeurIPS 2019 Solving a Class of Non-Convex Min-Max Games Using Iterative First Order Methods Maher Nouiehed, Maziar Sanjabi, Tianjian Huang, Jason Lee, Meisam Razaviyayn

NeurIPS 2018 Adding One Neuron Can Eliminate All Bad Local Minima Shiyu Liang, Ruoyu Sun, Jason Lee, R. Srikant

NeurIPS 2018 Algorithmic Regularization in Learning Deep Homogeneous Models: Layers Are Automatically Balanced Simon S Du, Wei Hu, Jason Lee

ICML 2018 Characterizing Implicit Bias in Terms of Optimization Geometry Suriya Gunasekar, Jason Lee, Daniel Soudry, Nathan Srebro

ICLR 2018 Emergent Translation in Multi-Agent Communication Jason Lee, Kyunghyun Cho, Jason Weston, Douwe Kiela

ICML 2018 Gradient Descent Learns One-Hidden-Layer CNN: Don’t Be Afraid of Spurious Local Minima Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, Barnabas Poczos

ICML 2018 Gradient Primal-Dual Algorithm Converges to Second-Order Stationary Solution for Nonconvex Distributed Optimization over Networks Mingyi Hong, Meisam Razaviyayn, Jason Lee

NeurIPS 2018 Implicit Bias of Gradient Descent on Linear Convolutional Networks Suriya Gunasekar, Jason Lee, Daniel Soudry, Nati Srebro

NeurIPS 2018 On the Convergence and Robustness of Training GANs with Regularized Optimal Transport Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, Jason Lee

ICML 2018 On the Power of Over-Parametrization in Neural Networks with Quadratic Activation Simon Du, Jason Lee

NeurIPS 2018 Provably Correct Automatic Sub-Differentiation for Qualified Programs Sham M. Kakade, Jason Lee

NeurIPS 2017 Gradient Descent Can Take Exponential Time to Escape Saddle Points Simon S Du, Chi Jin, Jason Lee, Michael I Jordan, Aarti Singh, Barnabas Poczos

ICML 2016 A Kernelized Stein Discrepancy for Goodness-of-Fit Tests Qiang Liu, Jason Lee, Michael Jordan

NeurIPS 2016 Matrix Completion Has No Spurious Local Minimum Rong Ge, Jason Lee, Tengyu Ma

NeurIPS 2015 Evaluating the Statistical Significance of Biclusters Jason Lee, Yuekai Sun, Jonathan E Taylor

NeurIPS 2014 Exact Post Model Selection Inference for Marginal Screening Jason Lee, Jonathan E Taylor

NeurIPS 2014 Scalable Methods for Nonnegative Matrix Factorizations of Near-Separable Tall-and-Skinny Matrices Austin R Benson, Jason Lee, Bartek Rajwa, David F Gleich

NeurIPS 2013 On Model Selection Consistency of Penalized M-Estimators: A Geometric Theory Jason Lee, Yuekai Sun, Jonathan E Taylor

NeurIPS 2013 Using Multiple Samples to Learn Mixture Models Jason Lee, Ran Gilad-Bachrach, Rich Caruana

NeurIPS 2012 Proximal Newton-Type Methods for Convex Optimization Jason Lee, Yuekai Sun, Michael Saunders

NeurIPS 2010 Practical Large-Scale Optimization for Max-Norm Regularization Jason Lee, Ben Recht, Nathan Srebro, Joel Tropp, Ruslan Salakhutdinov