Scientist Stories: Nauman Javed, Cross cell type regulatory predictions

My name is Nimaad. I'm in Brad Bernstein's group at the Broad. And I'm excited to tell you today about some of our efforts to use representation learning to make cross-cell type regulatory predictions. So, you know, we've heard at length by a number of different talks today about non-coding regulatory grammar. So just briefly again, what this refers to is the fact that gene regulation is actually governed by the context-specific activity of non-coding regions of the genome called enhancers. That typically lie in accessible chromatin, bind transcription factors or certain motifs. And this actual sequence specificity actually governs the activity of these non-coding regulatory elements. And so understanding how sequence governs this activity is really important for understanding things like non-coding disease-associated variation and understand— making better models of gene regulation at large. And so we've also seen today a number of really cool talks and more broadly in the field about how sequence-based deep neural networks are really effective tools for learning this grammar. And so briefly, the setup for this is, you know, you'll have genomic sequence, you'll one-hot encode it, you'll put it through some neural network and predict something like DNA accessibility or gene expression. And you're not actually super interested in the prediction task itself, but rather interpreting the network. And once you interpret the network, you can gain insights into things like motifs and motif orientation and spacing, et cetera. But one of the limitations of just using sequence actually is that you actually have to retrain separate networks for unobserved cell types. You can't actually generalize beyond what you've seen before. And so really the motivation for this, the work I'll be showing is how can we improve generalization and move towards cell-type agnostic models. And of course, the way to do this is, you know, to somehow incorporate cell-type-specific information. And we've already seen some work on how to do this effectively today. We wanted to, you know, do this in a way that makes sure that the cell-type-specific information is readily accessible and also preserves interpretability since sequence-based models are relatively more interpretable than models where you're incorporating, you know, many different modalities. And so what we introduce here is a model called GenFormer. So this is a multimodal transformer that basically integrates sequence and attack seq, and it incorporates, you know, the typical pre-training and fine-tuning paradigm that's become so common in, you know, language models at large. And so in the pre-training stage, what we do is we basically input to this transformer network genomic sequence as well as masked attack profile, and we tasked the network with predicting the masked accessibility signal. Once you've learned these representations over a sufficiently large corpus of ATAC-seq data, you can actually fine-tune it for downstream regulatory predictions, for example, gene expression prediction. And what I'll show is that masking actually improves the integration of the two input modalities and actually improves— allows, you know, preserves interpretability. You can now predict variant effects on both your input attack profile as well as your downstream target like gene expression. In terms of specific network architecture, so we're integrating a pretty large genomic bin, so about half a megabase of genomic sequence and the accompanying attack profile. We adopt, you know, a simple BERT-style masking approach, so 15% of the input with some peaks included, and we actually employ a span mask. So about 1.5 kilobases, in order to ensure that it's not really a trivial imputation task where the network's leveraging, you know, accessibility at neighboring adjacent bases to impute the missing signal. These two modalities are processed by separate convolutional layers, which also have a max pooling layer to reduce the dimension before they're concatenated and entered into a number of performer attention layers. So performer attention is basically similar to Transformer vanilla attention, but uses a random feature projection to decompose the attention matrix calculation and lets you reconstruct an approximation of the matrix in linear time rather than quadratic and significantly reduces the compute burden of training this type of model. And finally, once you've gone through the performer layers, we make a prediction over the masked tokens. It's a regression-based learning task. We're predicting at about 128— at 128 base pair resolution. And the loss function we're using is a sum of the total counts with a Poisson term, as well as a multinomial penalization of the shape, which is, you know, a BPNet-style loss, if you're familiar with that network. So what does the masking look— what does the masking task look like? So here we have a half-megabase locus at the LGR4 gene. And we've introduced 3 masks. So we can see for these holdout cell types and this holdout genomic region across a range of different signal values, we're able to accurately impute the missing signal and also pick up some cell type-specific peaks. So for example, the middle peak here. Looking more comprehensively just at the imputation task itself, here we've aggregated data from 14 holdout cell types. These are mostly cell lines and rare single cell populations that don't really have closely matching cell types in the original training corpus. So they are, you know, more likely to be true holdout cell types. We also are using holdout genomic regions. So these are the same regions in the informer test set. And we're just basically restricting the peaks as well as a number of negative matched background regions. And so we achieve a pretty, you know, good Pearson's correlation, Spearman's correlation about 0.8 here. And actually running the ablations for the two input modalities shows that, really the network is leveraging both in order to make the prediction, so both sequence and attack seq. The mass signal prediction actually seems to depend upon distal information. So we can see that in two different ways. The first is we can reconstruct the attention matrices from the performer attention layers. So here we are averaging across all 7 layers and heads and showing the same locus that I showed before, but for the two holdout cell types, K562 and PC3. And what you can see here is that, you know, many of— most of the tokens are attending to a number of ATAC-Seq peaks within the region. So what this shows is that perhaps, you know, distal information is important for imputing your missing signal. And another way to see this more quantitatively is to run the ablation. So if you, you know, go down to 65 kilobases, you can see, you know, decrease in the Pearson's correlation relative to if you're using the full input length at 524. Here's an example of interpretability within, you know, the masked regions. So you can also, you know, compute saliency scores. So ask, you know, what sequence features in the masked regions are driving accessibility? So here, for example, we can recover a CTCF motif, and we can, you know, also recover a GATA motif that's specific to K562. On the bottom-hand plot. Looking a little bit more comprehensively, we can look at CAQTL effect prediction. So these are single nucleotide variants that are, you know, impacting chromatin accessibility experimentally measured. So here on the left-hand side, we can see a DSQTL prediction with Genformer precision recall, you know, Average precision here is, you know, much better than a simple SVM-based baseline, GKM-SVM. And that's, you know, in turn significantly better than the actual positivity rate of about 0.02. Now we can also predict the effect sizes. So we're doing this in a sort of cell type manner. That's an important thing to emphasize. But, you know, across both LCL and bulk liver tissue, able to, you know, predict the effect sizes pretty well, achieving Spearman correlations of about 0.77 and 0.73 on the right-hand side. And finally, for the gene expression prediction task, what we're doing is taking the pre-trained network and replacing basically output head with an uninitialized new head, retraining that at a higher learning rate than the rest of the network. And we're doing this over a number of cell types in ENCODE that have paired ATAC and RAMPAGE-seq datasets. So RAMPAGE is an assay similar to CAGE-seq that measures 5' biased transcriptional output evaluated on a number of held-out cell types and genes. So here's an example of the profile predictions. So, you know, here we're accurately predicting the expression of ARMCX2 in lung relative to thoracic aorta. And, aggregating over TSSs, we can achieve a pretty good Spearman's correlation, again, of around 0.82. Looking more comprehensively, there's two sorts of tasks you might be interested in for gene expression prediction. So the first is across genes prediction, looking at how effectively you can predict the variance of genes within a single cell type, which is a little bit easier of a task given that networks are generally pretty good at discriminating highly and lowly expressed genes. Just based on some simple sequence features. The harder task here is predicting, you know, the variance of a single gene across cell types or tissues, because here you're having to learn something about tissue-specific regulation. And so for both of these tasks, actually GenFormer achieves, you know, outperforms the same network without pre-training, as well as the state-of-the-art model Informer, which is only sequence-based, use a multitask framework. We can ask the same question now for holdout cell types, so network, you know, cell types that the network has never seen. And so here again, we can achieve a good experiments correlation aggregated over TSSs. We can also break it down in the same way that we did on the left-hand side, and we see a little bit of a performance decline but still respectable correlations. So in terms of conclusions and next steps, so what I've shown here is masking-based pre-training is an effective way of integrating sequence and accessibility. The embeddings are interpretable for variant effect prediction. The pre-trained network actually leads to an improved gene expression prediction for held-out cell types. And in terms of future work, I think it would be very interesting to incorporate transregulatory information, also looking at eQTL prediction and distinguishing these from CAQTLs, as well as the extraction of enhancer-promoter and enhancer-enhancer interactions. And with that, I'd like to acknowledge my collaborators, Tom Weingarten and Arijit Sahanovish, Avi Dubi, Adam Roberts, Christoph, and Brad, as well as the TPU Research Cloud for giving me free compute. So, um, hi, really like the model. Um, I'm curious, within your ablation studies, are the actual Spearman metrics you're reporting before or after your kind of multinomial adjustment in the sense of how is the scale kind of, is that potentially leaking into your Spearman or your Pearson estimates? Like if you're giving it a TAC-seq only, kind of you're telling it the scale already. If you have any thoughts on that. Yeah, I think that's a good point. In order to sort of mitigate that, we do add like a random noise component to the input, but that won't sufficiently perhaps get rid of the scale, but it's not latching on directly to a certain ratio that's resulting from re-depth normalization or something like that by virtue of adding that noise. But yeah, that is a concern and a good point.

So my question was going to be, and you kind of alluded to it actually with the trans-regulation point. So when you're generalizing to other cell— unseen cell types, I guess the hope is that the model gets to see, oh, this chromatin region which had these transcription factor motifs is open. And so those transcription factors must be active. And therefore, in the masked region, if I see that motif, then I'm like, okay, it's more likely to be open. And so I guess, so I guess one question is, do you think that, are you able to see that that is actually what the model is learning? And the second question is, how much do you think there is to gain from including the trans regulation, which in a way is, I think in a way it's asking like, if you, if you had infinite compute, not just what TPU research cloud can give you, could you do this on like the whole genome? And then that would give you like the ultimate prediction?

Yeah, we actually haven't looked at that. I think that would be a really interesting thing to look at. So for example, like, just seeing which motifs are most accessible and how well that correlates to the presence of motif and the prediction accuracy. So that's definitely a great thing to look at. And yeah, ideally we would be doing this, you know, across all, like, active enhancers and a number of matched negative regulatory regions. But You know, like, for example, a screen from ENCODE has like more than a million elements. So it is difficult. Whereas, you know, on the gene expression side, like, you only have 20,000. That's still within the realm of like something reasonable to run a supervised self-supervised. Great talk. So I was wondering, have you benchmarked with CRISPR? Like enhancer perturbation datasets compared to other models to see if you can find, rank the enhancers better? Yeah, that's definitely like next steps. For some reason I can't change the slide, but that is one of the next steps that we want to do. So like for K562, I think there's a lot of data on, you know, validated EP interactions and now a number of other cell lines as well. So yeah.

All right, thanks everyone. So because we— so because we knew we're going to run over time, we created this fake session called Poster Highlights, which doesn't actually exist. But what we're going to do is walk over to the other building, which is you just get out of this building, turn left, and cross the street. There's the Allen Center. In the atrium, there will be the poster presentation. So there was about 70 posters that were accepted. About half of them will be presented here in person and half will be online, and we'll have a link for the online posters so that you can check out later. So see you guys across the street.