Learning the Language of Protein Structure

## Abstract

Abstract Representation learning and 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 Å. 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.

## 1 Introduction

The application of machine learning to large-scale biological data has ushered in a transformative era in computational biology, advancing both representation learning and *de novo* generation. Particularly, the integration of machine learning in molecular biology has led to significant breakthroughs (Sapoval et al., ; Chandra et al., ; Khakzad et al., ; Jänes and Beltrao, ), spanning many complex and inhomogenous data modalities, from sequences and structures through to functional descriptors and experimental assays, many of which are deeply interconnected and have attracted significant modelling efforts.

Currently, the deep learning landscape is increasingly converging towards a unified paradigm centered around attention-based architectures (Vaswani et al., ) and sequence modeling. This shift has been driven by the impressive performance and scalability of transformers, and even accelerated since treating non-standard data modalities as sequence-modeling problems has proven highly effective.
Indeed, transformers-based models are leading methods in many machine learning domains, including image representation (Radford et al., ), image (Chen et al., ; Chang et al., ) and audio generation (Ziv et al., ), and for Reinforcement Learning (Chen et al., ; Boige et al., ).
This trend has greatly benefited sequence-based biological models, allowing for the direct application of NLP methodologies like GPT (Radford and Narasimhan, ) and BERT (Devlin et al., ) with notable success (Ferruz et al., ; Lin et al., ).

In particular, large multi-modal models (LMMs), leveraging transformer backbones, are emerging as a key tool with applications in various fields such as: ubiquitous AI (GPT4 (Achiam et al., ), LLaVA (Liu et al., ), Gemini (Gemini et al., ), Flamingo (Alayrac et al., )); text-conditioned generation of images (Parti (Yu et al., ), Muse (Chang et al., )) or sounds (MusicGen (Copet et al., )); and even reinforcement learning (Reed et al., ).
LMMs are also instantiated in biological settings, such as medicine (MedLlama (Xie et al., ), Med-Gemini (Yang et al., ) Med-PaLM (Tu et al., )) and genomics (ChatNT (Richard et al., )).
Core to all these LMMs is the use of pre-trained encoders as a mechanism for combining modalities in sequence space.
For instance, LLaVa (Liu et al., ) and Flamingo (Alayrac et al., ) used pre-trained vision encoders, ViT-L/14 (Radford et al., ) and Normalizer Free ResNet (NFNet) (Alayrac et al., ) respectively, while Copet et al. () leverages the the codec of Défossez et al. ().
In general, training robust representations of specific modalities facilitates their incorporation into state-of-the-art LMMs.
However, to the best of our knowledge, there is no established methodology that readily allows the application of sequence-modelling to protein structures.

Foldseek (van Kempen et al., ) proposes a quantized autoencoder to encode *local* protein geometry which proved successful for database search and to some extent, in LMMs by augmenting the residue-based vocabulary (Su et al., ). However, since it only extracts local features at a residue level, it cannot provide global representation of the protein structures, preventing its use for instance for structure generation, or binding prediction for which including global information is necessary (Krapp et al., ).
Despite these advances in related domains, structure-based modeling of biological data such as proteins remains a formidable challenge. Unlike sequences, protein structures are inherently three-dimensional and continuous, which complicates the direct application of transformer models that primarily handle discrete data.
Instead, structure-based methods often design bespoke geometric deep learning methodologies to process Euclidean data; for example graph-neutral network encoders (Dauparas et al., ; Krapp et al., ) and structurally-aware modules such as those in AlphaFold (Jumper et al., ). Moreover, generative modelling of structures is typically performed with methods designed for continuous variables; such as diffusion (Watson et al., ; Yim et al., ) and flow matching (Bose et al., ), rather than the discrete-variable models that have proved so successful in sequence modelling.

In this work, we aim to address this gap by learning learning a quantized representations of protein structures enabling to efficiently leverage sequence-based language models.
The key objectives of this this work are:

- (i)
To convert protein structures into the discrete domain We propose the transformation of structural information of proteins into discrete sequential data, enabling seamless integration with sequence-based models.
- (ii)
To learn a discrete and potentially low-dimensional latent space By learning a discrete latent space through finite scalar quantization, we facilitate the mapping of continuous structures to a finite set of vectors. This effectively builds a vocabulary for protein structures, and can be pushed into low dimensions for applications with limited resources.
- (iii)
To achieve a low reconstruction error We aim to minimize the reconstruction error of the learned discrete representation, typically within the range of 1-5 Ångströms.

Our contributions are threefold. First, we introduce a series of quantized autoencoders that effectively discretize protein structures into sequences of tokens while preserving the necessary information for accurate reconstruction. Second, we validate our autoencoders through qualitative and quantitative analysis, and various ablation studies, supporting our design choices. Third, we demonstrate the efficacy and practicality of the learned representations with experimental results from a simple GPT model trained on our learned codebook, which successfully generates novel, diverse, and structurally viable protein structures.
All source code for experiments is publicly available online(^1^11Redacted to maintain anonymity).

## 2 Method

### 2.1 Protein Structure Autoencoder

Our objective is to train an autoencoder that maps protein structures to and from a discrete latent space of sequential codes. Following prior works (Yim et al., ; Wu et al., ), we consider the backbone atoms of a protein, $\mathtt{N}-\mathtt{C}_{\alpha}-\mathtt{C}-\mathtt{O}$, to define the overall structure.

For a protein consisting of $N$ residues, we seek to map its structure, represented by the tensor of the backbone atoms coordinates $\mathbf{p}\in\mathbb{R}^{N\times 4\times 3}$, to a latent representation $\mathbf{\tilde{z}}=[\tilde{\mathbf{z}}_{1},\ldots\tilde{\mathbf{z}}_{\frac{N}{
r}}]$, where a $r$ is a downsampling ratio, controlling the size of the representation.
Note that each element $\tilde{\mathbf{z}}_{i}$ can only take a finite number of values, with the collection of all possible values defining a codebook $\mathcal{C}$.

Figure: Figure 1: Schematic overview of our approach. The protein structure is first encoded as a graph to extract features from using a GNN. This embedding is then quantized before being fed to the decoder to estimate the positions of all backbone atoms.
Refer to caption: x1.png

A schematic overview of the our autoencoder is depicted in Figure 1.
In this section, we focus on the three components of the model; the *encoder* ${e_{\theta}}$ extracting
a set of $\frac{N}{r}$ embeddings of dimension $c$ denoted $\mathbf{z}\in\mathbb{R}^{\frac{N}{r}\times c}$ , the *quantizer* ${q_{\phi}}$ that discretizes $\mathbf{z}$ to obtain a quantized representation $\mathbf{\tilde{z}}$, and the *decoder* ${d_{\psi}}$ that predicts a structure $\mathbf{\tilde{p}}\in\mathbb{R}^{N\times 4\times 3}$ from $\tilde{\mathbf{z}}$.
The learnable parameters are respectively denoted $(\theta,\phi,\psi)$ and the learning setting summarizes as:

$$ $\mathbf{p}\overset{e_{\theta}}{\longmapsto}\mathbf{z}\overset{q_{\phi}}{ \longmapsto}\tilde{\mathbf{z}}\overset{d_{\psi}}{\longmapsto}\tilde{\mathbf{p}}.$ $$

##### Encoder

The encoder maps the backbone atoms positions $\mathbf{p}\in\mathbb{R}^{N\times 4\times 3}$, to a continuous downsampled continuous representation $\mathbf{z}\in\mathbb{R}^{\frac{N}{r}\times c}$ where $r$ is the downsampling ratio:

$$ $e_{\theta}:\mathbf{p}\in\mathbb{R}^{N\times 4\times 3}\mapsto\mathbf{z}\in \mathbb{R}^{\frac{N}{r}\times c}$ (1) $$

Note that when the downsampling ratio $r$ is set to 1, each component $\mathbf{z_{i}}\in\mathbb{R}^{c}$ can be interpreted as the encoding of residue $i$.

This representation learning task is similar to the traditional task of mapping point-clouds to sequences (Yang et al., ; Boget et al., ).
Inverse folding (Ingraham et al., ; Dauparas et al., ) is another example, that aims at estimating the sequence of amino acids corresponding to a given a protein structure.
Recently, ProteinMPNN has shown remarkable capacity at the inverse folding task. We follow their design choices and parameterize our encoder using a Message-Passing Neural Network (MPNN) (Dauparas et al., ).
In addition, we introduce a cross-attention mechanism, detailed in Section A.1 and Algorithm 2, allowing to effectively compress the structure representation by a downsampling ratio $r$.
We maintain locality using a custom attention masking scheme in the downsampling layer (see Section A.1), ensuring that each down-sampled node aggregates information from a small number of neighboring nodes in the original sequence space.

MPNN operates on a graph consisting of a set of vertices and edges. We follow Dauparas et al. (); Ganea et al. () and set as initial node features a positional encoding that reflects the residue’s ordering within the sequence, while for the edge features, we use a concatenation of pairwise distance features, relative orientation features and relative positional embeddings.
Note that as positional encoding, relative distance and orientation are invariant with respect to the frame of reference, the input data fed to the model are invariant to rotation and translation. This invariant encoding of the input structures guarantees the invariance of the learned representation regardless of the chosen downstream architecture. The encoding scheme is described in detail in Section A.2.

##### Quantization

The quantizer plays a crucial role in our work by discretizing all continuous latent representations into a sequence of discrete codes.
Traditional methods typically involve direct learning of the codebook (van den Oord et al., ; Razavi et al., ).
However, in line with the literature, we suffered several drawbacks associated with these approaches.
Indeed, explicit vector quantization is particularly expensive as it involves computation of pairwise distances, particularly problematic for long sequences and large codebooks.
Moreover, training instabilities arising from both the bias in the straight-through estimator and the under-utilization of the codebook capacity, often referred to as codebook collapse, make learning a discretized latent representation a hard optimization problem (Huh et al., ; Takida et al., ).
To address these challenges, we leverage the recent Finite Scalar Quantization (FSQ) framework (Mentzer et al., ), which effectively resolves the aforementioned issues, notably by reducing straight-through gradient estimation error.

FSQ learns a discrete latent space by rounding a bounded low-dimensional encoding of the latent representation.
Consider $\mathbf{z_{i}}\in\mathbb{R}^{c}$, the i^th element of the continuous encodings $\mathbf{z}$, and the quantization levels $\mathbf{L}=[L_{1},\ldots,L_{d}]\in\mathbb{N}^{d}$.
Its discretized counterpart, denoted as $\mathbf{\tilde{z}_{i}}\in\mathbb{Z}^{d}$, typically with $d\leq 8$, is defined by:

$$ $\mathbf{\tilde{z}_{i}}=\text{round}\left(\frac{\mathbf{L}}{2}\odot\text{tanh}( W\mathbf{z_{i}})\right)$ (2) $$

where $\odot$ the element-wise multiplication and $W\in\mathbb{R}^{d\times c}$ is a projection weight matrix (see Section A.1 for more details).
In doing so, each quantized vector $\mathbf{\tilde{z}_{i}}$ can be mapped to an index in $\{1,\ldots,\prod_{j=1}^{d}L_{j}\}$.
This implicitly defines a codebook by its indices, where each index is associated a unique combination of the each dimension values.
In our implementation, the quantized representation $\mathbf{\tilde{z}}=[\mathbf{\tilde{z}_{0}},\ldots,\mathbf{\tilde{z}_{\frac{N}{
r}}}]$ is then projected back to the latent space with dimension $c$ (line 5 of Algorithm 1).
Optimization is then conducted using a straight-through gradient estimator.
Equation 2 ensures that the approximation error introduced by the straight-through estimator is bounded.

##### Decoder

The decoder module of our framework is tasked with estimating a structure $\mathbf{\hat{\mathbf{p}}}$ from the latent quantized representation $\tilde{\mathbf{z}}$:

$$ $d_{\psi}:\tilde{\mathbf{z}}\in\mathbb{R}^{\frac{N}{r}\times c}\mapsto\tilde{ \mathbf{p}}\in\mathbb{R}^{N\times 4\times 3}$ (3) $$

The task our decoder addresses formulates more broadly as a sequence to point-cloud task.
A paradigmatic example of such a task in biology is the protein folding problem, where the confirmation of a protein is estimated given its primary sequence of amino-acids.
Jumper et al. () successfully tackled this task proposing a novel architecture for point cloud estimation from latent sequence of embeddings.

Specifically, the structure module of AlphaFold parameterizes a point cloud using frames.
A frame is defined by a tuple $\mathbf{T}=(\mathbf{R},\mathbf{t})$ where $\mathbf{R}$ is the frames’ orientation (*i.e.* a rotation matrix) and $\mathbf{t}$ is the frame center (*i.e.* a vector defining the translation of the center of the frame).
The origin of the frame is set to the $\mathtt{C}_{\alpha}$ carbon, and the orientation is defined using the nitrogen and
the other carbon atom.
For a thorough and mathematical description, we defer to Jumper et al. ().
However, note that AlphaFold’s structure module expects both a per-residue
representation $(\mathbf{s_{i}})_{i\leq N}$ and a pairwise representation $(\mathbf{k_{i,j}})_{i,j\leq N}$ between residues $i$ and $j$.
The per-residue representation $\mathbf{s_{i}}$ is constructed by mirroring the cross-attention mechanism described in Algorithm 2 used for downsampling in order to produce embeddings at the residue level by only varying the size of the initial input queries.
The process for constructing the pairwise representation $(\mathbf{k_{i,j}})_{i,j\leq N}$ used for reconstruction is described in Algorithms 3 and A.1.

### 2.2 Training

##### Objective

The Frame Align Point Error (FAPE) loss, introduced in Jumper et al. (), is a function that enables the comparison between point clouds.
Since there is no guarantee that the coordinates provided by the decoding module are expressed in the same basis as the input structure, direct comparison of the coordinates between the two point clouds becomes challenging.
The core concept behind the FAPE loss is to ensure that coordinates are expressed in the same global frame, thereby enabling the computation of the mean squared error.
To do so, given a ground-truth frame $\mathbf{T_{i}}$ and ground-truth atom position $\mathbf{x_{j}}$ expressed in $\mathbf{T_{i}}$, and their respective predictions $\mathbf{T^{p}_{i}}$ and $\mathbf{x^{p}_{j}}$,
the FAPE Loss is defined as:

$$ $\mathcal{L}_{\mathrm{FAPE}}=\|\mathbf{T^{p}_{i}}^{-1}(\mathbf{x^{p}_{j}})- \mathbf{T_{i}}^{-1}(\mathbf{x_{j}})\|$ (4) $$

In Equation 4, the predicted coordinates of atom $j$ ($\mathbf{x^{p}_{j}}$), expressed in the predicted local frame of residue $i$ ($\mathbf{T^{p}_{i}}$), are compared to the corresponding true atom positions relative to the true local frame, enabling the joint optimization of both the frames and the coordinates.
We found that clamping the FAPE loss with a threshold of $10$ improves the training stability.

##### Dataset

We use approximately 310000 entries available in the Protein Data Bank (PDB) (Berman et al., ) as training data.
The presence of large groups of proteins causes imbalances in the data set, with many protein from the same family sharing structural similarities.
To mitigate this issue, we sample the data inversely proportional to the size of the cluster it belongs to when clustering the data by sequence similarity using MMseqs2(^2^22The cluster size is readily available in the PDB data set: [https://www.rcsb.org/docs/grouping-structures/sequence-based-clustering](https://www.rcsb.org/docs/grouping-structures/sequence-based-clustering)).
We filter all chains shorter than 50 residues and crop the structures that have more than 512 residues by randomly selecting 512 consecutive amino acids in the corresponding sequences.
We randomly select $90\%$ for training and use the remaining as test set. Amongst these $10\%$ withhold protein structures, we retain $20\%$ for validation, the remaining $80\%$ being used for test.

##### Model Hyperparameters

For the encoder, we use a 3 layers message passing neural network following the architecture and implementation proposed in Dauparas et al. () and utilize the swish activation function. The graph sparsity is set to $50$ neighbours per residue. When the downsampling ratio is $r>1$, the resampling operation consists in a stack of 3 resampling layers as described in Algorithm 2, the initial queries being defined as the positional encodings. We strictly follow the implementation of AlphaFold (Jumper et al., ) regarding the structure module and use 6 structure layers.

##### Optimization and Training Details

The optimization is conducted using AdamW (Loshchilov and Hutter, ) with $\beta_{1}=0.9$, $\beta_{2}=0.95$ and a weight decay of $0.1$.
We use a learning rate warm-up scheduler, progressively increasing the learning rate from $10^{-6}$ to $10^{-3}$, and train the model for 250 epochs on 128 TPU v4-8 with a batch size 1024.

## 3 Experiments

The primary focus of our work is to develop an effective method to encode and quantize 3D structures with high fidelity. In this section we first evaluate, both qualitatively and quantitatively, our vector-quantized autoencoder by considering the compression and reconstruction performance and demonstrate that our tokenizer indeed permits high-accuracy reconstruction of protein sequences.
We then further highlight how this can be adapted to downstream tasks, by effectively training a *de novo* generative model for protein structures using a vanilla decoder-only transformer model trained on the next token prediction task.

### 3.1 Autoencoder Evaluation

##### Experiments

We trained six versions of the quantized autoencoder; with small ($K=4096$ codes) and large ($K=64000$) codebooks and increasing downsampling ratio ($r$) (from 1 to 4), and in doing so, varying the information bottleneck of our model.
For each codebook we evaluate the reconstruction performance achieved on the held out test set, to understand the trade-offs between compression and expressivity associated with these hyperparameter choices.

##### Metrics

To assess the reconstruction performances of the models, we rely on standard metrics commonly used in structural biology when comparing the similarity of two protein structures.
The *root mean square distance (RMSD)* between two structures is computed by calculating the square root of the average of the squared distances between corresponding atoms of the structures after the optimal superposition has been found. The *TM-score* (Y. and J., ) is a normalised measure of how similar two structures are, with a score of 1 denoting the structures are identical.
For context, two structures are considered to have similar fold when their TM-score exceeds 0.5 (Xu and Y., ) and a RMSD below 2Å is usually seen as approaching experimental resolution.

##### Results

Our results are summarised in Table 1. We find that with a codebook of $K=64000$ and a downsampling of $r=1$; our average reconstruction has $1.59$ Å  RSMD and a TM-score of $0.95$. For comparison, we also report the reconstruction performance of the exact same models without latent quantization. Whilst increasing the model capacity may allow these scores to be improved even further, this is already approaching the limit of experimental resolution (which is to say, the reconstruction errors are on par with the experimental errors in resolving the structures).

**Table 1: Test set reconstruction results of our discrete auto-encoding method for several down-sampling ratio and (implicit) codebook size. Note that a RMSD below $2$Å is considered of the order of experimental resolution and two proteins with a TM-score $>0.5$ are considered to have the same fold. The compression is defined as: the number of bits necessary to store the $N\times 4\times 3$ backbone positions divided by the number of bits necessary to store the $\frac{N}{r}$ tokens**
| Downsampling Ratio | Number of Codes | Compression Factor | Results |  |
| --- | --- | --- | --- | --- |
| RMSD ($\downarrow$) | TM-Score ($\uparrow$) |  |  |  |
| 1 | 4096 | 48 | 1.89 Å | 0.93 |
| 64000 | 48 | 1.59 Å | 0.95 |  |
| without quantization | $-$ | 0.97 Å | 0.98 |  |
| 2 | 4096 | 96 | 2.79 Å | 0.87 |
| 64000 | 96 | 2.31 Å | 0.91 |  |
| without quantization | $-$ | 1.45 Å | 0.95 |  |
| 4 | 4096 | 144 | 4.63 Å | 0.79 |
| 64000 | 144 | 3.23 Å | 0.85 |  |
| without quantization | $-$ | 2.19 Å | 0.91 |  |

Table 1 clearly indicates that increasing the downsampling factor or decreasing the codebook size correspondingly impacts the reconstruction accuracy.
This is expected as it essentially enforces greater compression in our autoencoder leading to a loss of information - nevertheless in all cases we find that the achievable reconstruction performance is still within a few angstrom with TM-scores clearly exceeding the 0.5 threshold on average.
The fact that the performances improve with increasing the codebook size also shows that our method does not suffer from codebook collapse (*i.e.* only a subset of the codes being used by the trained model), an issue well known and documented with other quantization methods (Huh et al., ; Takida et al., ).
Finally, comparing to continuous autoencoders, our learned quantization demonstrates significant information compression at the expense of only a small decrease in the reconstruction precision of the order of $0.5-1$ Å.

The reconstruction performance is illustrated in Figure 2, where examples of the model’s outputs are superimposed with their corresponding targets in the case of a downsampling factor of $r=2$ and $K=64000$ codes.
Visually, our model demonstrates its ability to capture the global structure of each protein.
Additionally, our model faithfully reconstructs local conformation of each protein preserving essential secondary structures elements of the protein such as the $\alpha$-helices and $\beta$-sheets.

Figure: Figure 2: Visualisation of the model reconstruction (blue) super-imposed with the original structures (green) for a downsampling factor of $r=2$ and $K=64000$ codes (fourth column of Table 1). Each row shows a different structures seen from a different rotation angle (column). The length and reconstruction RMSD are also given on the left of the most left column.
Refer to caption: extracted/5618651/images/reconstruction.png

### 3.2 De novo protein structure generation

##### Experiment

We now demonstrate that our learned discrete autoencoder can be effectively leveraged for downstream tasks.
In particular, we consider generation of protein structures from a model trained in our latent space as a paradigmatic demonstration of our tokenizers utility. This is not just because generative models for *de novo* protein design are of great interest for drug discovery – enabling rapid *in silco* exploration of the design space – but also as it directly leverages our compressed representation of protein structures using established sequence modelling architectures.

Specifically, we tokenize our dataset (defined in Section 2.2; full details on dataset preparation for this experiment can be found in Section A.3.1) using a downsampling factor of 1 and a codebook with 4096 codes (first line of Table 1).
This choice is motivated by the trade-off between reconstruction performance (this model achieves results close to experimental resolution), extensive datasets (GPT training benefits from large datasets, so we favor a low downsampling ratio, choosing $r=1$), and parameter efficiency (smaller codebooks imply fewer parameters , we set $K=4096$).
More specifically, we train an out-of-the-box decoder-only transformer model with 20 layers, 16-heads and an embedding dimension of 1024 (344M parameters) on a next-token-prediction task on the training split.
This decoder-only model is used to generate
new sequences of tokens which are then mapped back to 3D protein structures using our pre-trained decoder.

##### Metrics

We evaluate the generated structures on different aspects: *designability*, *novelty* and *diversity*.
*Designability*, or *self-consistency*, uses ProteinMPNN (Dauparas et al., ) to predict a potential sequence for the generated structure and refolds it using ESMFold (Lin et al., ) to report the structural similarity between the original and redesigned structure. We follow the literature and report the proportion of designed proteins with self-consistent TM-score (scTM) above 0.5 respectively the number of designed proteins with self-consistent RMSD (scRMSD) below 2Å.
The rationale is that generated structures should be sufficiently natural that established methodologies recognise them and agree on the fundamental biophysical properties.

Moreover, a useful generative model will provide samples that are varied and not simple replications of previously existing structures.
To assess this, we measure the *diversity* and *novelty* of our generation.
Diversity characterizes the structural similarities between the generated structures and is defined using structural clustering.
Specifically, we report the number of clusters obtained when using the TM-score as the similarity metric, normalised by the number of generated sequences.
Novelty compares the structural similarity of the generated structure with a dataset of reference. Here again, we use the TM-score as our similarity metric and consider a structure as novel if its maximum TM-score against the reference dataset is below 0.5. As is commonly done, we only report the proportion of novel samples in our generated dataset.
Overall, we follow Yim et al. () for the implementation of these metrics and refer the reader to Section A.3.2 for more details on the validation pipeline.

##### Baselines

To gauge how our structure generation GPT model fares against specifically designed protein design methods, we compare our model with FrameDiff (Yim et al., ) and RFDiffusion (Watson et al., ).
Both use bespoke SE(3) diffusion models specifically designed and trained the structure generation task.
Note that, the state-of-the-art RFDiffusion leverages the extensive pre-training of RoseTTAFold (Baek et al., ) on a large dataset and requires considerable computational cost. On the contrary, our generation model is a standard GPT and the objective here is to show how one can readily use the tokenized representation learned with our method for protein design.

##### Results

**Table 2: Structure generation metrics for our method alongside baselines (and nature)specifically designed for protein structure generation. Self-consistent TM-score (scTM) and self-consistent RMSD (scRMSD) are two different ways to asses the designability of the generated structure. Note that while high novelty score is desirable, structures that are too far from the reference dataset can also be a sign of unfeasible proteins.**
| Method | scTM ($\uparrow$) | scRMSD ($\uparrow$) | Novelty | Diversity |
| --- | --- | --- | --- | --- |
| Ours | 76.61% | 41.00% | 35.78% | 71.4 % |
| FrameDiff | 75.77% | 25.31% | 56.64% | 82.0% |
| RFDiffusion | 97.07% | 71.14% | 86.11% | 95.0% |
| Validation Set | 97.67% | 82.36 % | $-$ | 70.1% |

The results for the different metrics are given in Table 2 with the sampling strategy used for each method described in Section A.3.3.
It is noteworthy that while not specifically designed for the protein structure generation task, even our simple GPT model is able to generate protein structures of quality in par with a method specific to protein design such as FrameDiff.
For comparison, and to set an upper bound of what to expect in terms of the designability scores, we compute self-consistency and diversity metrics for 1600 randomly sampled structures from our validation and report them in Table 2.
While our method shows competitive performance at generative designable domains, the results are more nuanced for the novelty and diversity scores. Especially, our model seems to generate domains that are structurally closer to the reference data set than the ones generated by the baselines. One explanation for this can be found in the sampling method where the chosen parameters favored samples closer to the modes of the data distribution as discussed in Section A.3.3. Although generating structures that are novel and diverse is desirable, structures that differ too much from the natural proteins found in the reference dataset, can also indicate unrealistic structures, making the novelty and diversity metrics harder to interpret on their own.
Indeed, when visualizing novelty against designability (see Figure 7), we can see that while FrameDiff generates more novel but less designable structures (lower-right corner), our structures are more designable at the expense of lower novelty. In the light of Figure 7, we compute the number of structures with $\text{cathTM}<0.5$ (novel) and $\text{scRMSD}<2$Å(designable) and find that $9.01\,\%$ (70 out of 777) of our samples are designable and novel relative to $8.18\,\%$ (53 out of 648) for FrameDiff and $58.58\,\%$ (379 out of 647) for RFDiffusion.
We provide additional results and analyses in Section A.3.4.

In Figure 3, we visualize random samples from our model super-imposed with the ESM-predicted structure used in designability metric. We first notice that the generated sample exhibit diverse structure, with non trivial secondary secondary elements ($\alpha$-helices and $\beta$-sheets). We also note that the ESM-predicted structures are closely aligned with the original samples, showing the designability aspect of the generated structures.

Figure: Figure 3: Visualisation of generated samples (green) super-imposed with their self-consistent ESM-predicted structures (blue).
Refer to caption: extracted/5618651/images/samples.png

## 4 Related Works

##### Learning from Protein Structures.

Learning from protein structure is thriving research field that encompasses a wide variety of tasks crucial task.
For instance, inverse folding (Ingraham et al., ; Hsu et al., ; Dauparas et al., ; Jing et al., ) or binding site estimation (Krapp et al., ; Ganea et al., ) are critically enabling for drug design (Scott et al., ).
Others (Consortium, ; Robinson et al., ; Kucera et al., ; Gligorijević et al., ) focus on learning from protein structure to provide a better understanding of the role of the structure, hence pushing the knowledge frontier of biological processes.

##### Representation Learning of Protein Structures.

("Eguchi et al., ) learns a VAE using as inputs matrix distances and predicting 3D coordinates. The supervision is done by comparing distance matrices.
Despite the straightforward nature of sampling from a VAE latent space, it often yields subpar sample quality, see (Kingma et al., ).
Foldseek (van Kempen et al., ) propose to learn a VQ-VAE on small chunks of protein comparing true and reconstructed local geometrical features.
Because of the locality of the data used in Foldseek, decoding only small chunks of structures, this quantization defines an interesting tool for efficient structure database search but does not enable full backbone structure sampling.
Closer to our work, Gao et al. () propose a quantized autoencoder of protein structures, with however only limited reconstruction performances limiting its applicability.

##### Generation of Protein Structures.

A substantial body of literature addresses the challenge of sampling the protein structure space.
Numerous studies have advocated for the use of diffusion-based models Wu et al. (); Yim et al. (); Watson et al. () or flow matching techniques (Bose et al., ).
While many of these works, such as those by (Yim et al., ; Watson et al., ; Bose et al., ), employ complex architectures to ensure invariance to rigid-body transformations, (Wu et al., ) opted for an alternative parameterization directly preserving symmetries allowing the authors to capitalize on conventional architectures, yet working only on small protein crops.

## 5 Conclusion

This work demonstrates a methodology for learning a discrete representation of a protein geometry; allowing the mapping of structures into sequences of integers whilst still recovering near native conformations upon decoding.
This sequential representation not only simplifies the data format but also significantly compresses it compared to traditional 3D coordinates.
Our belief is that the primary contribution of our work lies in setting the stage for applying standard sequence-modelling techniques to protein structures.

A prerequisite for such development is the expressiveness of the tokenized representation, which must capture the necessary information to enable hight fidelity reconstructionf of the 3D structures.
Both empirical and visual inspection confirm this to be the case for our proposed methodology.
Indeed, as codebook collapse is effectively mitigated by the use of FSQ, larger codebook vocabularies and an increased tokens usageprovide straightforward recipes for improving the reconstruction accuracy by reducing the information bottleneck of the autoencoder

The first step towards sequence-based modelling of structures is the proof-of-concept GPT model, trained on tokenized PDB entries, that serves as a simple *de novo* generating protein backbones. That the achieved results are competitive with some recent diffusion-based model underlines the promise of this paradigm.
While a simple GPT does not yet match seminal approaches like RFDiffusion, it is important to recognize that the performances of the latter stem from extensive developments in 3D generative modeling.
Given the remarkable performance of sequence-modelling algorithms across a diversity of data modalities, equivalent efforts could also provide simpler and more powerful treatments of protein structure. This represents a promising direction for future research based on this work.

## Acknowledgments and Disclosure of Funding

This research was supported with Cloud TPUs from Google’s TPU Research Cloud (TRC).

## Appendix A Appendix

### A.1 Architectures

We provide in this section additional details on the autoencoder architecture.

##### Quantizer

For the FSQ quantizer, we use linear projections for encoding and decoding of the codes following the original work of [Mentzer et al., ].
For all the experiments, we fix the dimension of the codes to $d=6$. Then, we quantize each channel into $L$ unique values $L_{1},\ldots,L_{d}$ and refer to the quantization levels as $\mathbf{L}=[L_{1},\ldots,L_{d}]$. The size of the codebook $\mathcal{C}$ is given by the product of the quantization levels:
$|\mathcal{C}|=\prod_{j=1}^{d}L_{j}$. For the experiments with small codebooks with $|\mathcal{C}|=4096$, we use $\mathbf{L}=[4,4,4,4,4,4]$ and for large codebooks experiments, $|\mathcal{C}|=64000$, we take $\mathbf{L}=[8,8,8,5,5,5]$.

In more details the FSQ quantizer writes as Algorithm 1:

Figure: Algorithm 1 Finite Scalar Quantization

The product $W_{proj}\,\mathbf{z_{i}}$ facilitates scalar quantization within a lower-dimensional latent space, thereby defining a compact codebook. This is similar to the low dimensional space used for code index lookup in Yu et al. []. The subsequent up-projection operation then restores the quantized code to its original dimensionality.

##### Resampling

The cross-attention based resampling layer can be used for both down (resp. up) sampling, effectively reduces (resp. increases) the length of the sequence of embeddings is described in Algorithm 2.

Figure: Algorithm 2 Resampling Layer

##### Local Cross-Attention Masking

The encoder network used in this work preserves the notion of residue order, as defined in the primary structure of a protein (i.e. its ordered sequence of amino acids).
We do not provide our algorithm with information regarding the amino-acids.
However, we do include the order of the residues from which extract the atoms coordinates. When downsampling the encoder representations using a standard cross-attention operation, the resulting output virtually includes information from any other residue embeddings, irrespective of their relative position in the sequence.
In order to encourage the downsampled representation to carry local information, we propose to use local masks in the CrossAttention update of the resampling layer defined in Algorithm 2.
This will guide the network towards local positions (in the sequence) and prevent information from distant embeddings.
We illustrated the local masking in Figure 4, where only the direct neighbors are kept in the attention update.

Figure: Figure 4: Illustration of the local attention mechanism when using 2 neighbors for aggregation.
Refer to caption: x2.png

##### Decoder

For the decoder, we re-purpose the Structure Module (SM) of AlphaFold [Jumper et al., ]. In Jumper et al. [], a pair of 1D and 2D features (called the *single representation* and *pair representation* respectively) is extracted from the data by the Evoformer and fed to the SM to reconstruct the 3D structure.
Contrary to AlphaFold, we encode the structures with a set of 1D features - the sequence of discrete codes obtained after tokenization. Inspired by the OuterProductMean module of the Evoformer (Alg. 10 in SM of Jumper et al. []), we compute a pairwise representation of the structure by computing the outer product of the quantized sequence after projection, and concatenating the mean with the pair relative positional encoding, as defined in Algorithm 3.

Figure: Algorithm 3 Pairwise Module

Overall algorithm, the encoding and decoding processes write as described in Algorithm 4

Figure: Algorithm 4 Overall Algorithm Pseudo-Code

### A.2 Training Details and Metrics

##### Data Preprocessing

We consider the graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ consisting of a set of vertices - or nodes - $\mathcal{V}$ (the residues) with features $f_{V}^{1}\dots f_{V}^{|V|}$ and a set of edges $\mathcal{E}$ with features $f_{E}^{1}\dots f_{E}^{|E|}$.
For the node features, we use a sinusoidal encoding of sequence position such that for the $i$-th residue, the positional encoding is $\big{(}\phi(i,1)\ldots\phi(i,d)\big{)}$ where $d$ is the embedding size.
For the edge features we follow Ganea et al. []. More specifically, for each defined by residue $v_{i}$, a local coordinate system is formed by (a) the unit vector $t_{i}$ pointing from the $\alpha$-carbon atom to the nitrogen atom, (b) the unit vector $u_{i}$ pointing from the $\alpha$-carbon to the carbon atom of the carboxyl ($-\texttt{CO}-$) and (c) the normal of the plane defined by $t_{i}$ and $u_{i}$: $n_{i}=\frac{u_{i}\times t_{i}}{\|u_{i}\times t_{i}\|}$. Finally, setting: $q_{i}=n_{i}\times u_{i}$, the edge features are then defined as the concatenation of the following:

- •
relative positional edge features: $p_{j\rightarrow i}=(n_{i}^{T}u_{i}^{T}q_{i}^{T})(x_{j}-x_{i})$,
- •
relative orientation edge features: $q_{j\rightarrow i}=(n_{i}^{T}u_{i}^{T}q_{i}^{T})n_{j}$, $k_{j\rightarrow i}=(n_{i}^{T}u_{i}^{T}q_{i}^{T})u_{j}$, $t_{j\rightarrow i}=(n_{i}^{T}u_{i}^{T}q_{i}^{T})v_{j}$,
- •
distance-based edge features, defined as radial basis functions: $f_{j\rightarrow i,r}=e^{-\frac{\|x_{j}-x_{i}\|^{2}}{2\sigma_{r}^{2}}},r=1,2...R$ where $R=15$ and $\sigma_{r}=1.5$.

##### Compression Factor

We define the compression factor of Table 1 as the ratio between the number of bits necessary for encoding the backbone position atoms and the number of bits necessary to store the latent codes.
With positions stored as $64$-bit floats and latent codes stored as $16$-bit unsigned integers (since the highest number of latent codes is $64000<2^{16}$), we can write the compression factor as:

$$ $\text{Compression Factor}=\frac{N\times 4\times 3\times 64}{N\times 16/r}={r \times 48}$ (5) $$

### A.3 De novo Structure Generation

#### A.3.1 Training details

##### GPT hyperparameters

We use a standard decoder only transformer following the implementation of [Hoffmann et al., ] with pre-layer normalization and a dropout rate of 10% during training.
We follow Hoffmann et al. [] for the parameters choice with 20 layers, 16 heads per layers, a model dimension of 1024 and a query size of 64, resulting in a model with 344M parameters.

##### Training and Optimization

Given the tokenization of PDB training set, we respectively prepend a <bos> and append <eos> tokens and pad all sequences with <pad> token so that all sequences are of size $514$.
Hence, the maximum number of actual structural token per sequence is $512$.
The loss associated to <pad> tokens is masked out.

For the optimization, we utilize a ADAMw with $\beta_{1}=0.95$, $\beta_{2}=0.9$ and a weight decay of $0.1$. We also employ a learning rate scheduling linearly warming-up increasing the learning rate up to $5.10^{-5}$ for the first $1000$.
Following the seminal work of Radford and Narasimhan [], we employ embedding, residual and attention dropout with rate of $10\%$ rate.
We found the batch size to have a crucial importance on the optimization and adopt a batch size 65,792 tokens per batch.

The total number of actual structural tokens is $70$M.
In that perspective, and in line with works such as Hoffmann et al. [], we believe that leveraging large dataset of predicted structures such as AlphaFold (^3^33[https://alphafold.ebi.ac.uk/](https://alphafold.ebi.ac.uk/)) can provide significant improvements when training the latent generative model.

#### A.3.2 Structure generation metrics

##### Designability

We adopt the same framework than [Yim et al., , Trippe et al., , Wu et al., ] to compute the designability or *self-consistency* score:

- •
Compute 8 putative sequences from ProteinMPNN [Dauparas et al., ] with a temperature sampling of 0.1.
- •
Fold each of the 8 amino-acid sequences using ESMFold [Lin et al., ] without recycling, resulting in 8 folds per generated structure.
- •
Compare the 8 ESMFold-predicted structures with the original sample using either TM-score (scTM) or RMSD (scRMSD). The final score is taken to be the best score amongst the 8 reconstructed structures.

In Table 2, we report the proportion of generated structures that are said to be designable, *i.e* samples for which scTM>0.5 (or scRMSD$<2$ Å when using the RMSD).

##### Novelty

For the reference dataset, we use the s40 CATH dataset [Orengo et al., ], publicly available
[here](ftp://orengoftp.biochem.ucl.ac.uk/cath/releases/latest-release/non-redundant-data-sets/cath-dataset-nonredundant-S40.pdb.tgz).
In order to reduce the computation time, we first retrieve the top $k=1000$ hits using Foldseek [van Kempen et al., ]. We then perform TM-align [Y. and J., ] for each match against the targeted sample and report the TM-score corresponding to the best hit. A structure is then considered as novel if the maximum TM-score against CATH (cathTM) is lower than 0.5 and report the proportion of novel structures in Table 2

##### Diversity

Finally, we measure the diversity of the samples similarly to Watson et al. []. More specifically, the generated samples are clustered using an all-to-all pairwise TM-score as the clustering criterion and we observe the resulting number of structural clusters normalized by the number of generated samples. For a diverse set of generated samples, each cluster should be composed of only a few samples - or equivalently, the number of different clusters should be high. We use MaxCluster Herbert and Sternberg [] with a TM-score threshold of 0.6 as in Watson et al. [].

#### A.3.3 Sampling

##### Baselines

For each baseline methods, we follow a standardized process similar to that of Yim et al. [] to generate the testing dataset: we sample 8 backbones for every length between 100 and 500 with length step of 5: $[100,105,\ldots,500]$. We re-use the publicly available codes and use the parameters reported in Watson et al. [] and Yim et al. [] respectively.

##### Generating Structures with a Decoder-Only Transformer

Sampling for our method is a 2 steps process: first, sample a sequence of structural tokens from the trained prior described in Section A.3.1, then reconstruct the structures using the trained decoder. There are many ways to sample from a decoder-only transformer model [Vaswani et al., , Radford and Narasimhan, , Radford et al., , Holtzman et al., ]. We chose temperature sampling [Vaswani et al., ] with other alternative sampling strategies such as top-k [Radford et al., ] and top-p (or nucleus) sampling [Holtzman et al., ] resulting in little improvement at the cost of increased complexity. As showed in Vaswani et al. [], the temperature controls the trade-off between the confidence and the diversity of the samples. In order to tune the temperature, we sampled 2000 samples for each temperature between 0.2 and 0.8 in step of 0.2 and compute the designability score for these samples. As expected, the higher the temperature, the more varied the samples.
Indeed, the distribution of the proteins length, depicted in Figure 5, shows a greater diversity of higher temperatures.
On the other hand, with lower temperatures, the length distribution is more concentrated around few lengths. Similarly we can see than for lower temperatures, the samples closer to the modes of the length distribution (samples with length between 100 and approximately 300) achieve higher scTM-score (see Figure 5).
The results reported in Table 2 are obtained with a temperature of $0.6$, as it achieves a satisfying trade-off between designability and diversity.

Figure: Figure 5: Ablation of the temperature sampling. Left: Histogram of the generated structure length. Right: Designability score vs temperature sampling.
Refer to caption: extracted/5618651/images/temp_ablation_len_histo.png

Contrary to the baselines, our model learns the join distribution the length and the structures $p(s)=\int_{l}p(s,l)dl$ where the random variables $s$ and $l$ represent the structures and the length respectively.
Indeed, only the conditional distribution $p(s|l)$ is modeled by the diffusion-based baselines.
In Figure 5, we show the length distribution learned by the model.Since we can only sample from the joint distribution and not the conditional, we adopt the following approach: First, we sample 40,000 structures from the model (using a temperature of 0.6 as previously established).
We then bin the generated structures by length, with a bin width of 5 and bin centers uniformly distributed between 100 and 500, specifically: $[100,105,\ldots,500]$. Finally, we limit the maximum number of structures per bin to 10, randomly selecting 10 structures if a bin contains more than this number.

#### A.3.4 Additional Results

Figure: Figure 6: Designability score vs samples length. Left: scRMSD score for different structure lengths. Right: scTM score for different structure lengths.
Refer to caption: extracted/5618651/images/scrmsd_vs_len.png

Figure: Figure 7: Left: Novelty score for different structure lengths. Right: Novelty score *versus* designability score.
Refer to caption: extracted/5618651/images/novelty_vs_len.png

Figure: Figure 8: Distribution of the designability scores for novel domains (cathTM<0.5).
Refer to caption: extracted/5618651/images/scrmsd_distribution_cathtm_0.5.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Have you tried extending to side chain conformations?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 2
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.
soundness: 3
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
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.
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)

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
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.
soundness: 4
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.
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.
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.

---

confidence: 5
contribution: 1
initial_rating: 5
presentation: 2
questions: I'm wondering about the key advantages of this work compared to the listed works above. What's your unique contribution?
soundness: 2
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.
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.
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).

Meta-review: 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.

Decision: Reject