Paper: Liteformer: Lightweight Evoformer for Protein Structure Prediction (ICLR-2024)

Reviewer comment (Experiments):
Some of the performance figures that motivate this entire paper seem questionable. It is claimed in Figure 1 that AlphaFold2 OOMs on sequences of length 800. 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).

Author response:
**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.

Reviewer comment (Experiments):
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? Why mix and match these figures with AlphaFold2-Multimer evals?

Author response:
**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.

Reviewer comment (Reproducibility):
In general, details on the evaluation are very light (e.g. which CAMEO proteins were chosen? Why were all evaluation proteins filtered at 500? Which AlphaFold2 implementation served as the baseline?), to the point where it's difficult to know what is being run.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Writing):
The paper is a bit sloppily written, with unclear or incomplete sentences. For example: "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" is difficult to understand. Another example: "We trained 10,000 data for 5 days using 8 × DGX-A100 80G GPUs and inference on two multimeric datasets: VH-VL and DB5.5" contains no information about which 10,000 data points were used to train the model.

Author response:
**"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"

Reviewer comment (Experiments):
Regarding the evaluation datasets from CASP14 and CASP15, the authors state: "The 14th & 15th Critical Assessment of protein Structure Prediction, from which we respectively select 42 and 48 single protein targets with sequence lengths less than 500." Why filter in this way? The hardest proteins are the longest ones. Almost all entries at CASP15 did well on most of the proteins of length < 500 (excepting some orphans, etc.).

Author response:
**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.

Reviewer comment (Presentation):
Unclear motivation: The author emphasizes the huge computational complexity reduction brought by converting attention module into a flow network. However, how this method affects the computational complexity is only mentioned at Eeq 7 in Section 3. I think there should be a separate paragraph in the introduction detailing the motivations for using flow network.

Author response:
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. "

Reviewer comment (Evaluation):
The claim does not correspond to the experimental results. In Sec 2 and 3, the authors mainly claim that BFA can reduce computational complexity. However, the experimental results show that memory usage is the main advantage of BFA during training. The drop in computation time is less pronounced. The authors in sec 4 simply boil this down to reduced graph overhead when backpropagated. In my opinion, it is not helpful to prove the validity of this method.

Author response:
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.      

**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".

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Experiments):
The proposed method shows a great improvement in efficiency and memory cost. But there have been some more general works that play a similar role, such as memory-efficient attention and flash attention [Ref_1]. Can they jointly improve the training of the network? More ablation studies between them should be provided for further comparison.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.