Graph Machine Learning and Life Sciences with Hannes Stärk

Welcome to the Axial Podcast. Axial is an early-stage investment firm based in San Francisco. We partner with great founders and inventors investing in early-stage life science companies, often when they are no more than an idea. Axial is fanatical about helping the rare inventor who is compelled to build their own enduring business. Okay, I think we're put on record, Hans. Well, uh, it was great to chat up last 10 minutes or so about kind of just like debriefing everything. But, uh, how are you? Thanks for doing this.

I'm doing awesome, as always, I like to say.

Well, you've been— I'm really excited to talk because, uh, you've been at MIT for like a relatively short period of time and been hyperproductive. And so I just wanted to have this conversation and discuss your research around, you know, machine learning and molecular dynamics, and then kind of what, what else is on the horizon. But maybe to like backtrack it and just understand your research a little better, maybe you can talk about your story briefly. And you kind of started your career off at Munich and you've been at MIT for like 7 months, or maybe 2 months, but maybe all in all 7 months. Yeah, how did you get involved in like science and, you know, where did your journey begin?

So how I got involved into science and where my science journey began, that was basically in— so I started my master's at Technical University of Munich in like 2019, October 2019. And then in October 2020, I started my first like research project We're just, yeah, we need to do this interdisciplinary project at TUM as part of our courses. And then there I just reached out to her. So most people I think do it at some company, but I reached out to some lab at the university and asked them if they want to supervise me for a project. And I proposed I read a few papers and I was watching Yannick Kilcher's videos, in case you know him, and had some ideas on projects related to proteins. And yeah, they were a bioinformatics lab. I proposed a few projects to them and then in the end we ended up working on subcellular localization prediction, which was absolutely unrelated to the things I proposed. But it was, was interesting to me. So yeah, that was my first research project, and I guess I absolutely loved it and then decided that I want to do a PhD. And yeah, then I just finished that project in around, I guess, February, and then in March I had to write my master thesis, or I was looking for a project for a master thesis. And yeah, for that, I actually, in, in 2020, I went to my first conference. Okay, I went to— I just attended the online NeurIPS 2020, and Yeah, from that, or there, I met a bunch of people and asked them if they want to supervise my master thesis on this or that topic, and they agreed. And then I wrote my master thesis with them. That was the 3D Infomax paper, which is now also published at ICML. And yeah, after my master thesis, I or during my master's thesis, actually. So that was now the summer 2021. I had to— was reading all of these geometry processing papers, and one of those was a paper by Octavian Garnier, who I then— we got some ideas. While we were talking. I implemented those ideas and sent him the results, and he seemed to like that. And yeah, Octavian was here at MIT when we were talking about his papers or about his paper. And yeah, then, I mean, he seemed to like it that I just implemented some of those ideas, sent him the results. One of them, asked him for an internship in October. He said yes. And at the same time, yeah, then in October we started working on Equibind, this other paper that I guess we'll be talking about. And yeah, then I wrote my applications to a bunch of universities, and in, in February I was accepted to a bunch of universities and then decided that I'll join Tomi Yakula and Regina Basudai. Yeah, at MIT. And that's, I guess, where I am right now, because I started my PhD a little bit earlier actually than, um, than— yeah, usually people start in fall, right? But Tomi offered me to start on June already, and that's what I went for. Cool.

I think, yeah, I think it'd be great to dig in deep into Infomax as a way to kind of predict kind of small molecule features and using like a clever way to pre-train models and then Equibind as a way to kind of predict ligand binding to different targets and using a clever sampling, which I don't understand completely even after reading the paper. Yeah, I kept— also the labs you're working in, I can't pronounce your last names, so I'm gonna have to rely on you for that. It really— we'll post the labs in the description, but their last names have like 5 syllables each, so it's just— I'm not gonna try even. But what was the logic you know, you're trained as a computer scientist engineer and you're making it— you're at Munich and you're collaborating with people. What was the kind of spark to point your kind of expertise towards problems involved in chemistry, biology, life sciences overall? What was that kind of spark that kind of got you working on these problems?

So there are basically two, two main reasons. One is that I just think these areas offer very interesting types of data and are very methodologically rich. And if you work on problems in these areas, you can come up or you have to come up with super interesting algorithms or approaches to solve the problems. And you can bring ideas from physics, where a lot of statistics methods also are coming from. You can bring ideas from physics into the machine learning world and build interesting methods and just make machine learning methods themselves, general machine learning methods themselves better. And yeah, there you also have data types like graphs or 3D objects that you're dealing with. Instead of just the usual images, for example. And maybe then the priors or the prior knowledge of how these objects interact and things like that you can build into your models. And that is, yeah, I think very interesting. But the other, the second big reason is that I want my work to be impactful. Or like, this might sound, um, what do you call it, like saying it from a high horse maybe. But yeah, I want the work that I do to be helpful for humanity. Cool.

Yes.

What Joshua Elkington calls it, AI for good, I think. Yes, I want to do AI for good. Okay, cool. Okay.

You know, better than I thought.

So that's maybe not only AI, right? It is important to me that the work I do, like my professional life, my professional work has a positive impact on humanity, or it may be also as positive as I can possibly do it. There's also this 80,000 hours concept that, right, you spend 80,000 hours in your, in your life working, like on average, maybe. I'm not sure how they came up with this number, but those 80,000 hours you should, or maybe it is interesting to spend them as efficiently as possible for improving or for bringing humanity forward.

And so on that concept of spending time, how did you, on the machine learning side, I'm not an expert at all. I just kind of know a little bit and have friends. How did you determine to spend your time on graphs and graph methods and that kind of area of machine learning? Is there some sort of advantage there in terms of its relevance for chemistry biology, or is kind of graph-like methods, is that where is it a lot of kind of opportunities that are underexplored?

That's what I think. I am of the opinion that we should build as much prior knowledge as possible into our machine learning models if we want to solve a particular task or a particular problem. If you want to build AGI, then sure, maybe, or if you want to build a general reasoning robot, then maybe building in all prior knowledge that you have or as much prior knowledge as you can is not the optimum. But I think if you want to solve a particular problem that has, that is useful to in the real world, then we want to build in as much prior knowledge, reduce the space of functions that we have to model and bake in symmetries such that we can generalize to all symmetric cases. And yeah, using graphs is one way to do that, right? Or if we have some data that is, yeah, that has a specific graph structure, we can process the data ignoring that graph structure, or we can use the graph structure. And then there's the question of how to use the graph structure and that prior knowledge in the best possible way. And there you then get into interesting, yeah, interesting graph neural networks and message passing neural networks and different approaches. And I just think this is then also a very interesting data type in itself. Which makes it interesting to come up with new methods to process graphs better. Cool. Okay.

That's pretty interesting, actually. Okay. So graphs allow you to maybe model prior knowledge more efficiently, and especially in biology and chemistry, that is actually really useful because We still have a lot of the knowledge is implicit. Maybe we can use Inflamax as a case study of kind of that kind of the power of graphs. And so maybe you can talk about Inflamax like your first major preprint and paper piece of research, or maybe second after the first project. Can you talk about that story of Inflamax and how that project got started? Maybe to connect it to graphs and using prior knowledge, maybe that was kind of a theme of the project. And then maybe talk about some of the discoveries and we can just go from there.

So that's that project we started in like March 2021, right? And there I had my great supervisor or, well, my great advisors. Dominique Perrini, Prudencio Tso, Gabriele Corso, who's now my colleague because I joined them here at MIT, and Pietro Liu, also Stefan Günemann, Tom. And yeah, basically I reached out to Gabriele and Dominique to— and yeah, I, I read their very good, very interesting principal neighborhood aggregation paper and also talked about them, their directional graph network. So I knew that they also were interested in graphs or working on graphs. And yeah, I, my position was basically I'm interested in working on graphs and applications to molecules. And that's then how we came up, we brainstormed more or less, and we knew that we wanted to do some, bring the power of pre-training to the world of molecules, right? Because we have very often in the molecule world, we have very little data. Yeah, a small amount of data. But the pre-training approaches that are out there are often not very principled because they like just, we just mask out some atoms or something like that. And if you just mask out an atom and then you try to predict it with similar to masked language modeling, right? Then Yeah, you can probably get right almost, I don't know, 80% of the time by just predicting that the atom that you masked out is a carbon. So yeah, we just think that the pre-training approaches that were out there for molecules, they weren't very helpful. And yeah, what is it? And then the question was, what is it that we could have as a pre-training task that would help the model make better predictions and help reason about the molecule? Well, if the model knows about 3D structure, then it could make, plausibly make better predictions about those molecules for tasks where the 3D structure is important. So we devised this pre-training task where the, in which the small model, the model is trained to reason about the 3D structure of small molecules and then generates this latent 3D information, which it can then use to make better predictions for molecular properties. And there we have 3D-InfoMux.

So what was the kind of the why now in terms of there's people kind of know this is some of a problem where there's a lack of sparse data, 3D data on small molecules or certain molecules because it's expensive to generate at scale. And then, but there's a need to have better models. What was the kind of the breakthrough though train models to implicitly predict 3D structure and then improve predictive power? And then maybe you can talk about the things you predicted, because some of the properties were very, a little complicated in terms of like, you know, various quantum properties. Maybe you can discuss kind of like the readouts and what actually improved for the stuff you kind of were able to test on.

Yes, so for in our evaluation, so we found indeed that for quantum mechanical tasks, the improvements were quite significant or like sometimes, or well, what means sometimes, there were 22% improvements on average on the QM9 datasets, for example, for predicting its properties. If we did our pre-training, yeah, and that's, I would say, makes sense for quantum properties, right? Because it's an intrinsic property of a small molecule, and if you have information about the 3D structure, then we've also seen with models like DimeNet, also from Stefan Günnemann's group, or spherical message passing, or GEMNet, which use the 3D structure of small molecules in a smart way, there we've seen, but they need the ground truth 3D structures input, of course, which our method does not, there we've seen also large improvements on these quantum mechanical tasks. And I think if you have again intrinsic property of a molecule, then predicting it while only having, or additionally having the 3D structure available will help. But if you now just have a small molecule and you try to, you try to predict its toxicity without having any of the structures that it binds to, to have its toxic effect. If we only look at the small molecule and then if we additionally have the 3D structure, yeah, it might not help so much if we don't have the 3D structure of the molecule that it interacts with. And that is also what we found with 3D-InfoMax, where we then didn't see improvements with the 3D pre-training when evaluating on tasks like blood-brain barrier penetration or the TOX21 challenge or, well, actually then there were tasks like solubility prediction again or the lipophilicity where the 3D pre-training helped, right? Because these are intrinsic properties of the small molecule.

I think the paper was a really elegant way to convey the power of pre-training in models, and then, at least for intrinsic properties of these small molecules. And what kind of, from that, we published the research and we got accepted ICML, so congrats on that. And so what kind of, what did that unveil? Did Infomax— did you often, great research role, unveil new problems or new things to work on? Did that project set up Equibind, or was Equibind totally separate? Or what did you learn from Infomax that then like set you up for the research you did later on?

Oh well, what did I learn from Equi— uh, from the project?

I'm sure you have all these challenges and whatnot.

And yeah, it is just so much that I'm not sure what I should pick out now. I definitely learned a ton in that project, from writing a paper to, well, of course I've written another paper before, but Yeah, there I really got into, um, like the, the conference style of writing a paper, or the conference paper writing style, right? Because there was a conference submission, uh, but I, yeah, I also learned a lot about working together with experts from chemistry or industry partners. I learned a lot about interacting or reaching out to people and discussing papers with them, right? And like, because that's something I did a lot in the summer or when working on this project, I read a paper of some people. I had a bunch of questions about them. And then while previously maybe I would linger on those questions and fret about them and not find out the solution, there I really learned that most of the times researchers like talking about their papers and you can just reach out to them, ask them if they have time to discuss their paper and Oftentimes when you discuss their paper with them, they're super happy that someone's interested in their work. So yeah, I learned not to be afraid to ask people if they want to discuss their paper with me. What else did I—

Did you start— Han, if everyone— Han does a really great reading group. If you're a practitioner, you should partake in it. Did you start the reading group around this time? Or— Yes, exactly.

I had this experience that authors of papers are often very, very excited to talk about their research and very happy to spend time on that. And yeah, then I was looking for a reading group on graphs or about graphs. Because I knew some reading groups about learning theory, a new reading group about reinforcement learning, but yeah, that were online, right? Of course, there are some graph reading groups at some university that the people at the university can go to, but I wanted a reading group that I can go to as well, that's just online and a Zoom meeting. And we can discuss the paper there. So after a few months of not finding the reading group that I was looking for, I started it myself, I suppose. Cool.

Okay, congrats on that. Yeah, and it made me— sorry for the sidetrack, but yeah, if you work on graphs, you should probably look into Hans's reading group or something. We'll post it in the description. It's really cool. Any other lessons for Infomax? Or maybe we can go to Equibind.

Let's, let's go to Equibind. Cool.

You just have your names, your papers. They're like, in America, there's various credit bureaus that do credit scores. And so there's Equifax, there's like various names, there's like names of companies that like do credit scores. And so your papers are kind of like credit bureau names. So Equibind. And so Equibind was the paper you did at MIT, and really the breakthrough was really focused on ligand binding of targets. You described intrinsic properties of small molecules, Equibind was more about predicting when small molecules interact with proteins, and the key breakthrough is around sampling, which I still don't understand at all. So like, I don't have any insight on this paper, I've read it twice and I still don't get it, But maybe you could talk about Equibind, the story, and then kind of at least set up the paper, then we could discuss it.

Yeah, so the task that we tackle in Equibind, right, is predicting the structure in which a small molecule binds to a protein, right? And this is interesting for biologists or for drug discovery because sometimes if a molecule, small molecule, binds to a particular spot of protein, then the behavior of the protein is changed. Maybe some protein causes cancer, and if a small molecule binds to a particular location of that protein, then the abnormal or the, then the cancer-causing behavior of the protein is stopped. Yeah, and So this is a pretty helpful task for drug discovery if we can predict the location where a small molecule binds to a protein or attaches to a protein, right? And what we do in 3D-InfoMAX, or what previous approaches would always do, is they would try many, many different possible locations where the small molecule could go, and then score all of these locations or evaluate them and then use the best scoring location as their final prediction. But trying out all of these different positions takes a long time. And what we now do in, or what I think is the major takeaway of Equibind, is that this is also possible, or good 3D structure predictions are also possible in a single shot. If you build in the prior knowledge about how 3D structures could interact in a good way, yeah. And this new paradigm of predicting this in a single shot makes the method much faster, and this is very important for drug discovery where we sometimes want to make predictions for millions of molecules when we search through a library of molecules to find one that has the behavior that we would like.

So when you mean single shot, rather than scoring a bunch of locations on the target and then find the best one, single shot means you define the best one as the first result based on all the data you have, or maybe that's not the— maybe that's too simplistic.

No, that's based on what— based on the data that we train on, we have a a single-shot prediction, a single forward pass of our model that tells us, or that tells us a 3D structure in which the small molecule probably binds to the protein. But of course, this isn't a solution to all our problems, right? This, I would really see this as a, it's not the perfect method. I would really see this as a demonstration. Okay, well, you can also— it has also some practical real-world value, and actually some companies are using it in pipelines, and they are happy with some of the results. But the main thing I see from the paper is a demonstration that it is possible to make these 3D structure predictions in a much faster way if we're using machine learning methods, or with these machine learning methods, and that this is something that we should also invest time in and look at, and not only into improving the score functions with which we evaluate the different locations.

I think it's kind of another theme of like sampling better, but also having better, maybe on the back end, having better data.

Yeah, but that's not what we're doing, right? We're not sampling better, we're not sampling at all. And that's this new paradigm, which I mean, which we should also be looking at.

And so to not sample at all, how do you, what's the bar, at least for the protein structures you're working with? For the data? How many more different configurations do you need to map out for a given target to make single-shot possible?

We train with the same type of data that the previous methods train with, with which these scoring functions are also often trained with, because in the protein world, right, you just have data, or like there's, you have all the data that's out there in these data banks, like the Protein Data Bank, PDB, with like, I think around 180,000 protein structures, which are like the known protein structures as far as I know. And then we have PDBBind, which is a subset of that, which are, as far as I can tell, all the known structures of small molecules bound to proteins. And yeah, we use those for training because that's all there is out there.

Oh, interesting. And so is there like, you know, I'm assuming there's some limitations in terms of some targets are not as well studied as others. And so a single shot is not going to be possible for maybe proteins or intrinsically disordered domains or something like that. There's always going to be a limitation in terms of our data.

Yeah, let's talk about limitations of Equibind because there definitely are big limitations, I think. So first of all, I believe to make a good 3D structure prediction, what we really want is to have some iterative approach, right, that reasons about the distances between the small molecule and the protein as the small molecules may fit into the pocket or, yeah, where our model thinks it will end up. Also, we know that, right, this whole thing is not, or not in every case, it is completely 100% determined where the small molecule will bind up. It is a stochastic process and it is a Yeah, it's a probability distribution of where our small molecule is, and if we're able to model that, that would be maybe much better than having this single prediction. So if we have a generative model over where the small molecule could end up, this would be super valuable and maybe much more helpful in practice, and also reflecting reality much better. Yeah, and another thing is that we are not able to deal with symmetric proteins, right, because currently we are always, even our training data, we have a symmetric protein where the small molecule is bound on one side of the protein, then, and our model predicts the small molecule now to be on the other side, which would be a just as good prediction, we will still tell the model it was wrong because in the training data it was on the other side of the symmetric protein. And yeah, if you have a, it would be great to have a generative model that, or not necessarily a generative model, model, but for example, and that is able to tell us, yeah, with 50% probability it will land over here and with 50% probability it will bind over there. That would be a much better model of what's going on in reality. Cool.

Yeah, I think every new prompt, every limitation is an opportunity. Yeah, so maybe we can transition towards what you see long-term in the field of like machine learning, molecular dynamics, and more. Um, you know, you've worked on Infomax, Equibind. Um, you know, what's, what's on the horizon for you? What other projects you're exploring? Yeah, what's getting you excited?

Yeah, that's, that's also something. Um, right, I, I have been thinking about that a lot, and also because I've I've been asked that a few times before. I've been thinking about this question, and I, I mean, keep, keep in mind that I'm, uh, a 2-month PhD student, and I'm not a—

um, that's what makes you exceptional, Hans. You're 2 months in, you look at all the— how productive you've been. So Okay, thanks.

But still, I'm not like David Baker or Max Welling or whoever who have tons of experience and whose statements will be much more valuable about predicting the future, so to say, if you want to go with that. But yeah, maybe I'll just say what I'm really excited about. Out. And that is maybe in this direction of, right, we've spent a lot of time modeling individual molecules and coming up with better graph neural networks, for example, to process molecules and making predictions about a single molecule. But what I think is really exciting is to model molecular interactions. And the dynamic nature of molecules while they are interacting, which is of course core to the interactions of, in many cases. And why am I so excited about molecular interactions and the dynamics of molecules there? Well, that's because first of all, that is where a lot of impact is possible if we are able to model them. Because it's what's important in drug discovery and, right, in the end, in applications, we're not only interested in what a molecule looks like, so to say, but we're interested in what it does, so how it interacts with other molecules. Yeah, and the other thing, why the technical aspect of my why I find this very exciting is because when modeling interactions, you run into even more interesting problems and mathematical hurdles to solve than when only modeling a single molecule, because now you, for example, have two molecules floating around in space and maybe you want to be equivalent with respect to the initial placement of one of those molecules, yeah, and these, these symmetries that you can capture, they become a lot more intricate and even more interesting.

Cool, yeah, I think modeling small molecule protein interactions and protein-protein interactions are mind-boggling if you really think about it.

It does not only have to be protein-protein or small molecule-protein, which really maybe often happen in the drug discovery world. It can also be how a catalyst interacts with some material. If we're looking at the— there's this nice OC20 dataset from Open Catalyst dataset, that's where we were trying to predict the dynamics of a small molecule interacting or coming close to some material, yeah, and there are many, many different possible types of molecular interactions than just the ones in the drug discovery world, so to say.

Yeah, I'm a little biased, but I think, yeah, totally agree. And I think you'd get along with Katty Sarr. I'll email her. She's gonna love you. You'll get along really well with her. Yeah, any final thoughts? But this has been really great. I think great to have this initial conversation. And I know for a fact you're gonna do a bit more work. Yeah, it's gonna be cool. You're gonna continue to publish more research. I think we're all gonna read it. Uh, and then you're only 2 months in, so it'd be great to have a conversation maybe 2 years from now, and you'll be like, yeah, I have like 20 ICML papers, you know. I'm just, you know, you'd be a professor by then. Maybe you'll be a professor in 2 years. So, um, I hope—

maybe I would even say that I've done something wrong if I have 20 ICML papers by then. Because, yeah, I hope I produce impactful papers and maybe not that many, but really important papers. And this is also something, yeah, maybe that is a good point to finish with because that's something that my supervisors always tell me, right? Tom Yakula and Regina Bastille, but also which is to work on impactful or big and difficult problems instead of picking out a thought or an idea that one has, which can maybe lead to the next ICML paper. Of course, if that is one's goal, or like if one just does research for the sake of curiosity, that is absolutely fine, right? Many, many people do that just because you're curious. You work on a problem and you just are playful and explore the world of research out of curiosity. And oftentimes this also produces super impactful work, right? But yeah, for me, this is only Or what means only the curiosity and the fun of playing around is a huge part of why I love doing this research. But the other part that's important to me in the research is having impact in solving important hard problems. And that is something that my supervisors and Octavian also always gave me as advice is to, to find problems that not everyone could have solved and to work on problems that are actually important to solve. Cool.

I think that's a great way to end it. Hans, thank you for taking the time to do this. I'm gonna press stop record right now and then we can chat up. I really appreciate taking the time for this.

Like, quest.