Protein Structure Tokenization: Benchmarking and New Recipe

Thanks for your positive review and insightful comments! We respond to your concerns as below:

R1: Discrepancy observation that FoldSeek, which has the worst curve shape for the distinctiveness evaluation, has much better performance on efficiency benchmark. However, the distinctiveness and efficiency metrics should be related to each other. A more distinct codebook should lead to more efficiency utilization

We would like to humbly argue that the two distinctiveness and efficiency are related, but it is not necessarily the case that as one increases, and the other increases.

1. Importantly, the discrepancy stems from the different codebook size for different PST methods. And the reviewer’s intuition would be more sound when the codebooks share the same size: a more distinct codebook should lead to more efficiency utilization. For example, utilization is measured by the factor of #used codes and #total codes (i.e., the codebook size). Smaller the codebook size, higher the chance for utilization rate to be higher. FoldSeek has a very small codebook size K=20 (see Tab. 4), so it’s natural to have a high utilization rate, despite the observation that its used codes are not that distinct from each other.
2. We would like to clarify the concept of distinctiveness. The codebook has two parts of codes: used and unused ones. For distinctiveness, among the utilized codes in a codebook, the more distinct they are, the better chance for them to avoid ambiguous token-substructure mappings. We would like to emphasize the distinctiveness with used codes, and we show it in Fig. 3 (right panel). Because in reality, those used codes matter the most.

We’re open to hear more opinions on this topic from the reviewer, and are always eager to learn more to enrich our knowledge.

R2: Question on the meaning to have distinctiveness evaluation, given its discrepancy to “codebook utilization”.

1. The “Distinctiveness” evaluation was motivated from the observation in ESM3’s (Fig. S5) that ambiguous structural token mapping can hinder interpretability and harm model capability. “Distinctiveness” measures the similarity within codebook vectors, which serves as a proxy to assess token ambiguity.
2. We understand the reviewer’s concern that simply using cosine similarity for codebook vectors might be limited. **We would like to continue working on a better metric for measurement instead of simple cosine similarity. **
3. Please see R1 for more discussion on the “discrepancy between distinctiveness metric and codebook utilization metric”. In general, we would like humbly argue that the discrepancy mainly stems from the different codebook size used in different PST methods, but not the metrics themselves. We’re always open to hear more opinions on this topic from the reviewer, as we’re also in the exploration stage to find a good way to benchmark PST methods.

R3: “ESM3 is the leading model” should be rephrased to “the leading VQVAE-based model”

We agree and would modify this in paper.

R4: In sensitivity evaluation, TM-score is not a golden standard for structure similarity, but a state-of-the-art structure search algorithm. Acknowledgement of this limitation is useful.

Thanks for the valuable suggestion. We would acknowledge this limitation in both main text and appendix that TM-score is not a golden standard to measure structure similarity.

Also, commonly used metrics to measure structure similarity include TM-score (more globally) and RMSD (more locally). In App. F.2, we explored using RMSD instead of the TM-score in “Sensitivity” evaluation. As shown in Tab. 10, using RMSD is less effective than TM-scores. This might be because RMSD is sensitive to outliers and local structural discrepancies, while TM-score focuss on global structure and is less sensitive to local structure variations.

R5: It’s unclear whether the dynamic programming in sensitivity used global or local alignment, which should be stated both in main text and appendix

We used the global alignment algorithm by adopting the Biotite’s “align_optimal()” function. Thanks for pointing out that this essential detail is missing. We would add this content to both the main text and the appendix.

Thanks for invaluable suggestions. We respond to your questions below

R1: AminoAseed doesn’t perform well on binding site prediction in Fig. 4

We'd like to clarify our observation:

1. Fig. 4 shows that when using continuous structural representations, AminoAseed beats ESM3 on all three binding site prediction tasks, while being comparable to the continuous version of VanillaVQ on two of the tasks.
2. When using discrete structural tokens, which is the most important application setting for PSTs, AminoAseed prevails for most of the tasks, and gains improvement for averaged task perf. (in Tab. 2)

R2: Paper shows the benefits of scaling the model (e.g., increasing codebook size or encoder size) and the diminishing returns as model scales

We'd like to clarify this rephrased statement from the paper:

1. We show in paper the benefits of data-driven Pareto-optimal scaling (i.e., balancing) for the codebook size and codebook dimension. It’s important to notice that we kept the total codebook parameters the same
2. We show the diminishing sub-exponential returns in Fig. 7 when increasing encoder sizes, keeping the codebook configuration unchanged

R3: Pareto-optimal codebook configuration and codebook reparametrization, while innovative, add complexity to the model

We'd like to clarify that only codebook configuration adds complexity due to multiple runs of model training, while reparameterization shares the same complexity as the ablation model VanillaVQ

In general, for the codebook configuration, we leveraged the additional resources for exploration and provide insights for others who aim to develop better PSTs: under a fixed computational resources, balancing codebook sizes and dimensions could be useful. This is a discovery that’s previously unknown in the field

Simpler methods without using more computational resources are indeed favored, and we would like to explore it in our future work

R4: PST methods, including AminoAseed, tend to lose reliability under high noise levels in the input protein data

We would like to clarify that in Fig. 5, the noise is added to the induced structural representations, not on the raw protein structures. This is achieved by replacing them with the [Mask] representation under a given mask rate.

Essentially, it’s expected to observe large performance degradation at high noise levels (up to 90% tokens masked), since most structural information is removed

R5: AminoAseed does not do well in sequence-based tasks like remote homology detection

We would like to clarify that the remote homology detection task in StructTokenBench (i.e., “Homo” task) is not a sequence-based task, but a structure-based task. Namely, we use backbone structure as input, and do not leverage the residue sequences. As explained in App. A.1.1, our “Homo” task is from the same source TAPE as the popular sequence-based remote homology detection task, which might be the reason for the confusion.

Also, we appreciate the suggestion and would like to humbly point out this statement does not align with our observation. As shown in Tab. 2, on “Homo” task, our method AminoAseed performs the best, achieving 27.31% improvement over ESM3

R6: Do we plan on integrating AminoAseed with sequence models to improve its performance on remote homology detection? If so, how would we approach the integration?

1. We explored the setting of adding sequence information directly to PST (i.e., using structural representations + sequence token + positional encoding as input, and passing it to the probing layer). As shown in Tab. 5, adding sequence can mostly improve PSTs’ performance, as expected
2. Further explore using sequence models like ESM2 on this specific homology task may shifts from our main focus, i.e., to benchmark structural representations induced from PSTs. We would like to leave it for future work
3. For the integration method, one way is to combine the pretrained protein sequence representations from models like ESM2, with AminoAseed’s structural representations. Fusion attention layers can be added after the two models for better combination. Per-residue and per-protein contrastive learning can be used as training objectives.

R7: What’s the next step to scale up AminoAseed more effectively? Improvement to the architecture or optimization strategies could yield better performance with more compute?

1. We recommend data scaling: gathering more diverse pretraining datasets to improve AminoAseed, since its pretraining data includes only around 10% of the PDB database, with 48k protein single chains.
2. Improving the architecture design of VQVAE encoder might be useful, because the structure encoder is usually SE-(3) invariant and its capacity can be improved with better design.
3. Improving optimization would be helpful, which shares the same motivation as AminoAseed to use codebook reparameterization to enable better codebook gradient update.

We thank the reviewer for the insightful comments

R1: The claim “adding “a MLP” to reparametrize the codebook can alleviate codebook collapse” needs theoretical analysis and ablation experiments

1. We'd like to clarify that in Sec. 4.2 (L246), we add a simple linear layer (instead of a MLP) to reparameterize the codebook $C$ as $Q$=Linear($C$), and keep $C$ fixed during training

2. Ablation experiments: VanillaVQ is the ablation model, which shares the same pipeline and configurations as AminoAseed, except using $C$ instead of Q=Linear($C$). Main results can be found in Tab. 2/3/4/5 and Fig. 3/4/5/6/7. Generally, AminoAseed outperforms VanillaVQ across all four benchmarking perspectives

3. Theoretical analysis: We analyze the gradient flow on $Q$ in AminoAseed and $C$ in VanillaVQ, respectively. We will add this theoretical analysis in the appendix to enhance the method design justification

   (1) Setup: Assume $C$ of shape (K, D) is fixed, Linear() denotes a linear transformation with weight $W$ of shape (D, D) without a bias, and the reparameterized codebook $Q$=Linear(C)=$C * W$. We denote $L$ as loss, batch size as 1 for simplicity, and $p$ different codes are selected during tokenization with each code selected for once

   (2) Gradient on $Q$: For the gradient $\nabla_Q L$, only the $p$ selected rows of this gradient matrix are non-zero. The gradient updates $W$ via the chain rule: $\nabla_W L = C^T\nabla_Q L$, which is now a dense matrix with all non-zero values. Thus, all entries of $W$ are updated, enabling global adjustments to the entire codebook space through $Q=C * W$

   (3) Gradient on $C$: if $C$ is trainable, the gradient matrix $\nabla_C L$ only has the $p$ selected rows as non-zero. Thus, unused codes get no updates, leading to limited changes for the entire codebook space and eventually distribution shift (see Fig. 2)

R2: Why using the VQ-VAE framework when IF-based PSTs beat VQ-VAE-based methods on supervised tasks? And if AminoAseed sufficiently closes the perf. gap?

1. IF-based PSTs excel in supervised tasks but lack sensitivity to conformational changes, limiting their overall capability.

2. IF-based PSTs only support continuous tokens, while VQ-VAE-based PSTs produce both continuous and discrete tokens. As said in App. C, discrete tokens offer several advantages:

   (1) Multimodal LLM integration: enabling seamless fusion with sequence and functional text data, allowing direct use of NLP optimization techniques developed for LLMs.

   (2) Simplified structure modeling: eliminating the need to explicitly encode symmetry and physical constraints

   (3) Reduced overfitting risk: discrete representations may help mitigate overfitting compared to continuous features [a]

3. We acknowledge that there is still much more room for further improvement to close the performance gap between AminoAseed and IF-based PSTs on supervised tasks. AminoAseed proposes a simple yet effective strategy for improvement over vanilla VQ-VAE, and enjoys the benefits of VQ-VAE methods and shows promising results in “sensitivity” evaluation

[a]DeProt: Protein language modeling with quantized structure and disentangled attention

R3: Benchmarking AIDO.st

As suggested, we add AIDO.st and report the results below (some supervised tasks are not finished due to the rebuttal timeline). As shown, AIDO.st is less effective than AminoAseed on supervised downstreams, and also fall short in “sensitivity”. However, AIDO.st achieves very high codebook utilization

R4: Detailed description or diagram of the AminoAseed model architecture are not provided

1. We provided the detailed description for AminoAseed in App. E, including (1) overall pipeline of protein frame input, encoding, tokenization, decoding, and structure reconstruction objectives; (2) the geometric self-attention layer used in encoder; (3) explanation of straight through estimator of gradients in tokenization; (4) the standard Gram-Schmidt algorithm to create protein frames as input

2. We will add a diagram for AminoAseed in appendix for next edition

R5: The absence of complete code release undermines the reproducibility.

We apologize for not mentioning it in the paper. We will release all the code and processed data in the next months, because we’re currently in the process of code cleanup and preparation.

R6: Typo of in Sec. 2.3

Thanks for spotting the typo and we will modify it

Thanks for your valuable comments! We address concerns below

R1: Explain “StructTokenBench focuses on local(per-residue) over global(per-protein) structures”

PSTs tokenize per-residue substructures, matching StructTokenBench: supervised tasks and “Sensitivity” are at residue level. Current benchmarks(Sec. 6.1) cover per-protein tasks, justifying our per-residue focus. We’ll add in main text

R2: No citation for “Distinctiveness” motivation(L179)

We’ll cite ESM3 Fig. S5: ambiguous structural token mapping harms. This metric assess the token ambiguity via codebook similarity.

R3: Efficiency(Codebook Utilization) metric seem conflicting with Sec 5.7(L415): no correlation with reconstruction quality. Is it needed?

We rename “Efficiency” to “Codebook Utilization”(% of used codes). Underutilized codebooks waste resources and harm downstream perf.(Sec. 3.4), a well known issue in CV/NLP[a]. High utilization bring gains, like LLaMa3’s 128K-token vocabulary for efficient encoding[a]. Sec. 5.7 questions if reconstruction quality alone reliably indicates PST quality, given no correlation with both supervised task perf. and utilization(see Fig. 6,12) We welcome further discussion [a]LLama technical report

R4: Codebook utilization is distribution-dependent. PSTs' non-uniform token use may reflect accurate reconstruction(like BindGPT)

We agree on dependence. Current results hold across datasets. We’ll add more diverse datasets. We remain inconclusive on exact utilization-reconstruction relations. Shown in Sec. 5.7, we question if reconstruction quality alone reliably indicates PST quality, given no correlation with both supervised task perf. and utilization. We welcome further discussion. Thanks but we can’t test it as it lacks residue-level structural reprs.(see R8)

R5: Why use probing for supervised tasks? Why MLP not transformer?

Probing(fixed encoder+simple probing layer) is standard practice in CV[b]. Being simple ensures perf. reflects reprs. quality, not probing layer capacity. For BindInt, empirical perf. below show: linear layer lags 2-layer MLP, which matches 1-/2-layer transformer perf. Limited dataset sizes(Tab. 6,8) may explain transformers' marginal gains

R6: Why use positional encodings(PEs), as MLP isn’t permutation-invariant?

Our MLP is shared across structural tokens, so is permutation-equivariant. So it needs PEs for residue order. Empirically, removing PEs largely hurts perf.(see table below vs. Tab. 2)

R7: Test Cheap

It’s initially excluded as it models sequence and all-atom structure(37 atoms/residue), while our benchmark only considers backbone(4 atoms/residue, Sec. 2.1). Cheap’s perf.(vs. Tab. 2,3) is in https://anonymous.4open.science/r/cheap_perf-00EB limited by space: (1)it beats IF-based PSTs on supervised tasks (some results pending) (2)It largely lags AminoAseed in “sensitivity”. (3)No other evaluation as its has no codebook

R8: Test BindGPT, Geo2Seq

Thanks for referring to atom-wise coordinate tokenization methods. We’ll discuss in main text and appendix. BinGPT is untestable as: (1)It encodes atom types and xyz coordinates as strings. Unlike IF-/VQ-VAE PSTs, it doesn’t output residue-level structural reprs; (2)Being decoder-only, BindGPT also can’t produce residue-level reprs. Geo2Seq is not open-source

R9:Heuristic PSTs not tested

We’ll add them later due to deadline

R10:For fixed reconstruction quality, measure bits used

We use scatter plot: X-axis: reconstruction quality(RMSD or LDDT),Y-axis: Compression ratio(% byte reduction vs. XYZ coordinates).Plots in https://anonymous.4open.science/r/rebuttal_fig-3050. AminoAseed beats VanillaVQ in reconstruction but trails ESM3, while achieving better compression(smaller byte reduction %) than both

R11: Explicitly motivate PSTs

We’ll addR3/R4 discussions on distinctiveness and codebook utilization in main text

R12:State PSTs' modeling targets

We’ll add (1)Tab. 2 footnote: Cheap(sequence + all-atom structure) vs. others( backbone structure) and (2)appendix table: PST comparison (backbone/all-atom, sequence usage) including non-open-sourced

R13: Re-term

We change “efficiency” to “codebook utilization” and “effectiveness” to “downstream effectiveness” for clarity

R14:Computational efficiency as future work

Will leave for future

R15: Intention of noise analysis

Noise is added by randomly masking structural tokens. We reveal: (1) Structural reprs. are less robust than sequence reprs., so integrating sequence-structure may help (2) Practical guidance for masked language modeling on structural tokens

R16:Typo in L1029

Will modify