Liteformer: Lightweight Evoformer for Protein Structure Prediction

The claimed performance improvements are comparable to or less significant than those of other optimized versions of AlphaFold2, none of which are mentioned here. FastFold, UniFold, OpenFold have all already improved AlphaFold2 with e.g. FlashAttention.

Flash Attention is an implementation of the classical pairwise attention mechanism taking into account GPU details such as caches and communications. Flash Attention does not reduce the computational complexity in big $\mathrm{O}$ notation though. Such optimization is orthogonal to our work which focuses on deep learning architecture improvements to address computational complexity in big $\mathrm{O}$ notation. At the moment we did not find any public implementation of Flash Attention for Evoformer, especially the $O(L^3)$ component of the Evoformer.

Furthermore, Flash Attention solely focuses on unbiased attention computation and lacks the capability to expedite computation with a predetermined bias. As the communication between MSA representation and Pair representation relies on biased attention, Flash Attention isn't directly applicable for acceleration within this context.

Given these, we feel it will be unfair for us to invent a new Flash Attention implementation for Evoformer which is a new project on its own. In contrast, our contribution is to reduce the complexity of biased attention and achieve a similar accuracy for protein structure prediction tasks.

That said, we thank the reviewer for bringing up Flash Attention in this context. We think implementation and design ideas from Flash Attention are orthogonal to our Liteformer deep learning architecture design, and it will be a good future direction to incorporate ideas from Flash Attention into Liteformer to further improve its efficiency.

"However, since both the row and column dimensions of the pair representation, along with the row dimension of the MSA representation, are identical to the primary sequence length L." I'm not sure what the first sentence is saying. "We trained 10,000 data for 5 days using 8 × DGX-A100 80G GPUs" The second one contains no information about which 10,000 data points were used to train the model.

Thanks for your valuable suggestion. We have revised "However, since both the row and column dimensions of the pair representation, along with the row dimension of the MSA representation, are identical to the primary sequence length L." to "However, since both the row and column dimensions of the pair representation, along with the row dimension of the MSA representation, are the same as the primary sequence length L." and revised "We trained 10,000 data " to " We trained 10,000 randomly selected general protein data from PDB"

Table 2: Quantitative results comparison between Liteformer with ESM2-150M and Evoformer with ESM2-150M on CASP and CAMEO datasets.

Table 6: Quantitative results comparison between Liteformer with AlphaFold2 and Evoformer with AlphaFold2 on CASP and CAMEO datasets.

Thank you for your feedback. We enhanced the draft's clarity and updated the experimental results for both Liteformer with ESM2-150M and Liteformer with unmodified AlphaFold2 on the complete CASP and CAMEO datasets. Finally, we explained the contrast between Flash Attention and our model. We provide detailed answers to your comments below.

It is claimed in Figure 1 that AlphaFold2 OOMs on sequences of length 800, e.g. If the authors are running inference using the same 80GB A100s they use later in the paper, this simply cannot be true; one can get away with longer sequences even on ColabFold (Mirdita et al., 2022), which runs on free-tier Google Colab GPUs. The authors of OpenFold were able to run an unmodified version of the original AlphaFold2 on sequences of length 2000 on a 40GB A100 (Ahdritz et al., 2022).

We appreciate the opportunity to clarify a crucial aspect that might have caused confusion. In Figure 1 and throughout the context, our focus is on comparing training lengths, as opposed to the inference lengths highlighted in the studies mentioned by the reviewer. It's essential to differentiate between training and inference processes. During training, the significant memory consumption predominantly arises from backward gradient propagation, thereby necessitating shorter sequence lengths compared to inference tasks. We would like to emphasize that our comparison specifically highlights training sequence lengths, and we've made this distinction explicit in the text and Figure 1's caption. Our intention was to address the challenges in memory and computational complexity during the training phase due to the substantial memory requirements for gradient backpropagation.

The choice of ESM-2 150M as a baseline for monomeric structure prediction is extremely confusing. Why not use unmodified AlphaFold2? Why not use a larger ESM-2 protein language model?

We would like to provide clarification regarding the selection of ESM2-2 150M as a baseline for monomeric structure prediction in our experiment. Liteformer is a versatile protein structure module that seamlessly integrates with both unmodified AlphaFold2 and other Protein Language Models, such as ESM-2. Our objective was to assess Lightformer's performance in the context of single protein structure prediction tasks.

To further validate the effectiveness of our model, we included comparative experiments conducted with the unmodified AlphaFold2 framework. Due to the rebuttal's time limitation, we randomly select 10,000 general protein data from PDB and evaluate Liteformer and Evoformer on CASP and CAMEO datasets. Tab 6 shows that the accuracy predicted by Liteformer with AlphaFold surpasses Evoformer with AlphaFold on CASP and CAMEO.

Why mix and match these figures with AlphaFold2-Multimer evals?

Regarding the reviewer's concern about the integration of figures with AlphaFold2-Multimer evaluations, we would like to clarify that AlphaFold2-Multimer is tailored for complex structure prediction tasks. The assessment of its performance involves evaluating the accuracy of individual protein parts and docking performance. Both complex and single protein structure tasks necessitate assessing the predicted structure's accuracy using similar metrics. It's crucial to note that utilizing the same public metric doesn't imply a mix of results from two distinct tasks. Our experiment involved two tables explicitly demonstrating single and complex protein structure prediction performances separately, without amalgamating or analyzing them together. Figure 1 aims to show how our Liteformer achieves comparable results to the corresponding baseline model across different datasets. It's not about comparing two different tasks but rather demonstrating our model's performance across different datasets.

Which CAMEO proteins were chosen? Why were all evaluation proteins filtered at 500? Which AlphaFold2 implementation served as the baseline?

Thank you for your feedback. We inadvertently used a specific subset of data in our experiments due to an oversight in our experimental setup. Recognizing this error, we have re-conducted our experiments using the complete CASP and CAMEO datasets to ensure comprehensive and accurate results, as shown in Tab 2. It's important to note that during our initial experiments, a very small portion of excessively long-length data was filtered out. We apologize for any confusion caused by the previous use of filtering. The implementation of AlphaFold2 and AlphaFold-Multimer is from its original paper.

Dear Reviewer ffAR:

Thank you very much for your time spent on our submission and your questions. We have tried to address your concerns in the response and updated submission--any feedback from you would be appreciated. If you have further comments, please kindly let us know. We hope for the opportunity to respond to them.

Best wishes, Authors of paper 3633

Thank you for your thorough assessment and comments on our manuscript. We are excited that you appreciated the contribution and the novelty of our work. We have enhanced the clarity of our draft by indicating our motivation for using flow network in the Introduction section and revising the complexity analysis of our model in Section 2 and Section 3. We respond to your comments below.

I think there should be a separate paragraph in the introduction detailing the motivations for using flow network.

We have added a paragraph in the Introduction section to illustrate the motivation for utilizing flow network as below: " The diversity of protein sequences requires non-trivial attention to generality. However, previous linear attention mechanisms mainly use similarity decomposition to approximate Softmax($\mathrm{QK^T}$) with different kernel functions, which reintroduce inductive bias and devise linear complexity at the expense of generality. Therefore, we turn to reconstructing the biased attention mechanism from a new perspective, a flow network perspective. "

The claim does not correspond to the experimental results. In Sec 2 and 3, the authors mainly claim that BFA can reduce computational complexity at Equation 7. However, the experimental results show that memory usage is the main advantage of BFA during training.

We'd like to address the point you've raised, as we suspect there might be some misunderstanding. Our analysis in Section 2 and Section 3 have extensively delved into both memory and computational complexities within the manuscript, aligning closely with the experimental outcomes. In the context following Equations 1-10, we have intentionally provided insights into the tensor shapes involved in each operation, and we've analyzed the computational and memory complexities associated with each operation, albeit using the term "complexity." Following your advice, we updated the analysis in Section 2 and Section 3 using the term "memory and computational complexity".

Dear Reviewer WrTu:

Thank you very much for your time spent on our submission and your questions. We have tried to address your concerns in the response and updated submission--any feedback from you would be appreciated. If you have further comments, please kindly let us know. We hope for the opportunity to respond to them.

Best wishes, Authors of paper 3633

Dear Reviewer WrTu,

We extend our sincere gratitude for investing your time in reviewing our submission and raising pertinent questions. Our response aims to tackle your concerns, reflected in the updated submission. If you have further comments, please kindly let us know. We hope for the opportunity to respond to them.

Best regards, Authors of paper 3633

Dear reviewer,

Thank you very much for your work evaluating this review.

It is critical that you urgently address the author's responses, acknowledge their response, and eventually adjust your rating if warranted.

Best,

AC

Dear Reviewer,

Thank you for your thorough review. We are glad that you appreciate the contribution and novelty of our work. We improved the clarity of our draft, updated the method's figures, and indicated our model's tailoredness to protein tasks. Finally, we re-conducted the monomeric protein structure prediction on the complete CASP and CAMEO datasets and explained the contrast between Flash Attention and our model. We appreciate your insightful suggestions and provide detailed answers to your comments below.

In Fig. 2, the pipeline of the blocks in Evoformer and Liteformer is exactly the same. It's better to be combined into one figure. The difference between BFA and original attention is the key to which should be compared like this.

Following your advice, we have updated the Liteformer architecture figure in Figure 2 and the comparison between BFA and canonical attention in Figure 6. Figure 2 shows the distinction between Evoformer and our Liteformer architectures within a single panel. Through the implementation of Bias-aware Flow Attention, our Liteformer notably diminishes the memory and computational complexities from quadratic to linear in MSA row-wise updates and from cubic to quadratic in pair-wise triangle updates. Additionally, we have incorporated the attention computation process alongside our Bias-aware Flow Attention in Figure 6 to aid in contrasting our linear attention algorithm with the canonical attention method.

More ablation studies between Liteformer and Flash Attention should be provided for further comparison. Can they jointly improve the training of the network?

Flash Attention is an implementation of the classical pairwise attention mechanism taking into account GPU details such as caches and communications. Flash Attention does not reduce the computational complexity in big $\mathrm{O}$ notation though. Such optimization is orthogonal to our work which focuses on deep learning architecture improvements to address computational complexity in big $\mathrm{O}$ notation. At the moment we did not find any public implementation of Flash Attention for Evoformer, especially the $O(L^3)$ component of the Evoformer.

Furthermore, Flash Attention solely focuses on unbiased attention computation and lacks the capability to expedite computation with a predetermined bias. As the communication between MSA representation and Pair representation relies on biased attention, Flash Attention isn't directly applicable for acceleration within this context.

Given these, we feel it will be unfair for us to invent a new Flash Attention implementation for Evoformer which is a new project on its own. In contrast, our contribution is to reduce the complexity of biased attention and achieve a similar accuracy for protein structure prediction tasks.

That said, we thank the reviewer for bringing up Flash Attention in this context. We think implementation and design ideas from Flash Attention are orthogonal to our Liteformer deep learning architecture design, and it will be a good future direction to incorporate ideas from Flash Attention into Liteformer to further improve its efficiency.

The proposed BFA module seems more likely to be a general attention module, instead of a theme-related (protein-related) approach, which decreases the significance.

Regarding the perception that our approach is more of a general algorithm rather than tailored to a specific task, we would like to clarify that linearization of biased attention within our BFA method directly addresses the unique demands of protein modeling. In protein modeling, the exchange of information between MSA and Pair-wise representations crucially relies on biased attention mechanisms. Our focus on linearizing biased attention specifically caters to the requirements of this domain, ensuring an optimal strategy for enhancing the efficiency of protein structure prediction tasks.

What's the principle of selecting the targets from the CASP14, CASP15, CAMEO, and VH-VL, and why sequence lengths of them are restricted to be shorter than 500, since Liteformer has a better ability to handle the longer sequences?

Thank you for your feedback. We inadvertently used a specific subset of data in our experiments due to an oversight in our experimental setup. Recognizing this error, we have updated our experimental results using the entire CASP and CAMEO dataset to ensure comprehensive and accurate results, as shown in Tab 2 in the main context. It's important to note that during our initial experiments, a very small portion of excessively long-length data was filtered out. We apologize for any confusion caused by the previous use of filtering.

Table 2: Quantitative results comparison between Liteformer with ESM2-150M and Evoformer with ESM2-150M on CASP and CAMEO datasets.

Dear Reviewer 9JiD:

Thank you very much for your time spent on our submission and your questions. We have tried to address your concerns in the response and updated submission--any feedback from you would be appreciated. If you have further comments, please kindly let us know. We hope for the opportunity to respond to them.

Best wishes, Authors of paper 3633

Dear Reviewer 9JiD,

We extend our sincere gratitude for investing your time in reviewing our submission and raising pertinent questions. Our response aims to tackle your concerns, reflected in the updated submission. If you have further comments, please kindly let us know. We hope for the opportunity to respond to them.

Best regards, Authors of paper 3633

Dear reviewer,

Thank you very much for your work evaluating this review.

It is critical that you urgently address the author's responses, acknowledge their response, and eventually adjust your rating if warranted.

Best,

AC