Liteformer: Lightweight Evoformer for Protein Structure Prediction
Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Keywords: protein structure prediction, efficient transformer
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
Abstract: AlphaFold2 has achieved seminal success in predicting structures from amino acid sequences with remarkable atomic accuracy. However, its Evoformer module faces a critical challenge in terms of high memory consumption, particularly concerning the computational complexity associated with sequence length $L$ and the number of Multiple Sequence Alignments (MSA), denoted as $s$. This challenge arises from the attention mechanism involving third-order MSA and pair-wise tensors, leading to a complexity of $\mathcal{O}(L^3+sL^2)$.
This memory bottleneck poses difficulties when working with lengthy protein sequences. To tackle this problem, we introduce a novel and lightweight variant of Evoformer named Liteformer. Liteformer employs an innovative attention linearization mechanism, reducing complexity to $\mathcal{O}(L^2+sL)$ through the implementation of a bias-aware flow attention mechanism, which seamlessly integrates MSA sequences and pair-wise information. Our extensive experiments, conducted on both monomeric and multimeric benchmark datasets, showcase the efficiency gains of our framework. Specifically, compared with Evoformer, Liteformer achieves up to a 44\% reduction in memory usage and a 23\% acceleration in training speed, all while maintaining competitive accuracy in protein structure prediction.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 3633
Loading