Protein Structure Tokenization via Geometric Byte Pair Encoding

TL;DR

Protein structure tokenizers have struggled to balance interpretability, multi-scale control, and integration with language models. GeoBPE introduces a geometry-grounded, BPE-like tokenizer that discretizes backbone geometry into a hierarchical vocabulary of structural motifs by alternating k-medoids clustering with differentiable SE$(3)$ drift correction. The approach achieves a Pareto-optimal compression–distortion frontier, strong OOD generalization, and improved downstream transfer across diverse tasks, while providing interpretable motifs aligned with CATH functional families and enabling unconditional backbone generation when paired with transformers. This architecture-agnostic framework lays the groundwork for structure-native protein language models with interpretable, multi-resolution representations of folds.

Abstract

Protein structure is central to biological function, and enabling multimodal protein models requires joint reasoning over sequence, structure, and function. A key barrier is the lack of principled protein structure tokenizers (PSTs): existing approaches fix token size or rely on continuous vector codebooks, limiting interpretability, multi-scale control, and transfer across architectures. We introduce GeoBPE, a geometry-grounded PST that transforms continuous, noisy, multi-scale backbone conformations into discrete ``sentences'' of geometry while enforcing global constraints. Analogous to byte-pair encoding, GeoBPE generates a hierarchical vocabulary of geometric primitives by iteratively (i) clustering Geo-Pair occurrences with k-medoids to yield a resolution-controllable vocabulary; (ii) quantizing each Geo-Pair to its closest medoid prototype; and (iii) reducing drift through differentiable inverse kinematics that optimizes boundary glue angles under an $\mathrm{SE}(3)$ end-frame loss. GeoBPE offers compression ($>$10x reduction in bits-per-residue at similar distortion rate), data efficiency ($>$10x less training data), and generalization (maintains test/train distortion ratio of $1.0-1.1$). It is architecture-agnostic: (a) its hierarchical vocabulary provides a strong inductive bias for coarsening residue-level embeddings from large PLMs into motif- and protein-level representations, consistently outperforming leading PSTs across $12$ tasks and $24$ test splits; (b) paired with a transformer, GeoBPE supports unconditional backbone generation via language modeling; and (c) tokens align with CATH functional families and support expert-interpretable case studies, offering functional meaning absent in prior PSTs. Code is available at https://github.com/shiningsunnyday/PT-BPE/.

Figures (9)

• Figure 1: GeoBPE tokenizes a protein into discrete motifs linked by boundary glue angles and learns a hierarchical vocabulary of frequent structural primitives via k-medoids and recursively merging Geo-Pairs; at each step glue angles are optimized with differentiable inverse kinematics to preserve the global fold. Tokenization yields a merge tree that provides multi-resolution and interpretable representations of protein structure.
• Figure 2: (Top) GeoBPE tracks a Geo-Pair Encoding, a dictionary mapping Geo-Pair keys to occurrences at all times. Each step pops the most frequent Geo-Pair key, gathers the occurrences and selects $K$ prototypes to add to $\mathcal{V}$. All occurrences are quantized to the closest prototype. Glue angles are optimized to correct for the drift introduced. (Bottom) Toy example with two backbones; we initialize residue-orientation modes using two prototypes (green & yellow), pop the frequent Geo-Pair (blue), quantize occurrences (rounded $\rightarrow$ sharp corners), and optimize glue angles (red and orange).
• Figure 3: Plots of $(\operatorname{BPR}(\mathsf{T};\mathcal{D}),\ \Delta(\mathsf{T};\mathcal{D}_{\mathrm{test}}))$ across tokenizers for $\Delta\in \{\text{RMSD},\text{LDDT}\}$. We vary $|\mathcal{V}|\in \{128,256,512,1024\}$ for VQ-VAE/AminoASeed and $|\mathcal{V}|\in \{600,2500,6000,21000\}$ for GeoBPE to sample multiple points; we observe GeoBPE sweeps a smooth tradeoff curve.
• Figure 4: We plot the BPR (purple), length (orange), backbone distortion (RMSD, LDDT) as $|\mathcal{V}|$ across BPE steps. Ref. backbone RMSD/LDDT (dotted lines) uses random angle values for all internal angles, sampled from the empirical angle distribution.
• Figure 5: We ran an ablation for GeoBPE version with $|\mathcal{V}|=600$, keeping all parameters the same but toggling whether glue opt is skipped in Alg. \ref{['alg:res-init-tokens']}. We visualize the original (center), GeoBPE (left) and GeoBPE without glue opt (right) backbone states.