Scientist Stories: Zachary Sims, Masked Image Modeling Approach to CyCIF

Oregon Health and Science University. Hi. So yeah, my name's Zach. I'm from OHSU. I'm a research engineer, and I'm here with my PI, Youngwon Chang, here to talk to you about our efforts in producing a machine learning method that can address the limitations and potentially democratize multiplex tissue imaging by learning to generate tissue stains in silico. So to set up the problem, multiplex tissue imaging is this emerging family of technologies that is really powerful because it allows us to do this single-cell analysis within the spatial resolution of tissue. And so the technology that I'm going to be talking about today is cyclic immunofluorescence. And so this was a multiplex tissue imaging technology first proposed at Harvard Medical School. And so it involves this iterative process of 4-channel imaging along with the, the iterative cycle of 4-channel imaging with, and then followed by bleaching to remove the signal, and then repetitively doing this with, with different markers each time. And so you can, over successive rounds, build up this, this panel of dozens of different markers that measure protein expression within a single sample. And so in the bottom left here, I show an example panel. So we call the set of markers that we're imaging here a panel. And so this is an example panel done on a colorectal cancer sample here. And so this is a 25-marker panel. And so you can see here that we can do it in 8 rounds. And so the So in every round, one of the channels is reserved for DAPI, which is a nuclear marker. And so we need to have this marker in every round so that after we've done our imaging rounds, we can align and register the different rounds so that we can do our single-cell analysis. But then the other 3 channels are allowed for different markers of protein expression here. And so up here again is an example of an image of a 10-round cycle of CyCyF with the 3 markers in green, white, and red, and then DAPI in blue. And so the most common analysis that you do here after you've done the imaging is to segment each cell and then use the cell boundary for each individual cell to measure the average pixel value for every marker channel in your image. And that gives you a mean intensity value for every marker for every cell. And so you can gate that value, and that will give you a classification of a positive or negative cell for that marker. And that's how we can do cell phenotyping within a spatial resolution. And so this is a really powerful approach to investigating the tumor microenvironment and has a lot of potential applications for improving our ability to perform diagnosis in a pathology setting. But unfortunately, there are limitations that we have to address. So within Sisyphe, the main limitations here is time. So one of the main limitations is time. So the number of rounds that you can do is limited because one round takes up to one day to complete. And so additionally, as you start doing more and more rounds, the number of staining rounds that can be conducted is limited by tissue quality, which degrades in later rounds. So you can see up here, for example, In later rounds, you start to get tissue loss, and that means that certain cells are literally just not going to show up in these rounds. And so markers that were stained within this round, you essentially have a data missingness problem where you're not going to get markers for that cell, or those markers for that cell. Yeah. And so then in commercial multiplex imaging systems that are more commonly used, the number of channels is limited because you can only do single, single round staining. And so you can only do about 6 or 7. Markers in that setting. And so if you wanted to have a higher-plex system, like RareCyte Orion can do 17 or 18 markers per round, but those require special machinery and that is cost-prohibited. So in order to proliferate this technology, we want to be able to overcome these barriers. And so we look to a software solution to potentially address these issues. So to formulate this machine learning problem, we essentially proposed this two-step problem. So the first is, can we take our full panel of markers and identify a representative subset of markers that retains the full information of the full panel? And so we call this step the panel reduction step. And then so once we have a candidate reduced panel, then we want to be able to learn a way to impute or reconstruct the full panel and the markers that were withheld, essentially. So if we can do this, then we can have 3 immediate use cases. So the first is we can address the data missing this problem of where there's tissue loss in later rounds. We can impute those stains and address the data missing this issue. And then the second use case can be the potential to democratize higher-plex systems by taking lower-plex systems and synthetically expanding them to higher-plex systems by conducting the staining in silico. And then thirdly, we can also help to better optimize and design better panels by, if we can identify easier-to-predict markers, then we can leave those markers out of the actual staining and imaging, and we can replace them with more difficult experimental markers that can potentially identify more rare cell types. So this is emerging, this is an emerging problem. That has a few different solutions that have been proposed. So our lab was the first to publish a method for panel reduction and marker imputation done by a former grad student, Luke Turns. And so this method is a 2-step process. So first, we have to identify an optimal panel, optimal reduced panel by finding the subset of markers whose mean intensity are maximally correlated to the remaining withheld markers. The reduced set of markers is then fed into a multi-encoder VAE as a multi-channel single-cell image. And the— this image and the reduced set of markers is then used to reconstruct the missing channels that correspond to the withheld markers. And the next method was proposed by a group at Stanford that is a 3-step process. First, a concrete autoencoder is trained to learn a representative subset of 7 markers. Then a multiscale representation of each cell for each marker is fed into a ResNet to learn morphological features for each cell. Then those features are fed into a regression model to learn the mean intensity value of each withheld marker for each cell. The 3rd method in this family of methods is from Carnegie Mellon, and they propose an iterative approach to learning a reduced panel by training a U-Net to reconstruct image patches where a single marker is added to the reduced panel at each iteration, starting with 1 input marker and continuing until the target reduced panel size is reached. So the most recent method is from a group mostly from Harvard that uses a U-Net to reconstruct patches of a single marker using 6 input markers. So the first 2 methods that I mentioned are most comparable to our approach, and this Luke Turns implementation, it will serve as our baseline for my method that I'm discussing here. And so they're most common to us because they use single cell images as input. Whereas the second 2 methods use an image patch containing multiple cells as input. So all of these methods except the Sun et al. method suffer from the same limitation. And that is that the reduced panel is selected prior to and independent of marker imputation. So that results in a potentially asymmetric pair of problems where the identified reduced panel is not guaranteed to be the optimal panel for marker imputation. So, and additionally, all 4 methods suffer from the limitation that each trained model is only compatible with one reduced panel. So that limits out-of-the-box use cases. So the framework that we implement our panel reduction marker imputation process through is through using a clever implementation of the masked autoencoder. So we've talked a lot about over the past couple of days about masked language models. And so the masked autoencoder is this extension of masked token prediction to images, where the two key differences here is that we use an encoder-decoder framework where the encoder only sees the visible image patches, and the decoder then will take the latent representations of those visible patches along with the masked tokens and then use to regenerate the missing patches. And then it's trained with mean square error on the pixel values. And so if you have a standard RGB image, then you're going to be masking along the height-width axis of— so each token in your input is going to be like, for example, a 16 by 16 pixel by 3 color channel image. And so that's how you represent a sequence within this context. And so cleverly, what we can do here is simply mask along the channel axis for our multidimensional images. So what we feed in as input are these single-cell images cropped to a 32 by 32 pixel region, and then they have the 25 channels for each marker stain. And so conveniently, 25 can be— the 25 markers can be reformatted into this 5 by 5 grid where each Essentially, each patch now is a single image channel, and we can just mask channels the same way that we are masking patches in the original implementation. And so the original authors show, as opposed to BERT, that a high image masking ratio is effective here. And so they do about an 80% masking ratio, and that allows us to compress the information of this image, the semantic information of this image, by 80%. Which is effectively what we want to do here is compress the information of our multi-channel image into a subset of markers. So, to bring it to my implementation, how we use this. So, first, our dataset consists of a single breast cancer tissue microarray, which was collected under the auspices of the Human Tumor Atlas Network, and it contains 88 cores of 6 different breast cancer subtypes. And so we segment this TMA and we result in a dataset of over 700,000 cells. And so we crop to the 32 by 32 pixel ROI and zero out the background. And then we randomly mask 50% of the input channels during training and train the model to reconstruct the mask channels using mean square error as the loss function. Because we randomly mask channels during training, our model has generally learned to impute any marker from any other marker in our training set. So additionally, the transformer architecture allows us to utilize pre-trained weights on variable input sizes. So even though our model was trained to predict 13 markers from a panel of 12 randomly chosen markers, we find that during inference we can change the number of input markers and the better, and better prediction performance will result accordingly. So as we increase the amount of input markers, we'll get a better prediction for our remaining withheld markers. And so we seek to leverage this by using our trained marker to identify optimal reduced panels. And so this has the benefit of reducing the computational complexity by only having to train one model and also has the benefit of unifying the panel reduction and marker imputation steps into one model. So how do we leverage this? Is so it's still intractable to evaluate every subset of markers for every panel size that would result in over 30 million different combinations to experiment with. So instead, we propose a greedy selection process of finding the next best marker to add to the panel. And so to conduct the panel selection process, we can progressively increase the number of unmasked channels input to the trained model, starting with the best 1-marker panel and ending with the best N-1 marker panel. At each step, we find that the next best channel to be unmasked by, by choosing the channel that results in the best reconstruction of the masked channels, which we evaluate based off the Spearman correlation between the real and predicted mean intensity values. So this can be interpreted as choosing the most informative remaining marker. This process only requires us to evaluate 300 different candidate reduced panels for our dataset of 25 markers. And so this example on the bottom left here is showing a, uh, toy dataset of, uh, 4 markers. And so we first start by fixing the, uh, one marker panel to be DAPI because we always need DAPI, uh, to, uh, conduct the, uh, later downstream analyses. And so we then iterate through the remaining 3 markers to find the best 2-marker panel and select the next best marker. And then we include that into our reduced panel size and then iterate through the remaining 2 markers to find the best 3-marker panel. And so, yes, that clarifies. And so, This approach results in a significant performance increase over our lab's initial approach, as well as streamlines the process. Here we compare between the MEVAE model and the MAE model on the same breast cancer TMA dataset, and we evaluate the Spearman correlation between predicted mean intensity values for the withheld markers using 6 different reduced panel sizes. The yellow dashed line represents the performance of the MAE model on the— So on the panel size of 3, 6, 9, 12, 15, and 18, and the yellow violin plots here represent the performance of MAE using the same panels. And then the green violin plot here represents the performance using the reduced panel selected via iterative selection that I just demonstrated. So the other advantage of our MAE approach is the ability to evaluate any panel size with the same model. Since a new model must be trained for each panel, in our previous approach, analysis was only done on a fixed number of panel sizes. So I have, for example, here the optimal panel identified for 3 markers in our last implementation and then 6 markers here. And so the top table shows the reduced panel selected via correlation. You can see that the 6-marker panel consists of removing ECAD and adding CICLIN D1, VIM, and ER. The bottom table then shows our construct— our panels selected via iterative selection. And not only is it easier for us to evaluate more reduced panels, you can see that each reduced panel is a direct subset of the next panel, which greatly increases the interpretability of our panel selection procedure, which I will detail in the next slides here. So we can get a very interesting look at what's happening as we add markers to our reduced panel., here, as an example. So we first start by fixing DAPI as the first channel chosen on the right. The heat map indicates the Spearman correlation of each predicted mean intensity value for the marker, and the chart on the left shows the mean value of each column on the right. Using just DAPI as the reduced panel results in a poor reconstruction of the remaining 25 markers. We then iterate through the remaining 24 markers to find the best marker that can reconstruct the remaining 23 markers, and we find that EGFR is best. On the left, you can see that the prediction of the withheld markers has significantly increased. And on the right, the change in color of the pixels in each row demonstrates the improvement in the reconstruction for that marker when EGFR is added to the reduced panel. And then in the next step, this is where it gets interesting, is that we iterate through the remaining 23 markers and choose CK19. And then we can see that on the right, the reconstruction of PanCK has significantly increased with the inclusion of CK19. Into the reduced panel, which makes sense since PanCK and CK19 are highly coexpressed. So when we do this for every panel from one marker all the way up to N minus 1, then we get this nice little pyramid that with this red vertical line, we can slide it left or right based off our own personal threshold as to how acceptable of an error we would like to tolerate in choosing our reduced panel. And we, we personally find that a 9-marker panel works well and results in a 0.87 mean Spearman correlation on the 16 withheld markers. And so this scatter plot is showing the real, real versus predicted mean intensity values for each marker where each point on the scatter plot represents a cell here. And that's also shown here for 3, 6, 9, 12, 15 marker panel sizes. And as we increase the panel size, the prediction of the intensity values increases greatly. And so we also show this works on a colorectal cancer dataset, and we're still working on investigating multiple tissue types and working to integrate multiple panels from different tissue types as well as panels that have different markers on them. So that's our future work. And so I would like to thank my lab members at the Quantitative Bioimaging Lab at OHSU and our collaborators at Harvard