PaperlayerView exact revision on arXiv

Source-linked AI summary

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization

Romain Hardy, Xiaoman Zhang, Pranav Rajpurkar

arXiv:2610.18481v1cs.LG

TL;DR

Existing protein tokenizers encode sequence, structure, and function independently, leaving no modular vocabulary that jointly quantizes all three. ProTok uses modality-specific product-quantization groups and modality dropout to train one checkpoint for inference from any modality subset. Its structure tokens rank strongly across StructTokenBench while retaining reconstruction quality and supporting functional retrieval and generation.

  • Problem

    Existing tokenizers encode protein modalities independently, and joint sequence-structure approaches omit functional supervision.

    Source: §1 · pp. 1–2

  • Method

    ProTok quantizes structure, sequence, and function into one modular vocabulary using modality-specific product-quantization groups and modality dropout.

    Source: §1 · p. 2, §3.2 · p. 5, §3.4 · p. 6

  • Results

    ProTok places among the top two structure-only tokenizers on 27 of 28 StructTokenBench splits, with a mean rank of 1.4 across the splits.

    Source: §1 · p. 3

  • Takeaways & Limitations

    One ProTok vocabulary supports structure-token representation, reconstruction, and generation, including backbone-sequence co-generation.

    Source: §1 · p. 3, §6.3 · p. 11

    Generation results are preliminary: the prior is small, trained on CATH chains of at most 128 residues, and evaluated on 50 samples.

    Source: §6.2 · p. 11

Abstract

from arXiv · show

Protein sequence, structure, and functional annotations each capture complementary aspects of protein biology, yet existing tokenizers encode each modality independently, limiting the expressiveness of discrete protein representations. We present ProTok, a unified protein tokenizer that maps all three modalities into a modular discrete vocabulary via product quantization with per-modality codebook groups and modality dropout. A single ProTok checkpoint supports flexible inference from any modality subset without retraining. On seven StructTokenBench benchmarks spanning 14 tasks and 28 test splits, ProTok achieves an average rank of 1.4 across all structure-only tokenizers, placing first on 21 splits and among the top two on 27 of the 28. These gains come without sacrificing reconstruction: ProTok's tokens reach 0.939 lDDT on CATH 4.2, showing that multi-codebook quantization and multi-modal training yield tokens that are both reconstructable and informative. ProTok's structure tokens further lead all baseline structure tokenizers on zero-shot functional retrieval across CATH, Enzyme Commission, and Gene Ontology label hierarchies. The vocabulary also supports generation: a compact language model trained from scratch over ProTok tokens jointly generates backbone and sequence, with 52% of sampled backbones designable versus 21% for a parameter-matched GeoBPE baseline. ProTok shows that structure, sequence, and function can share one discrete vocabulary, motivating protein language models that reason over a unified token stream.

1 Introduction

Protein modeling spans sequence, structure, and function, but existing tokenizers encode these modalities independently and joint approaches omit functional supervision. ProTok unifies them in a modular vocabulary, and its structure tokens perform strongly across benchmarks while supporting reconstruction, retrieval, and generation.

1 Introduction

  • Existing tokenizers typically encode one modality, while joint sequence-structure approaches omit functional supervision; no existing tokenizer jointly quantizes structure, sequence, and function in a modular vocabulary.Single-codebook quantizers can suffer codebook collapse, shrinking the effective vocabulary; existing multi-codebook designs remain modality-specific.Source: §1 · pp. 1–2
  • ProTok quantizes representations from four pretrained encoders covering global structure, local geometry, evolutionary sequence features, and protein-level functional semantics.Each encoder output has a dedicated group of four codebooks, yielding 16 parallel codes per residue; orthogonality-regularized SimVQ projections keep each codebook full-rank.Source: §1 · p. 2
  • Modality dropout trains shared decoders to recover targets from any active modality subset, allowing one checkpoint to run in structure-only, sequence-only, or combined mode without retraining.During training, randomly masked groups are replaced with learnable mask tokens.Source: §1 · pp. 2–3, Figure 1 · p. 2
  • ProTok’s frozen structure tokens lead the evaluated baselines on zero-shot functional retrieval, while the same vocabulary supports reconstruction and backbone-sequence co-generation.The evaluation uses 311K protein chains for training and StructTokenBench’s seven benchmarks across functional-site prediction, flexibility regression, remote-homology classification, and conformational discrimination.Source: §1 · p. 3
  • 1.4 mean rank across 28 StructTokenBench splits places ProTok first on 21 and among the top two on 27, leading the evaluated structure-only tokenizers overall.The same tokens reach 0.939 lDDT on CATH 4.2; a language model over them yields 52% designable backbones versus 21% for parameter-matched GeoBPE.Source: §1 · p. 3

2 Related Work

Prior protein tokenizers vary in granularity, quantization, and modality coverage, but their tokenization pipelines do not jointly encode structure, sequence, and function. ProTok combines modular multi-codebook quantization with cross-modal training.

Protein structure tokenizers

  • Structure-tokenization approaches trade off reconstruction and downstream semantics or vary in quantization strategy, geometric features, and token granularity.Examples include learned per-residue codebooks, lookup-free quantization, frozen-feature clustering, a fixed 3Di alphabet, equivariant encoders, motif-level tokens, and all-atom tokenization.Source: §2 · p. 3
  • StructTokenBench evaluates seven benchmarks and serves as ProTok’s primary evaluation protocol.Prior tokenization pipelines remain modality-specific, with no existing tokenizer providing a unified structure-sequence-function encoding in one vocabulary.Source: §2 · p. 3

Multi-codebook and product quantization

  • Product quantization creates a large effective vocabulary from compact codebooks, while SimVQ moves codebook entries coherently to address dead codes.ProTok adds an orthogonality penalty on SimVQ projection matrices and assigns dedicated codebook groups to global structure, local geometry, sequence, and function.Source: §2 · p. 4

Multi-modal protein representations

  • Existing multi-modal systems use parallel structure and sequence streams or collapse both into one representation, but their tokenizers lack cross-modal incentives and functional supervision.Concurrent Yeti pairs a structure-only tokenizer with a masked-diffusion model for backbone-sequence co-generation, but does not model function.Source: §2 · p. 4

3 ProTok

ProTok uses frozen modality encoders, dedicated product-quantization groups, cross-group fusion, and shared decoders. Modality dropout and residue masking train the model to reconstruct from varied inputs and contexts.

3 ProTok

  • ProTok partitions quantization into modality-specific codebook groups, fuses their outputs through cross-group attention, and uses modality dropout to support reconstruction from any available input subset.Figure 1 shows the architecture and inference modes.Source: §3 · p. 4

3.1 Encoders

  • Four frozen encoders capture global structure, local backbone geometry, evolutionary sequence features, and protein-level function, with outputs projected into a shared 1024-dimensional space.The function embedding is quantized once per protein and broadcast across residues; its loss uses mean-pooled quantized representations.Source: §3.1 · pp. 4–5

3.2 Quantization

  • Each of four modality groups splits a 1024-dimensional residue representation into four 256-dimensional chunks, each assigned to one of 256 codebook entries.SimVQ projections move codebook entries coherently to discourage dead codes, while an orthogonality penalty helps prevent low-rank collapse; the model assigns 16 discrete codes per residue.Source: §3.2 · p. 5

3.3 Fusion and Reconstruction

  • A single-layer attention module attends across the four group representations at each residue, and the summed outputs feed decoders for backbone geometry, amino-acid identity, sequence features, and function embeddings.One checkpoint supports modes with eight, twelve, or sixteen active codebooks, without retraining.Source: §3.3 · p. 6

3.4 Training

  • Training combines structure reconstruction, sequence and function alignment, VQ, orthogonality, and entropy losses.Modality dropout samples one of 15 non-empty group combinations per protein; a separate 15% residue mask encourages learning from spatial context. Two stages train on 311K chains, with the second replacing the structure decoder with a 30-layer decoder and adding pLDDT and PAE prediction heads.Source: §3.4 · pp. 6–7

4 Experiments

The experiments assess structure-token quality through downstream prediction, intrinsic codebook and reconstruction measures, information probes, and generation. They compare ProTok with released discrete structure-tokenizer checkpoints under shared probe protocols.

Evaluating structural token quality

  • StructTokenBench evaluation covers seven benchmarks, 14 prediction tasks, and 28 held-out splits, using lightweight classifiers on frozen tokens and zero-shot evaluation for conformational tasks.ProTokstr reads out only structure groups, making its tokens directly comparable with structure-only tokenizers; intrinsic measures include codebook utilization, normalized perplexity, and CATH 4.2 reconstruction.Source: §4 · p. 7

Probing codebook information content

  • Three probes examine token information content through standardized inverse-folding and structure-prediction tasks, zero-shot functional retrieval, and nested inference modes.Functional retrieval spans nine CATH, Enzyme Commission, and Gene Ontology label hierarchies.Source: §4 · p. 7

Baselines

  • ProTok is compared with four discrete structure tokenizers using their official checkpoints and the same probe protocols and hyperparameter sweep.The baselines cover learned codebooks, lookup-free quantization, and k-means clustering of frozen geometric features.Source: §4 · p. 7

Assessing generation

  • Generation experiments train a masked prior over ProTokstr-seq tokens on CATH 4.2 chains up to 128 residues and assess samples by self-consistency.GeoBPE provides a parameter-matched designability comparison; larger models are excluded because their scale precludes a size-controlled comparison.Source: §4 · pp. 7–8

5 Results

ProTokstr performs consistently across StructTokenBench tasks while retaining high-fidelity reconstruction and supporting generation. Probing and ablations show that modality groups contribute complementary task-relevant information, while orthogonal regularization supports effective codebook diversity.

5.1 Representation, Reconstruction, and Generation

  • 0.939 lDDT on CATH 4.2 lets ProTokstr remain competitive downstream while reconstructing within 0.047 of ESM3’s native reconstruction.The paper contrasts this combination with ESM3’s near-perfect reconstruction but last-place downstream ranking, and AminoAseed’s opposite profile.Source: §5.1 · p. 8
  • 52% of sampled ProTok-prior backbones are designable, versus 21% for GeoBPE at a parameter-matched generator budget.Designability is defined by a redesigned sequence refolding to the sampled backbone with self-consistency TM-score (scTM) ≥0.5.Source: §5.1 · p. 8
  • ProTokstr combines high codebook utilization and normalized perplexity while encoding 64 bits per residue.Product quantization splits each residue across small codebooks, while orthogonal regularization keeps SimVQ projections well-conditioned.Source: §5.1 · p. 8
  • Mean rank 1.36 across 28 StructTokenBench splits, 1.5 positions ahead of AminoAseed (2.86), with top-two placement on 27 of 28 splits.ProTokstr has the highest category average across functional-site prediction, flexibility and dynamics, conformational discrimination, and remote homology.Source: §5.1 · p. 8, Table 1 · p. 9

5.2 Codebook Probing

  • 0.724 lDDT places ProTokstr’s structure tokens highest in matched-probe structure prediction, with inverse-folding recovery near the top, close behind DPLM-2.Matched probes isolate token information from each tokenizer’s decoder capacity.Source: §5.2 · p. 10
  • ProTokstr leads every structure tokenizer on all nine zero-shot label hierarchies, with retrieval margins suggesting codes carry signal beyond backbone reconstruction.The reported margins range from 0.012 P@1 at coarse CATH class to 0.185 at EC level 1.Source: §5.2 · p. 10, Table 2 · p. 10
  • 0.858→0.916→0.931 fold-level AUROC as catalytic-site inference adds sequence then function groups; on ATLAS RMSF, adding function regresses performance.ProTokstr-seq attains the best Spearman ρ on both ATLAS RMSF holdout splits.Source: §5.2 · p. 10

5.3 Ablation Study

  • 74.5% remote-homology macro F1 for the full model versus 64.1% for global structure only; removing orthogonal regularization drops projection effective rank from 80.4% to 33.3%.Without the regularizer, remote-homology performance returns to the structure-only baseline; the ablation table reports the evaluation scope and metrics.Source: §5.3 · pp. 10–11, Table 3 · p. 11

6 Discussion and Conclusion

Modality dropout trains active code groups to support reconstruction independently, enabling cross-modal transfer and inference from any modality subset using one checkpoint. ProTok unites representation, reconstruction, and generation, while its current evidence is bounded by data scale, preliminary generation tests, and task-dependent effects of adding modality groups.

6.1 Discussion

  • Modality dropout samples one of 15 non-empty subsets of four codebook groups, requiring active codes to support reconstruction and enabling any-subset inference from one checkpoint.The authors attribute cross-modal transfer to this training objective; the four 256-entry codebooks per group provide roughly 4 × 10^9 compositions per residue.Source: §6.1 · p. 11

6.3 Conclusion

  • ProTok’s single vocabulary serves representation, reconstruction, and generation, but does not establish superiority over specialized function predictors.The conclusion also reports top-two performance on 27 of 28 StructTokenBench splits and co-generation of backbones and sequences, with over half of sampled backbones designable.Source: §6.1 · p. 11, §6.3 · pp. 11–12

6.2 Limitations

  • The 311K-chain corpus trains the quantizer, fusion module, and decoders with frozen encoders, and is orders of magnitude smaller than production protein-language-model corpora.The authors suggest that longer training on more diverse structural data at PLM scale and broader functional annotation may yield further gains.Source: §6.2 · p. 11
  • Generation results are preliminary: the small prior was trained on CATH chains of at most 128 residues and evaluated on 50 samples.Scaling a ProTok-native language model to longer proteins and larger corpora, with conditional generation, is identified as future work.Source: §6.2 · p. 11
  • Adding modality groups is not universally beneficial: the protein-level function group helps residue-level functional-site prediction but can hurt ATLAS RMSF dynamics regression.Source: §6.2 · p. 11

A Training Details

Training uses distinct reconstruction targets for structure and sequence, confidence-weighted losses for predicted structures, and specified quantization and functional objectives. A second decoder fine-tuning stage adds confidence-signal prediction while keeping the relevant encoders and codebooks fixed.

Structure reconstruction targets

  • The Stage 1 structure decoder predicts pairwise backbone distances and directions, binned Cβ distances and direction cosines, and amino acid identity.Distances and directions use mean squared error; binned Cβ distances use 64-bin cross-entropy, direction cosines use 16-bin cross-entropy, and amino acid identity uses 20-class cross-entropy.Source: A Training Details · p. 16
  • AlphaFold-predicted chains weight each reconstruction term by per-residue pLDDT confidence, whereas PDB chains use uniform weighting.Source: A Training Details · p. 16

Loss coefficients

  • Equation 3 uses unit weights for reconstruction, VQ, sequence, and function losses, with λortho = 0.01 and λentropy = 0.1.The commitment and codebook alignment terms use λcommit = 0.25 and λquant = 1; functional InfoNCE uses temperature τ = 0.1, and proteins without UniProt annotations are excluded from its minibatch.Source: A Training Details · p. 16

Stage 2 decoder fine-tuning

  • Stage 2 fine-tunes a 30-layer ESM3-pretrained structure-token decoder on frozen Stage 1 structure and local-geometry tokens, adding pLDDT and PAE prediction heads.Sequence and function modalities are disabled; the encoder and structure and local-geometry codebooks remain fixed, and each auxiliary loss has weight 0.01.Source: A Training Details · p. 16

Optimization

  • Both stages train for 25 epochs on eight NVIDIA B200 GPUs, using AdamW at learning rate 5 × 10^-4 and effective batch size 256.Training adds σ = 0.1 Å coordinate noise to encoder inputs, uses clean reconstruction targets, and balances PDB and Swiss-Prot ratios per batch; Stage 1 and Stage 2 have approximately 146M and 632M trainable parameters, respectively.Source: A Training Details · p. 16

B Evaluation Setup

Evaluation covers seven StructTokenBench benchmarks, including supervised probes over frozen tokens and zero-shot conformational-switch comparisons. The appendix also specifies native reconstruction and retrieval protocols, reports seed variability, and lists data and model licenses.

StructTokenBench benchmarks and hyperparameter sweep

  • StructTokenBench spans seven benchmarks, 14 prediction tasks, and 28 test splits; 12 supervised tasks use a two-layer MLP over frozen tokens.Each tokenizer-task pair receives a 12-configuration sweep, and the best validation configuration is evaluated on every test split for that task.Source: B Evaluation Setup · pp. 16–17
  • The two conformational-switch tasks are evaluated zero-shot on fold-switching and apo-holo splits by aligning structure tokens and selecting the best temperature from five values.The reported criterion is the best Spearman correlation with the ground-truth labels.Source: B Evaluation Setup · p. 17

Seed variability

  • Across five-seed reruns, standard deviation is at most 2.7 macro F1 for remote-homology fold and at most 0.4 AUROC or Spearman ρ on other tasks.Variance is small relative to the best-baseline gap except ATLAS RMSF superfamily, where results are within seed variance; DPLM-2 leads on BioLIP2 catalytic fold across all seeds.Source: B Evaluation Setup · p. 17

Native reconstruction quality

  • Native reconstruction is measured as Cα lDDT on 1,120 CATH 4.2 test chains using each tokenizer’s native decoder, except ProTok uses its Stage 2 decoder.ProSST is excluded because it has no native decoder.Source: B Evaluation Setup · p. 17

Cross-modal probes

  • Zero-shot retrieval mean-pools frozen structure tokens and ranks candidates by cosine similarity, reporting top-1 precision across CATH, EC, and Gene Ontology label hierarchies.The same evaluation section describes inverse-folding amino-acid recovery and structure prediction probes trained over frozen structure tokens.Source: B Evaluation Setup · p. 18

Assets and licenses

  • The study uses publicly available data and models for non-commercial research under their respective terms, including CC0, CC BY, MIT, and Apache 2.0 licenses.The ESM3 open model uses the EvolutionaryScale Cambrian Non-Commercial License.Source: B Evaluation Setup · p. 18

C Codebook Analysis

ProTokstr combines broad codebook use and high normalized perplexity at high capacity, while utilization and fidelity vary across modality groups. Cross-group mutual information is uniformly small, with function codes essentially independent of the other groups.

Codebook analysis scope

  • Codebook diagnostics cover 1,120 CATH 4.2 test chains, with function-group statistics restricted to the 290 chains with annotations.Source: C Codebook Analysis · p. 18

Cross-tokenizer utilization and reconstruction

  • ProTokstr is the only compared tokenizer combining high utilization and normalized perplexity at 64 bits per residue, with UR 0.958 and PPL 0.769.ProSST nearly saturates its codebook (UR 0.997, PPL 0.871), but its k-means clustering populates every centroid by construction.Source: C Codebook Analysis · p. 18, Table S2 · p. 19

Per-group codebook utilization

  • The first three groups have UR 0.93–1.00 and PPL 0.75–0.94, while the protein-level function group has UR 0.40 and PPL 0.25.Function-group usage is bounded by the 290 annotated proteins, unlike the per-residue token counts of the other groups.Source: C Codebook Analysis · p. 18, Table S3 · p. 19

Per-group quantization fidelity

  • Global structure has the highest input-to-lookup fidelity at cosine similarity 0.354, versus 0.287 for local geometry, 0.276 for sequence, and 0.236 for function.The authors hypothesize that this ordering reflects differences in upstream encoder pretraining objectives.Source: C Codebook Analysis · pp. 18–19, Table S3 · p. 19

Cross-group mutual information

  • Cross-group normalized mutual information is uniformly small; global structure and local geometry share the most, while function is essentially independent of the other groups.The authors interpret this decoupling as indicating complementary signals under modality dropout, while noting that low NMI is necessary but not sufficient for complementarity.Source: C Codebook Analysis · p. 19

D Sequence Token Probing

ProTokseq improves on a fixed amino acid codebook across four StructTokenBench tasks, while remaining behind continuous ESM-2. This suggests the discrete bottleneck retains most, but not all, of the signal in continuous PLM states.

D Sequence Token Probing

  • ProTokseq beats the amino acid baseline on every task, while trailing continuous ESM-2 650M by 1.3 to 8.5 points across four tasks.The comparison uses the same 12-configuration hyperparameter sweep and selection protocol for all three representations.Source: D Sequence Token Probing · pp. 19–20, Table S4 · p. 20
  • The results motivate learned sequence codes over a fixed amino acid codebook, while indicating that discrete codes preserve most but not all continuous PLM signal.The interpretation is based on ProTokseq's consistent gain over the amino acid baseline and its remaining gap to ESM-2.Source: D Sequence Token Probing · pp. 19–20
  • Table S4 compares an amino acid baseline, ProTokseq, and continuous ESM-2 650M on four tasks, reporting AUROC and macro F1.All three representations use the same 12-configuration sweep and selection protocol.Source: D Sequence Token Probing · p. 19, Table S4 · p. 20

E Generation

ProTok tokens serve as a substrate for a masked generative prior that jointly samples residue-level product codes and decodes them into backbones and sequences. The reported ProTok priors are co-designable in 18% and 20% of samples, with comparison caveats including different folding oracles.

E Generation

  • A bidirectional masked transformer samples residue positions iteratively, jointly sampling per-group product codes at each committed position rather than flattening the 12-code stream.The prior is trained on CATH 4.2 chains of at most 128 residues, matching GeoBPE's reported generation regime.Source: E Generation · p. 20
  • Designability is assessed by redesigning each generated backbone with ProteinMPNN, refolding with ESMFold, and comparing the refolded structure by TM-score.A backbone is designable when its best scTM reaches 0.5; co-designability requires the prior's own generated sequence to refold to scTM ≥0.5.Source: E Generation · pp. 20–21
  • The 75M and 156M ProTok priors are co-designable 18% and 20% of the time, respectively; GeoBPE does not report co-designability.The 75M prior matches GeoBPE's generative model size, and the 156M prior performs comparably.Source: E Generation · p. 21, Table S5 · p. 21
  • The GeoBPE comparison is system-level rather than a controlled oracle match because it uses OmegaFold while ProTok uses ESMFold.The reported table compares designability, mean scTM, diversity, and uniqueness at a matched parameter budget.Source: E Generation · p. 21, Table S5 · p. 21

F Matched-Probe and Complementarity Tables

Matched probes show that ProTokstr has the strongest structure-prediction lDDT and nearly ties DPLM-2 in inverse-folding AAR. A separate table describes nested inference modes for evaluating codebook-group complementarity across five tasks.

F Matched-Probe and Complementarity Tables

  • The appendix reports full codebook-group complementarity and token-probe tables, while zero-shot functional retrieval appears in main-text Table 2.The appendix tables are summarized in Section 5.2.Source: F Matched-Probe and Complementarity… · p. 21
  • Table S7 evaluates three nested ProTok inference modes by varying input codebook count across five StructTokenBench tasks.It reports AUROC for functional sites and Spearman ρ for ATLAS RMSF.Source: Table S7 · p. 22
  • ProTokstr reaches 0.724 lDDT in structure prediction, the highest of the baselines, and 0.359 AAR in inverse folding, near-tying DPLM-2 at 0.369.Matched probes isolate token information content from each tokenizer's decoder capacity.Source: Table S6 · p. 21, F Matched-Probe and Complementarity… · p. 22

Related papers

Loading 2610.18481v1…