In this paper, we propose an end-to-end deep learning model, called E2Efold, for RNA secondary structure prediction which can effectively take into account the inherent constraints in the problem. The key idea of E2Efold is to directly predict the RNA base-pairing matrix, and use an unrolled algorithm for constrained programming as the template for deep architectures to enforce constraints. With comprehensive experiments on benchmark datasets, we demonstrate the superior performance of E2Efold: it predicts significantly better structures compared to previous SOTA (especially for pseudoknotted structures), while being as efficient as the fastest algorithms in terms of inference time.
RNA Secondary Structure Prediction by Learning Unrolled Algorithms
This paper proposes E2Efold, which is an RNA secondary structure prediction algorithm based on an unrolled algorithm. Previous methods rely on dynamic programming (which does not work for molecular configurations that do not factorize) or rely on energy-based models (which require a minimization step, e.g. by using MCMC to traverse the energy landscape and find minima). The former does not work for all molecules and the latter can be difficult to optimize. The method presented here is novel, shows strong SOTA performance, and would be of interest to the wider deep learning community.
The method is based on an unrolled algorithm, which is motivated by the inclusion of three inductive biases / constraints important underlying RNA folding. These constraints limit the wide RNA search space. The first component of the method is a “Deep Score Network” which uses a stack of Transformer encoders (with relative and exact positional embeddings) followed by 2D convolutional layers to output a L x L symmetric matrix describing the “scores” of base pairing. As these scores may not obey the rules of RNA folding, a second post-processing network is trained end-to-end together with the “Deep Score Network” to enforce constraints. This network starts with a transformation that symmetrizes the matrix and applies a constraint-enforcing mask. The problem is transformed into an unconstrained problem by using Lagrange multipliers; it is then solved using a proximal gradient. Finally, a recurrent cell is defined that implements this algorithm in a deep learning framework. This method is creative, could be applied to other tasks with constraints, and would be interesting to the wider deep learning community.
In addition to developing the deep score network and post-processing network, the authors also develop a differentiable F1 loss, so that the network can directly optimize for precision and recall on the task. The performance of this method significantly outperforms previous methods. There was a fruitful discussion on OpenReview regarding whether this was a result of overfitting on the task. Indeed, it is critical in deep learning applications to carefully construct train/test sets to avoid high performance by memorization alone. To address this, the authors train on RNAStralign and test on ArchiveII. As the original ArchiveII dataset contains subsequences of other RNA sequences, which can result in overfitting, the authors re-ran their experiment with that removed, and similar results were achieved. To support the hypothesis that ArchiveII and RNAStralign capture different distributions, they perform a permutation test on the unbiased empirical Maximum Mean Discrepancy estimator, finding that the distributions are different. I do wonder why they did not check if P(ArchiveII) = P(ArchiveII) as they do check if P(RNAStr_train) = P(RNAStr_train). On the specific task of pseudoknot prediction, the method also performs well (F1 is >0.23 over the baseline). On sequence length-weighted F1, the model does even better.
The paper is rich with ablations. The analysis of the number of unrolling iterations T helps support the use of an unrolled method and builds intuition for its importance - it would be useful to include this in the appendix of the paper. I also appreciated the visualizations, which are a good sanity check that the model correctly handles pseudoknots. The performance of the method is broken down by RNA family, which is also quite interesting -- the method outperforms LinearFold on all classes, besides 5S RNA, SRP, and Group I intron. Further analysis is required to better understand why the method is weaker on those datasets. Additionally, further work should explore training on one set of families and testing on a held-out set of families. This was pointed out by public comments on this paper. This is potentially a limitation of E2EFold (the authors do not seem to have tried this suggested experiment) and further exploration is required. Exploring this limitation (even if it is not overcome) would make this paper even more rich.
That said, I recommend acceptance of this paper due to the extensive experiments, polished writing, novel method, and strong results, which can inspire future research.