Learning the Language of Protein Structure

Authors: Jérémie DONA, Benoit Gaujac, Timothy Atkinson, Thomas PIERROT, Thomas D Barrett
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Strucutral Biology; Quantized Representation; Generative Modeling

Abstract:
Representation learning and \emph{de novo} generation of proteins are pivotal computational biology tasks. 
Whilst natural language processing (NLP) techniques have proven highly effective for protein sequence modelling, structure modelling presents a complex challenge, primarily due to its continuous and three-dimensional nature.  
Motivated by this discrepancy, we introduce an approach using a vector-quantized autoencoder that effectively tokenizes protein structures into discrete representations.  This method transforms the continuous, complex space of protein structures into a manageable, discrete format with a codebook ranging from 4096 to 64000 tokens, achieving high-fidelity reconstructions with backbone root mean square deviations (RMSD) of approximately 1-5 \AA. To demonstrate the efficacy of our learned representations, we show that a simple GPT model trained on our codebooks can generate novel, diverse, and designable protein structures. Our approach not only provides representations of protein structure, but also mitigates the challenges of disparate modal representations and sets a foundation for seamless, multi-modal integration, enhancing the capabilities of computational methods in protein design.

Decision: reject
Reviewer scores: 6, 5, 8, 3

Reviews and discussion:
<thread>
creation date:
1730081674639

thread_type:
official_review

<thread_object>
creator: reviewer_q6db

<content>
summary:
This paper, "Learning the Language of Protein Structure," proposes a new approach to protein structure modeling by representing the continuous 3D forms of proteins in a discrete, sequence-like format. This work tries to bridge protein sequence and structure modeling using techniques drawn from natural language. The authors introduce a vector-quantized autoencoder to translate 3D protein structures into discrete tokens. This tokenization leverages a codebook of 4096–64000 tokens, achieving effective compression while retaining high fidelity in structure reconstruction with root mean square deviations (RMSD) between 1–5 Å. The learned discrete representations are applied to generate novel and diverse protein structures using a simple GPT model, which demonstrated competitive performance in designing structurally viable proteins against established diffusion-based models.

strengths:
1. The paper is generally well-structured and clearly written.
2. The use of vector-quantized autoencoders combined with transformer-based NLP models is an interesting application of NLP methodologies to the domain of protein structure.
3. The application of finite scalar quantization (FSQ) to manage codebook stability issues shows creativity in overcoming known challenges in quantized representation learning.
4. The authors explore their model’s capacity with both qualitative and quantitative metrics. Ablation studies and variations in codebook size and downsampling ratios.

weaknesses:
1. 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.
2. The primary application explored in the paper is protein structure generation using a GPT model. While this is promising, the paper could be more impactful if it included additional downstream applications to show the broader utility of its tokenized representation.
3. 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. 

[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).

questions:
I'm wondering about the key advantages of this work compared to the listed works above. What's your unique contribution?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730237184958

thread_type:
official_review

<thread_object>
creator: reviewer_j5mn

<content>
summary:
This paper proposes to discretize protein backbone structures so that we can use the well-developed machinery for modeling sequences of discrete tokens to analyze and generate protein structure. The authors train use a ProteinMPNN-style embedder, a cross-attention based discretizer, and an AlphaFold folding-module based decoder to parametrize finite scalar quantization of protein backbones, showing that it is possible to achieve accurate reconstructions. Finally, they train a simple autoregressive transformer to generate protein backbone structures in this discrete latent space.

strengths:
The paper makes a clear and strong case for discretizing protein backbone structures in order to use discrete sequence models to analyze and generate protein backbones. While the individual parts (discretizing protein structure, generating protein structures, using sequence modeling architectures on discretized protein structures) are present in previous work, this work is unique in its focus, thoroughness, and attention to detail on the quality of the discretization and reconstruction. The main claims are well-supported by the experiments, which do a good job of following the current standards in the field. The paper generally flows very well, and the main points and contributions are clear.

weaknesses:
The two biggest areas for improvement are in explaining the paper's significance in light of the (missing) related work in this area and on the completeness of the ablations and metrics reported. I think this could be a really strong paper if most of these weaknesses are addressed. 

## Related work

The authors should mention and describe how their work is different from [ProTokens](https://www.biorxiv.org/content/10.1101/2023.11.27.568722v4.abstract), which also aims to discretize protein structures with high reconstruction accuracy, and [FoldToken2](https://www.biorxiv.org/content/10.1101/2024.06.11.598584v3.abstract), which also discretizes protein structures and trains (with very sparse details) flow matching and autoregressive generative models on those tokens. 

There are also other works that train sequence models on tokenized structures. [SaProt](https://www.biorxiv.org/content/10.1101/2023.10.01.560349v1) is a masked language model trained on both sequence and structure tokens, while [ProstT5](https://www.biorxiv.org/content/10.1101/2023.07.23.550085v1) uses tokenized structures to unify structure prediction and inverse folding into a single model. [ESM3](https://www.biorxiv.org/content/10.1101/2024.07.01.600583v1) trains their own structure tokenizer and claims to be able to co-generate sequence and structure. The paper should mention these related works and explain how it is distinct and/or builds on this prior work. 

To be clear, I believe that this paper does make a significant contribution even in the light of these related works. However, not all readers will know the related literature quite as closely, and the paper would be greatly improved by properly situating itself in this context. 

## Experiments

1. While Table 1 nicely summarizes reconstruction quality, we need to know more than the mean RMSD and TM. 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.
2. 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. 
3. 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. 
3. 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. 
4. The paper claims that we want to enforce locality with a special attention mechanism during tokenization. This should be ablated. 
5. It would be interesting to see how using a different-sized codebook affects the quality of the downstream generative model. Do we always just want the biggest codebook we can afford? 


## Other weaknesses

1. Algorithm 2 needs to be expanded so that the reader has some reasonable chance of being able to implement it. I do not see how the current description facilitates down- or up-sampling. 
2. The paper describes the training hardware for the quantization model but not for the generative model. The training time is not described for either model.
3. While the paper makes it clear later, the initial section describing the token decoder should explicitly state that it uses the AlphaFold2 structure module architecture.

questions:
In addition to the suggestions in the Weaknesses section:

1. How much compute was utilized to train the full model (tokenizer + generator)? For reference, in the United States, the White House recently issues an executive order saying that models trained on mostly biological sequences for more than 1e23 FLOPs need to be reported to the federal government. It's unclear whether protein structures count as biological sequences. By my estimate, 128 TPU v4-8s give a max of 140800 teraflops, so assuming a utilization of 50%, you'd get to 1e23 FLOPs after about 396 hours. 
2. It'd be interesting to know whether the generated length distribution matches that of the training set. It'd also be interesting to see how the designability/novelty/diversity metrics look if the baseline models were sampled from at exactly the same empirical distribution as the samples from the tokenized generator. 
3. I'm curious why the length vs scRMSD isn't monotonic for this model is for diffusion and flow matching backbone models.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693532737

thread_type:
official_review

<thread_object>
creator: reviewer_pvnp

<content>
summary:
This work introduces a method for tokenizing protein backbones. The structure is encoded as a graph to use a GNN as the encoder. Authors use the FSQ method for quantization and reconstructions are trained using the FAPE loss from AF2. To mitigate the fact that many protein families share similar structures, data is sampled inversely proportional to its MMSeqs-assigned cluster. At a codebook size of K=64000, the average reconstruction has 1.59A RMSD and TM score=0.95. Authors further examine the effects of downsampling and altering the number of codes (4096, 64000), finding that using 4096 codes reduces the performance, but not by too much.

strengths:
* Structure tokenization is an important area of research, setting the foundational work necessary for many downstream uses.
* Considering the clustered nature of biological data during training is a good practice that should be more commonly adopted in the field
* Investigating the optimal codebook size is useful both for understanding proteins and for building similar tools

weaknesses:
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). However, there are many open questions that authors could dig into, such as:
* Better understanding of how codebook size affects reconstruction, see [2]
* Investigating generation with more rigorous benchmarks; current results are sparse and do not perform great against baselines

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.

**Minor:**
* line 431: typo at ("Eguchi et al., 2022) ...

**References:**

[1] Bio2Token: All-atom tokenization of any biomolecular structure with Mamba (https://arxiv.org/abs/2410.19110)

[2] Tokenized and Continuous Embedding Compressions of Protein
Sequence and Structure (https://www.biorxiv.org/content/10.1101/2024.08.06.606920v1.full.pdf)

[3] FoldToken: Learning Protein Language via Vector Quantization and Beyond (https://arxiv.org/abs/2403.09673)

[4] ProTokens: A Machine-Learned Language for Compact and Informative Encoding of Protein 3D Structures (https://www.biorxiv.org/content/10.1101/2023.11.27.568722v2)

questions:
* It's a bit unclear from the writing what the architecture & output of the decoder is. Is it finetuning the structure module of AlphaFold? Or just using the structure module architecture? Or something different altogether (e.g. using a GNN like the encoder?) would be helpful to see this clarified.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697231273

thread_type:
official_review

<thread_object>
creator: reviewer_ai4u

<content>
summary:
The work proposes to use an FSQ-based auto-encoder to tokenize 3D protein structures at the residue level. They use a GNN encoder, IPA decoder, and a length compression scheme. They find that their autoencoder is able to achieve experimental precision in terms of RMSD and measure the RMSD and TM-score as a function of downsampling ratio/compression factor and codebook size. They then train a GPT model on the generated tokens and outperform FrameDiff but not RFDiffusion on scTM and scRMSD scores. They tend to lag behind in terms of novelty/diversity, but they note that there may be a tradeoff between novelty/diversity and validity.

strengths:
1. The problem is very well-motivated: Tokenizing 3D structures allows for multi-modal integration in downstream language models.
2. To the best of my knowledge, autoencoder-based tokenizers are pretty novel in the field of computational protein biology.
3. Experiments, though limited, are sufficient to demonstrate that their tokens can be used for both reconstruction and downstream language modeling.

weaknesses:
1. An ablation study comparing the FS quantization with VQ would help motivate the architectural choice.
2. 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. 
3. 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.
4. You may want to compare with a hand-crafted tokenizer (e.g. BPE-based tokenizer) in terms of reconstruction accuracy, downstream generation validity/diversity.

questions:
1. Have you tried extending to side chain conformations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732537284314

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their valuable comments and suggestions. 

As highlighted by several reviewers, protein structure quantization has recently gained significant momentum in the research community. Foldseek [1], one of the earliest and most established works in this area, introduces a structural vocabulary that encodes only local geometry. They demonstrate its utility for fast database searches, and subsequent works [2, 11] have leveraged this vocabulary to train structure-aware protein language models. However, as noted in our manuscript, these approaches lack the capacity to provide global representations of protein structures, as the vocabulary encodes only local structural features.
As the reviewers noted, several works are closely related to our method. Specifically, [12, 13, 14, 15] employ VQ-VAE or variations to train protein structure tokenizers. We would like to emphasize that these works are either concurrent with or subsequent to our own.
Similarly, many recent studies have explored the use of tokenized structures in training multimodal generative models, including [4, 7, 9, 10, 16].
We agree with the reviewer that including these references provides readers with a more comprehensive understanding of the field and its recent advancements. Accordingly, we have updated the manuscript to reflect these developments.


**Novel Results**

**Data Split and Novel Training**

As suggested by reviewer (j5Mn), due to our initial uniform sampling, our test performances may over-estimate the reconstruction power of our model for instance on OOD samples. To address the reviewers’ remark and to show the generalization power of our algorithm, we redesigned our training scheme: 
For our training, we randomly select 90% of the cluster and use the remaining clusters for validation and testing. Amongst these 10% withhold protein clusters, we retain 20% for validation, the remaining 80% being used for tests. This ensures that the test structures were never seen during training.
This has also been updated in the main paper.

**Metrics Reported**

This new training scheme enables us to further evidence the capacity of our model. We now report: 
1- The average reconstruction metrics on the test set, containing only structures belonging to left out clusters, i.e. from clusters carefully excluded from the training set.
2- The median reconstruction metrics (due to the limited dataset size) on CASP15 structures, known to be structurally complex
Also, with such a retraining, we highlight that we have been able to  improve our reconstruction metrics by more than 10% across all downsampling ratios and codebook sizes along with the ability to train our model on a smaller batch size drastically reducing the computation power necessary for training. 

**Novel Results, Ablations and Novel Codebook sizes!**


| Codebook Size | df | CASP15-RMSD | CASP15-TM | Test RMSD | Test-TM |
|---------------|----|-------------|-----------|-----------|---------|
| **4096**  	| 1  | 1.25 | 0.94 |1.55 |0.94|
| **64000** 	| 1  | 0.94| 0.97 |1.22|0.96|
|   |         	|       	|       	|     	|
| **4096**  	| 2  |  1.73 |0.89 |2.22 |0.90|
| **64000** 	| 2  | 1.82 | 0.90|1.95|0.92|
|  |         	|       	|       	|     	|
| **4096**  	| 4  |2.79|0.77|4.1|0.81|
| **64000** 	| 4  |2.55|0.80|2.96|0.86|


Table: 1 Updated Results for the main experiments with additional metrics on CASP-15



In order to also validate our proposition in restricted information settings,as required by reviewer (pVnp) we now include novel experiments, evidencing the capacity of our algorithm to train on restricted codebook size, effectively decreasing in the case of a downsampling ratio of 1, down to 1728 and even 432 codes.

| Codebook Size | df | CASP15-RMSD | CASP15-TM | Test RMSD | Test-TM |
|---------------|----|-------------|-----------|-----------|---------|
| **432**  	| 1  | 1.75 | 0.89 | 2.09 |0.91|
| **1728** 	| 1  |1.33  | 0.94 |1.79|0.93|

Table 2: Novel reconstruction results with limited codebook sizes
As required by reviewer (j5Mn), we now provide in appendix A.4 the boxplot of both RMSD and TM-score for the results of CASP15 and also both RMSD and TM-Scores distribution for  the test set.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732558773441

thread_type:
official_comment

<thread_object>
creator: reviewer_j5mn

<content>
title:
Reviewer discussion

comment:
The biggest discrepancy between me and the other reviewers seems to be whether the paper has sufficient technical novelty. In my view, while neither of the components (tokenizing structure and autoregressive generation) are novel, this is the first work I've seen combine those two. This was probably an obvious application from the moment FoldSeek was released, but it is nevertheless important to have a well-executed first attempt that tries to improve the tokenizer and also characterizes performance compared to the current standard diffusion methods.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732566295963

thread_type:
official_comment

<thread_object>
creator: area_chair_ozkm

<content>
title:
Last day for reviewers to ask questions to the authors!

comment:
Dear reviewers,

Tomorrow (Nov 26) is the last day for asking questions to the authors. With this in mind, please read the rebuttal provided by the authors earlier today, as well as the other reviews. Please explicitly acknowledge that you have read the rebuttal and reviews, provide your updated view _accompanied by a motivation_, and raise any outstanding questions for the authors. 

**Timeline**: As a reminder, the review timeline is as follows:

- November 26: Last day for reviewers to ask questions to authors.
- November 27: Last day for authors to respond to reviewers.
- November 28 - December 10: Reviewer and area chair discussion phase.

Thank you again for your hard work,

Your AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1733241617591

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Update: Novel results for generative modeling of protein structures

comment:
As requested by several reviewers [Ai4U, pVNP, J5MN], we explored the impact of codebook size on generative quality. 

To this end, we utilize the autoencoders trained with a downsampling factor of 1 and codebook sizes of 432, 1728, 4096, and 64,000. Using the corresponding tokenized datasets (generated with the respective encoder and codebook), we then trained four decoder-only transformer models for next-token prediction. We provide the designability and diversity results, computed on 256 structures generated using a decoder-only model with similar hyperparameters as those described in the manuscript. All decoder-only models were trained with the same batch size, learning rate and hyperparameters to ensure consistency.

| Codebook Size | SCTM (%) | SC-RMSD (%)   | Diversity |
| --------------------| ------------------- | ------------------|------------------|
| 432  |  89,5 % | 66.5 % | 63% |  
| 1728 |  88,6 % | 64.8 % | 66% |  
| 4096 | 87.2 %  | 61.8 % | 60 % |  
| 64,000 | 84.2 % | 56.2 % | 61.3 % | 
Table 1: Generative self consistency, novelty and diversity metrics 

*Key Insights*

The results presented here were achieved using a generator trained exclusively on the next-token prediction task, without any additional fine-tuning or refinement. Our results suggest that, given the available training data, a smaller vocabulary size can effectively produce high-quality protein structures. 
We managed to boost our self consistency results going from a scRMSD of 42% (for the 4096 codebook) to 61.8% using the same codebook size and even reaching 66.5% with a codebook size of 432.

*Significance of Findings*

We believe these additional experiments further validate our proposed framework, demonstrating its dual capability as:
- A robust autoencoder for protein structures, and

- A promising initial step toward de novo structure generation. These novel results suggest that our approach is closing the gap with tailored generative methods such as RF-Diffusion, offering a new and scalable pathway for protein design.
We will make sure to include these results in the camera-ready version of our manuscript and to make available the pretrained weights for the generation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission ended up with mixed reviews. While initially all reviewers except one were leaning towards recommending to reject this paper, the lively discussion with the authors led to a more split situation with two reviewers in favor of accepting and two reviewers against.

Several reviewers asked the authors to position their work better compared to (very recent) related work in their related work section. The authors have complied with this request and updated their related work section. 
Several reviewers, in particular reviewers Q6dB and pVnp, raised concerns about insufficient evaluation of the proposed method, indicating this could limit the potential community interest for this submission. Concerns were raised about the need to compare or ablate different tokenization choices. Furthermore, reviewer pVnp found there were missing metrics for the protein generation task that are used in other places in the literature (beta sheet percentage, sampling time, motif scaffolding, etc). During the discussion phase, the authors have put in significant effort to increase the thoroughness of their evaluation, for instance by looking into the effect of codebook sizes. 
The authors have not included the additional metrics that were requested by pVnp, but reviewer j5Mn was of the opinion that extending to these suggested metrics is a nice-to-have and not a must have for acceptance.
While I agree that the additional metrics are a-nice-to-have, and that this alone is not sufficient to reject the paper, I recommend rejecting this paper based on the shared concerns by reviewers Q6dB and pVnp that this paper has not presented sufficient empirical evidence of the value of their proposed tokenization approach on downstream tasks. I would like to thank the authors and reviewers for engaging in discussions and would strongly encourage the authors to resubmit an improved version of their work with more in-depth ablations to a future conference or venue.