Long-context Protein Language Model

• Also, we want to note that unlike ProtMamba, we propose to use BiMamba-S, which is the most proper way to realize bi-directionality in Mamba. BiMamba-S can be formulated in a more theoretical way as structured SSMs with quasi-separable matrices [Hwang et al.]. We did not apply quasi-separable mixers in our model since Mamba currently has a much better software-hardware interface environment and distributed training support in practical implementations that help us train a large foundation-level model feasibly and efficiently. We added this discussion in our paper as well. We also want to argue that, just like Transformer, which has been used in numerous impactful works such as GPT and ESM-2, BiMamba-S also serves as a versatile and effective architectural choice for sequence modeling. The reuse of proven architectures like Transformer or BiMamba does not diminish the novelty of a work. Instead, it allows researchers to focus on solving domain-specific challenges (in our case, exploring the capability of building up a new type of foundation pLM based on BiMamba-S). Our work follows this principle, leveraging this architecture to demonstrate its effectiveness in pushing the boundaries of foundation pLM.

• Besides the architectural choice based on Mamba, we have many other important contributions to this field as summarized in the Introduction: (1) method to encode biological interaction information into pLM, we propose a novel second-stage training based on random walks over graphs to extend the long-context capabilities of LC-PLM to leverage the PPI graph context; (2) we demonstrate that LC-PLM has improved length extrapolation capabilities, favorable scaling laws, and achieved a 7% to 34% improvement on downstream tasks (e.g. protein structure prediction (CASP15-multimers, CASP14, Benchmark2), tasks in TAPE and ProteinGym) compared to ESM-2, CARP, especially for longer proteins and protein complexes; (3) we demonstrate its effectiveness in capturing graph contextual information on remote homology detection (TAPE) in Table 3, protein function prediction (ogbn-proteins) in Table 14 and 16, and PPI link prediction (ogbl-ppa) in Table 15 and 17.

2. Significance.

We respectfully disagree with the assertion of “The significance of this work is also diminished by many false and tenuous claims”.

• "Such protein LMs cannot extrapolate to longer proteins [..]": this is not true, and cannot and has not been shown theoretically nor empirically. Transformer-based pLMs can extrapolate to longer contexts and can also handle relatively long contexts (e.g. 16k context size).

Since transformers are invariant to input position, they usually use positional encodings along with the sequence input. These encodings can be difficult to extend past the maximum length seen during training. Most transformer-based foundation pLMs limit the input length during pretraining such that they cannot extrapolate during inference. For example, ESM-2 has a maximum input length of 1024 residues; and it performs worse on both shorter and longer sequences [https://github.com/facebookresearch/esm/discussions/76]. We also provide evidence in Table 6 that ESM-2 fails to extrapolate on both longer (2k-8k) and shorter (0-128) sequences. It is worth noting that ESM-2 even used extrapolatable positional encodings [Zhao et al.], i.e. RoPE [Su et al.], which however performs poorly for protein sequences. CARP as a new type of attention-free foundation pLM increased the input sequence length to 4k but performs worse than ESM-2 due to its lack of expressiveness. Our LC-PLM built off BiMamba-S can take in infinite-length input sequences and extrapolate well on both shorter and longer sequences; and meanwhile provide even better performance on regular length proteins across various downstream tasks.

Table A. Protein structure prediction with LMFold. Structure prediction performance (TM score) are reported on different hold out datasets.

Table B. Evaluation on TAPE tasks in supervised fine-tuning setting. We report the top-1 accuracy for the Remote Homology fold-level test set; accuracy for the 3-class secondary structure prediction on the CB513 test set; Spearman’s correlation coefficients for the test sets for the Stability and Fluorescence prediction tasks. For the Jacobian contact map prediction task, we adopted the methods from [Zhang et al.] to use categorical Jacobian matrices computed from protein language models as the zero-shot prediction for protein contact maps and report the precision@2/L (L is the length of a protein sequence) on the validation set of ProteinNet dataset [AlQuraishi].

We thank the reviewer for their feedback on our paper. We performed multiple novel ablations and comparisons that we report in the main comment of the rebuttal. We also address the reviewer’s specific comments and concerns below:

1. Originality and comparison to ProtMamba.

We respectfully disagree with the assertion of “This method already exists as ProtMamba”. We would like to clarify several key distinctions:

• We have discussed the key differences between our work and ProtMamba [1] in lines 152-153 and 158-161 and Table 1. To summarize again, ProtMamba trained on concatenated homologous protein sequences with autoregressive causal language modeling and an infilling objective focusing on protein sequence generation. However, our LC-PLM focuses on learning a foundation-level long context protein language model (pLM) that can provide universal amino acid level protein representations for extremely long protein sequences; protein complexes, multimers, and heterodimers with encoded protein interaction context information. We would like to demonstrate that LC-PLM is a much better foundation pLM than the previous open-sourced SOTA foundation pLMs (i.e. ESM-2 [Lin et al.] and CARP [Yang et al.]). We additionally argue that, unlike ProtMamba [Sgarbossa et al.], which uses homologous sequences as context, our LC-PLM-G proposed a novel way to encode protein interaction graph information that contextualizes functionally related proteins rather than semantically similar ones.

• Here we provide new experimental results directly comparing our LC-PLM to ProtMamba. In our new experiments, we evaluate ProtMamba on the downstream tasks we used in our paper. The performance of ProtMamba is much lower than LC-PLM (even CARP and ESM-2 as shown in general response) across all tasks. This suggests that ProtMamba pretrained with concatenated homologous protein sequences potentially leads to degraded representations of individual protein sequences. After all, ProtMamba is trained to use homologous sequences as context for protein generation tasks (e.g. infilling and fitness prediction), rather than producing useful representations for protein sequences. Also it is worth noting that ProtMamba cannot extrapolate to sequence > 2048 since they use positional encodings with fixed length in training. We summarized the results in the tables below. We added these results to our manuscript. Also as a side note, ProtMamba is a concurrent submission to ICLR 2025 so other submissions can be excused for not comparing against them according to ICLR’s policy.

5. Can you clearly state which models have been trained in this work and for which models you just took pretrained versions?

We clearly indicated where we retrained models / used publicly pretrained models in our manuscripts, e.g. lines 338-341, lines 366-368, lines 433-434, the footnote in Table 2, etc. In summary, all ESM-2(-G) and LC-PLM(-G) models with 100B tokens are retrained; all models with “-public” postfix are taken from public pretrained checkpoints. We also provide training and dataset details in Appendices D and E.

• Also "they fail to account for the underlying biological mechanisms [...]" is just a tenuous claim without any evidence.

We provide much evidence throughout the paper in Tables 2, 3, 4 to show transformer-based SOTA pLM (i.e. ESM-2) fails to perform well when the task requires knowledge and information about biomolecular interactions and dynamics. For example, in Table 2, we show that ESM-2 performs much worse than LC-PLM on all folding benchmarks, especially on protein complexes, which need protein interaction information to better infer the structures; in Table 3, we show that ESM-2 also performs much worse on TAPE tasks, i.e. remote homology prediction and secondary structure prediction; in Table 4, we show that ESM-2 is also much worse on ProteinGym benchmark. We think there is sufficient evidence to demonstrate the suboptimality of transformer-based SOTA pLM (ESM-2) in capturing biomolecular interactions and dynamics and accounting for such underlying biological mechanisms.

• The claim that"our work [...] learns universal AA token level presentation" also bears any evidence and any definition what "universal" means.

We use the term “universal” since (1) this term has been widely-used in protein representation learning literature [Alley et al., Detlefsen et al.], where researchers define pLM is “learning universal, cross-family representations of protein space”; and (2) we pretrain our models on the Universal Protein Reference Clusters (UniRef) dataset which contains universal protein sequence resources. Given such learned high-quality protein embeddings, we can achieve decent performance across a variety of downstream tasks, which demonstrates the universality of such protein representations.

3. Why don’t we use homology-criterions like in ProtMamba?

Thank you for this question. We don’t consider homology-criterion because the context contained in homology representations, such as multiple sequence alignment (MSA) used for pretraining ProtMamba, are already contained in the sequences of individual proteins. After all, MSAs can be considered as performing clustering on a set of protein sequences. As such, we think homology does not bring additional information beyond protein sequences. On the other hand, we consider biological graphs such as PPI graphs, which contain functional interactions between protein sequences derived from additional sources such as annotations and wet lab biological experiments.

4. Train foundation pLMs with other Transformer architectures for comparison.

In our scaling law, length extrapolation, structure prediction, TAPE benchmark, and ProteinGym experiments, we pretrained our own Transformers/ESM-2 on the same dataset with the same number of tokens across different sizes from scratch as baseline comparisons. We clearly indicated where we retrained models / used publicly pretrained models in our manuscripts, e.g. lines 338-341, lines 366-368, lines 433-434, the footnote in Table 2, etc.

We want to note that our main goal is building up a foundation-level long context protein language model (pLM) that can provide universal amino acid level protein representations for extremely long protein sequences; protein complexes, multimers, and heterodimers with encoded protein interaction context information, not debating on selecting/tailoring a specific architecture, just like ESM-2 (there exist many other decent linear RNNs/Transformers architectures in the community but trying out all of them is beyond the scope of our paper). We would like to demonstrate that LC-PLM is a much better foundation pLM than the previous open-sourced SOTA foundation pLMs (i.e. ESM-2 and CARP). We compare our LC-PLM to them in Table 3. We provide more comparisons in the Table below.

Dear Reviewer jUT5,

As the discussion period is coming to an end tomorrow, we kindly ask you to review our response to your comments and let us know if you have any further queries. Alternatively, if you think we addressed your concerns properly and could raise the rating of the paper, we would be extremely grateful. We eagerly anticipate your response and are committed to addressing any remaining concerns before the discussion period concludes.

Best regards,

Authors

The differences to ProtMamba are just in the collection of training data and -- as the authors state -- in the focus on a foundation-level long context protein model. Also ProtMamba can be seen as a foundation level-protein model. The approach is still almost identical to ProtMamba (pre-print available since May). The adaption of a general architecture (such as Transformer, Mamba, etc) to a new application domain would be a relevant application paper, but this step has already been done. The rebuttal has again tenuous claims such as "ProtMamba cannot extrapolate to sequence > 2048", where it is clearly shown to perform well even for context sizes of 2^17. The claimed differences are overall too minor to represent an new ML application.

I appreciate the authors enthusiasm about their method and the effort that went into this rebuttal, but I decide to keep my score.

We thank the reviewer for their valuable feedback on our paper. We performed multiple novel ablations and comparisons that we report in the main comment of the rebuttal. We also address the reviewer’s specific comments and concerns below:

1.Weak motivation in using SSMs: no latency constraints for pLMs.

• We want to emphasize and restate our main motivation for using BiMamba-S as our architectural design choice in our work – our LC-PLM focuses on learning a long-context foundation pLM that can provide universal amino acid level protein representations for extremely long protein sequences; protein complexes, multimers, and heterodimers with encoded protein interaction context information. SSMs grant the model (1) long-context capability, (2) length extrapolation capability, and (3) high-resolution data modeling capability to realize our goal. In contrast, Transformers have to be retrained with adjusted positional encodings to extrapolate over length [Liu et al.] or can not perform well on high-resolution data [Wang et al., Schiff et al.].

• Along with such motivations, we empirically demonstrate that LC-PLM has improved length extrapolation capabilities, favorable scaling laws, and achieved a 7% to 34% improvement on downstream tasks (e.g. protein structure prediction (CASP15-multimers, CASP14, Benchmark2), tasks in TAPE and ProteinGym) compared to ESM-2 [Lin et al.] and CARP [Yang et al.], especially for longer proteins and protein complexes; (2) to encode biological interaction information, we propose a novel second-stage training based on random walks over graphs to extend the long-context capabilities of LC-PLM to leverage the protein interaction graph context; (3) we demonstrate its effectiveness in capturing graph contextual information on remote homology detection (TAPE) in Table 3, protein function prediction (ogbn-proteins) in Table 14 and 15, and PPI link prediction (ogbl-ppa) in Table 16 and 17.

• In addition to our main motivations above, we argue that the inference efficiency is important in computational protein design/drug discovery applications: one usually generate up to $10^6$ protein sequences as candidates [Adolf-Bryfogle et al.] and pLMs fine-tuned for scoring those designed protein sequences can be used for scoring, ranking, and filtering the designed protein sequences. The per-step constant time complexity of SSMs could be an advantage in accelerating this phase.

2.Weak motivation in long-context modeling of proteins.

• We thank the reviewer for this comment. We clarify the biological motivations and needs for long-context modeling of proteins into three perspectives: (1) functional, (2) structural, and (3) evolutionary. (1) many proteins function as part of multi-protein complexes (e.g. transcription factors) and physically or functionally interact with other proteins and molecules. The interaction information is often captured in protein-protein interaction graphs. Knowing the interacting partners of an individual protein is helpful in predicting the protein’s properties. We demonstrate this on protein function prediction tasks using the ogb-proteins graphs, as shown in Tables 14 and 16. With interaction information, the model shows performance gain. (2) Protein structure depends on global fold, which can involve residues and interactions across long distances and across multiple protein sequences. Modeling multi-protein systems captures distant dependencies critical for stability and function. Folding of multi-meric protein complexes relies on models capable of handling long contexts. We demonstrate this benefit in our LMFold experiments in Table 2. Our model outperforms ESM-2 across all folding benchmarks, especially on CASP15-multimers. (3) Proteins in the same pathway or family exhibit co-evolutionary patterns due to functional interdependencies. In fact, multiple sequence alignment (MSA) of homologous protein sequences is a common approach to increase the contexts for studying individual proteins. As other reviewers noted, ProtMamba [Sgarbossa et al.] is inspired by leveraging MSA as an individual protein’s context.

3. Compute-efficient Transformers.

• As we discussed in the above first and second points, our main goal is not to design a compute-efficient sequence modeling architecture. We emphasize other perspectives that SSMs / linear RNNs over Transformers in learning a pLM with better performances. This compute-efficient feature of SSMs is a gifted advantage that we naturally gain from using such a design choice. We also want to note that Flash Attention is a hardware-efficient implementation of Transformers that cannot reduce the training and inference time complexity of Transformers. It only provides faster implementation in terms of wall-clock time. Mamba currently also has decent hardware-efficient implementation [Dao & Gu,Megatron] such as parallelized associative scan [Gu & Dao] to make it more efficient in terms of wall-clock time. Other efficient Transformers that leverage approximate attention mechanisms (e.g. linear attention) are essentially a reformulation of linear RNNs/SSMs [Katharopoulos et al., Yang et al.].

• We also want to note that there exist many decent linear RNNs/Transformers architectures in the community but our work is not debating on selecting/tailoring a specific architecture. Just like how Transformers have been used in numerous impactful works such as GPT and ESM-2, BiMamba-S also serves as a versatile and effective architectural choice for sequence modeling. The reuse of proven architectures like Transformer or BiMamba does not diminish the novelty of a work. Instead, it allows researchers to focus on solving domain-specific challenges (in our case, exploring the capability of building up a new type of foundation pLM based on BiMamba-S). Our work follows this principle, leveraging this architecture to demonstrate its effectiveness in pushing the boundaries of foundation pLM.

• Although compute-efficient architecture is not our focus, we also want to argue that there is indeed practical demand for compute-efficient pLMs as we discussed above. The inference efficiency of pLMs is important in computational protein design/drug discovery applications: one usually generate up to $10^6$ protein sequences as candidates [Adolf-Bryfogle et al., 2018] and pLMs fine-tuned for scoring those designed protein sequences can be used for scoring, ranking, and filtering the designed protein sequences. The per-step constant time complexity of SSMs could be an advantage in accelerating this phase.

4. Why not other SOTAs?

• We want to highlight that our goal is to build up a foundation pLM that can learn meaningful universal amino acid level representations for pure protein sequences and generalize across various downstream tasks. As you pointed out, our well-trained pLM can be used as a drop-in replacement for ESM-2 to accommodate any of these existing methods on their tasks. In fact, we demonstrated this in practice with the LMFold experiments (a generalization of ESMFold), where a simple folding trunk is trained to predict protein structures based on pLM embeddings. We only show the comparison against existing SOTA open-sourced pLMs trained on protein sequences alone to avoid confounding gains resulting from ad-hoc methods built off pre-trained pLMs. For example, on the ProteinGym leaderboard, SOTA methods like PoET [Truong Jr et al.], TranceptEVE [Notin et al.] rely on combining family-specific models or alignment-based methods with a foundation pLM; SaProt [Su et al.] is a pretrained pLM with massive protein structure data. We added these SOTA methods to our table but we want to emphasize again that our work is to develop a foundation pLM (i.e. LC-PLM and LC-PLM-G) that can be utilized as a better pLM backbone in all these works. And we believe these methods built off our LC-PLM have a high probability of outperforming ESM-2 given all the evidence we provided in our work.

5. How did we hold out the test 250K? Do we retrain ESM-2 for Figure 4? What is “evaluation loss”?

• Similar to ESMFold [Lin et al.], the hold-out 250K sequences are randomly sampled from UniRef90. For the training set, we used mmseqs to filter out sequences in the UniRef50 data with >90% sequence identity, such that the remaining sequences in UniRef50 are no more than 90% similar to any of the sequences in the hold-out test set. We provide the details in Appendix B.
• Yes, we retrained ESM-2 following the official recipe for different sizes using the same train set and test set as we used for LC-PLM.
• Evaluation loss is the average cross-entropy across all tokens (equivalent to Perplexity (PPL)), which is the standard way that people use to evaluate language models . We added this description to our manuscript.

6. Align folding trunk in structure prediction.

In the folding evaluation presented in RQ4, Table2, we are essentially evaluating how well the residue-level embeddings from different pLMs can predict the 3D structures. We perform this evaluation by learning a folding trunk with its architecture adopted from ESMFold for each pLM on paired protein embeddings and ground truth structures, then validate the performance on hold out structure datasets (CASP14, CASP15-multimer, Benchmark2). It is worth noting that the folding chunks for each pLM are trained from scratch rather than transfer learning. This evaluation setting is similar to linear probing, where the parameters in the LMs are frozen and the probing head; in this case, a single folding trunk, is learned by stochastic gradient descent.

7. Robustness of downsampling.

Thanks for this great feedback! We add new results to assess the robustness of downsampling. We perform three 1.5% downsampling using three random seeds and retrain both our LC-PLM and ESM-2. The results are in the table below. We find that the downsampling is very robust to the performance with a small standard deviation that is close to (even smaller than) the standard deviation of using the same train set but training with different random seeds as we reported in our original table.

8. Usefulness of graph context.

Thanks for bringing up this concern! We want to refer to Tables 14, 15, 16, 17 to show that the encoded graph context information can help with protein function prediction and protein interaction prediction. But to make this claim much clearer, we also add more experiments (as shown in the table below) on downstream tasks to verify the effectiveness of the graph contextual training. LC-PLM-G outperforms its vanilla variant on 3/4 TAPE tasks, as shown in the table below. We also want to note that by comparing two LC-PLM-G variants trained on different PPI graphs, the performance also varied a lot, which indicates that the data quality of the PPI graph is also important for the performance boost. We think building up a high-quality PPI graph can be meaningful future work that makes the pretrained pLM better. Regarding the fact that incorporating PPI graphs drastically hurts the performance of ESM-2, this is potentially due to the poor length extrapolation capability (as shown in Figure 6) and the catastrophic forgetting issue [Kenneweg et al., Luo et al.] of Transformers. In Figure 6, we show that if we train ESM-2 on longer sequences, the model will fail to extrapolate on both shorter and longer sequences. Thus, after the second-phase graph context training, ESM-2 forgets the high-quality representations for regular-length protein sequences (shorter than graph-contextualized sequences) learned in the first-phase pretraining and fails to extrapolate on these shorter sequences. It can only provide degenerated representations of them.

9. Concrete modeling examples with biological significance that we have to turn to Mamba because the attention/Transformer sucks or even fails to handle the problem at all? Can the authors provide further evidence that the extended context window length improves the “related” downstream protein tasks beyond the results presented?

• We appreciate this comment and address reviewer’s concern in twofold:

--- (1) LC-PLM has favorable adaptability to long-context tuning: We performed three more downstream protein tasks to evaluate ESM2 and LC-PLM models. In our existing and new experiments (compiled in Table A), after performing the 2nd stage graph training, the performance of ESM2 (Transformers) degrades on many downstream tasks including protein fitness prediction in ProteinGym (Table 4, Spearman’s rho drop from 0.295 to 0.109), Contact map prediction (New table below, precision drops from 44.1 to 26.7), and TAPE stability prediction (New table below, Spearman’s rho drop from 0.763 to 0.750). On the other hand, our LC-PLM models maintain or slightly improve their performances on these tasks after the 2nd stage graph training. These results suggest it is difficult to tune Transformer models to adapt to longer contexts.

--- (2) Mamba-based models enjoy favorable inference efficiency in terms of both time and space complexity. This will satisfy the practical demands for in-silico protein design, where one needs to screen $10^6$ sequences using pLM based methods. The constant time complexity of Mamba/SSMs could be an advantage in accelerating this phase. There is also a GPU memory constraint in performing inference with the Transformer/ESM2 model on long protein sequences users of the ESM2 model have been facing [issue1, issue2].

Table A: Evaluation of pLMs before and after 2nd stage graph context training on downstream tasks. For the Jacobian contact map prediction task, we adopted the methods from [Zhang et al.] to use categorical Jacobian matrices computed from protein language models as the zero-shot prediction for protein contact maps and report the precision@2/L (L is the length of a protein sequence) on the validation set of ProteinNet dataset [AlQuraish]. We report the Spearman’s correlation coefficients for the test sets for the TAPE Stability prediction tasks and ProteinGym DMS substitutions benchmarks.

10. Weird y-ticks.

Thanks for this comment. The $2^1$ is due to the log scale on y-ticks. We fixed it in our manuscript to the normal scale.

Dear Reviewer 8QMc,

As the discussion period is coming to an end tomorrow, we kindly ask you to review our response to your comments and let us know if you have any further queries. Alternatively, if you think we addressed your concerns properly and could raise the rating of the paper, we would be extremely grateful. We eagerly anticipate your response and are committed to addressing any remaining concerns before the discussion period concludes.

Best regards,

Authors

We acknowledge and commend the reviewer's well-structured feedbacks for further improving our manuscript. We commit to incorporate those reasonable and constructive comments in the next version of our manuscript after the review period. While some of the comments are objective and reasonable, we do want to respond to a few points that are unfair and self-contradictory:

1. Problem Tackled
   However, across the tasks and results presented, LC-PLM does not demonstrate parity with or superiority over ESM-2.

Throughout the experiments presented in our manuscript, we benchmark LC-PLM against ESM-2 we pretrained with the exact same amount of tokens to demonstrate clear performance advantage of LC-PLM over ESM-2. This comparison is intended to rigorous control for the variabilities in the quality and quantity of pretraining dataset, such that the observed performance differences can be unbiasedly attributed to architectural advantages. On the other hand, it would be unfair to compare models pretrained with different datasets to draw conclusion about the architectural superiority.

2. Motivation
   PPI graph in my opinion is not a persuasive case.

It would be extremely helpful for the reviewer to articulate the reason why PPI graph is not a persuasive case for long-context capability, or perhaps suggest a more persuasive case for proteins' long-context use case/applications. From a systems biology perspective, PPI graphs capture functional contexts for individual proteins to help us understand how proteins function in cells/biological systems.

3. Support for Claims
   Retraining with limited training tokens and applying graph augmentation, which appear to disproportionately affect ESM performance, making the comparison feel unfair.

As mentioned in the response to the reviewer's point 1, the purpose of comparing LC-PLM and ESM2 with the same amount of training tokens is to control for the variabilities in pretraining data. The graph augmentation training for both LC-PLM and ESM2 were also performed under the exact same training data and procedure. The observed advantage of LC-PLM over ESM2 in controlled setting helps us tease out the confounding effects from differences in pretraining data. It would be helpful for the reviewer to clarify why such comparison "feel" unfair. We also respectively suggest the reviewer to objectively assess scientific works rather than using subjective views.

4. Significance.
   However, to establish significance, the authors could either propose new learning algorithms to address foundational protein modeling challenges or demonstrate LC-PLM’s advantages in specific yet biologically meaningful scenarios where ESM-2 underperforms.

In our manuscript, we demonstrate LC-PLM's advantage over ESM-2 on graph augmentation scenario, which is biological meaningful as PPI graphs, which can be generalized to biological knowledge graphs, are valuable tools to study how proteins function in actual biological systems rather than in isolation.

3. Criterions for Table 1 should be better specified (e.g., when is a method considered to be universal?)

We define universality as learning representations are learned to reflect general properties of all proteins, rather than a specific protein family or properties such as post-translational modifications. pLMs satisfying universality can be used as a foundation model to develop specialized methods such as PTM prediction and structure prediction. We added these clarifications to our manuscript.

We use the term “universal” since (1) this term has been widely-used in protein representation learning literature [Alley et al., Detlefsen et al.], where researchers define pLM is “learning universal, cross-family representations of protein space”; and (2) we pretrain our models on the Universal Protein Reference Clusters (UniRef) dataset which contains universal protein sequence resources. Given such learned high-quality protein embeddings, we can achieve decent performance across a variety of downstream tasks, which demonstrates the universality of such protein representations.

4. line 101: here also "Hochreiter, S. (1991). Untersuchungen zu dynamischen neuronalen Netzen" should be cited.

Thanks for pointing this out, we added this citation to this paper for their contribution to RNNs.

5. Which context size do you use for ESM-2? Could it have been increased for the experiments you carried out (i.e., did ESM2 and LC-PLM use approximately the same memory?)?

We use 1024 as the context size for the ESM-2 models we trained, which is the same context size used in the public ESM2 models. The sizes remain constant as the sizes of the positional embedding matrix can’t be changed.

6. What hyperparameters did you use for ESM-2? How was hyperparameter selection done for ESM-2 and LC-PLM(-G)?

We adopt the hyperparameter from the official ESM-2 paper [Lin et al. 2023] with some modifications: global batch size=0.5M tokens; peak learning rate = 2e-4, 2000 steps learning rate warm-up, followed by cosine decay schedule; weight decay=0.01; Adam beta1=0.9, beta2=0.98, eps=1e-8; We added those details in our Appendix E.

7. RQ3: In contrast to some other experiments UniRef50 is used for training. Evaluation is at UniRef90. Why is this the case?

Thanks for this question! We and others [Rives et al. 2019] found pretraining with UniRef50 leads to lower perplexity with the same training budget due to the lower-level of similar sequences in UniRef50. However, UniRef90 datasets have sizable proteins across all length bins, which is suitable for the length extrapolation evaluation experiments.

8. line 462-464: "As shown in Figure 7, the embeddings from LC-PLM-G captures the graph topology much better than LC-PLM, which aligns with the community detection results. What is the criterion to see that LC-PLM-G captures the graph topology much better?

We elaborate more on the criterion here. We first run community detection on the graph and label the nodes using its corresponding community labels. Since community detection helps us find local clusters in the graph, this labeling method provides enough information of how the topological structure presented in the graph. Thus, we do t-SNE dimensionality reduction on the learned protein embeddings and visualize the 2-D embeddings in the plot. If closer 2-D embeddings reveal the similar community labels, then our proposed graph context training can be demonstrated as capturing graph relational information well.

9. line 439-440: "This also suggests that, even for average-length protein sequences, long-range dependencies would be useful information and an important feature for protein structure prediction. What is the exact indication to come to this conclusion?"

This conclusion was made based on the performance improvement of LC-PLM on the CASP14 dataset (Table 2). CASP14 contains 37 mostly average-length protein sequences (mean=318.6; median=197).

10. line 533: "LC-PLM achieved 7% to 34% better performance on various downstream tasks." Which task do you exactly mean with 7% and which with 34%?

We found LC-PLM achieved 7% and 34% improvement over ESM-2 on TAPE Secondary Structure (85.07% vs 79.85% in accuracy), and Remote Homology (35.14% vs 26.57% in top-1 accuracy), respectively.

Lastly, we would like to sincerely ask you whether you are willing to increase your confidence score if you think we addressed your comments properly. Thank you so much for your support again!

2. Evaluation of pLMs on the other 3 TAPE tasks: Contact Prediction, Fluorescence and Stability.

We appreciate this suggestion and performed evaluations for all the models considered in this manuscript (LC-PLM, LC-PLM-G, ESM-2, ESM-2-G, CARP, ProtMamba). The results are shown in Table A and added to the revised version of our paper. It is worth noting that we adopt a different setting for the Contact Map prediction task for a fair comparison of attention-free models including CARP, LC-PLM, and ProtMamba. Instead of using the attention maps from Transformer-based pLMs to predict the contact maps [Rao et al. 2020], we followed [Zhang et al. (2024)] to use categorical Jacobian matrices computed from protein LMs as the zero-shot prediction for protein contact maps and report the precision@2/L (L is the length of a protein sequence) on the validation set of ProteinNet dataset [AlQuraishi 2019]. In Contact map and Stability prediction tasks, we noted a similar trend where LC-PLM outperforms ESM-2 with the same number of pretraining tokens, and that 2nd stage graph pretraining slight improve LC-PLM while hurting the performance of ESM-2, highlighting the difficulty of adopting ESM-2 to handle longer contexts. For Fluorescence tasks, we noted all models saturated at Spearman correlation coefficient around 0.69. This observation was also made by others [Rao et al. 2019; McDermott et al. 2023; Schmirler et al. 2024].

We sincerely thank the reviewer for the favorable rating and constructive suggestions! To address the weakness and questions:

1. Performance comparison with ProtMamba.

We evaluated the performance of ProtMamba model against ours across seven tasks to find ProtMamba significantly underperforms our models in all 5 tasks we evaluated (structure prediction, Table 2; Contact map, Remote Homology, Secondary Structure, Stability, Table 3;). Our results demonstrate ProtMamba is not good at producing good representations for protein sequences. Its training regime favors protein generation and contextual fitness prediction, which may degrade the representation quality of individual protein sequences in the meantime.

Table A. Protein structure prediction with LMFold. Structure prediction performance (TM score) are reported on different hold out datasets.

Table B. Evaluation on TAPE tasks in supervised fine-tuning setting. We report the top-1 accuracy for the Remote Homology fold-level test set; accuracy for the 3-class secondary structure prediction on the CB513 test set; Spearman’s correlation coefficients for the test sets for the Stability and Fluorescence prediction tasks. For the Jacobian contact map prediction task, we adopted the methods from [Zhang et al.] to use categorical Jacobian matrices computed from protein language models as the zero-shot prediction for protein contact maps and report the precision@2/L (L is the length of a protein sequence) on the validation set of ProteinNet dataset [AlQuraishi].

Dear Reviewer Z2Tu,

As the discussion period is coming to an end tomorrow, we kindly ask you to review our response to your comments and let us know if you have any further queries. Alternatively, if you think we addressed your concerns properly and could raise the rating of the paper, we would be extremely grateful. We eagerly anticipate your response and are committed to addressing any remaining concerns before the discussion period concludes.

Best regards,

Authors

We sincerely appreciate your advocacy on our behalf. Your support has been invaluable and deeply meaningful to us.

In our revised manuscript, we have incorporated the points we committed to addressing, making updates in both the main text and the appendix. Additionally, we have highlighted the new additions and adjustments, including the definitions and rationales of specific criterions used in Table 1 (see Appendix A from here given currently there is a technical issue on OpenReview for updating the PDF file). Please let us know if there is anything that you feel requires further clarification or adjustment.

Once again, we thank you for your insightful feedback. We kindly ask if you could consider increasing your confidence if you feel we have adequately addressed and clarified the concerns.

We thank the reviewer for their valuable feedback on our paper. We performed multiple novel ablations and comparisons that we report in the main comment of the rebuttal. We also address the reviewer’s specific comments and concerns below:

1. Little novelty and lack of comparison.

Thank you for this feedback. While we appreciate the recognition of our work’s relation to previous studies, we would like to elaborate on what we wanted to show from experimental comparisons. We also would like to clarify several key distinctions:

• We have discussed the key differences between our work and ProtMamba [Sgarbossa et al.] in lines 152-153 and 158-161 and Table 1. To summarize, ProtMamba trained on concatenated homologous protein sequences with autoregressive causal language modeling and an infilling objective focusing on protein sequence generation. However, our LC-PLM focuses on learning a foundation-level long context protein language model (pLM) that can provide universal amino acid level protein representations for extremely long protein sequences; protein complexes, multimers, and heterodimers with encoded protein interaction context information. We also note that LC-PLM is a much better foundation pLM than the previous open-sourced SOTA foundation pLMs (i.e. ESM-2 [Lin et al.] and CARP [Yang et al.]). We additionally indicate that, unlike ProtMamba, which uses homologous sequences as context, the proposed LC-PLM-G showed a novel way to encode protein interaction graph information that contextualizes functionally related proteins rather than semantically similar ones (i.e. proteins with similar sequences) .

• We thank the reviewer for suggesting a direct comparison with ProtMamba. We want to highlight again that our work aims to build a better protein foundation model instead of a task-specific protein model. We provide new experimental results comparing LC-PLM to ProtMamba. In these new experiments, we evaluate ProtMamba on the downstream tasks we used in our paper. The performance of ProtMamba is much lower than LC-PLM (even CARP and ESM-2 as shown in general response) across all tasks. This suggests that ProtMamba pretrained with concatenated homologous protein sequences potentially leads to degraded representations of individual proteins. After all, ProtMamba is trained to use homologous sequences as context for protein generation tasks (e.g., infilling and fitness prediction), rather than producing useful representations for proteins. Also it is worth noting that ProtMamba cannot extrapolate to sequence > 2048 since they use positional encodings with fixed length in training. We summarized the results in the tables below. We added these results to our manuscript. Also as a side note, ProtMamba (https://openreview.net/forum?id=BMfHO2lXGe) is a concurrent submission to ICLR 2025 so other submissions can be excused for not comparing against them according to ICLR’s policy (https://iclr.cc/Conferences/2025/FAQ).

Table A. Protein structure prediction with LMFold. Structure prediction performance (TM score) are reported on different hold out datasets.

Table B. Evaluation on TAPE tasks in supervised fine-tuning setting. We report the top-1 accuracy for the Remote Homology fold-level test set; accuracy for the 3-class secondary structure prediction on the CB513 test set; Spearman’s correlation coefficients for the test sets for the Stability and Fluorescence prediction tasks. For the Jacobian contact map prediction task, we adopted the methods from [Zhang et al.] to use categorical Jacobian matrices computed from protein language models as the zero-shot prediction for protein contact maps and report the precision@2/L (L is the length of a protein sequence) on the validation set of ProteinNet dataset [AlQuraishi].

• We agreed that Cadeucus [Schiff et al.] used a similar architectural design choice and we did cite them and other papers using BiMamba in lines 219-221. However, we want to note that we used BiMamba-S since it is the most proper way to realize bi-directionality in Mamba. BiMamba-S can be formulated in a more theoretical way as structured SSMs with quasi-separable matrices [Hwang et al.]. We did not apply quasi-separable mixers in our model since Mamba currently has a much better software-hardware interface environment and distributed training support in practical implementations that help us train a large foundation-level model feasibly and efficiently. We added this discussion in our paper as well.

• We also want to note that, just like Transformer, which has been used in numerous impactful works such as GPT and ESM-2, BiMamba-S also serves as a versatile and effective architectural choice for long-context sequence modeling. The reuse of proven architectures like Transformer or BiMamba does not diminish the novelty of a work. Instead, it allows researchers to focus on solving domain-specific challenges (in our case, exploring the capability of building up a new type of foundation pLM based on BiMamba-S). Our work follows this principle, leveraging this architecture to demonstrate its effectiveness in pushing the boundaries of foundation pLM.

• Besides the architectural choice based on Mamba, we have many other important contributions to this field as summarized in the Introduction: (1) method to encode biological interaction information into pLM, we propose a novel second-stage training based on random walks over graphs to extend the long-context capabilities of LC-PLM to leverage the PPI graph context; (2) we demonstrate that LC-PLM has improved length extrapolation capabilities, favorable scaling laws, and achieved a 7% to 34% improvement on downstream tasks (e.g. protein structure prediction (CASP15-multimers, CASP14, Benchmark2), tasks in TAPE and ProteinGym) compared to ESM-2, CARP, especially for longer proteins and protein complexes; (3) we demonstrate its effectiveness in capturing graph contextual information on remote homology detection (TAPE) in Table 3, protein function prediction (ogbn-proteins) in Table 14 and 16, and PPI link prediction (ogbl-ppa) in Table 15 and 17.

• We want to highlight that our goal is to build up a foundation pLM that can learn meaningful universal amino acid level representations for pure protein sequences and generalize across various downstream tasks. That said, for ProteinGym and etc., it is not fair to compare our method to the other ad-hoc methods built off pre-trained pLMs. For example, on the ProteinGym leaderboard, SOTA methods like PoET [Truong Jr et al.], TranceptEVE [Notin et al.] rely on combining family-specific models or alignment-based methods with a foundation pLM; SaProt [Su et al.] is a pretrained pLM with massive protein structure data. We added these SOTA methods to our table but we want to emphasize again that our work is to develop a foundation pLM (i.e. LC-PLM and LC-PLM-G) that can be utilized as a better pLM backbone in all these works. Such ad-hoc design built off the base pLM to improve on some specific task is beyond the scope of our paper. Our original table only compares foundation pLMs trained on protein sequences alone.

2.Why don't we simply train a GNN?

• We want to note that, by training with graph context, we are not aiming to only perform graph tasks (e.g. node prediction, link prediction, etc.), but also want to show that a good representation with encoded protein interaction context information helps many other downstream tasks (shown in Tables 3, 4, and 12) like remote homology prediction. We also trained GNNs on other graph-related tasks like ogbn-proteins and ogbl-ppa and conducted ablation studies with and without our learned protein embeddings, as shown in Tables 16 and 17. We show that with better protein representations and improved graph context-aware representations, the performance can be further boosted.

3.Wrong direction in the figure.

• Thanks for the feedback and pointing this out. We have fixed it in our paper.

Dear Reviewer sY9s,

As the discussion period is coming to an end tomorrow, we kindly ask you to review our response to your comments and let us know if you have any further queries. Alternatively, if you think we addressed your concerns properly and could raise the rating of the paper, we would be extremely grateful. We eagerly anticipate your response and are committed to addressing any remaining concerns before the discussion period concludes.

Best regards,

Authors

Dear Authors,

Thank you for your detailed answers. In the following you will find some follow-up thoughts:

Novelty: To improve clarity regarding the fact that BiMamba-S is not a novel architectural component but was introduced previously and applied in this work to another domain, you might consider moving the BiMamba-S section to the Preliminaries section.

Goal of Learning a Long-Context Foundation pLM: Thank you for clarifying the focus of your work. With this context in mind, I believe the manuscript requires improvements in clarity, motivation, and focus:

a) Clarity and Motivation: You should define what properties should a protein foundation model have and how could you test them. I guess, showing downstream task performances to demonstrate representation learning capabilities makes sense, just the motivation is missing.

b) Universal Representation Claim: The manuscript states that LC-pLM “learns universal AA token-level protein representations” (line 161). However, this claim is vague, and the term “universal” is not well-defined.

c) Long-Context Motivation: I agree with reviewer 8QMc’s initial comment that the necessity for long-context capabilities isn’t well motivated in the manuscript, despite the fact that Figure 6 is relevant and interesting, and implicitly shows the foundation models’ need to be able to adapt to larger sequences.

d) Multi-Modal Foundation Models: The authors state in one of the review answers that they want to focus on pure sequence foundation models. Will trained models be relevant though? Wouldn’t people just use multi-model models like ESM-3? I guess, one cannot expect the proposed method matches ESM-3 performance. However, it would be important to know how significant differences are.

I appreciate the performance gains demonstrated by LC-pLM in the experiments. However, my concerns regarding architectural novelty remain. If the primary focus is to develop “protein foundation models,” the manuscript would benefit from significant restructuring and rewriting to reflect that focus more clearly.

For these reasons, I will stick to my initial rating.

Thank you for your feedback. Here we provide further discussion.

1. No novelty of BiMamba-S; Moving it to Prelim.

We respectfully disagreed.

(1) Note that, as we mentioned in our rebuttal, Cadeucus has similar architectural design choice that is not exactly same. Here are several key differences:

• "shared layers" is not qual to "tied weights". (a) "Shared layers" are more efficient since they reuse the same layer object, whereas weight tying tracks separate layers with identical weights, introducing minor overhead. (b) "Shared layers" simplify optimization as they operate on a single computational graph, while weight tying requires additional constraints to enforce equality during training.
• We have another design choice "untied input/output embedding layers" and we carefully study it. This can help alleviate collapsed embedding space and improve the uniformity of learned embeddings, which make the model more expressive and avoid anisotropic issue.
• We perform the reverse operation after a normalization layer and we provide a residual connection at the end of the block to help gradient flow.

(2) Also note that Caduceus is not the first one proposing to use bidirectional Mamba (BiMamba). There has been a lot of works built off BiMamba, just as we discussed in our manuscript -- "... time-series forecasting (Liang et al., 2024), audio representation learning (Erol et al., 2024), visual representation learning (Zhu et al., 2024), DNA modeling (Schiff et al., 2024), and graph learning (Behrouz & Hashemi, 2024)." These works came out on various dates but all discussed their BiMamba architecture designs in their model/methodology. We think it's worth to also discuss our own architectural designs in the method part to indicate the similar insights and key differences in a totally different field.

2. Clarity and motivation.

Note that we're not proposing a protein foundation model; instead, we're introducing another foundation-level protein language model (pLM). We strongly recommend the reviewer to read ESM-2 (Lin et al.) to get the motivation of having a foundation-level pLM if they're not familiar with this field and think the motivation is missing.

Also, the community hasn't had a clear definition of "foundation model" and there has been a long debate on this. Thus, it's also not easy for us to define its necessary properties and claim we achieved "foundation" in our work. Given these, we choose to avoid this term but focus on demonstrating we have a better pLM compared to existing works. As we introduced in the manuscript, we think we achieved this by the following evidence:

• We demonstrate that LC-PLM has improved length extrapolation capabilities, favorable scaling laws, and achieved a 7% to 34% improvement on downstream tasks (e.g. protein structure prediction (CASP15-multimers, CASP14, Benchmark2), tasks in TAPE and ProteinGym) compared to ESM-2, especially for longer proteins and protein complexes.
• To encode biological interaction information, we propose a novel second-stage training based on random walks over graphs to extend the long-context capabilities of LC-PLM to leverage the PPI graph context. We demonstrate its effectiveness in capturing graph-contextual information on remote homology detection (TAPE), protein function prediction (ogbn-proteins), and PPI link prediction (ogbl-ppa).

3. "Universal Representation" is not well-defined.

We already addressed this concern in our reply to Review jUT5. For your convenience, we restate here again:

We use the term “universal” since (1) this term has been widely-used in protein representation learning literature [Alley et al., Detlefsen et al.], where researchers define pLM is “learning universal, cross-family representations of protein space”; and (2) we pretrain our models on the Universal Protein Reference Clusters (UniRef) dataset which contains universal protein sequence resources. Given such learned high-quality protein embeddings, we can achieve decent performance across a variety of downstream tasks, which demonstrates the universality of such protein representations.

5. Why don't people just use multi-modal models like ESM-3? Why do we still work on pure pLM?

We elaborate the reasons as follow:

• We kindly note that ESM-3 is also a pLM but with training data tokenized from multiple sources including protein sequences, structures, and functions. So if we can build up a better pLM, we can also make a better ("multi-modal") version of LC-PLM in the future with the same training data, which can outperform ESM-3.
• Also, note that ESM-3 still suffers the same intrinsic issue that it has weak long-context capability of modeling protein complexes, heterodimers, interaction contexts, and etc. Therefore, if our proposed LC-PLM can alleviate this weakness, we will have a better model to adapt to these mentioned scenarios/tasks.
• Developing a better base model is orthogonal to training a multi-modal model. For example, the researchers are still putting efforts on developing new language models that have other features like highly compute-efficient, capture long-context, etc. Why don't they just use GPT-4o and stop researching on new LMs? Our work follows the same motivation and insights.

Lastly, we kindly ask the reviewer to carefully read the related literature in this domain and look into our response and manuscript before making arguments. Thank you.

3. Distinctions with between LC-PLM and ProtMamba

There are two key distinctions between our method and ProtMamba:

• Definition of long contexts for proteins: ProtMamba exclusively use evolutionary contexts of individual proteins in the forms of flattened MSAs, whereas LC-PLM-G primarily use functional and structural contexts of proteins stored in PPI graphs. In fact, LC-PLM-G’s graph contextual method is more generalizable: one can feed the evolutionary contexts in the form of sequence graphs, an alternative representation of MSA [Benedict et al. 2014], to train or inference with our LC-PLM-G model. We added these discussions to our manuscript and consider these new directions as future works.
• Foundation model vs specialized model: our goal is to build a foundation pLM that can learn meaningful representations for protein sequences and generalize across various downstream tasks, including prediction of protein structures, functions, fitness, and interactions. On the other hand, ProtMamba is trained to use homologous sequences as context for protein generation tasks (e.g. infilling and fitness prediction), rather than producing useful representations for protein sequences.

4. Benefit of SSM/Mamba-based architecture

• SSM/Mamba-based models can adapt to long-context training better than Transformer-based architectures

We demonstrate our LC-PLM has favorable adaptability to long-context tuning: We performed three more downstream protein tasks to evaluate ESM2 and LC-PLM models. In our existing and new experiments (compiled in Table B), after performing the 2nd stage graph training, the performance of ESM2 (Transformers) degrades on many downstream tasks including protein fitness prediction in ProteinGym, Contact map prediction, and TAPE stability prediction. On the other hand, our LC-PLM models maintain or slightly improve their performances on these tasks after the 2nd stage graph training. These results suggest it is difficult to tune Transformer models to adapt to longer contexts.

Table B: Evaluation of pLMs before and after 2nd stage graph context training on downstream tasks.

• Practical applications that would benefit from efficient pLMs

Mamba-based models enjoy favorable inference efficiency in terms of both time and space complexity [Gu & Dao, 2024]. The inference efficiency is important in computational protein design/drug discovery applications: one usually generate up to 10^6 protein sequences as candidates [Adolf-Bryfogle et al., 2018] and pLMs fine-tuned for scoring those designed protein sequences can be used for scoring, ranking, and filtering the designed protein sequences. The per-step constant time complexity of SSMs could be an advantage in accelerating this phase.