Learning the Language of Protein Structure
Authors: Jérémie DONA
Benoit Gaujac
Timothy Atkinson
Liviu Copoiu
Thomas Pierrot
Thomas D Barrett
Venue: Accepted by TMLR ()

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-4 \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.

Under review as submission to TMLR
Learning the Language of Protein Structure
Benoit Gaujac 11 Jérémie Donå
Liviu Copoiu', Timothy Atkinson', Thomas Pierrot' and Thomas D. Barrett
"Equal contributions: {b (b.gaujac, j.dona)@insttddeppcom 'InstaDeep
Reviewed on OpenReview: http:///ppanrree...aeee////tttet
Abstract
Representation learning and de noto generation of proteins pivotal computational biology lasks. Whilst natural language processing (NTP) techniques have proven highly effective for protein sequence modelling, structure modelling presents complex challenge, primarily due to its continuous and three-dimensional nature. Motivated by this discrepancy, we introduce an approach using a quuttr-quantied 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 codebook ranging from 432 to 64000 tokens, achieving high-fidelity reconstructions with backbone root mean square deviations (RMSD) of approximately 1-4 A. To demonstrate the efficacy of our learned representations, we show that 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 protein design.
1 Introduction
The application of machine learning to large-scale biological data has ushered in transformative era in computational biology, advancing both representation learning and de novo generation. Particularly, the inlegration of machine learning in molecular biology has led to significant breakthroughs (Sapoval et al. 2022 Chandra et al. 2023 Khakzad et al. 2023 Jänes and Beltrao 2024). 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 siguificant modelling elforts.
Currently, the deep learning landscape is increasingly converging towards a unified paradigm centered around attention-based architectures (Vaswani et al. 2017 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 sequence-moddeling problems has proven highly effective. Indeed, transformer-based models are leading methods in many machine learning domains, including image representation (Radford et al. 2021),
image (Chen et al. 2020 Chang et al.] 2022) and audio generation (Ziv et al. 2024). and for reinforcement learning (Chen et al.] 2021: Boige et al.) 2023). This trend has greatly benefited sequence- segueece-bbsed biological models, allowing for the direct application of NLP mothodologics like GPT (Radford and Narasimhan, 2018
and BERT (Devlin et al. 2019) with notable success (Ferruz el. al. 2022 Lin et al. 2023b).
In particular, large multi-modal models (LMMs), leveraging transformer backbones, are emerging as key tool with applications in various fields such as: ubiquitous AI (GPT4 (Achiam et al. 2023 LLaVA (Liu
el al. 2024b), Gemini (Gemini et al. (202)), Flamingo (Alayrac et al. 2022 lexl-conditioned generation of images (Parti (Yu et al., 2022b), Muse [Chang et al. 2023)) or sounds (MusicGen (Copet et al.
2023 ); and
even reinforcenent learning (Reed et al. 2022). LMMs are also instantiated in biological settings, such as
medicine (MedLlama (Xie et al.] 22244) Med-Gemini (Yang al. (024bb Med-PaLM (Tu et al. 2024)) and
genomics (ChatNT (Richard et al. 2024 ) Core to all these L.MSs is the use of pre-trained encoders

---

Under review as submission to TMLR
mechanism for combining modalities in sequence space. For instance, LLaVA (Liu al. 2024b and Flamingo 2022) used pre-trained vision encoders, VIT-I./14 (Radford et al. 2021 and Normalizer Free ResNet (NFNet) (Alayrac et al. 2020) respectively, while Copet et al. (2023) leverages the codec of Défossez ict al. 2023 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.
[Alayrac et al.,
Despite these advances in related domains, structure-based modeling of biological data such as proteins remains formidable challenge. Unlike sequences, protein structures are inherently three-dimensional and contimuous, 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-neural network encoders (Dauparas et al. 2022 Krapp et al. 2023) and structurally aware modules such as those AlphaFold (Jumper et al. 2021). Morcover, generative modelling of structures is typically performed with methods designed for continuous variables; such as diffusion (Watson eet al. 2023 Yim et. al., 2023 and flow matching (Bose al. 2024). 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 quantized reprosentations of protein structures enabling to efficiently leverage sequence-based language models. The key objectives of 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 scamless intogration 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 contimuous 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 learned discrete representation, typically within the range of 1-4 Angströms.
We aim to minimize the reconstruction error of the
Our contributions are threefold. First, we introduce a series of quantized autocncoders 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 gemerates novel, diverse, and structurally viable protein structures. We release all code at model at htt.gs//wwgggugggeo..co///aaaapp hrtpap/pee.ppppooaooocccocoooonr
wtwpr/ttcctttttttteettree//////e /|nd the trained
2 Method
2.1
Protein Structure Autoencoder
Our objective is train an autoencoder that maps protein structures to and from a discrete latent space of sequential codes. Following prior works (Yim al. 2023 Wu al. 2024), we consider the backbone atoms of protein, N C, C 0, 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 p E IR x4x3 to latent representation ž = [2,.... xx where r is downsampling ratio, controlling the size of the representation. Note that each element 2; can only take a finite number of values, with the collection of all possible values defining a codebook C.
A schematic overview of the our autoencoder is depicted in Figure In this section, we focus on the three components of the model; the encoder eg extracting a set of N embeddings of dimension denoted z R
2

---

Under review as submission to TMLR
ENCODS
QUANTZZE
64
ASCOVETRUCT
15
eg
S
dæ
28
FAPE Loss
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.
the quantizer ds that discretizes to obtain quantized representation ž, and the decoder dy that prodicts structure P E R x1x3 from 2. The learnable parameters are respectively denoted (0, o, 15) and the learning setting summarizes as:
2 P.
Encoder The encoder maps the backbone atoms positions E Nx4x3, to a continuous downsampled continuous representation z IR. xe where r is the downsampling ratio:
co p E Nx4x3 + E R
X
(1)
Note that when the downsampling ratio is set to 1, each component 2 IR 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 YYngg
et al. 2024a Boget et al. 2024h. Inverse folding (Ingraham et al. 2019 Dauparas et al. 2022 is another example, that aims at estimating the sequence of amino acids corresponding Recently, ProteinMPNN has shown remarkable capacity at the inverse folding task. We follow their design Paassag-Passing Neural Network (MPNN) (Dauparas al.
given a protein structure.
choices and parameterize our encoder using 2022). In addition, introduce a cross-attention mechanism, detailed in Appendix A. and Algorithm 2 enabling us to effectively compress the representation of the structure by downsampling ratio r. We maintain locality using a custom attention masking scheme in the downsampling layer (see Appendix A. ensuring that each downsampled node aggregates information from a small number of neighboring nodes in the original sequence space.
MPNN operates on a graph consisting of set of vertices and edges. We follow Dauparas et al. (2022 Ganea et al. (2022) and set as initial node features a positional encoding that reflects the residue's is ordering within the sequence, while for the edge features, we use a concatenation of pairwise distance leatures, 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 rolation 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 Appendix A.2
The quantizer plays crucial role in our work by discretizing all continuous latent representa- tions into a sequence of discrete codes. Traditional methods typically involve direct learning of the codebook (van den Oord el al.; 2017 Razavi et al. 2019). However, in line with the literalure,, we suffered several drawbacks associated with these approaches. Indeed, explicit vector quantization particularly expensive as it involves computation of pairwisc distances, particularly problematic for long sequences and large codebooks. Morcover, we faced training instabilities that arise from both the bias in the straight-througe estimator and the under-utilization of the codebook capacity, often referred to codebook collapse, make learning
Quantization
3

---

Under review as submission to TMLR
discretized latent representation a hard optimization problem (Huh al 2023 Takida et al. 2024 To
address these challenges, we leverage the recent. Finite Scalar Quantization (FSQ) framework (Menizer et al. 2024), which effectively resolves the aforementioned issues, notably by reducing straight-through gradient estimation crror.
FSQ learns discrete latent space by rounding a bounded low-dimensional encoding of the latent representation. Consider ž R,, the jt element of the continuous encodings z, and the quantization levels L [L....., La] E N Its discretized counterpart, denoted as 2 2d, typically with d d 8, is defined by:
L (Le
3 = round
@ tanh(W2z)
(2)
where O the element-wise multiplication and W E _gdxe is a projection weight matrix (see Appendix A.1 for more details). In doing so, each quantized vector 2 can be mapped to an index in (1..... IT L This implicitly defines a codebook by indices, where each index associated unique combination of the each dimension values. In our implementation, the quantized representation ž = [žo..... ................................ 2N] is then projected back to the latent space with dimension e (line of Algorithm D Optimization is then conducted using a ttraight gradient estimator. Equation (2) ensures that the approximation error introduced by the straight-throug estimator is bounded.
Decoder The decoder module of our framework is tasked with estimating structure from the latent quantized representation 2:
de : ã E TR xe p R Nx4x3
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 conformation of protein is estimated given its primary sequence of amino-acids. Juupper et al. (2021) successfully tackled this task proposing a novel architecture for point cloud estimation from latent sequence of embeddings. Therefore, we use AlphaFold-2 structure module to parameterize our decoder and learn its parameters from scratch.
Specifically, the structure module of AlphaFold-2 parameterizes a point cloud using frames. A frame is defined by a tuple = (R,t) where R is the frames' orientation (i.e. rotation matrix) and the frame center (í.e. vector defining the translation of the center of the frame). The origin of the frame is set to the C. carbon, and the orientation defined using the nitrogen and the other carbon atom. For thorough and mathematical description, we defer to Jumper et al. 22021). However, note that AlphaFold-2's structure modulc expects both a per-residue representation (s))ign and a pairwise representation (kıj)a.sgN between residues and j. The per-residue representation S; is constructed by mirroring the cross-allention 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 (kij) used for reconstruction described in Algorithm 3 and Appendix A.1
2.2 Training
Objective The Frame Align Point Error (FAPE) loss, introduced in Jumper et al. (2021). 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 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 T; and ground-truth atom position Xj expressod in Ti, and their respective predictions TP and xf. the FAPE Loss is defined as:
LEAPE = |TP (XX'((()) - T-- (xx)]l
(4)
In Equation 4 the predicted coordinates of atom j (xf), expressed in the predicted local frame of residue (TP), 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.
4

---

Under review as submission to TMLR
Dataset We use approximately 310000 entries available in the Protein Data Bank (PDB) (Berman et al. 2000) as training data. The presence of large groups of proteins causes imbalances in the data set, with many proteins 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 2 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% of the clusters for training and use the remaining as test set. Amongst these 10% withheld protein-structure clusters, 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 el. al. (2022) and utilize the swish activation function. The graph sparsity set to 50 neighbors per residue. When the downsampling ratio is r > 1, the resampling operation consists of a stack of 3 resampling layers as described in Algorithm 2 the initial queries being defined as positional encodings. We strictly follow the implementation of AlphaFold-2 (Jumper el. al., 2021 regarding the structure module and use 6 structure layers.
Optimization and Training Details The optimization is carried out using AdamW (Loshchilov and Hutter 2019) with 81 = 0.9, B2 = 0.95 and a weight, decay of 0.1. We use a learning rate warm-up scheduler, progressively increasing the learning rate from 10 to 10-3 over the first 1000 steps, and train the model for 100 epochs on 8 TPU v4-8 with a batch size of 128. With such hyperparameters, the autoencoder model has 4.5. parameters, and the training lasts 32 hours on TPU v4-8, which amounts to total number of FLOP comprised between 7 x 1029 and 1020.
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 downstream tasks, by effectively training de novo generative model for protein structures using 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. To further assess the compression capacity of our model, we also train two additional quantized autoencoders with smaller codebook sizes, K = 432 and K = 1728.
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 (Zhang and Skolnick 2005) is a normalised measure of how similar two structures are, with a score denoting the structures are identical. For context, two structures are considered to have similar fold when their TM-score exceeds 0.5 (Xu and Zhang 2010) and a RMSD below 2 usually seen as approaching experimental resolution.
Results Our results are summarised in Table 1 We find that with codebook of K 64000 and a downsampling of r = 1; our average reconstruction has 1.22 Å RMSD and TM-score of 0.96. For comparison,
The cluster size is readily available in the PDB data set: hwtp://wwwrcbboo. http://www.rrrbooorr//////////// grouping structurre
sequence-based-ee based-clustering
5

---

Under review as submission to TMLR
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: Avcrage Test set reconstruction results of our discrete auto-cncoding method for several down- sampling ratios and (implicit) codebook sizes. For CASP-15 we report the median of the metrics due to the limited dataset size. Note that RMSD below 2 is considered of the order of experimental resolution and two proteins with TM-score > 0.5 considered to have the same fold. The compression defined the number of bits necessary to store the N 4 x 3 backbone positions divided by the number of bits necessary to store the N tokens multiplied by log, # Codes)
Results
CASP15
Downsampling Ratio
Number of Codes
Compression Factor
RMSD (4)
TM-Score (1)
RMSD (L)
TM-Score (1)
2.09 A
1.75 À
0.89
1 728
71
1.79 A
0.93
1.33 Á
0.94
4096
64
1.55 A
0.94
1,25 Á
0.94
64000
1.22 A
0.96
0.94 Á
0.97
without quantization
0.97 A
0.98
1.07 A
0.93
4096
128
2.22 A
0.90
1.73 A
0.89
64000
1.95 A
0.92
1.82 À
0.90
2
without quantization
1.45 A
0.95
1.44 A
0.90
4096
256
4.10 A
0.81
2.79 A
0.77
64000
192
2.96 Á
0.86
2.55 Á
0.80
4
without quantization
2.19 A
0.91
1.98 Á
0.84
Table 1 clearly indicates that increasing the downsampling factor decreasing the codebook size correspond- ingly impacts the reconstruction accuracy. This is expected as it essentially enforces greater compression in our autoencoder leading to loss of information nevertheless in all cases we find that the achievable reconstruction performance is still within a few angstrom with TM-scores clearly exceoding 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 subset of the codes being used by the trained model), an issue well known and documented with other quantization methods (Huh et al. 2023 Takida et al., 2024h. This is also noticeable in Figure 2 that shows that, given downsampling ratio, larger codebooks effectively decrease the reconstruction errors. Finally, comparing to continuous autoencoders, our learned quantization demonstrates significant information compression at the expense of only small decrease in the reconstruction precision of the order of 0.5 - 1 Å. Additionally, we provide in Appendix A.3 detailed distribution of the RMSD and TM-Scores for both CASP-15 and the held-out test set. Notably, we can see that increasing the downsampling ratio, or decrcasing the codebook size tends thickens the right tail (resp. the left) of the distribution for the RMSD (resp. the TM score).
The reconstruction performance is illustrated in Figure 3 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 the protein such as the a-helices and 8-sheets.
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 model trained in our latent space as a paradigmatic demonstration of our tokenizers utility. This is just because generative models for de novo protein design are of great interest for drug discovery enabling rapid in silico exploration

---

Under review as submission to TMLR
5
0.8
3-
0.6
0.4
1
codes 432
codes 1728
cados 4096
codes 64000
rodes 432
codes - 1728
codes 4096 codes 64000
Figure Evolution of the RMSD (left) and TM-score (right) distribution with the codebook size for a downsampling ratio of 1 CASP-15 data.
Length: RMSD: 1.343
Length: 409 RMSD: L804
Length: RMSD: 1.398
Reconstruction
Target
(y-axis)
(y-axis)
90 (y-axis)
Figure 3: Visualisation of the model reconstruction (blue) super-imposed with the original structures (green) for a downsampling factor of r = 2 and = 64000 codes (see Table for detailed results). 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.
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 datasct (defined in Scction 2.2 full details on datasct preparation for this experiment can be found in Appendix A.4.1) using downsampling factor of and codebook with 4096 codes (first line of T This choice is motivated by the trade-off wetn reconstruction performance (this

---

Under review as submission to TMLR
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 assess the designability of the generated structure. Note that while high novelty score is desirable, structures that are too far from the reference dataset also be sign of unfeasible proteins.
Method
scTM (t)
scRMSD (t)
Novelty
Diversity
Ours
87.22 %
61.83%
23.3%
60.11 %
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%
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- of-the-box decoder-only transformer model with 20 layers, 16-heads and an embeddding dimension of 1024 (344M parameters) on a next- next-token-predee trenc-criiininn task on the training split. This decoder-only model is used to generate new soquences 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: designabilily, novelly and diversily. Designability, or self-consistency, uses ProteinMPNN (Dauparas et al 2022 to predict potential sequence for the generated structure and refolds it using ESMFold (Lin et al. 2022 to report the structural similarity between the original and redesigned structure. We follow the literature and report the proportion of designed proteins with celf-consistent TM-score (scTM) above 0.5 respectively the number of designed proteins with self-consistent RMSD (scRMSD) below 2A. The rationale is that generated structures should be sufficiently natural that establishod methodologies recognisc them and agroe on the fundamental biophysical properties.
Moreover, 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 structure as novel if its maximum TM-score against the reference dataset below 0.5. As is commonly done, only report the proportion of novel samples in our generated dataset. Overall, we follow Yim et al. (2023) for the implementation of these metrics and refer the reader to Appendix A.4.2 for more details on the validation pipeline.
Basclines To gange how our structure generation GPT model fares against specifically designed protcin design methods, we compare our model with Fifm (Yim et al. 2023) and RFDiffusion (Watson et al. 2023). Both use bespoke SE(3) diffusion models specifically designed and trained the structure generation task. Note that, the state-offthe-art RFDiffusion leverages the extensive pre-training of RoseTTAFold (Baek et al., 2021) on 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 the tokenized representation learned with our method for protein design.
The results for the different metrics given in Table 2 with the sampling strategy used for each method described in Appendix A.4.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 on par with a method specific to protein design such as FifmeDiff. For comparison, and to set an upper bound of what to expect in terms of the designability scores, we compute self-consistency and diversity metries 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
Results
8

---

Under review as submission to TMLR
L: 427 / scRMSD: 1.38
L: 364 / scRMSD: 1.27
L: 463 / scRMSD: 1.45
ESM-prediction Sample
Figure 4: Visualisation of generated samples (green) super-imposed with their sell-consistent ESM-predicted structures (blue).
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 discussed in Appendix [A.4.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 20 we can see that while Fifm 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 20 we compute the number of structures with cathTM < 0.5 (novel) and 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 Fiam and 58.58 % (379 out of 647) for RFDiffusion. We provide additional results and analyses in Appendix A.4.4
In Figure 4 we visualize random samples from our model superimposed with the predicted ESM structure used in the designability metric. We first notice that the generated samples exhibit diverse structure, with non trivial secondary elements (a-helices and B-sheets). We also note that the ESM-predicted structures are closely aligned with the original samples, showing the designability aspect of the generated structures.
4 Related Works
Learning from Protein Structures. Learning from protein structure is thriving research field that
encompasses wide variety of tasks crucial task. For instance, inverse folding (Ingraham et al. 2019 Hsu
ct al.
2022
Dauparas et al. 2022 Jing et al.] 2021) or binding site estimation (rapp et al. 2023 Ganea
et al.
2022
are critically enabling for drug design (Scott et al. 2016). Others (Consortium] 2006 Robinson
et al.
2023
Kucera et al. 2023; Gligorijovíć al. (202]) 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. (2022) 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 VAE latent space, it often yields subpar sample quality,
see Kingma et al. (016)).
Discrete representation learning for protein struct has recently garnered increasing altention. Foldseek van Kempen et al.) 2024) introduces a quantized autoencoder to encode local protein geometry, demonstrating success in database search tasks. However, as it focuses solely on local features at the residue level, it lacks the capacity to provide global representation of protein structures. This limitation restricts its application in tasks like structure generation or binding prediction, where global information is critical (Krapp al.
9

---

Under review as submission to TMLR
2023).
Building on the 3Di-alphabet introduced by Foldseck, Su et al 2024 Heinzinger et al (2023) propose
structure-aware protein language models that integrate structure tokens with sequence tokens. Additionally, Li et al. (2024) combines a structural autoencoder with K-means clustering applied to the latent representation of a fixed reference dataset.
2017) for protein Lin al. (2023a) explores discrete structural representation learned by such VQ-VAEs (van den Oord et al.
More closely related to our work, Gao et al. (2023) adapts VQ-VAF (van den Oord al.,
structures. However, its limited reconstruction performance constrains its applica
2017p, while Liu
et al. (023) trains diffusion model on the discrete latent space derived from this approach. Similarly, and
concurrently with our work, Liu et al. (2024a) combines finite scalar quantization (Mentzer et al. 2024) with a specialized transormer-baseee autoencoder for proteins, RNA, and small molecules.
Very recently, efforts have emerged to combine quantized structural representation with discrete sequence representation, enabling multimodal generative models trained on a joint discrete latent space. FoldToken (Gao et al. 2024a b), is a concurrent approach that shares conceptual similarities with our work but differs in key methodological and application-oriented aspects. FoldToken employs joint quantization of sequence and structure, enabling integration across modalities, whereas our method focuses exclusively on structural information. This decoupling allows for mode-specific pretraining, aligning with strategies from subsequent works (Hayes et al., 2024 Lu et al.) (2224). Methodologically, FoldToken introduces a series of improvements to existing VQ methods (van den Oord et al. 2017) aimed at enhancing reconstruction accuracy; whereas gradient estimation gap inherent to VQ methods while improving codebook utilization. Furthermore, while Fold'Token primarily emphasizes backbone inpainting and antibody design (Gao et al.. 2024b our work considers de novo generation of complete structures.
we adopt the FSQ framework, which reduces the straight-throut
Generation of Protein Structures. A substantial body of literature addresses the challenge of sampling the protcin structure space. Numerous studics have advocated for the use of diffusion- models (Wu
et al. 2024 Yim et al. 2023 Walson et al. 2023 or flow matching techniques (Bose et al. 2024 While
many of these works, such as those by Yim et al. (2023); Watson et al. (2023) Bose et al. 2024). employ
complex architectures to ensure invariance to rigid-body transformations, Wu et al. (2024) opted for an
alternative parameterization directly preserving symmetries allowing the authors to capitalize on conventional
architectures, yet working only on small protein crops. Recently, [Wang et al. (2024b) employs lookup-free quantizer ( (Yu et al. 2024) as a structure tokenizer and trains a diffusiin protein language model @Wang
et al. 2024a) on the concatenated sequence and structure tokens. Other works (Hayes et al. 2024 Lu el al.
2024 simply uses VQ-VAEs (van den Oord et al. 2017) to tokenize the structures and train large language models (LLM) on the combination of sequence and structure tokens.
5
Conclusion
This work demonstrates a methodology for learning 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-modellingg techniques to protein structures.
A prerequisite for such development the expressiveness of the tokenized representation, which must capture the necessary information to enable high-fidelity reconstruction 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 increased tokens usage provide straightforward recipes for improving the reconstruction accuracy by reducing the information bottleneck of the autoencoder.
The first step towards sequence-basee 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- model underlines the promise of this paradigm. While a simple GPT does not yet match seminal approaches like RFDiffusion, it is important to recognize
10

---

Under review as submission to TMLR
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.
6 Broader Impact
The discrete protein-structur tokenizer and accompanying generative models presented in this work can advance computational protein engincering in several ways. By converting atomic coordinates into unified token vocabulary (432 64 000 tokens), the framework allows researchers to apply the full range of transformer- based language-model techniqus-pre-trriiingg pre-training, fine-tuning, and prompt-driven structural data. Because the models and weights are openly released in reproducible configurations, laboratories with limited computational resources can replicate our baselines, while larger centres can exploit the highest- fidelity variants. shared structural vocabulary also facilitates interoperability: into existing sequence predictors, docking pipelines, function classifiers without redesigning model inputs, thereby lowering experimental (urnaround times for enzyme, antibody, and other protein-based therapeutic design. Potential risks must also be acknowledged. Any generative protein model could, in principle, be misapplied to design toxic or immuno-evasive molecules. In silico structures may be over-interpreted if they are not validated experimentally, leading to misplaced confidence in downstream applications.
to
can be incorporated
References
J. Achiam, S. Adler, Agarwal, L. Ahmad, Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv, 2023.
J.-B. Alayrac, A. Recasens, Schneider, Araudjelović, J. Ramapuram, J. De Fauw, Smaira, Dieleman, and A. Zisserman. Self-super multimodal versatile networks. In Advances in Neural Information Processing Systems, 2020.
J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, el al. Flamingo: visual language model for few-shot learning. In Advances in Neural Information Processing Systems, 2022.
M. Back, F. DiMaio, I. Anishchenko, J. Dauparas, S. Ovchinnikov, G. R. Lee, J. Wang, Q. Cong, L. N. Kinch, R. D. Schaeffer, C. Millán, H. Park, C. Adamss C. R. Glassman, A. DeGiovanni, J. H. Pereira, A. V. Rodrigues, A. A. van Dijk, A. C. Ebrecht, D. J. Opperman, T. Sagmeister, C. Buhlheller, T. Pavkov-Keller, M. K. Rathinaswany, U. Dalwadi, C. K. Yip, J. E. Burke, K. C. Garcia, N. V. Grishin, P. D. Adams, R. J. Read, and D. Baker. Accurate prediction of protein structures and interactions using a threc-track ncural network. Science, 2021.
M. Berman, J. Westbrook, Z. Feng, G. Gilliland, T. N. Bhat, H. Weissig, Shindyalov, and P. E. Bourne. The protein data bank. Nucleic Acids Research, 2000.
Y. Boget, M. Gregorova, and A. Kalousis. Discrete graph auto-encoder. Transactions on Machine Learning Research, 2024.
R. Boige, Y. Flet- Berliac, A. Flajolet, G. Richard, and T. Pierrot, PASTA: Pretrained action-state transformer agents. In NeurIPS 2023 Foundation Models for Decision Making Workshop, 2023.
J. Bose, Skdound--aadegh, Fatras, G. Huguet, J. Rector-Brooks, C.-H. Liu, A. Nica, M. Korablyov, M. M. Bronstein, and A. Tong. -stocsstcc flow matching for protein backbone generation. In International Conference on Learning Representations, 2024.
A. Chandra, L. Tünnermann, T. Löfstedt, and R. Gratz. Transformer basod deep learning for predicting protein properties in the life sciences. eLifee, 2023.
11

---

Under review as submission to TMLR
H. Chang, H. Zhang, L. Jiang, C. Liu, and W. T. Freeman. MaskGIT: Masked generative image transformer. In Conference Computer Vision and Pallern Recognilion, 2022.
H. Chang, H. Zhang, J. Barber, A. Maschinot, J. Lezama, L. Jiang, M.-H. Yang, K. Murphy, W. T. Freeman, M. Rubinstein, Y. Li, and D. Krishnan. Muse: Text-to-image generation via masked generative transformers. In Proceedings of the 40th International Conference on Machine Learning. 2023.
L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling, arXiv, 2021.
M. Chen, A. Radford, R. Child, Wu, H. Jun, D. Luan, and I. Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning, 2020.
Gene Ontology Consortium. The gene ontology (GO) project in 2006. Nucleic Acids Research, 2006.
J. Copet, F. Kreuk, I. Gat, T. Remez, D. Kant, G. Synnaeve, Y. music genoration. In Advances in Neural Information Processing Systems, 2023.
Adi, and A. Defossez. Simple and controllable
J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. Ragotte, L. Milles, B. Wicky, A. Courbet, R. Haas, N. Bethel, P. Leung, T. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. Bera, N. P. King, and D. Baker. Robust deep learning-bass based protein sequence design using ProtcinMPNN. Science, 2022.
Défossez, Copet, G. Synnaeve, and Y. Adi. High fidelity neural audio compression. Transactions on Machine Learning Research, 2023.
J. Devlin, M.W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. Tn North American Chopler of the Associalion for Compulational Linguistics, 2019.
R. Eguchi, C. Choe, and P.-S. Huang. Ig- VAE: Generative modeling of protein structure by direct 3D coordinate generation. PLoS Computational Biology, 2022.
N. Ferruz, S. Schmiddt, and B. Höcker. Protgpt2 is deep unsupervised language model for protein design. Nature Communications, 2022.
O.-E. Ganea, X. Huang, C. Bunne, Y. Bian, R. Barzilay, T. S. Jaakkola, and A. Krause. Independent models for end-to-end rigid protein docking. In International Conference on Learning Representations, 2022.
Z. Gao, C. Tan, and S. Z. Li. VQPL: Vector quantized protein language. arXiv, 2023.
Z. Gao, C. Tan, and S. Z. Li. Foldtoken2: Learning compact, invariant and generative protein structure language. bioR 2024a.
Z. Gao, C. Tan, J. Wang, Y. Huang, L. Wu, and S. Z. Li. Foldtoken: Learning protein language via vector quantization and beyond. arxiv, 2024b.
Team Gemini, R. Anil, S. Borgeaud, Y. Wu, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv, 2023.
V. Gligorijevié, P. D. Renfrew, T. Kosciolek, J. K. Leman, D. Berenberg, T. Vatanen, C. Chandler, B. C. Taylor, I. M. Fisk, H. Vlamakis, R. J. Xavier, R. Knight, K. Cho, and R. Bonneau. Structure-basee protein function prediction using graph convolutional networks. Nature Communications, 2021.
T. Hayes, R. Rao, H. Akin, N. J. Sofroniew, D. Oktay, Z. Lin, R. Verkuil, V. Q. Tran, Deaton, M. Wiggert, R. Badkundri, I. Shafkat, Gong, A. Derry, R. Molina, N. Thomas, Y. Khan, C. Mishra, Kim, L. J. Bartie, M. Nemeth, D. Hsu, T. Sercu, S. Candido, and A. Rives. Simulating 500 million years of evolution with language model. bioRriv, 2024.
12

---

Under review as submission to TMLR
M. Heinzinger, K. Weissenow, J. Gomez Sanchez, A. Henkel, M. Steinegger, and B. Rost. Prostt5: Bilingual language model for protein sequence and structure. bioRrin, 2023.
A. Herbert and M. Sternberg, MaxChuster: a tool for protein structure comparison and chustering, arXiv, 2008.
J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. Rae, and L. Sifre. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems, 2022.
A. Holtzman, J. Buys, L. Du, M. Forbes, and Y. Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020.
C. Hsu, R. Verkuil, Liu, Z. Lin, B. Hie, T. Sercu, A. Lerer, and A. Rives. Learning inverse folding from millions of predicted structures. In International Conference on Machine Learning, 2022.
M. Huh, B. Cheung, P. Agrawal, and P. Tsola. Straightening out the straight-through estimator: overcoming optimization challenges in vector quantized networks. In International Conference on Machine Learning, 2023.
J. Ingraham, V. Garg, R. Barzilay, and T. Jaakkola. Generative models for graph-based protein design. In Advances in Neural Information Processing Systems, 2019.
J. Jänes and P. Beltrao. Deep learning for protein structure prediction and design progress and applications. Moleculor Systems Biology, 2024.
B. Jing, S. Eismann, P. Suriana, R. J. L. Townshend, and R. Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2021.
J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakoo,, R. Bates, A. Zídck, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, A. J. Ballard, A. Cowie, B. Romera-Paredes, S. Nikolov, R. Jain, J. Adler, T. Back, S. Petersen, D. Reiman, E. Clancy, M. Zielinski, M. Steinegger, M. Pacholska, T. Berghammer, S. Bodenstein, D. Silver, O. Vinyals, A. W. Senior, K. Kavukcuoglu, P. Kohli, and D. Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 2021.
H. Khakzad, 1. Igashov, A. Schneuing, C. Goverde, M. Bronstein, and B. Correia. A new age in protein design empowered by deep learning. Cell Systems, 2023.
D. P. Kingma, T. Salimans, R. Jozefowiez, X. Chen, I. Sulskever, and M. Welling. Improved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems, 2016.
L. F. Krapp, L. A. Abriata, F. Cortés Rodriguez, and M. Dal Peraro. Pesto: parameter-free geometric deep learning for accurate prediction of protein binding interfaces. Nature Communications, 2023.
T. Kucera, C. Oliver, D. Chen, and K. Borgwardt. Proteinshake: Building datasets and benchmarks for deep learning protein structures. In Advances in Newral Information Processing Systems Datasets and Benchmarks Track, 2023.
M. Li, Y. Tan, X. Ma, B. Zhong, Z. Zhou, H. Yu, W. Ouyang, L. Hong, B. Zhou, and P. Tan. Prosst: Protein language modeling with quantized structure and disentangled attention. Rio 2024.
X. Lin, 7. Chen, Y. Li, X. Lu, C. Fan, Z. Cao, S. Feng, Y. Q. Gao, and J. Zhang, Protokens: A machine-learned language for compact and informative encoding of protein 3d structures. bioRriv, 2023a.
Z. Lin, H. Akin, R. Rao, B. Hic, Z. Zhu, W. Lu, A. dos Santos Costa, M. Zazal-Zarandi, T. Sercu, S. Candido, al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRriv 2022.
13

---

Under review as submission to TMLR
Z. Lin, H. Akin, R. Rao, B. Hic, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, and A. Rives. Rvolutionary-scale prediction of atomic-level protein structure with a language model. Science, 2023b.
A. Liu, A. Elaldi, N. Russell, and O. Viessmann. Bio2token: All-atom tokenization of any biomolecular structure with mamba. arXiu, 2024a.
H. Liu, C. Li, Q. Wu, and Y. J. Lee. Visual instruction tuning. In Advances in Neural Information Processing Systems, 2024b.
Y. Liu, L. Chen, and H. Liu. Diffusion in a quantized vector space generates non-idealized protein structures and predicts conformational distributions. bioRriv, 2023.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
J. Lu, X. Chen, S. Z. Lu, C. Shi, Н. Guo, Y. Bengio, and J. Tang. Structure language models for protein conformation generation. arXie, 2024.
F. Mentzer, D. Minnen, E. Agustsson, and M. Tschannen. Finite scalar quantization: VQ- VAE made simple.
In International Conference on Learning Represenlations, 2024.
C. A. Orengo, A. D. Michic, S. Jones, D. T. Jones, M. B. Swindells, and J. M. Thornoon. CATH hierarchic classification of protein domain structures. Structure, 1997.
A. Radford and K. Narasimhan. Improving language understanding by generative pre-training, arXiv, 2018.
A. Radford, Wu, R. Child, D. Luan, D. Amodei, and Sutskever. Better language models and their implications. Technical report, OpenAI, 2019.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021.
A. Razavi, A. van den Oord, and O. Vinyals. Generating diverse high-fidelity images with VQ-VAF-2. In Advances in Newral Information Processing Systems, 2019.
S. Reed, K. Zolna, E. Parisotto, S. Gomez Cohmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, T. Eccles, J. Bruce, A. Razavi, A. Edwards, N. Heess, Y. Chen, R. Hadsell, O. Vinyals, M. Bordbar, and N. de Freitas. generalist agent. arXiv, 2022.
G. Richard, B. P. de Almeida, H. Molle,Toree C. Blum, L. Hexemer, P. Pandey, S. Laurent, M. Lopez, A. Laterre, M. Lang, et al. Chatnt: A multimodal conversational agent for dna, rna and protein tasks. RioR,ie, 2024.
L. Robinson, T. Atkinson, L. Copoiu, P. Bordes, T. Pierrot, and T. Barrett. Contrasting sequence with structure: Pre-training graph representations with PLMs. In NewrIPS 2023 Al for Science Workshop, 2023.
N. Sapoval, A. Aghazadeh, M. G. Nute, D. A. Antunes, A. Balaji, R. Baraniuk, C. J. Barberan, R. Dannenfelser, C. Dun, M. Edrisi, et al. Current progress and open challenges for applying deep learning across the biosciences. Nature Communications, 2022.
D. F. Scott, A. R. Bayly, C. Abell, and J. Skidmore. Small molecules, big targets: drug discovery faces the protein protein interaction challenge. Nature Reviews Drug Discovery, 2016.
J. Su, C. Han, Y. Zhou, J. Shan, X. Zhou, and F. Yuan. Saprot: Protein language modeling with structure- aware vocabulary. In International Conference on Learning Representations, 2024.
Y. Takida, Y. Ikemiya, T. Shibuya, K. Shimada, W. Choi, C. H. Lai, N. Murata, T. Uesaka, K. Uchida, W.-H. Liao, and Y. Mitsufuji. HQ-VAE: Hierarchical discrete representation learning with variational bayes. Transactions on Machine Learning Research, 2024.
14

---

Under review as submission to TMLR
B. Trippe, J. Yim, D. Tischer, D. Baker, T. Broderick, R. Barzilay, and T. S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem. In International Conference Learning Representations, 2023.
T. Tu, S. Azizi, D. Driess, M. Schackermann, M. Amin, Chang, A. Carroll, C. Lau, R. Tanno, Ktena, al. Towards generalist biomedical AI. NEJM AI, 2024.
A. van den Oord, O. Vinyals, and K. Kavukcuoglu. Neural discrete representation learning. In Advances in Nearal Information Processing Systems, 2017.
M. van Kempen, S. S. Kim, C. Tumescheit, M. Mirdita, J. Lee, C. L. M. Gilchrist, and J. Söding. Fast and accurate protein structure search with foldseek. Nature Biotechnology, 2024.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu. Diffusion language models are versatile protein learners. In International Conference Machine Learning, 2024a.
X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu. DPLM-2: A multimodal diffusion protein language model. arXiv, 2024b.
J. L. Watson, D. Juergens, N. R. Bennelt, B. L. Trippe, J. Yim, H. E. Eisenach, W. Abern, A. J. Borst, R. J. Ragotte, L. F. Milles, B. L. M. Wicky, N. Hanikel, S. J. Pellock, A. Courbet, W. Sheffler, J. Wang, P. Venkatesh, I. Sappington, S. Vázquez Torres, A. Lauko, V. De Bortoli, E. Mathicu, S. Ovchinnikov, R. Barzilay, T. S. Jaakkola, F. DiMaio, M. Back, and D. Baker. De novo design of protein structure and function with RFdiffusion. Nature, 2023.
K. E. Wu, K. K. Yang, R. van den Berg, S. Alamdari, Y. Zou, A. X. Lu, and A. Amini. Protein structure generation via folding diffusion. Nature Communications, 2024.
Q Xie, Q. Chen, A. Chen, C. Peng, Y. Hu, F. Lin, X. Peng, J. Huang, J. Zhang, V. Keloth, X. Zhou, H. He, L. Ohno-Machado, Y. Wu, H. Xu, and J. Bian. Me Ilama: Foundation large language models for medical applications. arx, 2024.
J. Xu and Y. Zhang. How significant is protein structure similarity with TM-score 0.5? Biotnformatics, 2010.
L. Yang, Y. Tian, M. Xu, Z. Liu, S. Hong, W. Qu, W. Zhang, B. Cui, M. Zhang, and J. Leskovec. VQGraph: Rethinking graph reprosentation space for bridging GNNs and MLPs. In International Conference on Learning Representations, 2024a.
L. Yang, S. Xu, A. Sellergren, T. Kohlberger, Y. Zhou, I. Ktena, A. Kiraly, F. Ahmed, F. Hormozdiari, T. Jaroensri, E. Wang, E. Wulczyn, et al. Advancing multimodal medical capabilities of gemini. arXie, 2024b.
J. Yim, B. Trippe, V. De Bortoli, E. Mathieu, A. Doucet, R. Barzilay, and T. Jaakkola. SE(3) diffusion model with application to protein backbone generation. In International Conference Machine Learning, 2023.
J. Yu, X. Li, J. Y. Koh, H. Zhang, R. Pang, J. Qin, A. Ku, Y. Xu, J. Baldridge, and Y. Wu. Vector-quantized image modeling with improved VQGAN. In International Conference on Learning Representations, 2022a.
J. Yu, Y. Xu, J. Y. Koh, T. Luong, G. Baid, Z. Wang, V. Vasudevan, A. Ku, Y. Yang, B. K. Ayan, B. Hutchinson, W. Han, 7. Parekh, X. Li, H. Zhang, J. Baldridge, and Y. Wu. Scaling autoregressive models for content-rich text-to-image generation. Transactions on Machine Learning Research, 2022b.
L. Yu, J. Lezama, N. B. Gundavarapu, L. Versari, K. Sohn, D. Minnen, Y. Cheng, A. Gupta, X. Gu, A. G. Hauptmann, B. Gong, M.-H. Yang, I. Essa, D. A. Ross, and L. Jiang. Language model beats diffusion tokenizer is key to visual generation. In International Conference on Learning Representations, 2024.
15

---

Under review as submission to TMLR
Y. Zhang and J. Skolnick. TM-align: protein structure alignment algorithm based on the TM-score. Nucleic Acids Research, 2005.
A. Ziv, I. Gat, G. Le Lan, T. Remez, F. Kreuk, J. Copet, A. Défossez, G. Synnaeve, and Y. Adi. Masked audio generation using single non-auttoregressive transformer. In International Conference on Learning Representations, 2024.
16

---

L number of Under review as submission to TMLR
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 al., 2024p. For all the experiments, we fix the dimension of the codes to d = 6. Then, we quantize each channel into L unique values L L,,...; La and refer to the quantization levels as L Li..... La]. The size of the codebook C is given by the product of the quantization levels: Ly. For the experiments with small codebooks with |C| 4096, we use L [4, 4, 4, 4, 4, 4] and for large codebooks experiments, [C] 64000, we take L [8, 8, 8, 5,,,5,
Il -
In more details the FSQ quantizer writes Algorithm
Algorithm Finite Scalar Quantization
1: Input: z₁ E R (residue embedding), Wpros E Raxe W Wap-pros E Rexd (weight matrices), levels
2: Output: & (Quantized output)
// Compute low dimensional embedding
2i W proj Z₁
// Bound cach element Zij between L//2
4: Zy = 2 tanh(2ij)
// Round cach element to the nearest integer
W Wep-proj round(2))
return 2;
The product Wrooj 2₁ facilitates scalar quantization within a lower-dimensional latent space, thereby defining a compact codebook. This similar to the low dimensional space used for code index lookup in Yu el al. (202aa) ) 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
Algorithm Resampling Layer with Positional Encoding
Input: Inputs E Rxxd Mask, target size: p, features dim: d Queries: [Optional]
Output:
Queries, Inputs
3: If Queries = None then:
4: Queries +- Sin.ossitiooallEoooinnggg... E Rxxd
5: Queries, Keys, Values +- Linear (Queries) Linear (Inputs), Linear (Inputs)
AttentionWeighhts Softmax QueriessReye' Mask) RPxT
Vd
Output AttentiinNeightts Values E Roxd Queries MLP(Dut)ut)
Inputs +- MIP(Inputs)
10: return (Queries, Inputs)
Local Cross-Attentes Masking The encoder network used in this work preserves the notion of residue order, as defined in the primary struct ure 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 residucs 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
17

---

Under review as submission to TMLR
representation to carry local information, 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 where only the direct neighbors are kept in the attention update.
Enbeddings
Attention Maak
Downeampled
Enbedding
Figure 5: Illustration the local attention mechanism when using 2 neighbors for aggregation.
Decoder For the decoder, we re-purpose the Structure Module (SM) of AlphaFold (Jumper et al. 2021 In Jumper al. (2021), 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,
encode the structures with set of 1D features the sequence of discrete codes obtained after lokenization. Inspired by the OuterProductMean module of the Evoformer (Alg. 10 in SM of Jumper et al. (2021)), ) we compute 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 Algorithm
Algorithm Pairwise Module
Input: (8:)igN
Output: = (ki))issN
linear transforms of the initial embedding
s left Wieft-s, gright 8
// (n k): protein length, d: embedding dim 4: k = (nds kd -> nkd, sieft, griht)
= MLP(Kgg,
return (ki))ssN
Overall algorithm, the encoding and decoding processes write as described in Algorithm 4
A.2 Training Details and Metrics
Data Preprocessing We consider the graph G (V,E) consisting of set of vertices or nodes V (the residues) with features f .....lV and a set of edges & with features f/ f For the node features, we use a sinusoidal encoding of sequence position such that for the i-th residuc, the positional encoding is (1, ((,,d) d)) where d the embedding size. For the edge features we follow Ganea et al. 2202 More specifically, for each defined by residue v,, a local coordinate system is formed by (a) the unit vector t, pointing from the a-carbon atom to the nitrogen atom, (b) the unit vector u, pointing from the a-carbon to the carbon atom of the carboxyl (-co-) and (c) the normal of the plane defined by 4 and uy: n4 ]a, Finally, setting: qi n4 X uy, the edge features are then defined as the concatenation of the following:
relative positional edge features: Pj-vi (n u, 9: )xz xz),
relative orientation edge features: O5-r (nf uz of )ng, Ro (nf un of Juy, li (nr of )
18

---

Under review as submission to TMLR
Algorithm 4 Overall Algorithm Pseudo-Code
Input: p E Nx4x3 (0, ())
Output: 2.
// Compute embedding at the residno-level
X GNN(P)
// Downsample N + N/r
4: 2 = Resampling(z)
// Quantize
2 (z)
// Upsample N/r N
Resampling(2)
// Make pairwise representation for decoding
7: k PairWisseMduue(())
// Decode
8: p StructureeMouulo(ss
return 2, P (Quantized output)
lky x111 , 1, 2...R where
distance-based edge features, defined as radial basis functions: so-2s = R = 15 and or = 1.5.
207
Regularization We found that introducing scheduled commitment loss, similar to the original VQ-VAE approach van den Oord et al. 2017), significantly improves late-stage stability in training. However, imposing this penalty too early can harm encoder expressivity. To address this, we delay the onset of the commitment loss until step 20000 and then linearly increase its weight from 0 step 20000 a maximun value of Xmax= 0.2. Concretely, the auxiliary loss takes the form:
(stxpstep, z:) X(step) 2 stopgrad(roundedl2
(5)
where the schedule function X(step) is given by
0,
if step < To:
stop To
A(step)
Amax
if To S step S Tı,
(6)
Ti - To
Amax:
if step > Ti.
To is the step at which the ramp-up starts and Ti is the step at which the ramp-up finishes. We empirically set To = 20000 and Ti = 40000 which we found work well in practice. This delayed, gradually increasing penalty ensures that the encoder can initially learn expressive representations, then become more stable in later training.
Compression Factor We define the compression factor of Table as the ratio between the number of bits necessary to encode the atoms of the position of the backbone and the number of latent codes store multiplied by log( (oodss. With positions stored as 64-bit floats, we write the compression factor as:
x 4 x 64
768
Compression Factor
N log, #Codes)/r
log,(#
A.3 Additional Results: Structures Autoencoding
In Figures 6 and 7 we dotail the evolution of the RMSD and TM-scores on CASP-15 dataset when varying the downsampling ratio with fixed codebook size 4096. The results indicate that the average error increases alongside the standard deviation of the error distribution.
19

---

Under review as submission to TMLR
RMSD on CASP-15; Experiment Resu is for Codebook Size 4096
doveroa
domnsamoi mg
Figure 6: Evolution of the RMSD distribution on CASP-15 dataset with the downsampling ratio, with a fixed codebook size 4096
TM-Score on CASP-15: Experiment Results for Codebook Size 4096
0.7
0.8
dectnampling
downtarpping mio
Figure 7: Evolution of the TM-Score distribution on CASP-15 dataset with the downsampling ratio, with a fixed codebook size of 4096
A.4 De novo Structure Generation
A.4.1 Training details
GPT hyperparameters We use standard decoder only transformer following the implementation of Hoffmann et al.) 2022) with pre-layer normalization and a dropout rate of 10% during training. We follow Hoffmann et al. (2022) for the parameters choice with 20 layers, heads per layers, model dimension of 1024 and query size of 64, resulting in 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 AdamW with 31 0.95, 3n = 0.9 and a weight decay of 0.1. The learning rate follows a linear warm-up schedule, increasing linearly to 10 over the first 1 000 training steps. Following the seminal work of Radford and Narasimhan (2018), we employ embedding, residual, and attention dropout with rate of 10%. We found that batch size has crucial importance optimizing and adopt batch size of 65,792 tokens per batch.
20

---

Under review as submission to TMLR
CASM15 IMSC: Experiment Results for Codebook Size 64 .000
dowera
dourrums -
domnsamci ng rats
Figure 8: Evolution of the RMSD distribution on CASP-15 dataset with the downsampling ratio, with a fixed codebook size 4096
TM Score on CASP 15: Experiment Results for Cacebook Size 64,000
00
0.8
0.4
Anwnarraang rioo
Figure 9: Evolution of the TM-Score distribution on CASP-15 dataset with the downsampling ratio, with a fixed codebook size of 4096
The total number of actual structural tokens 70M. In that perspective, and line with work such as Hoffmann al. 20022.. we believe that leveraging large dataset of predicted structures such AlphaFold provide significant improvements when training the latent generative model.
A.4.2 Structure generation metrics
Designability We adopt the same framework than Yim et al. 2023 Trippe et al. 2023 Wu et al. 2024 to compute the designability or self-consistency score:
Compute putative sequences from ProteinMPNN DDuparas et al. sampling of 0.1.
2022 with temperature
Fold each of the 8 amino-acid sequences using ESMFold (Lin et al. in 8 folds per generated structure.
2022 without recycling, resulting
Compare the 8 ESMFold- 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.
htttps///alphafoll
21

---

Under review as submission to TMLR
0.8
0.16
0.7
0.14
0.6
0.12
0.5
Frequency
Frequency
0.08
0.04
0.1
0.0
0.00
2
6
10
12
40
60
70
80
90
100
RMSD
TM-Scores
Figure 10: RMSD (left) and TM-Score (right) distribution on the held out test set for the codebook size of 432 and downsampling ratio of
0.20
0.8
0.6
0.15
Frequency
Frequency
0.4
0.10
0.2
0.05
40
100
RMSD
TM-Scores
Figure 11: RMSD (left) and TM-Score (right) distribution on the held out test set for the codebook size of 1728 and downsampling ratio of 1.
In Table report the proportion of generated structures that are said to be designable, i.e samples for which se scM5>0.5 (or scRMSD< 2 Á when using the RMSD).
Novelty For the reference dataset, we use the s40 CATH dataset (Orengo et al. 1997): publicly available at fttp://rregoffii biochem. acl.ac. uk/ cath httth--ddtast--noeeddddaaaa
sets/ pdb. tgz. order to reduce the computation time, we first retrieve the top 1000 hits using Foldseck (van Kempen al. 20244). We then perform TM-align (Zhang and Skolnick 2005) for each match against. the targeted sample and report the TM-score corresponding to the best hit. A
22

---

Under review as submission to TMLR
1.2
0.25
10
0.20
0.8
Frequency
Frequency
0.15
0.10
0.2
0.05
0.0
0.00
2
6
10
12
50
60
70
90
100
RMSD
TM-Scores
Figure 12: RMSD (left) and TM-Score (right) distribution on the held out test set for the codebook size of 4096 and downsampling ratio of 1.
0.40
14
0.30
Frequency
Frequency
0.20
0.15
0.10
0.2
0.05
Th.m
0.0
0.00
6
40
50
70
80
90
100
RMSD
TM-Scores
Figure 13: RMSD (left) and TM-Score (right) distribution on the held out test set the codebook size of 64,000 and downsampling ratio of 1.
structure is then considered as novel if the maximum TM-score against CATH (cathTM) 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. (2023). 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 diverse set of generated samples, each cluster should be composed of only few samples
23

---

Under review as submission to TMLR
0.7
0.14
Frequency
Frequency
0.4
0.08
0.3
0.06
0.04
0.02
0.0
0.00
2
6
10
12
30
40
50
60
70
80
90
100
RMSD
TM-Scores
Figure 14: RMSD (left) and TM-Sccre (right) distribution on the held out test set for the codebook size of 4096 and downsampling ratio 2.
0.8
0.175
0.150
0.6
0.125
Frequency
Frequency
0.100
0.4
0.075
0.2
0.050
0.025
0.0
0.000
80
90
100
RMSD
TM-Scores
Figure 15: RMSD (left) and TM-Score (right) distribution on the held out test for the codebook size of 64,000 and downsampling ratio of 2.
equivalently, the number of different. clusters should be high. We use MaxCluster Herbert and Sternberg (2008) with TM-score threshold of 0.6 as in Watson et al. (2023)
A.4.3 Sampling
Baselines For each baseline methods, we follow a standardized process similar to that of Yim et al. (2023 to generate the testing datasot: we sample backbones for every length between 100 and 500 with length stop of 5: [100, 105....... 5000). We re-use the publicly available codes and use the parameters reported in Watson et al. 2023) and Yim et al. (2023) respectively.
24

---

Under review as submission to TMLR
0.07
0.35
0.06
0.30
0.05
0.25
Frequency
Frequency
0.04
0.20
0.03
0.15
0.02
0.10
0.05
0.01
0.00
0.00
2
6
10
12
20
40
60
80
100
RMSD
TM-Scores
Figure 16: RMSD (left) and TM-Sccre (right) distribution on the held out test set for the codebook size of 4096 and downsampling ratio of 4.
0.10
0.08
0.4
Frequency
Frequency
0.06
0.1
0.02
D:
0.0
0.00
2
6
8
10
12
30
40
50
60
70
80
90
100
RMSD
TM-Scores
Figure 17: RMSD (Jeft) and TM-Score (right) distribution on the held out test set for the codebook size of 64,000 and downsampling ratio of 4.
Generating Structures with Decoder-Only Transformer Sampling for method is 2 steps
process: first, sample sequence of structural tokens from the trained prior described in Appendix A.4.1 then reconstruct the structures using the trained decoder. There are many ways to sample from decoder-only transformer model (Vaswani et al. 2017 Radford and Narasinhan 2018 Radford et. al.) 2019 Holtzman et al. 2020 We chose temperature sampling (Vaswani et al.) (0177) with other alternative sampling strategies such as top-k (Radford al. 2019) and top-p (or nuclous) sampling (Holtzman et al., 2020) resulting in little improvement at the cost of increased complexity. As showed in Vaswani et al. (017). the temperature controls the trade-off between the confidence and the diversity of the samples. In order to tune the temperature,
25

---

Under review as submission to TMLR
sampled 2000 samples for each temperature between 0.2 and 0.8 in steps 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 [8] shows à 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 sc-M-scree (sce Figure 18]. The results reported in Table 2 are obtained with a temperature of 0.6, it achieves satisfying trade-off between designability and diversity.
10
Temperature
0.8
0.2
0.6
10"
0.4
Temperature
Count
0.6
0.2
0.4
0.8
0.4
10
0.2
0.6
0.8
0.0
10
100
150
200
250
300
350
400
450 500
100
200
300
400
500
Length
Length
Figure 18: Ablation of the temperature sampling. Left: Histogram of the generated structure length. Right:
Designability score vs lemperature sampling.
Contrary the baselines, our model learns the joint distribution the length and the structures = p(s, l)dl where the random variables and 1 represent the structures and the length respectively. Indeed, only the conditional distribution p(s[() is modeled by the diffusion-based baselines. In Figure [8] we show the length distribution learned by the model. 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 temperature of 0.6 as previously established). We then bin the generated structures by length, with bin width of 5 and bin centers uniformly distributed between 100 and 500, specifically: [100, 105,..... 500). Finally, we limit the maximum number of structures per bin to 10, randomly selecting 10 structures if bin contains more than this number.
26

---

Under review as submission to TMLR
A.4.4
Additional Results
1.0
model
6
vq3d
0.8
framediff
ScRMSD
rfdiffusion
ScTM
0.6
model
0.4
vq3d
framediff rfdiffusion
0.2
0.0
100
150
200
250
300
350
400
450
500
100
150 200 250 300 350
400
450
500
Length
Length
Designability score samples length. Left: scRMSD score for different structure lengths. Right: scTM score for different structure lengths.
Figure 19:
1.0
model vq3d
0.8
0.9
framediff ridiffusion
0.6
cathTM
TM
160
0.4
model
0.6
240
320
vq3d
400
480
0.2
framediff rfdiffusion
0.0
100
150 200 250 300 350 400 450 Length
500
3
5
6
8
scRMSD
Figure 20: Left: Novelty score for different structure lengths. Right: Novelty score versus designability score.
27

---

Under review as submission to TMLR
model
0.125
model
0.3
vq3d
vq3d
0.100
framediff rfdiffusion
framediff
Probability
Probability
rfdiffusion
0.075
0.2
0.050
0.1
0.025
0.000
0.0
2
4 scRMSD
0.4
0.6 ScTM
0.8
1.0
Figure 21: Distribution of the dosignability scores for novel domains (cathTM<0.5).
Negative Log-Likelihood Scores per Residue vs. Sequence Length
2.5
NLL Score (Log Scale)
Random Baseline (-luu(1((C)
Mean SEM
250
300
Sequance Length
Figure 22: Evolution with the length the generated sequence of the per-residue negative Log-Likelihood of the selected amino acids provided by ProteinMPNN
28


Reviews:
[{"review_id": "SowD9ryogY", "reviewer": "Reviewer_fkWK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper demonstrates that continuous protein structures can be effectively converted into discrete sequences using Finite Scalar Quantization (FSQ), enabling direct application of transformer-based language models to structural data.\n\n- The proposed method constructs a flexible discrete latent space with adjustable codebook size and downsampling ratio, achieving significant compression (up to 192×) while maintaining structural fidelity.\n\n- The autoencoder achieves low reconstruction errors, with RMSD in the range of 1–4 Å and TM-scores above 0.9, validating that the learned token representations preserve detailed geometric information.\n\nWeakness:\n\n- Limited novelty: The method primarily combines existing modules (FSQ, MPNN, AlphaFold decoder) without introducing fundamentally new algorithmic ideas.\n\n- Missing comparisons with concurrent work: The paper does not compare against recent related methods like FoldToken or Liu et al. (2024), which also use FSQ and structure-tokenization approaches.\n\n- Insufficient downstream validation: The paper only evaluates reconstruction and generation; no downstream tasks such as fold classification or sequence-conditioned generation are tested.\n\n- No analysis of token semantics: The authors do not analyze token usage distributions, clustering patterns, or whether the tokens encode meaningful geometric or functional features.\n\n- Decoder heavily relies on AlphaFold: The use of a pretrained structure module raises concerns about whether the reconstruction accuracy stems from the token representation or the strong decoder.\n\n[1] Gao, Zhangyang, Tan, Cheng, Wang, Jue, Huang, Yufei, Wu, Lirong, & Li, Stan Z. FoldToken: Learning Protein Language via Vector Quantization and Beyond.\n\n[2] Gao, Zhangyang, Tan, Cheng, & Li, Stan Z. FoldToken2: Learning Compact, Invariant, and Generative Protein Structure Language.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a framework for learning discrete representations of protein 3D structures using a vector-quantized autoencoder. The authors utilize Finite Scalar Quantization (FSQ) to convert continuous backbone atom coordinates into discrete tokens. These tokens enable the use of transformer-based sequence models, such as GPT, for de novo protein structure generation. The method is evaluated on structure reconstruction and generative design, demonstrating competitive reconstruction accuracy and moderate generative performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add comparisons with concurrent methods such as FoldToken. These works also explore protein structure tokenization via vector quantization (including FSQ) and provide strong baselines for both reconstruction and generation tasks. A quantitative comparison — including reconstruction metrics (RMSD, TM-score), codebook utilization, and model size — is essential to situate the contribution of this work relative to the state of the art.\n\n- Clarify the source of reconstruction performance (decoder vs token). The current decoder relies heavily on AlphaFold's structure module, which may obscure how much structural fidelity stems from the discrete token representation versus the powerful decoder. The authors should clarify whether the decoder was trained from scratch or frozen, and ideally provide an ablation with a simpler decoder to evaluate the informativeness of the learned tokens.\n\n-Improve diversity of generated structure. While designability is high, the novelty of generated structures is relatively low (23.3%). Please consider exploring alternate sampling strategies (e.g., higher temperature, nucleus sampling) or training regimes that encourage greater diversity in the generated token sequences.", "broader_impact_concerns": "The manuscript does not currently include a Broader Impact Statement, which is required given the potential real-world implications of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ACwHCqBOOW", "reviewer": "Reviewer_wvd5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n++ This paper presents a promising method for transforming complex and continuous space of protein structures into discrete tokens, which enables the training of a protein generation model using a GPT-like architecture.\n\n++ The paper is relatively easy to follow.\n\nWeaknesses:\n-- The values of the reported metrics in Section 3.1 (Results, page 5) do not match those presented in Table 1.\n\n-- The experiments on tokenization and reconstruction should include a comparison between the proposed method and the VQ-VAEs approach.\n\n-- The generated proteins exhibit lower novelty and diversity compared to those produced by diffusion-based methods, as shown in Table 2.\n\n-- It would be better to include some metrics such as training and testing time, FLOPs, or other measures of computational cost to provide a clearer comparison of the method's efficiency.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a method that uses a vector-quantized autoencoder to tokenize protein structures into discrete representations. The resulting codebook contains ranges from 4096 to 64000 tokens, enabling high-fidelity reconstructions. These discrete tokens can be used to train a simple sequence-modelling-based GPT model for generating novel, diverse and designable protein structures.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-- The experiments on tokenization and reconstruction should include a comparison between the proposed method and the VQ-VAEs approach.  (critical)\n\n-- It would be better to include some metrics such as training and testing time, FLOPs, or other measures of computational cost to provide a clearer comparison of the method's efficiency. (critical)\n\nMinor\n\n-- Section 3.1: \"we also train two additional quantized autoencoders with smaller codebook sizes, K = 1728 and K = 1728\", however, both values are the same, which may be unintended.\n\n-- In the caption of Figure 3: \"fourth column of Table 1\" does not match the content described in the caption; \"a different structures\"; \"on the left of the most left column\".\n\n-- A full stop is missing at the end of the second paragraph in Section 5.\n\n-- Page 19: the last paragraph has some issues: \"a ADAMw\"; \"5.10^{-5}\"; \"for the first 1000\"; \"with rate of 10% rate\"", "broader_impact_concerns": "The authors did not present a discussion about broader impacts. Please add some discussions.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jy97zfjDNE", "reviewer": "Reviewer_XHLD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n​1. Combines FSQ with geometric deep learning to address the challenge of representing continuous 3D structures in a discrete, sequence-friendly format.\n\n​2.  Achieves near-experimental-resolution reconstruction and demonstrates generative capabilities with a vanilla GPT model (87% designability).\n\n​3. Shows that a simple GPT model trained on discrete codes achieves designability comparable to diffusion-based methods, reducing computational costs.\n\nWeaknesses:\nMore empirical comparisons (e.g., against other VQ-based protein models like FoldToken or VQPL) would strengthen the position..", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a method for transforming continuous protein structure data into discrete token sequences through a vector-quantized autoencoder. By doing so, it enables the application of sequence-based language models, like GPT, to protein structure modeling—previously hindered by the continuous and 3D nature of such data. The contributions of the paper are as follows,\n\n1. Leveraging Finite Scalar Quantization, the authors develop a structural tokenizer that converts backbone atom coordinates into sequences of discrete tokens while maintaining low reconstruction error.\n\n2.  The model supports a range of codebook sizes and downsampling factors, achieving compression rates up to 256× while preserving structural fidelity, with TM-scores typically >0.9.\n\n3. Using a vanilla GPT model trained on the token sequences, the method is able to generate novel, diverse, and designable protein backbones, with quality competitive to methods like FrameDiff.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is encouraged to extend evaluation to include at least one functional or predictive task to showcase the utility of discrete structure tokens beyond generation.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)