CaTran: ultra-light neural network for predicting gene-gene interactions from single-cell data

Artemy Bakulin

Keywords: Self-attention, gene embeddings, causality guided attention

TL;DR: Making DCDI simpler with learnable gene embeddings

Abstract: Part of the difficulty of learning a gene-regulatory network from expression data is related to the fact that edges in such a network represent different interactions with a different effect size. Therefore modeling gene associations requires learning an individual function for each pair of interacting genes. This may greatly inflate the number of parameters in a model and lead to insufficient generalization. In this paper we propose a method for gene regulatory network inference, called CaTran (Causal Transformer), which avoids explicitly learning pairwise relation between genes, which allows it to significantly reduce the size of the model. The key feature of this approach is learning for each gene a low dimensional embedding and then using a self-attention mechanism to estimate its relation to other genes. Our method is applicable for both observational data and data with interventions. For the latter it implements a differentiable gene importance test and forces attention values to be in accordance with it. Because the gene regulatory network in CaTran is learned as a soft adjacency matrix, it allows sampling graphs with arbitrary number of edges based on a set threshold. Comparison of these graphs with the gene networks from databases showed that even for large graphs the edges are predicted with high precision.