Learning the Language of Protein Structure

We thank the reviewer for acknowledging interest in our work and are grateful he appreciates the narrative and experimental part of our work.

We address the reviewers' remarks point by point and hope to be able to discuss.

Tokenization for protein structures is not entirely novel; works like FoldSeek and several others [1-10] have experimented with tokenized representations. Although the paper claims novelty in its discrete, sequence-based approach, it could better establish its unique contributions by clearly differentiating its tokenization strategy from these previous methods. To strengthen the paper’s contribution, the authors should clarify and empirically demonstrate the advantages of their approach over existing methods.

Indeed, structure tokenization is not entirely novel and some of the methods reported in your references [1-10] such as Foldseek are already discussed in our work and we explain why our approach is complementary to [1], as they only treat protein sub-structures.

Moreover, while we concur with the reviewer that we missed [2] which indeed performs a tokenization of structures, we highlight that all subsequent works [4-10] were not pre-printed at the time of writing of the manuscript.

Nonetheless, we agree with the reviewer that providing the reader with a fair overview of all methods is essential, in that sense we updated the manuscript to include all references to our related works on protein structure tokenization.

The technical novelty is quite limited. Almost all the components of this approach are built on existing techniques, like MPNN, FSQ quantization, AlphaFold's structure module, and FAPE loss.

Our primary objective in this study is to propose an efficient structure representation and decoding, to do so we indeed leverage existing techniques, combining them in a novel way that addresses the specific challenges in protein structure modeling.

By carefully integrating these components, we were able to propose a model that achieves significant advances in tokenizing, reconstructing, and generating protein structures with high fidelity and efficiency.

Our carefully designed experiments show the relevance of our work. Therefore, we firmly believe that the use of efficient, principled and well tested approaches is in that sens valuable to the community.

I appreciate your response, but I must admit I’m not fully satisfied with it. To clarify, I’m not asking you to compare these approaches in your experiment section; my request specifically pertains to the related work section. It’s important that you highlight the unique advantages of your method in relation to prior work.

Regarding your statement that “all subsequent works [4-10] were not pre-printed at the time of writing of the manuscript,” I believe referencing the time of submission would be more appropriate and precise than referencing the time of writing.

To be clear, I’m fine with the manuscript being accepted or rejected, but my main concern is that the authors clearly and transparently present their unique contributions. It’s important to acknowledge that your work is not, in fact, the first in this area. The published FoldSeek and the unpublished FoldToken both precede your work. If they are there earlier than yours, do you have any special contributions in this field? Ultimately, the completeness and depth of a contribution may hold more value than its chronological timing.

I believe that explicitly situating your contribution in this context while highlighting its distinctiveness would strengthen the manuscript significantly.

We first would like to thank the reviewer for its thorough and insightful review. We address its remarks point by point and hope to engage in a fruitful discussion.

On the related work

Thank you for the constructive feedback on the related work. We agree that situating our work within the latest developments is essential for reader clarity, and we appreciate the reviewer’s acknowledgment of our contribution. We also address this remark in our comment for all reviewers. In response, we expand our related work section to include ProstT5 as an example of using FoldSeek local structure tokens to incorporate structural information in language modeling. We also clarify that SaProt is already discussed in our current submission. The rapid pace of progress in this field makes it challenging to propose a comprehensive overview. At the time of writing, newer models such as FoldTokens2, ProTokens, and ESM-3 had only recently emerged and are still unreviewed preprints. Nevertheless, we commit to incorporating these recent works to provide readers with a fuller context, while underscoring that our contributions remain significant and complementary in light of these advances.

Relatedly, the authors use a random split for training and evaluating the quantization and reconstruction. This means that there are likely test structures that are very similar to training structures. The authors should evaluate how well the reconstruction generalizes for structures that are more/less similar to those seen in training.

We concur with the reviewer on its view. To address the reviewers’ concern, we retrain all our algorithms on novel splits leaving out some full clusters for validation and test. That way we would better understand the generation capacity of the proposed model. The detailed results are available in the response to all reviewers.

Note that, to further assess the capacity of our model, we also now include reconstruction results on the CASP15 dataset. This dataset, recognized by the community as including complex proteins, includes structures that are not included in the dataset used for training.

The detailed results are now available in the main paper and in the answer to all reviewers. Put simply, the results obtained on the CASP15 dataset are inline with the results obtained on the test. For instance we report for all training with a downsampling factor below two a median RMSD below two Angstroms.

At a minimum, we need to see distributions of reconstruction quality. It would be even nicer to have an analysis of what kinds of structures are reconstructed well and poorly.

We concur with the reviewer. We update the manuscript with a distribution of the errors on both CASP15 and the newly built dataset split and discuss them both in the main body of the paper as well as in the appendices.

While it's true that backbone diffusion models generally list the proportion designable, then the proportion of the designable backbones that are novel, then the proportion that are diverse, this isn't what we really care about for unconditional generation. What we actually want is the number of designable backbones that are distinct from each other and from natural backbones, in other words, the number (or proportion) of total backbones generated that are designable and novel and diverse. Given that the models are trained on the PDB, it would be better to do the novelty assessment against the PDB than against CATH.

Indeed, the proportion of both novel and designable is indeed a crucial metric when addressing de-novo generation, we want to highlight that the purpose of our generative experiments is to showcase that our learned tokenizer provides an informative description of proteins structures, not to build a new state-of-the-art generative model for protein structures. Our goal is to show that a generative model learned only using simple pretraining methods conducted on our learned representation is indeed already competitive with existing benchmarks. While we concur that the novelty would be better estimated on the full PDB dataset rather than only on the CATH, this raises a significant computational challenge running a one on one comparison with such a large database.

The authors have done a great job addressing most of my concerns and strengthening the paper. The paper is now properly situated in the existing literature, the methods are clearer, they've backed up the claim that local attention helps, and I'm much more confident that their tokenization method is robust to backbones that aren't in the training set. The improved reconstruction performance also strengthens the submission.

Some points I'd like to address from the discussion:

Another metric for reconstruction quality especially relevant for protein design is whether the reconstructed backbones are predicted to be realized by the original sequence. This could be evaluated, for example, by obtaining the perplexity of the original sequence when designing the reconstructed backbone using ProteinMPNN.

To be clear, I wasn't talking about doing this to evaluate the backbones generated from the autoregressive model. I agree that that would be redundant with the existing self-consistency metrics. I meant to do this as an additional metric for measuring the quality of reconstructions from the tokenizer.

Indeed, the proportion of both novel and designable is indeed a crucial metric when addressing de-novo generation, we want to highlight that the purpose of our generative experiments is to showcase that our learned tokenizer provides an informative description of proteins structures, not to build a new state-of-the-art generative model for protein structures. Our goal is to show that a generative model learned only using simple pretraining methods conducted on our learned representation is indeed already competitive with existing benchmarks.

I agree that the point of the paper is not to show SOTA generative performance and that the existing performance is strong. I'm still curious though, about the proportion of generations that are novel and designable. I think this should be pretty easy to compute given the underlying data used to compute the current metrics. Selfishly, I'd love for the authors to include this in the paper in order to move the field towards reporting this.

While we concur that the novelty would be better estimated on the full PDB dataset rather than only on the CATH, this raises a significant computational challenge running a one on one comparison with such a large database.

This is fine as long as they run generations from every model against the same database.

I'm also still curious about train time/compute!

There have been a number of structure tokenization methods introduced, including all-atom ones [1,2] and backbone-only [3,4]. These should at minimum be cited, and preferably include a baseline comparison to. It's certainly not fair to discount the novelty of a piece of work because other works have been since preprinted (I believe a version of this work was preprinted earlier than some of the others listed below).

We agree with the reviewer and strongly believe in the relevance of our work.
We want to highlight that at the time of writing most the literature were not even preprinted. Nonetheless, as the reviewers correctly noticed, the field is rapidly evolving. As stated in the response to all reviewers, we have taken this remark into consideration, significantly expanding our related work section and updated our manuscript accordingly.

Better understanding of how codebook size affects reconstruction, see [2]

In the main text, we present a detailed ablation study on codebook size and refer the reviewer to Table 1 in the main paper for comprehensive results. Here, we summarize the key findings. At a fixed downsampling ratio, the validation RMSD decreases as the number of codes in the codebook increases. This trend is particularly pronounced at a high downsampling ratio (4), where reducing the codebook size by a factor of 16 results in a 40% decrease in reconstruction performance. In contrast, when the downsampling ratio is 1, this decrease is only 20%.

To address your question more thoroughly, we retrained our model with additional codebook sizes of 432 and 1728 to achieve further compression. The resulting median RMSD and TM-scores on CASP-15 are 1.75 and 0.90, and 1.33 and 0.94, respectively. These results demonstrate that protein structures can indeed be effectively compressed within a more constrained codebook. Additionally, the manuscript provides an in-depth analysis, including error distributions and visualizations, to illustrate the interplay between downsampling ratios and codebook sizes. This helps better characterize the probability distribution in this challenging benchmark.

We hope these new results adequately address the reviewer’s concerns.

Investigating generation with more rigorous benchmarks; & current results are sparse and do not perform great against baselines Thank you again for this thoughtful feedback, which will be beneficial as we continue to refine our evaluation methods and would be glad to include any actual suggestions in an updated version of the paper. As stated in the main paper, generative results are illustrative of the representation power of the method. We recall the presented generation results are conducted with a decoder-only model, learned solely on a next-token-prediction task using our learned tokenizer. We motivate our choice to use our tokenizer as such to illustrate that even a simple sequence generation method could provide competitive results without any sort of refinements.

There are numerous ways to improve the generation results and we defer them as future work.

At the current state of the field, for this work to be of interest to the ICLR 2025 audience, I think the authors should consider probing further into their results. There's a ton of interesting findings, but introducing a tokenizer in itself (especially given the limited backbone-only setting) might not be a significant contribution.

We strongly disagree with the reviewer. Indeed, as the reviewer suggested, several tokenization methods are now being developed, for both efficient information compression and downstream usage. Therefore our work very much aligns with the interest of the community. We believe that proposing a well tested model, based on grounded methods for efficient structure encoding, decoding and tokenization is not only relevant to the community but also lays the foundations for a thriving line of research.

"I agree and am aware that this was one of the first works to be preprinted on this subject, and it's rather unfair to discount the novelty of a work because the ML conference circuit comes around too infrequently to catch up to the pace of research in this space."

We thank the reviewer for acknowledging this point, however, we strongly disagree with the statement :

"Nonetheless, by the time April 2025 comes around, I'm just not too sure if this paper in its current form will be a significant contribution to readers."

We evidenced a strong case showing that our proposition could be of interest for the practitioner thanks to:

• Information encoding and reconstruction: Our tokenizer achieves reconstruction performance at the level of experimental resolution using only backbone structural information, unlike prior works [2, 3], which incorporate sequence data. This capability highlights the potential of our method to support multimodal protein language models.
• Generation: when training "a simple off-the shelf decoder-only transformer" without any refinement, we provided generation results that have shown to be competitive with diffusion models specifically tailored to excel at that task.

"...will this be the standard tokenizer that people should go to because its comprehensive evaluation has convinced us that that should be the case? Or is this uncovering some fascinating insights about how to better do this type of work?"

Our previous answer stresses the methodological and technical merits of our approach. Moreover, thanks to the reviewers' insightful feedback we have been able to further study how the codebook size and downsampling ratio affects reconstruction.

In conclusion, our present study provides the reader thanks to extensive analyses with practical guidelines on how to train and utilize our structure-only tokenizer effectively, providing clarity on expected performance for reconstruction and downstream applications.

"Since other tokenizers exist, unless there are benchmarks against other tokenizers, I'm not sure if I would choose to use this one in particular."

As you previously highlighted, indeed, we assume that a comprehensive evaluation that convinced reviewers of its merits along with open source code (that can already be accessed but that will not be disclosed here for anonymity purpose, but will be for the camera ready version) constitutes a strong case for considering its use.

"This is exacerbated by the fact that it is backbone only, and all-atom ones now exist. I'm also not too sure what the takeaways are about what to do when I am building my own tokenizer."

We highlight that backbone-only is not a precursor to all-atom but rather is a deliberate design choice as it allows the tokenisation model to run without context of the sequence (which would otherwise be implicit in the side-chain atoms). For our purposes this held two primary benefits. (1) The ability to have modality specific pre-training of tokenisation models. (2) The use of our tokenizer in protein design tasks where the sequence is not a priori known - for example, a common workflow for ML-based protein engineering is to first generate the protein backbone (e.g. with RFDiffusion), then compute the corresponding sequence (e.g. with ProteinMPNN).

We will make sure to further highlight in our manuscript the benefits of this specific design choice.

"The additional codebook sizes are helpful (and I will raise my ratings to address this)"

We are grateful for the reviewers comments. We hope that the above discussion will allow further re-consideration of the reviewers score; but note that at present the suggested ratings increase has not happened.

We would like to make sure if the reviewer have had the opportunity to review the additional elements and analysis provided notably regarding the your latest comments, further highlighting the merits and differences of our approach.

We have addressed in this rebuttal many of the reviewer's concerns, including additional analysis on the codebook size, and in depth analysis on the errors as they requested.

Also, the reviewer stated we would increase its ratings, and we have been unable to notice such a change.

We'd appreciate to hear your feedback.

An ablation study comparing the FS quantization with VQ would help motivate the architectural choice.

We began by exploring traditional VQ-VAE-based architecture. However, training the encoder within this framework proved challenging, necessitating numerous adjustments—such as k-means initialization, EMA updates, and periodic resets—to maintain stability. Despite these efforts, the approach remained unstable, frequently suffering from latent collapse or poor expressivity. In contrast, our current FSQ-based autoencoder yields results that are significantly more robust and stable, outperforming the VQ-VAE approach by several orders of magnitude.

You can include an ablation study demonstrating the importance of invariance in your encoder architecture. Consider trying a non-invariant architecture and see how that impacts downstream language modeling.

The use of invariant / equivariant features in structure modeling is an open debate in the field, whilst we concur with the reviewer that ablating this component also provides application specific elements to the debate, it is not the primary goal of our work. Nevertheless, we hypothetize the invariant encoder to be important in our work, maintaining desirable properties of the latent space. Indeed, it ensures that two structures differing only by rotation or a translation have similar latent representations. We also incorporate as a benchmark the use of a non equi/invariant transformer encoder so as to provide the reader with a comparison with a natural benchmark. We report the metrics for the autoencoding.

The downsampling of the sequence makes it harder for the tokenizer to reconstruct the structure, but also makes it easier for the language model to learn. It would be great to demonstrate this tradeoff in your experiments. Also, I think the same can be said about the invariant feature representation.

Our rationale behind these choices is grounded in the observed training behavior of language models, as documented by Hoffmann et al. [1]. A common rule of thumb when training decoder-only models is to ensure there are approximately 20 observed tokens per parameter, adjusted by the size of the vocabulary. Achieving this balance typically requires avoiding downsampling, as downsampling would significantly reduce the number of observed tokens, limiting model performance. We agree with the reviewer that exploring the training of decoder-only models with higher downsampling ratios is promising, we must point out that it would introduce challenges: increased downsampling would reduce the amount of training data, impairing model training capacity. Additionally, the hyperparameters required for training with different downsampling rates would likely diverge significantly, complicating direct comparisons between models.

In that sense, this paves the way to another line of research studying both the optimal tokenizer given a generation strategy (here through next token prediction).

[1]: Hoffmann et al, Training Compute-Optimal Large Language Models

You may want to compare with a hand-crafted tokenizer (e.g. BPE-based tokenizer) in terms of reconstruction accuracy, downstream generation validity/diversity.

Our language model is learned on a next-token-prediction task using our learned tokenizer. We motivate our choice to use our tokenizer as such to illustrate that even a simple sequence generation method could provide competitive results without any sort of refinements. Naturally we could apply further refinements either in the tokenization, leveraging BPE, or in the generation, for instance applying some preference fine tuning methods, however, we believe that evidencing the “vanilla” generation power is more useful to practitioners and opens the way to further works.

Have you tried extending to side chain conformations?

We thank the reviewer for this comment. Including side chain information can indeed provide the model with further valuable information and as suggested, we plan to include it in future works.

Ablations

Notably as asked by reviewers (j5Mn, Ai4U), we included for the 4096 codebook size and downsampling ratio of 1, two ablations:

1. Removing the local attention.
2. A non equivariant encoder (replacing the GNN with a non-equivariant graph transformer). To ensure it has the same information as our original encoder, we modify the regular transformer architecture:

• Initial embeddings: the flattened edges features concatenated with the residues’ backbone positions (making the encoder non equivariant).
• Bias the attention score: using linear transformation of the original edges features effectively providing information for the computation of the embedding -Key offsetting: we make the key being the average of the original key and a linear transformation of the node’s flattened edges features.

Table 3: Ablation Studies

We can notice that our local attention scheme enables very consistent training, and removing it significantly hampers training. Interestingly, the graph transformer architecture we propose is very consistent with our original encoder proposition.

Main paper corrections

All corrections, additions and new results are now included in the main paper, notably for the related work and the experimental sections.

[1] van Kempen, Michel, Stephanie S. Kim, Charlotte Tumescheit, Milot Mirdita, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. "Foldseek: fast and accurate protein structure search." Biorxiv (2022): 2022-02.

[2] Trinquier, Jeanne, Samantha Petti, Shihao Feng, Johannes Söding, Martin Steinegger, and Sergey Ovchinnikov. "SWAMPNN: End-to-end protein structures alignment." In Machine Learning for Structural Biology Workshop, NeurIPS. 2022.

[3] Lin, Xiaohan, Zhenyu Chen, Yanheng Li, Zicheng Ma, Chuanliu Fan, Ziqiang Cao, Shihao Feng, Yi Qin Gao, and Jun Zhang. "Tokenizing Foldable Protein Structures with Machine-Learned Artificial Amino-Acid Vocabulary." bioRxiv (2023): 2023-11.

[4] Gao, Zhangyang, Cheng Tan, Jue Wang, Yufei Huang, Lirong Wu, and Stan Z. Li. "Foldtoken: Learning protein language via vector quantization and beyond." arXiv preprint arXiv:2403.09673 (2024).

[5] Li, Mingchen, Yang Tan, Xinzhu Ma, Bozitao Zhong, Huiqun Yu, Ziyi Zhou, Wanli Ouyang, Bingxin Zhou, Liang Hong, and Pan Tan. "ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention." bioRxiv (2024): 2024-04.

[6] Gaujac, Benoit, Jérémie Donà, Liviu Copoiu, Timothy Atkinson, Thomas Pierrot, and Thomas D. Barrett. "Learning the Language of Protein Structure." arXiv preprint arXiv:2405.15840 (2024).

[7] Hayes, Tomas, Roshan Rao, Halil Akin, Nicholas J. Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil et al. "Simulating 500 million years of evolution with a language model." bioRxiv (2024): 2024-07.

[8] Lu, Amy X., Wilson Yan, Kevin K. Yang, Vladimir Gligorijevic, Kyunghyun Cho, Pieter Abbeel, Richard Bonneau, and Nathan Frey. "Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure." bioRxiv (2024): 2024-08.

[9] Wang, Xinyou, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. "DPLM-2: A Multimodal Diffusion Protein Language Model." arXiv preprint arXiv:2410.13782 (2024).

[10] Lu, Jiarui, Xiaoyin Chen, Stephen Zhewen Lu, Chence Shi, Hongyu Guo, Yoshua Bengio, and Jian Tang. "Structure Language Models for Protein Conformation Generation." arXiv preprint arXiv:2410.18403 (2024)

[11] Jin Su, Chenchen Han, Yuyang Zhou, Junjie Shan, Xibin Zhou and Fajie Yuan “SaProt: Protein Language Modeling with Structure-aware Vocabulary”, ICLR 2024

[12] Z. Gao, C. Tan, and S. Z. Li. “VQPL: Vector quantized protein language”. arXiv, 2023.

[13] X. Lin, Z. Chen, Y. Li, X. Lu, C. Fan, Z. Cao, S. Feng, Y. Q. Gao, and J. Zhang. “Protokens: A machine-learned language for compact and informative encoding of protein 3d structures.” bioRxiv, 2023.

[14] Y. Liu, L. Chen, and H. Liu. “Diffusion in a quantized vector space generates non-idealized protein structures and predicts conformational distributions.” bioRxiv, 2023.

[15] A. Liu, A. Elaldi, N. Russell, and O. Viessmann. “Bio2token: All-atom tokenization of any biomolecular structure with mamba.” arXiv, 2024.

[16] Z. Gao, C. Tan, and S. Z. Li. “Foldtoken2: Learning compact, invariant and generative protein structure language.” bioRxiv, 2024.