Q&A with Senior Scientist Jason Lee
This month, we sat down for a conversation with the newest member of our leadership team, Senior Scientist Jason Lee. Jason is an Associate Professor of EECS and Statistics at UC Berkeley. He completed a PhD from Stanford and was a postdoc at Berkeley, before joining the faculty at USC and then Princeton, and working at Google DeepMind. His research interests are in the theory of machine learning, optimization, and statistics. Jason is the recipient of the Samsung AI Researcher of the Year award, an NSF Career Award, the ONR Young Investigator Award in Mathematical Data Science, a Sloan Research Fellowship, and the NeurIPS Best Student Paper Award, and was a finalist for the Best Paper Prize for Young Researchers in Continuous Optimization. He has been participating in Simons Institute programs since we opened our doors in Fall 2013, when he was a graduate student.
Simons Institute: How did you become interested in theoretical computer science?
Jason Lee: So I don’t consider myself really a theoretical computer scientist. I don’t have any papers in FOCS or STOC, but I’ve always been sort of interested in it. I work mostly in learning theory or statistics, so machine learning but on the theoretical side. A lot of times, the problems we look at are closely connected to things that people have looked at in TCS. Some conference communities like COLT are sort of half TCS people, half machine learning people. So it’s sort of this diffusion.
Simons Institute: And how did you first get interested in the kinds of problems that you work on, and your field generally?
Jason Lee: I started trying to understand what happens in machine learning when I was a postdoc, around 2015. That was when deep learning had clearly already taken off. For practitioners in computer vision, deep learning had started replacing their standard feature-based schemes. And natural language processing hadn’t taken off yet, but every few weeks, you would see some deep learning system that had started replacing some older system, and the results were getting better.
And it was very mysterious, because I was trained in this viewpoint that algorithms should be simple linear things, or things that result in convex optimization. So then the algorithms had guarantees and could be more easily analyzed. But deep learning was much more complicated. And so I started around then to look at simple cases of nonconvex models, simple neural nets. And that naturally transitioned when transformers came out, and natural language started being revolutionized, into looking at language modeling and the current systems.
Simons Institute: Tell us a little bit more about your research. What research questions are you interested in right now?
Jason Lee: So I’m most interested in understanding how AI systems work, and whether you can have any guarantees, and then how that leads back into making better AI systems. Because right now, they’re very black box, in the sense that you don’t write the code and then the code implements the algorithm. You write a code that trains something based on data, and that gives back the algorithm. So there’s a layer of indirection that makes it very hard to know what the training does. So a lot of my work has been focused on how we can provably train models that work, and understanding this.
Simons Institute: And is that possible?
Jason Lee: So, in some sense, I think it is very behind the empirical success, but a lot of times, we can find simple toy models and simpler situations that exhibit similar characteristics and that you can analyze. So, for example, instead of analyzing very deep hundred-layer networks, we’ll analyze simple two- or three-layer networks that have some of the same properties, and that it’s tractable to analyze. Or instead of looking at the case with general data written by all humans, we’ll think of simple mechanisms that generate data, like Markov chains, or other hidden Markov models. And then under this case, we try to analyze what you learn when you know where the ground truth comes from.
Simons Institute: And will this ultimately translate? Will we ultimately be able to have models that are guaranteed not to hallucinate, for example?
Jason Lee: Something that categorically does not hallucinate is probably not possible. But maybe you can aim for things like helping design optimizers, or the future generation of training architectures, or a sandbox of tasks that you know all architectures should be able to do, because they’re very common. You can test your future architectures on these toy problems. It’s very hard to say this will definitely lead to more accurate models or something like that. I think that would be very hard.
Simons Institute: What research result are you most proud of?
Jason Lee: The paper I’m probably most proud of is from when I was a postdoc here. I wrote a very short paper showing that analyzing gradient descent would always escape saddle points and converge to local minimizers. And the reason that this was significant back then was that we had a lot of models where we ran gradient descent or variants, and they would have very good performance. But we didn’t really know why. And the theoretical state was that you think of gradient descent as just sort of local search. And then there were some models, such as the matrix completion problems and many other similar problems, where we knew that if you could reach a local minimizer, then the performance would be good statistically. But there was this kind of corner case where you might be trapped in a saddle point, but intuitively, it should be impossible. And this was a very short and clean argument that showed this is, in fact, impossible. It grew out of one line from a very classical result in math that hadn’t been written down, or it was folklore. And the result was just very surprising, I think. The whole paper was maybe, I don’t know, eight pages. It was very short.
Simons Institute: Can you tell us about your past experience at the Simons Institute?
Jason Lee: The first time I came to the Simons Institute, I remember driving over from Stanford to Berkeley as a graduate student. I think it was the Big Data program in Fall 2013. I wasn’t an official participant. I was just driving or carpooling in for the workshops, and there was a huge amount of activity. I’ve definitely spent a lot of time here over the years. I’ve probably spent half my time in the last several years here, potentially.
I’ve really liked being at the Simons Institute. There’s a large community of people kind of similar to me here, interested in understanding how machine learning works, but coming from a theoretical lens — computer science and statistics and electrical engineering. So that’s been really good, and a lot of my collaborations have resulted from being at the Institute. Even when I was a graduate student, the Big Data program already led me to work a lot more on machine learning and optimization. And now I work with a lot of the fellows at the Simons Institute, especially in the Machine Learning Pod. We have a weekly reading group with the pod postdocs, and a lot of our students, and some faculty — Peter Bartlett, Nikita Zhivotovskiy, Song Mei, and Jingfeng Wu, who was a postdoc and is now a new faculty member. It’s a group of us that have pretty similar-ish interests, and there’s been a lot of collaborations between people in the pod. I’ve had a lot of collaborators in the programs as well. And some of the students I met through the Simons Institute later became my postdocs.
Simons Institute: How is the emergence of powerful AI influencing TCS as a field?
Jason Lee: We’re living through how AI is influencing TCS, math, and all theoretical fields — especially in the past month. We’re kind of in the middle of it, so it’s very hard to say right now. At the moment, it’s an extremely powerful tool, where if you give it a list of 20 problems, maybe it’ll solve two or three. It’s not perfect, but it’s definitely very, very good. It can do some things that are very difficult, that no human could do. And it really is changing things. Everyone’s just still developing their workflow with AI. You’re seeing that sometimes, even people who are nonexperts can just prompt and solve a problem, but then it’s not clear they understood the solution or even the problem. There’s a very large burden on reading all these solutions now. And reading solutions is very, very difficult. I think maybe writing an exposition will become much more important going forward. People are putting more emphasis on formal verification, like Lean. But even if you verify, it’s not clear you understand. So somehow it has to be that we place more emphasis or give more credit to explaining.
Simons Institute: And how has TCS benefited machine learning in general, and in the generative AI era in particular?
Jason Lee: I think probably the earliest examples would come from boosting, in the ’90s, where I think boosting was solved in the COLT community, and then this led to AdaBoost and gradient boosting. That’s an older example. More recently, the class of diffusion models and diffusion processes for generative AI generating images — mostly images and video, not text — is a clear place where the algorithm was developed with a very theoretical mindset. The earliest paper was in 2015 by a group of statistical physicists, but this was later reexplained and generalized considerably in 2019, and there are many books about it now. This is a very clear example where theoretical reasoning led to the algorithm.
Within the last two years, in optimizer research, I think you can point to Muon as well. Although it wasn’t ever a published paper, there was a blog post on it, and there were earlier papers in the optimization literature about spectral steepest descent. And it seems like a lot of places use it in place of the AdamW optimizer. And it’s at least, I think, very theoretically motivated.
Simons Institute: And what do you envision as the trajectory of your own research in the next five or 10 years?
Jason Lee: I’m hoping that I can continue to try to understand phenomena that arise in AI, and what makes this boundary between AI and theory. The nice thing is there are always new phenomena, in the sense that, it seems the practice of AI is ahead of the theory — at least it has been in the recent 10, 20 years. There are a lot more practitioners than theorists. So one thing I would want to look at, for example, is test-time inference or test-time compute, which is an entirely new axis, where previously we just had training and test time. And now we have this new axis, where you can spend compute resources on a test time, and the model seems to get better.
Some basic questions are how you can trade off data, training time, and other resources. This also relates to scaling up models — which dimension should you scale up on, and when should you scale up which one? There are so many dimensions now: model size, data size, how long you train for, how long you let the test-time inference go for, whether it should be sequential or parallel. And I feel like all of these can be distilled into some simple questions, which if you can answer somewhat cleanly, yield a lot of good guidance.
Simons Institute: Do you have any instinct for how the availability of LLMs as research assistants may impact how you approach these problems?
Jason Lee: Yeah, I personally use them very, very heavily [laughs]. In the last two or three months, whenever I have something that I can write down as a concrete problem, my first instinct now is to type it into the LLM and see what it says. I’ve recently started trying to brainstorm with an LLM. I think it’s less good at that, but it’s still very useful because it just knows so much more than you. You don’t have to go to Wikipedia and read things that you don’t really know, or maybe you sort of knew from an undergrad class, or you have some very vague recollection of. If you just say a bunch of keywords, it’ll kind of remind you of everything. So I use it super heavily. I think it’s very different for different people. Some of my students told me they don’t use it at all, and I was shocked.
Everyone has been thinking a lot about how you can cowork with the AI productively. It does seem that there’s much less value in just entering a well-defined problem, and then waiting for the model to come back with an answer, and then just saying, “Continue five times.” Then I’m not really needed. So I have been trying to think more about how to do this. I don’t have a very good answer yet, but I’ve been proactively trying to shift toward less well-defined questions.