Scientist Stories: Ellen Zhong, Reconstructing continuous distributions of 3D protein structure

Alright, thank you to the organizers. I'll be talking about our work on fully unsupervised reconstruction of three-dimensional protein structures from two-dimensional cryo-EM images. And specifically, we'd like to reconstruct heterogeneous dynamical protein structures. So as the main biological motivation for this work, as you all know, large macromolecular complexes underlie all essential biological processes. Of course, we'd like to obtain the structures of these complexes in order to understand how they work.. As a first pass, it would be nice to obtain a static structure, but what we would really like to know is how they move in relation to how they function. And so in pursuit of this goal, we can use cryo-EM, which stands for cryo-electron microscopy. Over the last decade or so, there have been dramatic advances in this technology, mostly driven by better hardware, so these direct electron detectors, and by better image processing algorithms, so these probabilistic probabilistic methods that are typically implemented on GPUs. And so just as a quick background of cryo-EM image formation, so in a cryo-EM imaging experiment, you have a purified sample of your target molecule of interest whose structure you want to determine. This is frozen in a very thin layer of ice, and it's imaged with an electron beam, and your recorded image is typically modeled as the integral of this particle's three-dimensional electron density function along the imaging axis. And so these molecules are randomly oriented as you flash freeze them, so there's an orientation, unknown orientation in each of these particles, and they're also unlike the schematic, extremely noisy due to the low electron dose required to limit radiation damage. And so that's the forward model at a very high level. The task of reconstruction is to solve the inverse problem of going from images, so a typical number being in the hundreds of thousands, to the 3D electron density. So here's a real data set of the ADS ribosome. So a couple of challenges, as I mentioned, there are the unknown particle poses, so you don't know the orientation of each of these volumes within the image. There's a low signal-to-noise ratio. You have to take into account image degrading filters and microscopy such as the point spread function of the microscope, discretization of measurements, but then the problem that we're particularly interested in is the so-called heterogeneity problem. So this is where you break the standard assumption that each of these structures in each of these volumes within each image is the exact same static structure. But before I talk about heterogeneous reconstruction, just one more quick piece of background for the homogeneous case. So cryo-EM— traditional cryo-EM algorithms all use what's called the Fourier slice theorem to relate these two-dimensional projection images to the three-dimensional volumes. And so the Fourier slice theorem states that the Fourier transform of these 2D projection images is a central slice out of the 3-dimensional Fourier transform of the volume perpendicular to the projection direction. And so this suggests a method for reconstruction. Given a bunch of images, if we knew their orientations, we could just take their Fourier transforms, insert these slices at their particular orientation in the 3-dimensional Fourier transform, build up this 3-dimensional Fourier transform, and then take the inverse FFT to get your real space volume. And so this works well in the homogeneous case. Traditional reconstruction algorithms use the Fourier slice theorem. However, the so-called heterogeneity problem I think is actually a major opportunity in cryo-EM. So this is definitely, like, throws a wrench in the reconstruction task. Each image contains a unique molecule, but because we actually capture the data for these heterogeneous states, we can actually interrogate the diverse functional states of these different proteins. So standard approaches for heterogeneous reconstruction, other than just discarding your heterogeneous data, is something called multiclass reconstruction, which models a discrete mixture model of k independent volumes. You don't know the nature or the number of states a priori. For example, your proteins might lie along a continuum of states, and so this, in general, the identification analysis of heterogeneity, especially its continuous forms, is an open problem for cryo-EM. And so what we would like to do is learn a continuous low-dimensional manifold of structures in a fully unsupervised fashion. So our particular method is called CryoDRaGN, which stands for Deep Reconstructing Generative Networks. We developed, in brief, coordinate-based neural networks to directly approximate the three-dimensional structure. We use a Fourier space image encoder volume decoder architecture based on the VAE to learn a latent variable describing the protein's structural heterogeneity. We also employ an exact inference— we also employ exact inference for the pose variables and variational inference for the heterogeneity. So our coordinate-based neural networks can be best understood in contrast to traditional algorithms. So instead of representing our volume, our structure, as discrete points in a 3-dimensional voxel array, we instead use a neural network to approximate the continuous function which maps from 3D Cartesian coordinates to the electron density at that particular location in 3-dimensional space. Instead of directly feeding in the 3-dimensional coordinates, we use a fixed positional encoding function to featurize the coordinates in terms of sinusoids of varying frequency. And then the reason why we do this is to model heterogeneity. So in this traditional multi-class reconstruction framework, we would model a mixture of k independent separate voxel arrays, whereas in our approach we can simply extend the input dimension of our neural network to include a latent variable z denoting heterogeneity. So that becomes a generative model in our framework. That's the volume decoder. This is a schematic of the overall architecture. We have an image encoder which takes an input image and produces a distribution in latent space describing this particular image's heterogeneity. And then we know that each image by the image formation model is a slice out of a 3-dimensional volume. So we can take those 3-dimensional coordinates and a sample from that approximate posterior, feed that into the volume decoder, and reconstruct the image pixel by pixel. Pixel. We don't know the orientation of the slice a priori, so instead we do this search over all possible— or over discretization of all possible orientations and in-plane shifts using a branch-and-bound global optimization procedure. And so I also want to note that this coordinate-based representation is really nice. It enforces geometric consistency between these two-dimensional views. Each of these two-dimensional projections should share a common line— should share a common line in common. So this is all enforced by our positional-based folding decoder. And so we want to evaluate this first on synthetic data. So in the simplest case, we create this heterogeneous complex— heterogeneous dataset containing a model protein complex with 1 degree of freedom. And so we have this atomic model here, and I rotate one of the dihedral angles connecting this blue domain of the protein to the rest of the complex. I sample 100 atomic models along this reaction coordinate, and generate 500 randomly oriented and translated projections for each model, yielding a total dataset of 50,000 images. Um, and so the question for method is, can we reconstruct this entire continuum of states? And so if we train a one-dimensional latent variable model, here I'm showing the predicted latent code for each image, uh, produced by the encoder along the true reaction coordinate, which we know in this case because it's synthetic data, and we can see that it correlates very nicely. We can also sample or generate the volumes corresponding to the different points in the length code, and so here are 10 different structures shown at high isosurface and low isosurface representation for the volume. You can see that we reconstruct the ground truth structure. We can generate movies, which look kind of silly in this very synthetic dataset. And in contrast, CryoSPARC, which is an existing tool for heterogeneous or for traditional cryo-EM reconstruction, Its multi-class reconstruction is— with 3 classes is only able to reconstruct 1, maybe 2 volumes that resemble the ground truth, with a third having very weird alignment artifacts. And in general, you also lose the underlying heterogeneity in the data. And so there are a couple of additional datasets we've generated that have more complex latent structure. In this linear 2D dataset, we have 2 independent domains of the protein complex that are allowed to move independently. In the circular 1D dataset, we have one domain of the protein which rotates 360 degrees. We have a discrete mixture of 10 different classes. And what I'm showing here is the dataset embedded in latent space. And what's really cool is that the topology of our data distribution in latent space matches the topology of the underlying heterogeneity. And we also quantify the accuracy of our volumes using an FSC resolution metric, which is a standard in the field. And then finally, I just wanted to show some results on some real cryo-EM reconstructions. So this is a dataset containing a mixture of ribosome assembly intermediates. The biological question was trying to understand how the ribosome is formed, and this is my advisor Joey's postdoctoral work. And so through many manual iterations of this multi-class reconstruction, they identified 4 major classes of the large ribosomal subunit. From this multi-class reconstruction approach. And here what I'm showing is our method training a 10-dimensional latent variable model, and this is a distribution of data in latent space colored by the class assignments from the publication. And so you can see that the 4 clusters for the major classes come out really nicely. And here are the structures generated at these points at the cluster centers, and the resolutions of these structures match the resolution of the publication. We can also sample additional volumes from latent space. This blue cluster all the way on the right actually corresponds to an impurity during the sample preparation, which is the 70S ribosome, and so this comes out automatically in our reconstruction method. And here are just 4 additional structures sampled from this particular class, this particular blue cluster— purple cluster, showing the heterogeneity which corresponds to the known intermediate states for the ribosome during assembly. And then finally, I just have one last dataset of the pre-catalytic spliceosome. So this structure of the spliceosome was the first structure published in 2017, and you can see that there's a low-resolution core— or high-resolution core with low-resolution distal elements. And this structure was actually resolved separately, or the different domains of the structure were resolved separately and then pieced together just to form this composite map because there was so much heterogeneity in this dataset. And so in contrast, this is the reconstruction that we get using our method. This is just one particular latent traversal in latent space, but we can see that there's this stable core and lots of motion in the distal elements of the complex. And what we can do is not just do one latent traversal, we can do a random walk along latent space. So this is 132 sampled structures. We do this random walk and we can see complex heterogeneity at the domain interfaces. And so in contrast to the traditional approaches, which are able not even to get one structure, we can sample this entire conformational ensemble that really, to me, resembles a lot of dynamics trajectory. Almost. And so in conclusion, this is our method for fully unsupervised reconstruction of a continuous low-dimensional random folder protein structures from cryo-EM data. It's the first method to apply neural networks to the cryo-EM reconstruction task and the first that can learn a continually heterogeneous representation of protein structure from cryo-EM data. And we started sharing this method to different research groups who were unable to use the traditional algorithms for reconstruction of their protein structures, and we hope to identify new structures from real cryo-EM data. So with that, I'd like to thank both my labs as well as other people who have provided good feedback, and thank you. Yes, any more questions?

This is really beautiful work.

There's been some recent advances with differentiable ray tracers. Differentiable, pardon?

Like ray tracing for images, right, where you have a volume but then you generate a 2D image, you can differentiate a bit.

Have you thought about any [INAUDIBLE] So the image formation model is slightly different. The images are actually these projection images, so like integrals through the density along the imaging axis. Whereas in other like domains of 3D shape reconstruction and computer vision, they're like surfaces where you have this ray tracing. So I'm not exactly sure how that— like we already directly model the image formation process by treating our images in Fourier space as slices of the volume.

Yeah. That's really cool. I guess two sort of unrelated questions. One of them just has to do with that ribosome example, and the red cluster seems to actually be in two clusters.

And if you look at the other one—

Yes. These kind of methods obviously benefit from lots of data. How easy have you found it to— I mean, we'd like it in principle just to scale up and automatically learn and data that's all out there?

Do you have ways to do that in this building? Yeah, so I guess like answering the second question first, one thing about this method is that we train a separate model per experimental dataset. So for each new, I guess, cryo-EM dataset, we're training a separate model and we don't really expect it to generalize to unseen protein structures or proteins that aren't in the particular dataset. Although it is kind of an interesting question of, can we take advantage of all the structures, all the protein structures that we know now through cryo-EM. And then to answer your first question, I think— so the answer is yes. The actual published dataset had 13 different structures of the ribosome. So within these 4 major classes, you have minor variation in classes C, D, and E. And then if you're particularly curious— oh wait, it's not in here. But the different clusters within the— the different clusters within this red cluster, the red major cluster, correspond to these minor states of ribosome assembly. Because we know it's this complex parallel process. Yeah.

Yeah? [INAUDIBLE] Correct. Yeah.

Um, is there like guidance on how to choose that or—

Mm-hmm.

And also how does it change for different— Mm-hmm. Yeah. So in all the examples that I've shown here, other than that first synthetic example with the 1-dimensional motion, they're all trained with 10-dimensional latent variables and that's just a hyperparameter that you have to select in your architecture before you train. Um, you could cross-validate the latent dimension, but so far I haven't done that. And what we've seen is that as long as the dimensionality of your latent variable exceeds the dimensionality of, the heterogeneity you're trying to model, you're so far okay. But that's definitely something that maybe could be tuned.

Yeah. [Speaker:AUDIENCE] Yeah, that's pretty cool stuff.

I was just wondering, when you're correlating this with like a new simulation or other methods, how do you— [INAUDIBLE] [Speaker:KAREN] Mm-hmm. Yes. So that would be the next step. So after the output of cryo-EM reconstruction is this electron density map. But then as a practitioner or as a biologist, you want the atomic coordinates. So that's like the next step in the cryo-EM pipeline, and maybe there's some ways of incorporating the known sequence information in how we model the structures. Right now our volume representation is just an MLP, so it's a very naive representation for the volume, which is very flexible, but maybe we can do something better by incorporating the atomic positions in some way.

Yeah.

Cool.

Thank you.