Jason D. Lee
 |
Jason D. Lee |
About Me
I am an associate professor of Electrical Engineering and Computer Science (secondary) at Princeton University.
Previously, I was a research scientist at Google Deepmind, member of the IAS, and an assistant professor at USC for three years. Before that, I was a postdoc in the Computer Science Department at UC Berkeley working with Michael I. Jordan, and also collaborated with Ben Recht. I received my PhD in Applied Math advised by Trevor Hastie and Jonathan Taylor. I received a BS in Mathematics from Duke University advised by Mauro Maggioni.
I am a native of Cupertino, CA.
My research interests are broadly in
Any email from prospective students, postdoc, or visitors, please include the string “filter_student” in the subject, else I will not receive your email. I advise students in Computer Science, Electrical Engineering, Math, ORFE, and PACM. I am recruiting PhD students and postdoctoral scholars starting in 2023 at Princeton University, please email me a CV apply. Princeton PhD students interested in machine learning, statistics, or optimization research, please contact me.
My current focus is on machine learning with a focus on foundations of AI, deep learning, representation learning, and deep reinforcement learning. I have lectured on the Foudations of Deep Learning at MIT Video and Slides; my tutorial at the Simons Institute: Slides and Video; and my tutorial at Machine Learning Summer School (MLSS 2021): Video and Slides.
I have also given tutorials on Representation Learning at the Johns Hopkins Winter School and Beijing AI Institute; Slides and Video.
I am also happy to host visitors. Summer visitors please contact me around February to schedule your visit. See a list of past visitors at here.
Awards
Samsung AI Researcher of the Year Award 2023
NSF Career Award 2022
ONR Young Investigator Award 2021
Sloan Research Fellow in Computer Science 2019
NIPS 2016 Best Student Paper Award for ‘‘Matrix Completion has no Spurious Local Minima"
Finalist for Best Paper Prize for Young Researchers in Continuous Optimization
Princeton Commendation for Outstanding Teaching for ELE538B
Selected Publications
Computational-Statistical Gaps in Gaussian Single-Index Models.
Alex Damian, Loucas Pillaud-Vivien, Jason D. Lee, Joan Bruna.
COLT 2024.
Scaling Laws in Linear Regression: Compute, Parameters, and Data.
Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter L. Bartlett, Jason D. Lee.
NeurIPS 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.
How Transformers Learn Causal Structure with Gradient Descent.
Eshaan Nichani, Alex Damian, and Jason D. Lee.
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, and Tri Dao.
code and blog.
ICML 2024.
Optimal Multi-Distribution Learning.
Zihan Zhang, Wenhao Zhan, Yuxin Chen, Simon S. Du, and Jason D. Lee.
COLT 2024.
Settling the Sample Complexity of Online Reinforcement Learning.
Zihan Zhang, Yuxin Chen, Jason D. Lee, and Simon S. Du.
COLT 2024.
Fine-Tuning Language Models with Just Forward Passes.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora.
NeurIPS 2023.
Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee.
NeurIPS 2023.
Looped Transformers as Programmable Computers.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos.
ICML 2023.
Neural Networks can Learn Representations with Gradient Descent.
Alex Damian, Jason D. Lee, and Mahdi Soltanolkotabi.
COLT 2022.
Offline Reinforcement Learning with Realizability and Single-policy Concentrability.
Wenhao Zhan, Baihe Huang, Audrey Huang, Nan Jiang, and Jason D. Lee.
COLT 2022. Video and Slides.
Label Noise SGD Provably Prefers Flat Global Minimizers.
Alex Damian, Tengyu Ma, and Jason D. Lee.
NeurIPS 2021.
Few-Shot Learning via Learning the Representation, Provably.
Simon S. Du, Wei Hu, Sham M. Kakade, Jason D. Lee, and Qi Lei.
ICLR 2021. Video and Slides.
Predicting What You Already Know Helps: Provable Self-Supervised Learning.
Jason D. Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo.
NeurIPS 2021. Video and Slides.
Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks .
Yu Bai and Jason D. Lee.
ICLR 2020. Video and Slides.
On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift
Alekh Agarwal, Sham M. Kakade, Jason D. Lee, and Gaurav Mahajan.
JMLR. Video 1) and Video 2 and Slides by Sham Kakade.
Gradient Descent Finds Global Minima of Deep Neural Networks
Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai.
ICML 2019.
Gradient Descent Converges to Minimizers.
Jason D. Lee, Max Simchowitz, Michael I. Jordan, and Benjamin Recht.
COLT 2016
Matrix Completion has No Spurious Local Minimum.
Rong Ge, Jason D. Lee, and Tengyu Ma.
Best Student Paper Award at NeurIPS 2016.
Exact Post-Selection Inference with the Lasso.
Jason D. Lee, Dennis L Sun, Yuekai Sun, and Jonathan Taylor.
Annals of Statistics 2016.