E NSEMBITS: an alphabet of protein conformational ensembles Kaiwen Shi1,2 Carlos Oliver1,2,3 Department of Computer Science, Vanderbilt University 2 Center for AI in Protein Dynamics, Vanderbilt University 3 Department of Molecular Physiology and Biophysics, Vanderbilt University [email protected] [email protected]
arXiv:2605.13789v1 [cs.LG] 13 May 2026
1
Abstract Protein structure tokenizers (PSTs) are workhorses in protein language modeling, function prediction, and evolutionary analysis. However, existing PSTs only capture local geometry of static structures, and miss the correlated motions and alternative conformational states revealed by protein ensembles. Here we introduce E NSEMBITS, the first tokenizer of protein conformational ensembles. E NSEMBITS address challenges inherent to tokenizing dynamics: deriving informative geometric descriptors across conformations, permutation-invariance encoding of variable-size ensembles, and conquering sparsity in dynamics data. Trained with a Residual VQ-VAE using a frame distillation objective on a large molecular dynamics corpus, E NSEMBITS outperforms all related methods on RMSF prediction, and is the strongest standalone structural tokenizer on an token-conditioned ANOVA test on per-residue motion amplitude. E NSEMBITS further matches or exceeds static tokenizers on EC, GO, binding site/affinity prediction, and zero-shot mutation-effect prediction despite using far less pretraining data. Notably, the distillation objective enables E NSEMBITS to predict dynamics token from one single predicted structure, which alleviates dynamics data sparsity. As the field moves from static structure prediction toward ensemble generation, E NSEMBITS offer the discrete vocabulary needed to bring dynamics into protein language modeling and design. Our codes can be found at: https://github.com/OliverLaboratory/Ensembits_release.
1
Introduction
The effort to decode the language of protein function has progressed steadily from mapping DNA sequences to solving 3D structure, and is currently facing the challenge of predicting and capturing conformational dynamics. The ensemble of conformations explored by a protein interfaces directly with biological function: catalytic loops in enzymes such as adenylate kinase open and close to gate substrate access, and G-protein coupled receptors transit between inactive and active states to relay extracellular signals. In each case, function is encoded not in any single snapshot but in the distribution of conformations a protein samples and the correlated motions that connect them. It stands to reason that such dynamical behaviours form the protein language that is shaped by evolution, akin to DNA codes, can be exploited by ML tools such as protein language models (PLMs) if properly captured. Grasping such a code will lead to improved protein structure and function prediction, search and alignment, and ultimately protein design. In ML terms, the objects we allude to are tokens. Tokens are the substrate upon which PLMs learn to reason, and as our grasp of protein behaviour improves, these tokens have improved in granularity and information richness from simple amino-acid sequences to structure-aware alphabets [van Kempen et al., 2022, Lin et al., 2023b, Sun et al., 2026], and to all-atom tokenization which enables modeling Preprint.
beyond proteins [Liu et al., 2025a]. Each step has unlocked downstream possibilities: 3Di tokens drive Foldseek’s billion-fold speedup in structural search [van Kempen et al., 2022]; ESM3’s joint sequence-structure-function vocabulary enables generative design [Lin et al., 2023b]; and discrete structural tokens have been used as drop-in inputs for function prediction, remote homology detection, and inverse folding [Heinzinger et al., 2023, Liu et al., 2025b, 2024]. A more expressive alphabet lets downstream models reason about a richer slice of protein biology. Access to rich dynamics data is growing rapidly. Large-scale MD repositories such as mdCATH [Mirarchi et al., 2024], ATLAS [Vander Meersche et al., 2023], and MISATO [Siebenmorgen et al., 2024] now provide thousands of all-atom trajectories spanning diverse fold space, replicas, and temperatures, offering an unprecedented supervisory signal for learning dynamical priors directly from physics. Concurrently, the next generation of protein structure prediction is focusing on ensemble prediction: BioEmu [Lewis et al., 2025] emulates equilibrium ensembles at scale, AFsample2 [Kalakoti and Wallner, 2024] perturbs AlphaFold’s MSA [Jumper et al., 2021] and dropouts to sample alternative conformations, and methods such as AlphaFlow,ESMDiff, and Distributional Graphormer [Jing et al., 2024, Lu et al., 2025, Zheng et al., 2024] recast structure prediction as a generative problem over conformational distributions. Together, these advances mean that ensemble data will not remain the bottleneck for long; instead, the bottleneck will shift to representation. Efforts to represent the high-dimensional ensemble data digestable by transformer architectures have recently begun. Portal et al. [Portal et al., 2026] and the [Guo et al., 2025] show enhanced function prediction through graph-based representations using distograms from Boltz2 Passaro et al. [2025] and MD-derived residue correlations respectively. PETIMOT [Lombard et al., 2025] represents motions at the residue level by encoding the principal angles of backbone motion and is able to reconstruct some dynamic signatures from sequence. ProtProfileMD [Lüth et al., 2026] encodes distributions of structural states as histograms built on 3Di tokens [van Kempen et al., 2022] (thus not a tokenizer) to enhance language models, but doing so over an alphabet trained on static structure forfeits joint modeling of correlated motion across residues. While these methods encode various dynamical features, none have yet described a full dynamics-aware token alphabet that jointly compresses local correlated motion patterns into a discrete vocabulary. We propose E NSEMBITS, the first tokenizer of protein conformational ensembles. E NSEMBITS maps an unordered set of structural frames (protein conformers) to a discrete token through three coupled design choices: (i) SE(3)-invariant descriptors computed across frames whose neighbour identities are themselves frame-dependent, so contact formation and breakage along the trajectory enter the descriptors directly; (ii) a permutation-invariant set encoder so that variable-size ensembles map deterministically to a single latent regardless of frame order; and (iii) a single-frame-to-token distillation objective that pulls the embedding of a sub-sampled ensemble—down to a single frame— toward that of the full ensemble, enabling the tokenizer to be queried at test time even when only a static structure is available. The result is a residual-VQ codebook over local dynamical motifs that can be plugged into any downstream PLM in place of, or alongside, existing static structural tokens. Our contributions are: • Formulation, model, and training pipeline. We formulate ensemble tokenization as a multiset-to-token compression problem, design SE(3)-invariant descriptors of protein frames, and train an end-to-end tokenizer of protein dynamics. • Single-frame-to-token distillation (SFTD). We devise a training objective that aligns sub-ensemble embeddings with their full-ensemble counterparts, letting one tokenizer serve both full/partial trajectories and single static frames at inference time. • Empirical validation across dynamics, function, and mutation. We empirically validate that E NSEMBITS leads RMSF prediction, dominates ANOVA tests on token-dynamics correlation, and matches or exceeds existing static structural tokenizers on EC, GO, bindingsite prediction, and zero-shot mutation-effect prediction with far less pretraining data.
2
2
Method
2.1
Problem formulation and Ensembits
Let P = {x1 , . . . , xP } be an unordered multiset of P frames of a single protein, where each frame xp ∈ RL×d collects the (per-atom) coordinates of the L residues in conformation p. Here d = 3A encodes the A atoms retained per residue — e.g. A = 1 for Cα-only. We write xpr ∈ Rd for the stacked atom coordinates of residue r in frame p. An SE(3)-invariant descriptor function ϕ : RL×d × {1, . . . , L} −→ RDf ,
ϕ(xp , r) = frp
maps a frame and a residue index to a per-residue, per-frame descriptor of dimension Df . A dynamics tokenizer is then a function [ P T : RD f −→ C = {C1 , . . . , CM } P ≥1
that sends the per-residue descriptor multiset {fr1 , . . . , frP } to a token in a codebook C of M entries, where each token Cm carries an embedding qCm ∈ RD . Crucially, T is P -agnostic: the union over P ≥ 1 in the domain expresses that it accepts a multiset of any positive cardinality. E NSEMBITS realises this map with a Residual Vector-Quantized Variational Autoencoder (RVQVAE) [Zeghidour et al., 2021] (Figure 1). The per-residue descriptor multiset {fr1 , . . . , frP } is passed through a permutation-invariant set encoder, which produces a single latent vector z. A residual vector quantizer with K levels discretizes z into a tuple of codebook indices (the residue’s dynamic token) and outputs the corresponding quantized embedding q. A decoder then maps q back to P descriptor vectors, supervised by a permutation-invariant (Hungarian-matched [Kuhn, 1955]) reconstruction loss against the input descriptor multiset. Training jointly minimizes this multiset-reconstruction loss and the standard VQ commitment loss.
Figure 1: The Ensembits tokenization pipeline. The P -frame per-residue descriptor multiset is encoded by a permutation-invariant set encoder, residually quantized into a token tuple (C1 , . . . , CM ) with token embedding q, and decoded back to P descriptors under a Hungarian-matched reconstruction loss. Codebooks are updated by EMA; the alternative gradient-based codebook loss is shown for completeness. 2.2
Descriptors
Ensembits captures protein dynamics by reconstructing per-residue SE(3)-invariant descriptor multisets. We consider two routes, both feeding the same multiset objective. The 3Di-style route extends Foldseek [van Kempen et al., 2022]: for each residue we gather its k nearest neighbors and compute a 3Di-style pairwise descriptor with a backbone ψ dihedral. We use 3
k-NN rather than Foldseek’s single neighbor because conformational motion cannot be summarized by one pair-interaction alone. Between consecutive neighbors we insert a “glue” descriptor inspired by GeoBPE [Sun et al., 2026], though not identical. The ESM3-style route replaces the 3Di + ψ + glue construction with relative-frame descriptors [Hayes et al., 2025, Lu et al., 2025]: each neighbor contributes its relative SE(3) frame to the anchor’s N/Cα/C triple. This is our production setting; we acknowledge it as a design caveat we cannot prove, and we chose it because relative frames are nearly an information upper bound on 3Di-style descriptors and performed empirically better. A complete list of components for both routes is in Appendix A. On top of this descriptor set, we propose three ways to capture the intrinsic dynamical changes within a conformational ensemble: modes we call fixed, dynamical, and fused, described in Appendix A.4. For our final tokenizer we adopt the dynamical mode: it yields a smaller per-residue descriptor than the fused mode, does not require a canonical ordering of the frames like the other two, and enables a distillation trick that makes single-frame to token-sequence prediction trivial. We expand on this last point in Section 2.4. For convenience, we describe the dynamical mode here. Dynamical mode. For each frame p ∈ {1, . . . , P } and each residue r ∈ {1, . . . , L}, we (i) compute residue r’s K nearest neighbors using the Cα positions in frame p, obtaining a frame-specific neighbor set Nrp = {npr,1 , . . . , npr,K } ⊆ {1, . . . , L} \ {r}, where npr,i is the index of residue r’s i-th nearest neighbor in frame p (sorted closest to farthest), and (ii) compute the per-residue, per-frame feature vector frp = ϕ(xp , r) ∈ RDf from the per-pair feature block; the dynamical mode is style-agnostic and works with either descriptor family from Appendix A. Each residue’s descriptor is therefore the multiset {fr1 , . . . , frP } of P per-frame feature vectors, which the set encoder ingests directly. Because the neighbor identities Nrp are allowed to vary with p, the descriptor naturally encodes contact formation and rupture along the trajectory. The per-frame dimension Df is independent of P , so the same encoder ingests ensembles of any size, which is critical for SFTD in Section 2.4. 2.3
Tokenizer
E NSEMBITS relies on RVQ-VAE [Zeghidour et al., 2021] for quantization in the latent space. We detail the individual model components in the following section. Set encoder. The set encoder maps the per-residue descriptor multiset {fr1 , . . . , frP } ⊂ RDf to a single latent vector z ∈ RD . We adopt a PerceiverIO-style design [Jaegle et al., 2022] that is permutation-invariant by construction, removing any need to learn invariance from data. A shared per-element MLP first embeds each descriptor frp independently, after which a small set of learnable query tokens cross-attends to the P embeddings, taking a softmax-weighted average over them; the queries are then refined by a stack of self-attention plus feed-forward blocks operating only among themselves, and finally concatenated and linearly projected to z. Permutation invariance follows because (i) the per-element MLP is shared across descriptors, (ii) the softmax-weighted average in cross-attention is invariant to the order of the key-value pairs, and (iii) all subsequent operations act on the queries, whose ordering is fixed by the architecture and independent of the input. Residual quantizer VQ-VAE [van den Oord et al., 2018] discretizes continuous embeddings via a learnable neural net and has been applied in the structural biology field to tokenize protein structures [van Kempen et al., 2022]. RVQ-VAE [Zeghidour et al., 2021] is a multi-stage extension [Vasuki and Vanathi, 2006] in which a continuous embedding is approximated by a sum of K codewords drawn from K independently learned codebooks C1 , . . . , CK . Given the latent z from the set encoder, the residual quantizer Q outputs a token index tuple (c1 , . . . , cK ) and the corresponding summed quantized embedding q =
K X
Ccℓℓ ,
cℓ =
arg min i∈{1,...,Mℓ }
ℓ=1
4
ρℓ−1 − Ciℓ ,
with the following terms: ℓ • Cℓ = {C1ℓ , . . . , CM } ⊂ RD is the level-ℓ codebook, a learned set of Mℓ codewords in the ℓ encoder’s latent space. • cℓ ∈ {1, . . . , Mℓ } is the discrete codebook index selected at level ℓ, the residue’s level-ℓ token. Each cℓ is the argmin of the distance between the current residual ρℓ−1 and the candidate codewords in Cℓ .1 • Ccℓℓ ∈ RD is the codeword at that index, i.e., the level-ℓ centroid embedding selected from Cℓ and contributed additively to q. • ρ0 = z is the encoder latent, and ρℓ = ρℓ−1 − Ccℓℓ is the residual at level ℓ, the portion of z not yet explained by levels 1, . . . , ℓ.
To ensure stable training, we have implemented E NSEMBITS with EMA codebook update [van den Oord et al., 2018], dead-code revival [Dhariwal et al., 2020], and cosine schedule [Loshchilov and Hutter, 2017]. The full algorithm and training setup are described in Appendix B.1. A practical concern with RVQ-VAE is that the total addressable codebook size grows multiplicatively: QK the token tuple (c1 , . . . , cK ) ranges over ℓ=1 Mℓ possibilities, which explodes with even two or three levels, provided that each Mℓ is moderately large. In practice we found that training on all levels but using only the first-level tokens is sufficient for good downstream representations and works better than a plain single-codebook VQ-VAE at the same primary vocabulary size M1 = 2048. We suspect that subsequent levels absorb fine-grained, residue-specific variation, freeing C1 to encode the dominant dynamical mode in a denoised form. Decoder and reconstruction loss The decoder maps the quantized embedding q back to a multiset of P reconstructed per-frame descriptors F̂r = {f̂r1 , . . . , f̂rP } for each residue r. Similar to the set-encoder philosophy, we want the decoder to output a bag of descriptors with no pre-defined order. To this end, we run a Hungarian optimal matching [Kuhn, 1955] between the predicted multiset F̂r and the target multiset Fr = {fr1 , . . . , frP }: P
Lrecon F̂r , Fr
1 X p 2 f̂r − frπ(p) 2 , π∈SP P p=1
= min
(1)
where SP denotes the symmetric group on P elements (i.e., the set of all permutations of {1, . . . , P }). The Hungarian assignment is not differentiable, so at each training step we solve for the optimal matching π ⋆ and treat it as fixed during backpropagation, differentiating only through the matched MSE. The total reconstruction loss is the average of Lrecon (F̂r , Fr ) over all residues r in the protein. 2.4
Single-frame-to-token distillation
Next, we treat the question of how to handle situations where conformation frames are not available. Can we still recover the information learned from multi-conformation dynamics by just looking at one frame? To address that, we propose single-frame-to-token distillation (SFTD), which biases single-frame inputs towards their multi-frame counterparts. Since the dynamical mode (Section 2.2) admits an arbitrary number of frames in arbitrary order, in each training step we construct two ensembles P1 and P2 from the same protein, where P1 contains all Pmax available frames and P2 ⊆ P1 is a randomly sampled sub-multiset of size |P2 | = P ′ ∈ {1, . . . , Pmax }. The two ensembles are encoded by the same set encoder to produce latent vectors z1 and z2 , respectively. We pull z2 towards z1 under a stop-gradient on the full-ensemble side: 2 Ltotal = Lrecon + β · Lcommit + λ · z2 − sg[z1 ] 2 , (2) where sg[·] is the stop-gradient operator, Lrecon and Lcommit are the mean of the per-branch RVQVAE reconstruction and commitment losses across the two forward passes, and β, λ are tunable hyper-parameters. The stop-gradient ensures that z1 acts as a fixed teacher target while only the sub-ensemble branch z2 receives the distillation gradient. ℓ ℓ should be distinguished from the upper-case Ci , which denotes the i-th codeword (a continuous vector D in R ), and from the calligraphic Cℓ , which denotes the entire codebook (the set of Mℓ codewords at level ℓ). 1 The lower-case c
5
3
E NSEMBITS Validation & Results
3.1
Datasets
We train E NSEMBITS on a combined MD corpus that fuses two complementary all-atom sources: mdCATH-div, a homology-deduplicated derivative of the mdCATH dataset [Mirarchi et al., 2024] (∼5,400 CATH protein domains, 5 temperatures, 5 replicas per tuple), and MISATO [Siebenmorgen et al., 2024], an MD dataset of ∼17,000 protein–ligand complexes derived from PDBbind. Training directly on every (domain, temperature, replica) ensemble of mdCATH would introduce strong homology leakage between train and evaluation proteins, since CATH groups many sequenceredundant domains under the same family; we deduplicate at the H-superfamily level to mitigate this. We defer dataset-curation details and the combined-corpus construction to Appendix B.2. In case of data deficiency, one can use generative models such as BioEmu [Lewis et al., 2025] or CFRandom [Lee et al., 2025] to synthesize multiple frames per protein and treat that as a surrogate ensemble. We tested an augmentation of the training of E NSEMBITS (denoted E NSEMBITS _aug) on AFSample2 generated ensembles [Kalakoti and Wallner, 2024] and document the probe results in Section 3.3. However, the quality of the learned tokens then becomes a function of the generation quality. The production tokenizer is trained on the mdCATH-div + MISATO combined corpus only. 3.2
Inspecting the learned alphabet
To test whether individual codebook entries capture coherent local structural motifs, we visualize a few representative tokens by selecting three distinct-protein exemplar residues per token and overlaying them in a common reference frame. The detailed alignment algorithms for exemplar selection and additional token representatives can be found in Appendix C, including the t-SNE analysis of the codebook.
Figure 2: Token 1699 — a near-stationary local motif. Three distinct-protein exemplars (1e6vC00:99, 1acoA01:19, 3DCT:26), ranked left-to-right by encoder-latent distance dz to the codebook centroid. Each panel overlays five MD frames Kabsch-aligned on all Cα atoms with the highlighted regions excluded. The frames superimpose tightly both within each panel and across exemplars, with mean motion amplitude, the first PCA singular value, ⟨s1 ⟩ ≈ 2.2. Token 1699 (Fig. 2) illustrates a stationary motif: the three highlighted regions superimpose tightly both across the five frames within each panel and across the three protein exemplars, evidence that this token represents a rigid, well-defined local conformation. The exemplars are drawn from both training sources (mdCATH and MISATO), demonstrating that the codebook generalizes the same geometric pattern across heterogeneous corpora. Token 1063 (Fig. 3) illustrates the opposite regime: the anchor sits in a sequence-local fold context, and its sixteen canonical k-NN are almost entirely residues at small sequence offsets along the same chain segment (|i − j| ≤ 15 for nearly all sixteen). Across the five overlaid frames the highlighted fan out visibly, showing that this token encodes a flexible local fragment whose conformation varies substantially across the ensemble while its 3D-neighbour pattern is preserved. 6
Figure 3: Token 1063 — a flexible local motif. Three exemplars (1c6rA00:23, 1gytL01:140, 1b9iA02:103) ordered left-to-right by increasing dz (2.57, 2.79, 3.13). Same alignment and coloring as Fig. 2; mean motion amplitude ⟨s1 ⟩ ≈ 8.9, roughly four times higher than Token 1699. 3.3
Does E NSEMBITS encode flexibility?
We compare E NSEMBITS against a set of tokenizer baselines on the task of predicting Root Mean Square Fluctuation (RMSF), the average displacement per residue across frames. Specifically, SFTD empowers us to test against single-frame tokenizers, since often times multi-frame ensembles are not available. We describe the implementations of each baseline in Appendix D, and prove that our tokenizer captures the dynamics through an RMSF probe test. We probe the flexibility signal carried by each token vocabulary on two datasets: mdCATH-div (CATHH structure split, P = 10 frames) and MISATO (sequence/structure split, P = 8 frames) [Siebenmorgen et al., 2024]. For every protein, the label is the single-frame Cα root-mean-square fluctuation (RMSF) at 320 K, averaged across MD replicas. Single-frame features are obtained as described in Appendix D, and a single-frame MLP head is fit on the pooled-residue training set. Performance is measured by the Spearman rank correlation between predicted and ground-truth RMSF on held-out residues, averaged over 10 random seeds (Table 1). E NSEMBITS dominates the RMSF benchmark in every setting. Among multi-frame methods, E NSEMBITS, P=full is best on every column, beating the strongest baseline (ProtProfileMD_K) by ∼ 2.5 Spearman points on mdCATH-div, ∼ 11 on MISATO sequence, and ∼ 18 on MISATO structure. More strikingly, E NSEMBITS, P=1 is also the best single-frame tokenizer on every column, ahead of the second-best single-frame method (AminoAseed) by ∼ 5 points on mdCATH-div, ∼ 1.3 on MISATO sequence, and ∼ 4.3 on MISATO structure. The gain from P =1 to P =full (+14.6 Spearman points on mdCATH-div, and +1.9 on each MISATO split) quantifies the marginal value of feeding the full P -frame ensemble through our permutation-invariant set encoder at inference, yet the strength of E NSEMBITS, P=1 on its own confirms that SFTD transfers dynamics signal into serving single-frame tokenization. We also perform an ANOVA test to further show that meaningful dynamics signals are learned by the tokens in Appendix E. In this test, E NSEMBITS dominate and lead by a large margin. This ANOVA test serves as another proof of our claim that protein dynamics is encoded into E NSEMBITS. 3.4
Does E NSEMBITS alphabet carry functional information?
The preceding experiments validated that E NSEMBITS carry quantized dynamics information, and SFTD enables predicting dynamics from a single, stationary frame. In this section, we test the E NSEMBITS tokens on multiple downstream tasks to investigate whether learned dynamics can help with function prediction, even when the pretraining data is sparse. Mutation Effect Prediction We evaluate zero-shot mutational fitness prediction on the P ROTE IN G YM substitution benchmark [Notin et al., 2023], restricted to the 96 deep-mutational-scanning assays whose wild-type sequence length is at most 200 residues. For each assay we fold the wild-type sequence and every unique mutant sequence with ESMFold [Lin et al., 2023b] and tokenize each 7
Table 1: RMSF Spearman across two datasets, mean ± std over 10 seeds. mdCATH-div uses the CATH-H structure split with P = 10 frames; MISATO is reported on both the sequence (UniProtclustered) and structure (CATH-H clustered) splits with P = 8 frames. The best tokenizer is bold, and the second best underlined. Model
mdCATH-div
MISATO (seq)
MISATO (struct)
AA 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P = 1
0.125 ± 0.000 0.337 ± 0.000 0.009 ± 0.002 0.408 ± 0.001 0.031 ± 0.001 0.365 ± 0.001 0.458 ± 0.000
0.212 ± 0.001 0.360 ± 0.000 0.006 ± 0.001 0.500 ± 0.000 0.059 ± 0.001 0.461 ± 0.000 0.513 ± 0.000
0.235 ± 0.002 0.302 ± 0.001 0.004 ± 0.001 0.488 ± 0.002 0.024 ± 0.002 0.411 ± 0.001 0.531 ± 0.002
ProtProfileMD_K Vote_3Di E NSEMBITS _aug, P = full E NSEMBITS, P = full
0.579 ± 0.001 0.429 ± 0.000 0.592 ± 0.000 0.604 ± 0.001
0.420 ± 0.001 0.366 ± 0.000 0.523 ± 0.000 0.532 ± 0.001
0.367 ± 0.001 0.323 ± 0.001 0.545 ± 0.001 0.550 ± 0.001
predicted single-frame structure with our model at P =1. Following the rest of the paper, we score each variant from the L1 codeword only: we take theP sum, across positions, of the Euclidean distance between the wild-type and mutant L1 codewords, i ∥Cc1wt − Cc1mut ∥2 , and negate so that larger 1,i 1,i values predict more wild-type-like (and presumed more fit) variants. We do not expect a structural tokenizer to beat the ESM2-650M baseline on a benchmark designed for sequence-fitness scoring; we instead test whether our token-distance signal carries complementary information by combining it with ESM2 via a within-assay z-score blend α · z(E NSEMBITS) + (1 − α) · z(ESM2) at α=0.3 (the grid-best α for 4/5 tokenizers, including ours) and comparing to ESM2 alone. The same protocol is applied to four baseline structural tokenizers; results are reported in Table 2. Judging from Table 2, one can conclude that E NSEMBITS, P=1, even as a distilled model serving one single predicted structure, outperforms every baseline structural tokenizer when blended with ESM2. When standalone, it only trails ESM3struct by a small margin, which is predictable given that ESM3struct was trained with inverse folding loss that aids sequence-related tasks. Table 2: Zero-shot mutational fitness on P ROTEIN G YM (96 DMS assays, L ≤ 200). Per-assay Spearman correlation between predicted score and experimental DMS score, averaged across assays. All tokenizers consume ESMFold-predicted atom14 backbones (real N/Cα/C/O/Cβ); no idealgeometry reconstruction. The “+ ESM2” rows show the within-assay z-score blend at α=0.3 (grid-best for four of five tokenizers); % gain is computed against ESM2 alone. The best result in each block is bold, the second-best is underlined. Mean Spearman
∆ vs ESM2
% gain
3di_tokens (Foldseek/mini3di) AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.327 0.366 0.358 0.384 0.381
−0.157 −0.118 −0.126 −0.100 −0.104
−32.4% −24.4% −26.0% −20.7% −21.4%
ESM2-650M alone
0.484
—
—
+ ESM2 blend (α=0.3): 3di_tokens + ESM2 AminoAseed + ESM2 ProToken + ESM2 ESM3struct + ESM2 E NSEMBITS, P=1 + ESM2
0.507 0.517 0.508 0.515 0.518
+0.023 +0.033 +0.023 +0.030 +0.033
+4.6% +6.8% +4.8% +6.3% +6.9%
Model
Enzyme Commission, Gene Ontology, and binding site/affinity prediction. We evaluate E NSEMBITS tokens on the MISATO downstream suite using the ProteinShake [Kucera et al., 2023] sequence, structure, and random-clustered splits, on Enzyme Commission classification at depths 1–3, 8
top-50 Gene Ontology term prediction, per-residue binding-site classification, and a complementary binding-affinity regression. Full per-task and per-split tables are in Appendix F.
Figure 4: Skill-score radar across MISATO downstream tasks. Each panel is one ProteinShake split; each axis is one (task, metric) pair. Skill is computed per axis and per split as skill(m) = (score(m) − score(R ANDOM))/(score(ESM3 STRUCT) − score(R ANDOM)), so the inner dashed polygon (0) is the random-token floor and the outer solid polygon (1) is the ESM3struct performance on the same split. Polygons further from the centre recover more of the available functional signal. AA and Random rows are excluded from the polygons for clarity. Seed variance is high. E NSEMBITS, P=full and E NSEMBITS, P=1 match or exceed every baseline on the binding-site and binding-affinity tasks across all three splits, with the largest margins on the structure split (skill scores up to 2.08 on binding-site AP and 1.41 on AUROC). On EC and GO they are competitive across splits. E NSEMBITS ties or beats ESM3struct on 7 of 11 structure-split axes, losing only on EC depth-1 and depth-3 where fold-correlated function and ESM3’s inverse-folding co-training prior dominate. On the sequence split, E NSEMBITS trails ESM3struct on several EC/GO axes, but on the random splits, E NSEMBITS, P=full/1 almost always maintain a top place. We remind the reader, however, that functional performance is not the main contribution of E NSEMBITS: our contribution is to quantize ensemble dynamics into a discrete vocabulary.
4
Discussion
In this work we formalized the problem of protein-dynamics tokenization, propose a concrete architecture for solving it, verify that the learned tokens faithfully capture ensemble-level dynamics, and demonstrate that the quantized representation transfers to downstream functional tasks. Limitations. Because E NSEMBITS tokens are learned end-to-end from observed conformational ensembles, their quality is bounded by the coverage and fidelity of the underlying trajectories. Current ensemble generators are either costly (MD) or not yet mature enough (ML) to serve as a reliable source of augmentation, although our experiments suggest this is a viable direction once the generators improve. We emphasize that the goal of this work is to introduce a methodology for instilling dynamical information into discrete tokens, not to deliver a universal alphabet covering all protein dynamics. Future directions. We highlight two directions we find particularly compelling. First, given access to substantially larger and higher-quality ensemble corpora, it would be valuable to scale the present architecture and study how the token alphabet evolves with model and data size. Second, because our tokens carry explicit dynamical information, they are a natural input modality for protein language models; we are especially interested in whether augmenting pLMs with dynamics tokens yields representations better suited to multi-frame structure generation. 9
5
Acknowledgment
We want to thank Benjamin P. Brown and Hassane Mchaourab for constructive discussions and feedback.
10
References Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music, 2020. URL https://arxiv.org/abs/2005.00341. Vladimir Gligorijević, P. Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C. Taylor, Ian M. Fisk, Hera Vlamakis, Ramnik J. Xavier, Rob Knight, Kyunghyun Cho, and Richard Bonneau. Structure-based protein function prediction using graph convolutional networks. Nature Communications, 12(1), May 2021. ISSN 2041-1723. doi: 10.1038/s41467-021-23303-9. URL http://dx.doi.org/10.1038/ s41467-021-23303-9. Pengkang Guo, Bruno Correia, Pierre Vandergheynst, and Daniel Probst. Boosting protein graph representations through static-dynamic fusion. bioRxiv, pages 2025–02, 2025. Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J. Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q. Tran, Jonathan Deaton, Marius Wiggert, Rohil Badkundri, Irhum Shafkat, Jun Gong, Alexander Derry, Raul S. Molina, Neil Thomas, Yousuf A. Khan, Chetan Mishra, Carolyn Kim, Liam J. Bartie, Matthew Nemeth, Patrick D. Hsu, Tom Sercu, Salvatore Candido, and Alexander Rives. Simulating 500 million years of evolution with a language model. Science, 387(6736):850–858, February 2025. ISSN 1095-9203. doi: 10.1126/science.ads0018. URL http://dx.doi.org/10.1126/science.ads0018. Michael Heinzinger, Konstantin Weissenow, Joaquin Gomez Sanchez, Adrian Henkel, Martin Steinegger, and Burkhard Rost. Prostt5: Bilingual language model for protein sequence and structure. bioRxiv, 2023. doi: 10.1101/2023.07.23.550085. URL https://www.biorxiv.org/content/ early/2023/07/25/2023.07.23.550085. Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, and Joāo Carreira. Perceiver io: A general architecture for structured inputs i& outputs, 2022. URL https://arxiv.org/abs/2107.14795. Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles, 2024. URL https://arxiv.org/abs/2402.04845. John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino RomeraParedes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583–589, July 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-03819-2. URL http://dx.doi.org/10.1038/s41586-021-03819-2. W. Kabsch. A discussion of the solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A, 34(5):827–828, 1978. doi: https://doi.org/10.1107/S0567739478001680. URL https://onlinelibrary.wiley.com/doi/abs/10.1107/S0567739478001680. Yogesh Kalakoti and Björn Wallner. Afsample2: Predicting multiple conformations and ensembles with alphafold2. bioRxiv, 2024. doi: 10.1101/2024.05.28.596195. URL https://www.biorxiv. org/content/early/2024/06/02/2024.05.28.596195. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL https://arxiv.org/abs/1412.6980. Tim Kucera, Carlos Oliver, Dexiong Chen, and Karsten Borgwardt. Proteinshake: Building datasets and benchmarks for deep learning on protein structures. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 58277–58289. Curran Associates, Inc., 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ b6167294ed3d6fc61e11e1592ce5cb77-Paper-Datasets_and_Benchmarks.pdf. 11
H. W. Kuhn. The hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2):83–97, 1955. doi: https://doi.org/10.1002/nav.3800020109. URL https: //onlinelibrary.wiley.com/doi/abs/10.1002/nav.3800020109. Myeongsang Lee, Joseph W. Schafer, Jeshuwin Prabakaran, Devlina Chakravarty, Madeleine F. Clore, and Lauren L. Porter. Large-scale predictions of alternative protein conformations by alphafold2based sequence association. Nature Communications, 16(1), July 2025. ISSN 2041-1723. doi: 10.1038/s41467-025-60759-5. URL http://dx.doi.org/10.1038/s41467-025-60759-5. Sarah Lewis, Tim Hempel, José Jiménez-Luna, Michael Gastegger, Yu Xie, Andrew Y. K. Foong, Victor García Satorras, Osama Abdin, Bastiaan S. Veeling, Iryna Zaporozhets, Yaoyi Chen, Soojung Yang, Adam E. Foster, Arne Schneuing, Jigyasa Nigam, Federico Barbero, Vincent Stimper, Andrew Campbell, Jason Yim, Marten Lienen, Yu Shi, Shuxin Zheng, Hannes Schulz, Usman Munir, Roberto Sordillo, Ryota Tomioka, Cecilia Clementi, and Frank Noé. Scalable emulation of protein equilibrium ensembles with generative deep learning. Science, 389(6761): eadv9817, 2025. doi: 10.1126/science.adv9817. URL https://www.science.org/doi/abs/ 10.1126/science.adv9817. Xiaohan Lin, Zhenyu Chen, Yanheng Li, Xingyu Lu, Chuanliu Fan, Ziqiang Cao, Shihao Feng, Yi Qin Gao, and Jun Zhang. Protokens: A machine-learned language for compact and informative encoding of protein 3d structures. bioRxiv, 2023a. doi: 10.1101/2023.11.27.568722. URL https://www.biorxiv.org/content/early/2023/11/27/2023.11.27.568722. Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637):1123–1130, 2023b. doi: 10.1126/ science.ade2574. URL https://www.science.org/doi/abs/10.1126/science.ade2574. Andrew Liu, Axel Elaldi, Nathan Russell, and Olivia Viessmann. Bio2token: All-atom tokenization of any biomolecular structure with mamba, 2025a. URL https://arxiv.org/abs/2410.19110. Dan Liu, Francesca Young, Kieran D. Lamb, Adalberto Claudio Quiros, Alexandrina Pancheva, Crispin J. Miller, Craig Macdonald, David L. Robertson, and Ke Yuan. Plm-interact: extending protein language models to predict protein-protein interactions. Nature Communications, 16(1), October 2025b. ISSN 2041-1723. doi: 10.1038/s41467-025-64512-w. URL http://dx.doi. org/10.1038/s41467-025-64512-w. Wei Liu, Ziye Wang, Ronghui You, Chenghan Xie, Hong Wei, Yi Xiong, Jianyi Yang, and Shanfeng Zhu. Plmsearch: Protein language model powers accurate and fast sequence search for remote homology. Nature Communications, 15(1), March 2024. ISSN 2041-1723. doi: 10.1038/ s41467-024-46808-5. URL http://dx.doi.org/10.1038/s41467-024-46808-5. Valentin Lombard, Sergei Grudinin, and Elodie Laine. Petimot: A novel framework for inferring protein motions from sparse data using se(3)-equivariant graph neural networks, 2025. URL https://arxiv.org/abs/2504.02839. Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts, 2017. URL https://arxiv.org/abs/1608.03983. Jiarui Lu, Xiaoyin Chen, Stephen Zhewen Lu, Chence Shi, Hongyu Guo, Yoshua Bengio, and Jian Tang. Structure language models for protein conformation generation, 2025. URL https: //arxiv.org/abs/2410.18403. Finn H. Lüth, Victor Mihaila, Milot Mirdita, Martin Steinegger, Burkhard Rost, and Michael Heinzinger. Protein language modeling beyond static folds reveals sequence-encoded flexibility. bioRxiv, 2026. doi: 10.64898/2026.01.21.700698. URL https://www.biorxiv.org/content/ early/2026/01/22/2026.01.21.700698. Antonio Mirarchi, Toni Giorgino, and Gianni De Fabritiis. mdcath: A large-scale md dataset for datadriven computational biophysics. Scientific Data, 11(1), November 2024. ISSN 2052-4463. doi: 10.1038/s41597-024-04140-z. URL http://dx.doi.org/10.1038/s41597-024-04140-z. 12
Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, Jonathan Frazer, Mafalda Dias, Dinko Franceschi, Yarin Gal, and Debora Marks. Proteingym: Largescale benchmarks for protein fitness prediction and design. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 64331–64379. Curran Associates, Inc., 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ cac723e5ff29f65e3fcbb0739ae91bee-Paper-Datasets_and_Benchmarks.pdf. Saro Passaro, Gabriele Corso, Jeremy Wohlwend, Mateo Reveiz, Stephan Thaler, Vignesh Ram Somnath, Noah Getz, Tally Portnoi, Julien Roy, Hannes Stark, David Kwabi-Addo, Dominique Beaini, Tommi Jaakkola, and Regina Barzilay. Boltz-2: Towards accurate and efficient binding affinity prediction. bioRxiv, 2025. doi: 10.1101/2025.06.14.659707. URL https://www.biorxiv.org/content/early/2025/06/18/2025.06.14.659707. Nicolas Portal, Wissam Karroucha, Vincent Mallet, and Massimiliano Bonomi. Learning dynamic protein representations at scale with distograms. February 2026. doi: 10.64898/2026.01.29.702509. URL http://dx.doi.org/10.64898/2026.01.29.702509. Till Siebenmorgen, Filipe Menezes, Sabrina Benassou, Erinc Merdivan, Kieran Didi, André Santos Dias Mourão, Radosław Kitel, Pietro Liò, Stefan Kesselheim, Marie Piraud, Fabian J. Theis, Michael Sattler, and Grzegorz M. Popowicz. Misato: machine learning dataset of protein–ligand complexes for structure-based drug discovery. Nature Computational Science, 4(5):367–378, May 2024. ISSN 2662-8457. doi: 10.1038/s43588-024-00627-2. URL http://dx.doi.org/10.1038/s43588-024-00627-2. Michael Sun, Weize Yuan, Gang Liu, Wojciech Matusik, and Marinka Zitnik. Protein structure tokenization via geometric byte pair encoding, 2026. URL https://arxiv.org/abs/2511. 11758. Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning, 2018. URL https://arxiv.org/abs/1711.00937. Michel van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Foldseek: fast and accurate protein structure search. Biorxiv, pages 2022–02, 2022. Yann Vander Meersche, Gabriel Cretin, Aria Gheeraert, Jean-Christophe Gelly, and Tatiana Galochkina. Atlas: protein flexibility description from atomistic molecular dynamics simulations. Nucleic Acids Research, 52(D1):D384–D392, November 2023. ISSN 1362-4962. doi: 10.1093/nar/gkad1084. URL http://dx.doi.org/10.1093/nar/gkad1084. A. Vasuki and Ponnusamy Thangapandian Vanathi. A review of vector quantization techniques. IEEE Potentials, 25:39–47, 2006. URL https://api.semanticscholar.org/CorpusID: 40842602. Xinyu Yuan, Zichen Wang, Marcus Collins, and Huzefa Rangwala. Protein structure tokenization: Benchmarking and new recipe, 2025. URL https://arxiv.org/abs/2503.00089. Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Soundstream: An end-to-end neural audio codec, 2021. URL https://arxiv.org/abs/2107.03312. Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, He Zhang, Shidi Tang, Hongxia Hao, Peiran Jin, Chi Chen, Frank Noé, Haiguang Liu, and Tie-Yan Liu. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, 6(5):558–567, May 2024. ISSN 2522-5839. doi: 10.1038/s42256-024-00837-3. URL http://dx.doi.org/10.1038/s42256-024-00837-3.
13
A
Descriptor specification
This appendix gives the full definition of the E NSEMBITS descriptors used as input to the RVQ-VAE. We describe two descriptor families sharing the same neighbour-selection machinery but differing in their per-neighbour feature blocks: a 3Di-style descriptor built from Cα unit vectors, and an ESM3-style relative-frame descriptor built from full backbone N/Cα/C atoms. A.1
Conventions
Inputs are the per-frame coordinates xp ∈ RL×d for p ∈ {1, . . . , P }, with d = 3A encoding the A atoms retained per residue. The 3Di-style descriptor uses only Cα, optionally augmented with backbone N, C (and Cβ) when those are available; the ESM3-style descriptor requires real backbone N, Cα, C. Where backbone N/C are absent, they are reconstructed from consecutive Cα positions using ideal peptide geometry (canonical |N−CA| = 1.46 Å, |CA−C| = 1.52 Å, and the N–CA–C angle ≈ 111◦ ). We write Cαr for the Cα position of residue r in the current frame, and ua→b = unit(Cαb − Cαa ) for the unit vector between two Cα atoms. Neighbours are found via Foldseek-style virtual centers (3Di descriptor) or via ESM3’s kNN-onCα-distance routine (ESM3 descriptor). Within a frame, the k nearest neighbours of residue i are selected by neighbour-set distance; the 3Di descriptor enforces a minimum sequence separation of |i − j| > 3, whereas the ESM3 descriptor places no sequence-gap filter. A.2
3Di-style descriptor
Per-neighbour 3Di features (10D). For an ordered residue pair (i, j) we compute ten SE(3)invariant features from Cα-derived unit vectors: #
Feature
Definition
0 1 2 3 4 5 6 7 8 9
dij a1 a2 a3 a4 a5 a6 a7 seq1 seq2
∥Cαi − Cαj ∥ ⟨ui−1→i , ui→i+1 ⟩ (bend at i) ⟨uj−1→j , uj→j+1 ⟩ (bend at j) ⟨ui−1→i , ui→j ⟩ ⟨uj−1→j , ui→j ⟩ ⟨ui−1→i , uj→j+1 ⟩ ⟨ui→i+1 , uj−1→j ⟩ ⟨ui−1→i , uj−1→j ⟩ (backbone alignment) sign(i − j) · min(|i − j|, 4) sign(i − j) · log(|i − j| + 1)
Backbone ψ dihedral (4D, optional). The backbone ψ dihedral ψr at residue r is the torsion angle defined by the four consecutive backbone atoms (Nr , Cαr , Cr , Nr+1 ), encoded as (sin ψr , cos ψr ) to avoid the ±180◦ wrap. For a neighbour pair (i, j) we append a 4D block (sin ψi , cos ψi , sin ψj , cos ψj ). Termini and residues for which the next residue’s N is unavailable are zero-padded. When real backbone N/C are absent, ψr is computed from N/C atoms reconstructed via ideal peptide geometry; the descriptor module falls back transparently in this case. The dihedral is non-redundant with the 3Di block: the 3Di features involve at most three consecutive Cα per residue and so capture local bend, whereas ψr captures the backbone twist around the Cα–C bond — the dominant local degree of freedom in conformational change, and a strong single-value discriminator of secondary structure. Inter-neighbour glue (4D). Between two consecutive neighbours jm and jm+1 in the (ordered) neighbour list of residue i, we add a 4D block describing their relative geometry. With local backbone tangent dirr = unit(Cαr−1 → Cαr+1 ) and contact direction cd = unit(Cαjm → Cαjm+1 ): 14
#
Feature
Definition
0 1 2 3
Distance Alignment Approach Twist
∥Cαjm − Cαjm+1 ∥ ⟨dirjm , dirjm+1 ⟩ ⟨dirjm , cd⟩ ⟨dirjm+1 , cd⟩
The first neighbour in the list has no preceding neighbour and its glue block is omitted (not zeropadded). Per-residue layout. concatenation
Let n be the number of neighbour slots per residue. The descriptor is the
[ 3Di1 , dih1 , glue1→2 , 3Di2 , dih2 , glue2→3 , . . . , 3Din , dihn ], with dihm blocks present only when the ψ dihedral is enabled. With D3Di = 10, Dψ = 4, Dglue = 4: ψ dihedral
Dfirst
Drest
Df
on off
14 10
18 14
14 + (n − 1) · 18 10 + (n − 1) · 14
The total Df is the 3Di-family per-frame descriptor dimension referenced in the methods section. A.3
ESM3-style relative-frame descriptor
The ESM3-style descriptor encodes a residue’s local backbone environment as the SE(3) transformations between its own backbone frame and those of its K spatial neighbours, using the pretrained ESM3 structure-encoder utilities of Hayes et al. [2025], Lu et al. [2025] for frame construction and neighbour-finding (no learnable parameters inside the descriptor itself). Per-residue local frame. Given real backbone atoms (Nr , Cαr , Cr ) in frame p, we build a perresidue local frame Trp ∈ SE(3) via Gram–Schmidt orthonormalisation of (Nr − Cαr , Cr − Cαr ) (ESM3’s build_affine3d_from_coordinates). Per-neighbour relative frame (K · 12D). We select the K = 16 spatial nearest neighbours of r in Cα distance via ESM3’s find_knn_edges — with no sequence-gap filter, so chain-adjacent residues are included — and express each neighbour’s frame in r’s local coordinates: p,rel Tr→j = (Trp )−1 ◦ Tjp ,
j ∈ Nrp , |Nrp | = K.
Each relative transform is flattened to its 12-dimensional Affine3D.tensor representation (a 3 × 3 rotation matrix plus a 3-vector translation, all in r’s frame). The slate of K = 16 relative frames is concatenated to give a K · 12 = 192-dimensional per-residue, per-frame descriptor: p,rel p,rel frp = Tr→n ∈ R192 . p , ..., T r→np r,1
r,K
The descriptor is SE(3)-invariant by construction: a global rigid-body motion of the entire structure left-multiplies every Tjp by the same group element, which then cancels under the (Trp )−1 ◦ Tjp composition. Real backbone N/Cα/C are required and are sourced per dataset (mdCATH h5 trajectories, MISATO MD); when only Cα is available, N and C are reconstructed via ideal peptide geometry before frame construction. We use this descriptor for the shipped tokenizer; the 3Di-family descriptor remains as an option in the codebase. A.4
Neighbour selection by mode
The three modes below describe how the slate of n neighbour slots is populated; they apply identically to the 3Di and ESM3-style descriptors — only the per-pair feature block downstream of the slate differs. 15
Fixed.
For each residue r, we select a reference frame p⋆ (r) = arg max Rgloc (r, p), p∈{1,...,P }
where Rgloc (r, p) denotes the local radius of gyration of a small backbone window centered at r in frame p — equivalently, the conformation in which the local environment is most expanded. This gives us a permutation-invariant ordering. We compute the k nearest neighbours of r in this reference frame, obtaining the neighbour set Nr = {nr,1 , . . . , nr,k } ⊆ {1, . . . , L} \ {r}, and hold this set fixed across all frames. The per-frame feature vector frp ∈ RDf is then computed from the per-pair feature block (3Di + ψ + glue, or ESM3 relative frame) between r and Nr evaluated in frame p. The dimension is independent of P . Dynamical.
Rather than anchoring to a single frame, we recompute the k-nearest-neighbour set Nrp = {npr,1 , . . . , npr,k } ⊆ {1, . . . , L} \ {r}
independently in every frame p. The per-frame feature vector frp is computed between r and Nrp , so neighbour identities are themselves frame-dependent. The descriptor therefore captures contact formation and breakage along the ensemble, a signal that is invisible to fixed-neighbour variants. The dimension again matches the fixed mode and is independent of P . The shipped tokenizer uses this mode for the ESM3-style descriptor with K = 16 and no sequence-gap filter. Fused. The fused mode retains every neighbour that appears in any frame. For each residue r, we sort the P frames in decreasing order of Rgloc (r, p), then concatenate the per-frame k-nearestneighbour lists in this canonical order to form a fused neighbour list of length P · k. Duplicate occurrences are kept, so every residue receives a slate of fixed cardinality. Per-frame feature vectors frp are computed against this fused list in every frame. This mode retains the maximum information available from the ensemble at the cost of a P -dependent dimension (linear in P for both descriptor families), and won’t permit SFTD. A.5
Resulting dimensions
For the 3Di family with ψ dihedral on (the setting used in all reported 3Di-family experiments) and the ESM3-style family with the shipped K = 16: Family
Mode
k=1
k=2
k=3
k = 3, P = 5
k = 3, P = 10
3Di 3Di 3Di
Fixed Dynamical Fused
14 14 14
32 32 32
50 50 50
50 50 266
50 50 536
ESM3 (K = 16, fixed or dynamical) ESM3 (K = 16, fused, never used)
Df = 192 (P-independent) Df = 192 · P (e.g. 960 at P = 5, 1920 at P = 10)
16
B
Model Details
B.1
RVQ-VAE
RVQ-VAE [Zeghidour et al., 2021] is a multi-stage extension [Vasuki and Vanathi, 2006] in which a continuous embedding is approximated by a sum of K codebook entries, one drawn from each of K independently learned codebooks C1 , . . . , CK , with associated nearest-neighbor quantization operators Qℓ (ρ) = arg min∥ρ − e∥. e∈Cℓ
Concretely, given the latent z produced by the set encoder, the quantized embedding q =
K X
Qℓ (ρℓ−1 ),
ρ0 = z,
ρℓ = ρℓ−1 − Qℓ (ρℓ−1 ),
(3)
ℓ=1
is built up additively: each Qℓ returns its nearest entry from Cℓ to the input residual ρℓ−1 , and the codebook indices (c1 , . . . , cK ) produced along the way constitute the discrete token tuple. The full procedure is summarized in Algorithm 1. The straight-through estimator is applied to the summed quantized embedding q, so the encoder receives a single clean reconstruction gradient through the bottleneck while the codebooks themselves are updated separately (discussed below). Algorithm 1 Residual Vector Quantization Require: z, the latent produced by the set encoder; codebooks Cℓ with quantization operators Qℓ for ℓ = 1, . . . , K Ensure: the quantized embedding q 1: q ← 0 2: ρ ← z ▷ residual 3: for ℓ = 1 to K do 4: q ← q + Qℓ (ρ) 5: ρ ← ρ − Qℓ (ρ) 6: end for 7: return q The total addressable codebook size grows multiplicatively: with Mℓ = |Cℓ | codes in the ℓ-th QK codebook, the token tuple (c1 , . . . , cK ) ranges over ℓ=1 Mℓ possibilities, which explodes easily with even two to three levels. In practice, we found that using only the first-level codes is sufficient for good representations and works better at downstream tasks than a plain VQ-VAE. We suspect that subsequent levels absorb fine-grained, residue-specific variation, freeing the first level to encode the dominant dynamical mode in a denoised form — a luxury a plain single-codebook VQ-VAE does not have. We also evaluated the basis-reparameterization trick of AminoAseed [Yuan et al., 2025], which constrains each codeword to a fixed orthogonal basis composed with a learnable projection. We found it relatively ineffective in our setting as the explained dynamics variance dropped (see Section 3.3). We posit that such behavior is caused by our data scale. Since our training data is relatively small, the update per round uses almost all nodes, thus bypassing the distribution drift phenomenon in their paper. Beyond the standard RVQ-VAE objective, several training choices proved essential for stable convergence in our setting. We apply the EMA codebook update independently within each codebook Cℓ : with decay γ = 0.99 [van den Oord et al., 2018], for each code i ∈ {1, . . . , Mℓ } the running cluster mean of level-ℓ inputs assigned to that code is updated as (t)
(t)
Ni
(t−1)
= γ Ni
(t)
+ (1 − γ) ni ,
(t)
(t−1)
mi = γ mi
+ (1 − γ)
X
ρ′ ,
(t)
ei =
(t) ρ′ ∈Ai
(t)
(t)
mi
(t)
, (4)
Ni
(t)
where Ai is the multiset of level-ℓ residuals ρ′ assigned to code i at training step t, ni = |Ai |, and (t) ei is the updated codeword. EMA avoids the slow, gradient-driven oscillations we observed under 17
the loss-auxiliary update mode and lets the codebook track the encoder distribution directly. To guard against codebook collapse — where a small subset of codes captures all assignments and the rest stop receiving updates — we adopt the random-restart heuristic of [Dhariwal et al., 2020]: any code whose (t) EMA usage count Ni falls below one is reseeded to a uniformly sampled encoder output from the current batch, giving “dead” codes a fresh chance to participate. Finally, we anneal the learning rate from 10−3 down to 10−6 along a cosine schedule [Loshchilov and Hutter, 2017] over the full training horizon, which we found important once the encoder is wide and deep — without it, late-stage updates were large enough to perturb a converged codebook and cause utilization to oscillate. With these training modifications, we observed stable loss descent and near-100% code utilization at every level. A more detailed description on model specifics can be found in Appendix B.3. B.2
Dataset
The production E NSEMBITS tokenizer is trained on a combined corpus that fuses two complementary MD datasets: mdCATH-div, a homology-deduplicated derivative of the mdCATH all-atom MD corpus, and MISATO, an MD dataset of ∼17,000 protein–ligand complexes derived from PDBbind. Combining the two sources exposes the model to CATH-balanced single-domain dynamics from mdCATH-div together with the protein–ligand complex dynamics of MISATO, under a single descriptor and a single training recipe. The two source corpora and their union are described below. mdCATH-div. We start from the mdCATH dataset [Mirarchi et al., 2024], a large-scale all-atom MD corpus of CATH protein domains spanning ∼5,400 domains, 5 temperatures, and 5 replicas per tuple. Training directly on every (domain, temperature, replica) ensemble would introduce strong homology leakage between train and evaluation proteins, since CATH groups many sequenceredundant domains under the same family. To avoid this, we build a homology-deduplicated derivative corpus we call mdCATH-div as follows. For each H-superfamily (the finest level of the canonical CATH classification, grouping domains by structural and evolutionary homology), we select a single representative domain, yielding 2,442 ensembles drawn from 2,447 H-groups in total. For every representative we draw frames from the 320 K trajectory across all five replicas, stride-sampling every 10th saved frame, and apply Farthest-Point Sampling (FPS) on the Cα RMSD between candidate frames to select maximally diverse subsets of size K ∈ {3, 5, 10}. The resulting multi-frame file stores the selected Cα tensors, the corresponding pairwise RMSD matrices, and the CATH metadata (topology, architecture, homology). Splits are 80/10/10 at the H-superfamily level (1,953 train / 244 val / 245 test) so the tokenizer’s own validation and test partitions share no H-superfamily with its training set. MISATO. mdCATH covers CATH single-chain domains. To broaden coverage we add MISATO [Siebenmorgen et al., 2024], an MD dataset of 16,972 protein–ligand complexes derived from PDBbind, each simulated for 10 ns in explicit water following semi-empirical QM refinement of the starting PDB structure. For each entry we apply the same FPS procedure on the all-atom trajectory to obtain 8- and 10-frame ensembles, stored in the same per-protein layout as mdCATH-div so the two corpora can be consumed by a single descriptor pipeline. Combined training corpus. The training set is the union of mdCATH-div and MISATO — 19,414 protein ensembles in total (2,442 from mdCATH-div and 16,972 from MISATO), yielding ∼6.6 M training residues and ∼720 K validation residues per epoch. The two corpora are concatenated in manifest order with MISATO indices offset by the mdCATH-div size, and the training split is the union of each corpus’s own H-superfamily-disjoint train partition; validation and test residues are drawn from the corresponding held-out splits. Augmentation. It should be noted that in case of data deficiency, one can use generative models, such as BioEmu [Lewis et al., 2025], CFRandom [Lee et al., 2025], and etc., to generate multiple frames for the same protein and treat that as an ensemble. However, the quality of learned tokens will be then a function of the generation quality. For our production model, we only trained on mdCATH-div and MISATO. We experimented with this promising data augmentation, see Section 3.3. 18
B.3
Implementation details
We summarize the final architectural and training hyperparameters used to train the production tokenizer (combined mdCATH-div + MISATO corpus, ESM3-style nearest-neighbor descriptors at K=16, P =10, descriptor dimension D = 192). Architecture. • Set encoder (PerceiverIO-style): per-element MLP with hidden size 256, followed by one cross-attention block from nq = 8 learnable queries to the P descriptor embeddings, and a stack of nenc = 4 self-attention plus FFN blocks operating on the queries. All attention layers use h = 4 heads. Queries are concatenated and linearly projected to a latent of dimension dz = 128. • Residual quantizer: K = 3 levels with codebook sizes [L1 , L2 , L3 ] = [2048, 128, 128], giving an addressable token space of L1 · L2 · L3 ≈ 3.4 × 107 . At inference, only L1 tokens (M1 = 2048) are used. Codebooks are updated by EMA with decay γ = 0.99. Codes whose EMA usage count drops below 1 are reseeded to a uniformly sampled encoder output from the current batch. • Decoder: ndec = 3-layer MLP with hidden size 256 and GELU activations, mapping the quantized latent ŷ ∈ R128 to P = 10 descriptor vectors. • Total trainable parameters: ≈ 3.4M. Loss. The training objective is Ltotal = Lrecon + β Lcommit + λ MSE stop_grad(zfull ), zsub with commitment cost β = 0.5 and SFTD weight λ = 0.1, where Lrecon is the Hungarian-matched MSE of Eq. equation 1 averaged over the two SFTD branches.
Figure 5: Training curve. Optimization. AdamW [Kingma and Ba, 2017] with initial learning rate 10−3 , weight decay 10−5 , and a 1000-step linear warm-up followed by a cosine schedule decaying to 10−6 over the full training horizon. Batch size 4096, gradient clipping at norm 1.0. We train for at most 1000 epochs with early stopping on validation reconstruction loss (patience 40 epochs); the final run converged at epoch 195. Descriptors are standardized to zero mean and unit variance per feature using statistics computed on the training split; the same (µ, σ) are bundled with the model checkpoint for downstream inference. During training, the number of input frames is sampled uniformly from peff ∼ U{1, . . . , 10} at each step (the variable-P schedule used by SFTD), so the encoder sees every P ∈ [1, 10]. Hardware and runtime. Training was performed on a single NVIDIA H200 and converged in approximately 7.2 hours (∼ 25,964 seconds for 195 epochs over 6,557,466 training residues / 719,507 validation residues). Final codebook utilization. At convergence, the L1 = 2048-code primary codebook reaches 96.3% utilization on validation set (1,973 unique codes assigned at least once; perplexity ≈ 1114); both L2 = L3 = 128 refinement codebooks reach 100% utilization (all 128 codes used; perplexities ≈ 126.7 and ≈ 126.3).
19
C
Token visualizations
C.1
Tokens
For each token, exemplars are picked by ranking all of its assigned residues by Euclidean distance dz between the encoder’s pre-quantization latent and the token’s codebook centroid; panels are arranged left-to-right by increasing dz so the leftmost exemplar is the most token-central and the rightmost the most peripheral. For each exemplar residue (anchor), we identify the sixteen canonical k-nearest neighbours (defined as the residues most frequently chosen as descriptor neighbours across the P = 10 ensemble frames, matching the K = 16 kNN that the ESM3-style descriptor uses at training time; no sequence-gap filter is applied) and highlight the local 3-mer (j−1, j, j+1) around each. To reveal token-level dynamics rather than global rigid-body motion, we Kabsch-align [Kabsch, 1978] all P frames using all Cα atoms of the protein with the 3-mer residues excluded, applying the same rigid-body transform to the all-atom coordinates from the source trajectory (mdCATH at T = 348 K, replica 0, or MISATO MD; five evenly-spaced frames in both cases). We here show representative tokens from the learned codebook, each illustrated with five distinctprotein exemplars drawn from the combined mdCATH-div + MISATO corpus on which the tokenizer was trained. Rows are ordered from top to bottom by increasing average per-residue motion amplitude ⟨s1 ⟩2 , the leading singular value of the locally-aligned Cα frame matrix, so the top rows correspond to compact, near-stationary local motifs and the bottom rows to more dynamical motifs whose 3-mer fragments visibly fan out across the five overlaid frames. Within each row, columns are ordered left to right by increasing encoder-latent distance dz from the token’s codebook centroid, so the leftmost panel is the most “token-central” exemplar and the rightmost the most peripheral. The same local geometry recurs across unrelated proteins drawn from both source corpora: the anchor (red sphere) sits at a consistent position relative to its sixteen canonical k-NN 3-mers, demonstrating that the codebook captures protein-agnostic structural motifs rather than protein-specific or dataset-specific peculiarities. Most of the sixteen neighbours fall within a sequence-local window of |i − j| ≤ 15 (because nearby chain residues dominate the 3D-nearest-neighbour set for typical fold geometries); the remaining neighbours, when present, are long-range contacts from other secondary structure elements packed against the anchor’s local environment. Selection, alignment, and rendering follow the procedure described above.
2 See Appendix E for a more detailed definition of s . 1
20
Figure 6: Five representative tokens from the codebook (rows), ordered top to bottom by increasing average residue motion amplitude ⟨s1 ⟩ (compact-stationary → dynamical), and within each row by increasing encoder-latent distance dz to the codebook centroid (most → least token-central, left to right). Each cell shows one distinct-protein exemplar of the row’s token: the anchor Cα (red sphere), its sixteen canonical kNN 3-mers (cartoons coloured by frame index along the plasma colormap, five overlaid frames from the source trajectory – mdCATH at T = 348 K, replica 0, or MISATO MD), and the surrounding local backbone (gray, single frame). The same local geometry recurs consistently across proteins drawn from both mdCATH and MISATO; moving down the rows, the 3-mers spread out across the overlaid frames, indicating progressively larger conformational fluctuations encoded by the token.
C.2
Codebook structure
To inspect the geometry of the learned token vocabulary, we project the L1 primary codebook (M = 2048 entries, each in R128 ) to two dimensions via t-SNE on the cosine distance, and color each point by how often the code is used as a per-residue token assignment across the combined mdCATH-div + MISATO corpus on which the tokenizer is trained (∼7.3M residue-token assignments over 19,406 ensembles: 2,447 mdCATH-div domains and 16,972 MISATO trajectories); see Fig. 7. Two observations are notable. First, the codebook achieves near-complete utilization: 2045/2048 = 99.85% of codes are assigned at least one residue across the corpus, with no dead-code clusters — the three unused codes (gray) appear as isolated points scattered among the live ones, not as a connected region. This confirms that the EMA codebook update combined with random-restart dead-code revival succeeds at populating the entire vocabulary. Second, the per-code usage distribution remains close to uniform: the most-used token receives 15,371 residues while the mean over used codes is ≈ 3,573, a top-to-mean ratio of only 4.3× (and a p99/p50 ratio of 3.66×). In typical VQ-VAE training, this ratio can exceed 10× when the codebook is over-capacity — a few attractor codes 21
absorb most assignments and the rest atrophy. The fact that high- and low-usage codes are visually intermixed in the t-SNE further reinforces this: there is no pocket of attractor codes carving out a high-density region. Together these properties indicate that the codebook is neither collapsed nor over-allocated; it tiles the encoder’s latent manifold approximately uniformly, with each code contributing meaningfully to the tokenization.
Figure 7: t-SNE projection of the L1 primary codebook (M = 2048 entries, d = 128) for the production tokenizer (combined mdCATH + MISATO training, k = 16, P = 10). Each point is one codebook entry; the color of used codes encodes log10 (1 + usage count) across the combined corpus, while the 3 unused codes are drawn in light gray. The twelve most-used token IDs are annotated in red. Codebook utilization is 99.85% with a near-uniform usage distribution (top-to-mean ratio ≈ 4.3×), and the t-SNE shows no separation between high- and low-usage codes — the codebook tiles the latent manifold approximately uniformly.
22
D
Baselines
We describe the implementations of each baseline and E NSEMBITS here. When N/C atoms are unavailable in the source structure, missing backbone atoms are reconstructed with ideal geometry (canonical |N−CA| = 1.46 Å, |CA−C| = 1.52 Å, and the N–CA–C angle ≈ 111◦ ). For datasets we used in this paper, such approximation is not used. Single-frame methods: • aa: 20-dimensional one-hot encoding of the amino-acid identity at each residue. • 3Di_tokens: the Foldseek 3Di alphabet [van Kempen et al., 2022] applied to the single representative conformation of each protein. Each residue is mapped to one of 20 structural states by Foldseek’s pretrained k-means quantizer over its pairwise descriptor; we treat the resulting integer sequence as a token stream of vocabulary size 20 with their frozen centroid embeddings obtained from mini3di. • Random: a sanity-check baseline in which every residue is assigned a random integer in {0, . . . , K − 1}, with K matched to the codebook size of the corresponding E NSEMBITS run. This isolates the contribution of the learned vocabulary from the downstream classifier capacity. • E NSEMBITS, P=1: the single-frame E NSEMBITS tokenizer. Tokens are obtained by computing the same per-residue descriptor used at training time – 192-D ESM3-style K=16 nearest-neighbor relative SE(3) features – from a single conformation, and passing it through the multi-frame set encoder at P =1. Because the encoder is permutation-invariant over the frame axis, the P =1 inference mode requires no architectural change. The vocabulary is the primary RVQ codebook (K1 = 2048 in our final configuration). • Aminoaseed: a single-frame structural VQ tokenizer from StructTokenBench [Yuan et al., 2025]; the codebook_512x1024-linear-fixed-last checkpoint trained on PDB. Per residue, the encoder (consuming only N/Cα/C atoms) produces a VQ index in [0, 512) which we look up in the frozen 1024-D codebook. • ProToken: ProToken-1.0 [Lin et al., 2023a], a single-frame structural VQ tokenizer trained on PDB. Per residue, the encoder (consuming a minimal N/Cα/C/O backbone) emits a VQ index in [0, 512), which we look up in its frozen 32-D codebook. • ESM3struct: the structure-track VQ tokenizer of ESM3 [Lin et al., 2023b], specifically the StructureTokenEncoder of esm3-sm-open-v1. Per residue, the encoder (consuming N/Cα/C atoms) emits a VQ index in [0, 4096), which we look up in the frozen 128-D EMA codebook. Multi-frame methods: • ProtProfileMD_K: per-residue empirical distribution over the 20 mini3di structural states. For each residue we tokenize K MD frames independently and record the fraction of frames assigned to each state (the construction of [Lüth et al., 2026]); K=8 on misato and K=10 on MD-CATH. We additionally interleave each bin’s probability with the fixed 2-D mini3di centroid of that state, yielding a 60-D per-residue feature. • Vote_3di: majority-vote 3Di tokens across the available frames per residue. In case of a tie, mean-pooled to a 21-D protein feature. Captures the most-frequent local structural state. • E NSEMBITS, P=full: the multi-frame E NSEMBITS tokenizer run on the full P -frame ensemble of each protein. Identical model and codebook as E NSEMBITS, P=1; differs only in the number of conformations fed to the permutation-invariant set encoder (P =10 on mdCATH-div, P =8 on MISATO).
23
E
Tokens encode distinguishable local dynamics: an ANOVA test
Throughout the paper we claim that E NSEMBITS tokens capture ensemble dynamics rather than purely static structure. The clean way to test this claim is to ask: does knowing a residue’s token identity reduce uncertainty about its local fluctuation amplitude in MD? We answer this with a one-way analysis of variance (ANOVA) on the per-residue Cα motion amplitude conditioned on the primary token. Setup. For every residue r in the mdCATH-div corpus we have a P =10 Cα trajectory and a primary token assignment t(r) ∈ {1, . . . , M } with M = 2048. To extract a scalar measure of local fluctuation amplitude that is independent of global rigid-body motion, we Kabsch-align the 10 frames using the local Cα ball of radius 10 Å around residue r in frame 0, then compute the top singular value s1 (r) of the resulting (10, 3) aligned-coordinate matrix. Intuitively, s1 (r) measures how much the central Cα atom fluctuates along its dominant local motion direction relative to its neighbourhood; small s1 means a rigid residue, large s1 means a flexible one. We restrict the analysis to tokens with at least 80 residues assigned (852 of 2,048 tokens), giving n = 301,979 residues spanning 2,447 proteins. One-way ANOVA: variance decomposition. ANOVA decomposes the total variance of s1 into “token-explained” and “within-token-residual” components: X X XX 2 2 (t) (t) 2 = s1 (r) − s̄1 nt s̄1 − s̄1 + s1 (r) − s̄1 , (5) r
|
t
{z
SStotal
}
|
t
{z
SSbetween
}
|
r∈t
{z
SSwithin
}
(t)
where s̄1 is the corpus mean of s1 , s̄1 is the mean within token t, and nt is the count of residues assigned to t. The decomposition is exact and makes no parametric assumption on the distribution of s1 . Replacing each residue’s value by its token’s mean would leave SSbetween as the remaining variance — this is the variance that “token id alone” recovers. The headline summary is the proportion of variance explained, η2 =
SSbetween ∈ [0, 1]. SStotal
(6)
Equivalently, η 2 is the R2 of the one-hot-token regression s1 ∼ onehot(t). Two extreme cases anchor the interpretation: if all tokens have identical mean s1 then η 2 = 0 (token id is uninformative about motion amplitude); if every residue with the same token has the same s1 then η 2 = 1 (token id deterministically predicts amplitude). F-statistic and parametric p-value. The classical inferential question — is η 2 distinguishable from zero? — is answered by the F -statistic F =
SSbetween /(M − 1) ∼ F (M − 1, N − M ) SSwithin /(N − M )
under H0 ,
(7)
where H0 asserts that all tokens have the same true mean. Under H0 , both numerator and denominator estimate the same residual variance and one expects F ≈ 1. Values much larger than 1 indicate that the between-token variance exceeds what residual scatter alone could produce. Permutation null: a distribution-free check. The F -test assumes within-group normality and homoscedasticity, both of which are mildly violated for s1 (it is positive and right-skewed). We additionally run a non-parametric permutation test: shuffle the 852 token labels uniformly at random across the 301,979 residues, recompute η 2 on the shuffled data, and repeat 1,000 times. This produces an empirical null distribution against which the observed η 2 can be compared without any distributional assumption. The permutation null also calibrates a known finite-sample bias of η 2 : with 852 groups, even random labels can absorb a small amount of variance simply because every group gets its own mean parameter. 24
Feature motion amplitude s1 motion shape s1 /s2
η 2 (observed) 0.371 0.040
F -test
permutation null mean −300
F (851, 301,127) = 208.5, p < 10 F (851, 301,127) = 14.5, p < 10−300
0.0028 —
Table 3: One-way ANOVA on per-residue motion features grouped by primary token, n = 301,979 residues, M = 852 tokens with ≥ 80 residues each.
Results. Table 3 reports the test for two motion features: the dominant amplitude s1 (above) and the anisotropy ratio s1 /s2 (a scalar measure of how directional the local motion is). Token identity explains η 2 = 0.371 of the variance in motion amplitude: knowing the token reduces the uncertainty about a residue’s local fluctuation by 37%. The associated F -statistic of 208.5 is enormous — between-token variance is two orders of magnitude larger than what within-token scatter alone would produce. Even more telling, the permutation null places η 2 in a tight neighbourhood of 0.0028; not one of 1,000 random shuffles came close to the observed value. The real signal is approximately 131× the null mean, giving pperm < 10−3 and confirming that the F-test result is not an artifact of distributional violations. The motion-shape result η 2 = 0.040 is more modest but still highly significant. This makes sense: motion direction (encoded in s1 /s2 ) is geometrically harder to align across residues with different local frames, while motion amplitude is a scalar quantity that the codebook can carve into clean strata. Interpretation. The codebook is not interchangeable on dynamics: each token corresponds to a statistically distinguishable distribution of local motion. A useful comparison: a perfect oracle that mapped each residue to its own state would have η 2 = 1, a random codebook at the same vocabulary size has η 2 ≈ 0.0028, and E NSEMBITS sits at η 2 = 0.371 — roughly a third of the way to the oracle and ∼ 131× above a random codebook of the same size. The remaining 63% of variance is residue-specific noise the token does not capture (sequence context, position in the fold, side-chain identity), which is consistent with the codebook being a deliberately compact summary of geometry rather than a per-residue identifier.
Figure 8: ANOVA test that different tokens encode different dynamics. Per-residue motion amplitude s1 is the top singular value of the locally Kabsch-aligned P =10 Cα-frame matrix around each residue, computed on n = 301,979 residues across 852 tokens with ≥ 80 residues. (a) Histogram of (t) s1 values across all residues (blue), with the per-token mean s̄1 shown as a red tick on the lower axis for each token. Per-token means span a wide range of the corpus distribution, indicating that tokens map onto distinct motion-amplitude strata. The inset reports the ANOVA summary: η 2 = 0.371, F (851, 301,127) = 208.5, p < 10−300 . (b) Permutation null distribution of η 2 obtained by shuffling token labels 1,000 times (gray histogram), with the observed η 2 marked in red. None of the 1,000 random shuffles came within an order of magnitude of the observed value; the real signal is ∼ 131× the null mean.
25
Negative controls. The above η 2 value rules out the null hypothesis that token id is uninformative about motion amplitude, but it leaves open a second concern: are the tokens merely recovering a coarse confounder such as fold class, chain position, or protein size? To rule this out we re-run the same one-way ANOVA on s1 with three alternative grouping variables that should explain at most a small share of the amplitude variance: CATH structural class (4 broad fold types plus special bucket), residue position quintile within the chain (5 groups), and protein-length quintile (5 groups). Table 4 reports the resulting η 2 . None of the controls comes close to the token signal: CATH class and protein length each capture η 2 ≈ 0.054 and residue position only η 2 = 0.020, each about 7×–18× smaller than the token result. The codebook is therefore distinguishing residues at a much finer granularity than fold class or chain position, and is not reducible to either confounder. # groups M
η 2 on s1
E NSEMBITS
852
0.371
CATH structural class Protein-length quintile Residue position quintile (within chain)
5 5 5
0.054 0.054 0.020
Grouping variable
Table 4: One-way ANOVA on per-residue motion amplitude s1 across 301,979 residues, with grouping by E NSEMBITS primary token and three negative controls. Tokens explain ∼ 7× more variance than the strongest non-token baseline, ruling out fold-class or position-along-chain as alternative explanations of the result in Fig. 8.
Comparison to existing structural tokenizers. We run the same η 2 test using each baseline tokenizer’s per-residue token assignment as the grouping variable, on the same ∼302k mdCATHdiv residues (Table 5). E NSEMBITS has the strongest amplitude conditioning by a wide margin: η 2 = 0.371 against ≤ 0.128 for every other tokenizer. The gap separates into four regimes: (i) ProToken sits at η 2 = 0.006 — essentially indistinguishable from a random codebook of the same size, consistent with its training objective being single-conformation structural reconstruction with no flexibility signal; (ii) the single-frame structure tokenizers 3Di_tokens and ESM3struct cluster at η 2 ≈ 0.05 (∼ 7× less dynamics signal than E NSEMBITS), despite ESM3struct’s PDBscale pretraining; (iii) AminoAseed, a single-frame structural VQ-VAE with a learned 512-code vocabulary, sits at η 2 = 0.079 (about ∼ 5× below E NSEMBITS), suggesting its codebook has organically discovered some flexibility-correlated structural classes beyond the fixed-vocabulary 3Di alphabet; (iv) multi-frame aggregations of mini3di tokens (ProtProfileMD_K10 and Vote_3Di, both summarising 10-frame mini3di token distributions) recover the most, plateauing at η 2 ≈ 0.12–0.13, still ∼ 3× below E NSEMBITS. The closest baseline is within ∼ 3× of E NSEMBITS and the next tier within ∼ 5×, so the gap is not "an order of magnitude" once the multi-frame mini3di aggregations are included; nevertheless every non-E NSEMBITS row sits well below the half-height of the E NSEMBITS result. Tokenizer
M active codes
n residues
η 2 / R2 on s1
E NSEMBITS
852
301,979
0.371
Vote_3Di (mini3di plurality, K=10) ProtProfileMD_K10 (mini3di histogram, R2 ) AminoAseed ESM3struct 3Di_tokens (mini3di, single-frame) ProToken
20 20 384 1,369 20 412
308,293 308,293 307,578 193,613 308,713 304,931
0.128 0.121 0.079 0.054 0.051 0.006
Table 5: One-way ANOVA on per-residue motion amplitude s1 (locally Kabsch-aligned P =10 Cα frames) across mdCATH-div residues, grouped by primary token from each tokenizer. Active codes M counts only tokens with at least 80 residues assigned. For the continuous ProtProfileMD baseline we report the linear-regression R2 of s1 ∼ histogram, which is the natural analog of η 2 for continuous predictors. E NSEMBITS captures ∼ 3× the variance of the strongest baselines (Vote_3Di and ProtProfileMD_K10, both multi-frame mini3di aggregations), ∼ 5× that of singleframe AminoAseed, and ∼ 7× that of the single-frame structural tokenizers ESM3struct and 3Di, consistent with its multi-frame, dynamics-aware training objective.
26
F
Downstream task experiments
This appendix expands the downstream summary in Section 3.4. We give the full task setup, evaluation protocol, and per-task tables for EC classification, GO term prediction, binding-site identification, and binding-affinity regression on the misato structure split, the most pertinent of the sequence, structure, and random splits. The complete experiments on all splits are documented in Appendix G. The reader should note that the seed variance is high on most of the function task, and read the results accordingly. EC Function Prediction We evaluate E NSEMBITS tokens on Enzyme Commission (EC) class prediction [Gligorijević et al., 2021, Kucera et al., 2023] under the misato structure split, reporting mean ± std over 10 seeds. Each protein is annotated with one or more EC numbers; we collapse them to depth-1 (the top-level enzyme class) for the headline table and report depths 2 and 3 in Appendix G. Per-residue features are obtained as token embeddings, and a 1D-CNN classifier is fit on top: a stack of three Conv1D layers, masked mean+max pooling across residues, and a small MLP head producing per-class logits trained with multi-label binary cross-entropy. We report five metrics on the test split: top-1 hit — the fraction of test proteins whose highest-scored class is among the ground-truth labels; mAP (macro-averaged AP) and µAP (micro-averaged AP) — average-precision computed per class then averaged, and over all (protein, class) pairs pooled, respectively; and the macro / micro F1 at the 0.5 decision threshold ([email protected], µ[email protected]). Table 6 reports the depth-1 results; full per-depth breakdowns are deferred to Appendix G. Among the single-frame methods, ESM3struct leads every metric (top-1 0.882, µAP 0.943), with AminoAseed a close second on every metric (µAP 0.912, top-1 0.841). Two caveats apply here. First, both ESM3struct and AminoAseed were pretrained on PDB-scale structural corpora that almost certainly overlap with the misato test proteins, so EC depth-1 (a 7-class fold-correlated task) plays to their strengths. Second, ESM3struct’s tokenizer is trained jointly with an inverse-folding objective, so each token implicitly carries sequence information; the tokens are therefore not purely geometric, unlike the descriptors used by every other entry. Both caveats recur throughout this appendix. In the multi-frame block, E NSEMBITS, P=full is the best entry on every metric, ahead of Vote_3Di by ∼ 6–8 pp on top-1 and AP-style measures (top-1 0.839 vs. 0.774; µAP 0.871 vs. 0.818). However, E NSEMBITS, P=1 does not match ESM3struct at this depth, and E NSEMBITS, P=full also trails ESM3struct by a clear margin (Welch’s t-test on top-1: t=3.6, p=0.003). We attribute this gap to the two caveats above: ESM3struct’s PDB-scale pretraining and inverse-folding co-training both prime it for the 7-class fold-correlated EC depth-1 task, which is known for sequence-leaky. Table 6: EC depth-1 top-1 / mAP / µAP / F1 — misato structure split, mean ± std over 10 seeds. Bolding and underlining are separate for the single-frame (top) and multi-frame (bottom) groups. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P = 1
0.333 ± 0.000 0.548 ± 0.240 0.360 ± 0.042 0.841 ± 0.032 0.376 ± 0.005 0.882 ± 0.019 0.674 ± 0.239
0.346 ± 0.052 0.429 ± 0.084 0.212 ± 0.011 0.569 ± 0.061 0.265 ± 0.059 0.615 ± 0.042 0.462 ± 0.115
0.350 ± 0.050 0.533 ± 0.274 0.305 ± 0.028 0.912 ± 0.025 0.385 ± 0.061 0.943 ± 0.020 0.714 ± 0.260
0.080 ± 0.042 0.288 ± 0.141 0.108 ± 0.005 0.514 ± 0.060 0.118 ± 0.034 0.548 ± 0.059 0.365 ± 0.157
0.266 ± 0.140 0.520 ± 0.256 0.347 ± 0.048 0.830 ± 0.039 0.374 ± 0.032 0.869 ± 0.024 0.643 ± 0.284
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P = f ull
0.374 ± 0.000 0.774 ± 0.136 0.839 ± 0.033
0.247 ± 0.043 0.504 ± 0.061 0.543 ± 0.045
0.330 ± 0.039 0.818 ± 0.199 0.871 ± 0.045
0.087 ± 0.046 0.422 ± 0.089 0.456 ± 0.082
0.295 ± 0.155 0.765 ± 0.144 0.800 ± 0.125
Model
GO Label Prediction We evaluate top-50 Gene Ontology term prediction [Gligorijević et al., 2021, Kucera et al., 2023] under the misato structure split (Table 7). Each protein is annotated with the union of its three GO branches (molecular function, biological process, cellular component); we keep the 50 most frequent terms in the training set and treat the task as multi-label classification. The architecture, optimizer, loss, and metrics (top-1, mAP, µAP, [email protected], µ[email protected]) are identical to EC. In the multi-frame block, E NSEMBITS, P=full is best, with Vote_3Di a near second (top-1 0.692 vs. 0.676; µAP 0.846 vs. 0.849 — Vote_3Di edges ahead by 0.3 pp on this one metric; E NSEMBITS, 27
P=full wins the other four including the threshold-sensitive F1 scores). ProtProfileMD_K8 sits far below the other multi-frame methods on the F1 metrics, indicating that its histogram-of-3Di representation, while informative on top-1, leaves a less linearly separable decision surface for the per-class threshold head. In the single-frame block, AminoAseed leads every metric (top-1 0.704, µAP 0.859), with ESM3struct second (top-1 0.665, µAP 0.835). E NSEMBITS, P=1 is third (top-1 0.619), ∼ 5–9 pp behind the leaders — the gap is larger than on EC d1, suggesting GO’s 50-class setting separates the methods that have explicitly seen PDB-scale structural priors from those that have not. Notably, E NSEMBITS, P=full (multi-frame block) outperforms ESM3struct (single-frame block) on 4 of 5 metrics, indicating that ensemble dynamics close most of the gap to PDB-pretrained structural tokenizers on the harder 50-class problem. Table 7: GO top-50 — misato structure split, mean ± std over 10 seeds. Bolding and underlining are separate for the single-frame (top) and multi-frame (bottom) groups. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.354 ± 0.348 0.641 ± 0.052 0.294 ± 0.163 0.704 ± 0.048 0.468 ± 0.224 0.665 ± 0.073 0.619 ± 0.057
0.300 ± 0.141 0.400 ± 0.021 0.114 ± 0.005 0.444 ± 0.021 0.208 ± 0.108 0.428 ± 0.040 0.397 ± 0.033
0.500 ± 0.301 0.837 ± 0.026 0.247 ± 0.058 0.859 ± 0.018 0.492 ± 0.243 0.835 ± 0.101 0.781 ± 0.071
0.116 ± 0.186 0.304 ± 0.040 0.000 ± 0.000 0.374 ± 0.028 0.084 ± 0.127 0.342 ± 0.083 0.283 ± 0.073
0.299 ± 0.411 0.777 ± 0.072 0.000 ± 0.001 0.814 ± 0.022 0.257 ± 0.359 0.747 ± 0.220 0.681 ± 0.127
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.610 ± 0.073 0.676 ± 0.073 0.692 ± 0.073
0.199 ± 0.074 0.423 ± 0.020 0.428 ± 0.033
0.386 ± 0.166 0.849 ± 0.030 0.846 ± 0.047
0.040 ± 0.108 0.335 ± 0.038 0.347 ± 0.052
0.087 ± 0.252 0.790 ± 0.043 0.798 ± 0.065
Model
Binding Site We evaluate per-residue binding-site prediction under the misato structure split (Table 8). Each residue carries a binary label indicating whether it lies within 5 Å of any heavy atom of the bound ligand. Per-residue features are obtained as for EC and GO, and a per-residue 1D-CNN head with the same conv stack (no pooling — the prediction is per residue) outputs a binding-site logit, trained with masked binary cross-entropy. We report two metrics, both pooled over all test-set residues: AUROC (area under the ROC curve) measures ranking quality independent of decision threshold, and AP (average precision = area under the precision-recall curve) is the threshold-free metric of choice for this strongly imbalanced task (positive rate ≈ 0.16). Binding-site is the task where dynamics-aware tokenization shows its largest margin. E NSEMBITS, P=full achieves AUROC 0.750 and AP 0.471 on the structure split — a ∼ 16 pp AUROC and ∼ 28 pp AP gap over Vote_3Di, the next-best multi-frame baseline (AUROC 0.594, AP 0.191). More strikingly, E NSEMBITS, P=1 also leads the single-frame block (AUROC 0.719, AP 0.418), beating ESM3struct (0.678, 0.291) and AminoAseed (0.694, 0.318) by clear margins. This reverses the EC / GO pattern where ESM3struct and AminoAseed lead the single-frame group: binding-site detection rewards information about local conformational flexibility (interface residues are often the more flexible parts of a fold), which static PDB-pretrained tokenizers cannot encode. The dynamics-aware multi-frame ensemble extends the margin further, with E NSEMBITS, P=full’s structure-split AP nearly 2.5× that of every static baseline. Binding Affinity We additionally evaluate ligand-affinity regression under the misato structure split (Table 9). The label is − log Kd /Ki for the bound ligand of each complex, and the head conditions on a 167-bit MACCS structural-key fingerprint of the ligand (resolved from the misato ligand_id via the PDB Chemical Component Dictionary) concatenated to a mean+max-pooled protein representation; the resulting joint vector is mapped to the affinity scalar by a small MLP. We report mean ± std over 10 seeds on three metrics: R2 (variance-explained, higher is better), Spearman ρS (rank correlation, higher is better), and MSE (mean squared error in log-affinity units, lower is better). Affinity on the structure split is the hardest residue-level task in this appendix. The test proteins are CATH-H-disjoint from training, the protein-only features cannot resolve which ligand is bound, and the predictive signal collapses to a narrow band around the constant-mean baseline (R2 = 0): every 28
Table 8: Binding-site — misato structure split, mean ± std over 10 seeds. Bolding and underlining are separate for the single-frame (top) and multi-frame (bottom) groups. Model
AUROC
AP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.649 ± 0.054 0.594 ± 0.012 0.500 ± 0.005 0.694 ± 0.032 0.626 ± 0.007 0.678 ± 0.043 0.719 ± 0.019
0.300 ± 0.095 0.184 ± 0.020 0.124 ± 0.002 0.318 ± 0.067 0.208 ± 0.013 0.291 ± 0.110 0.418 ± 0.032
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.583 ± 0.004 0.594 ± 0.015 0.750 ± 0.011
0.169 ± 0.005 0.191 ± 0.025 0.471 ± 0.010
entry sits in the range R2 ∈ [0.04, 0.12] and no method materially exceeds R2 = 0.12. Within the multi-frame half, E NSEMBITS, P=full is best on every metric, edging Vote_3Di on R2 (0.090 vs. 0.075), Spearman (0.415 vs. 0.400), and MSE (2.966 vs. 3.016); the margins are within seed-to-seed standard deviation so we do not claim a robust separation. Within the single-frame half, aa is best on R2 and MSE while AminoAseed is best on Spearman (0.465); the geometry-only single-frame tokenizers (3Di_tokens, ProToken, E NSEMBITS, P=1) all sit at or below the AA one-hot baseline on this split. The structure-split affinity result is best read as a diagnostic of which input is carrying the signal, not as a tokenizer ranking. A token-shuffling control Random — random integer tokens per residue — lands at R2 = 0.103, statistically indistinguishable from every learned representation in Table 9. The MLP head with random protein features hits the same ceiling as the MLP with ESM3struct, ProToken, Vote_3Di, or E NSEMBITS, because all of these protein representations are out-of-distribution on the CATH-H-disjoint test split and the MLP cannot extract a generalising signal from them. What remains is the MACCS-167 ligand fingerprint, which is in-distribution by construction (no ligand-disjoint constraint on this split) and accounts for essentially all of the ∼ 0.10 R² floor observed here. In short, on the structure split the ligand is the driving force and the protein representation contributes no measurable additional signal. The sequence- and random-split numbers in Appendix G show that the protein representation does carry signal once the test distribution is in-domain. Table 9: Binding affinity (− log Kd /Ki ) — misato structure split, ligand-aware head (MACCS-167), mean ± std over 10 seeds. Bolding and underlining are separate for the single-frame (top) and multi-frame (bottom) groups. MSE: lower is better. R2
Spearman
MSE
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.113 ± 0.082 0.080 ± 0.063 0.103 ± 0.053 0.039 ± 0.126 0.091 ± 0.052 0.077 ± 0.074 0.092 ± 0.082
0.411 ± 0.034 0.378 ± 0.041 0.389 ± 0.032 0.465 ± 0.034 0.376 ± 0.034 0.432 ± 0.060 0.412 ± 0.045
2.862 ± 0.265 3.000 ± 0.207 2.924 ± 0.173 3.133 ± 0.411 2.965 ± 0.169 3.009 ± 0.240 2.959 ± 0.267
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.058 ± 0.056 0.075 ± 0.060 0.090 ± 0.082
0.364 ± 0.042 0.400 ± 0.045 0.415 ± 0.052
3.070 ± 0.183 3.016 ± 0.194 2.966 ± 0.267
Model
29
G
All Experiments
This appendix reports the full per-split results for each downstream benchmark. Each task is evaluated under three splits — structure, sequence, and random. Within every table the upper half lists the per-residue (single-frame) methods and the lower half lists the per-protein (multi-frame) methods; bold marks the best and underline the second-best result within each half per column. All entries are mean ± std over 10 random seeds. Before reading the section, we would like to remind the reader that the purpose of E NSEMBITS is to capture protein dynamics, not to improve downstream performance. We include the study here just to show that Ensembit can have performance on par with other tokenziers. G.1
EC Function Prediction
We report Enzyme Commission classification at three taxonomy depths (top-level enzyme class, subclass, and full EC number) under all three splits. structure split The hardest setting: train and test proteins share neither sequence identity nor structural class. Table 10: EC depth-1 — structure split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.333 ± 0.000 0.548 ± 0.240 0.360 ± 0.042 0.841 ± 0.032 0.376 ± 0.005 0.882 ± 0.019 0.674 ± 0.239
0.346 ± 0.052 0.429 ± 0.084 0.212 ± 0.011 0.569 ± 0.061 0.265 ± 0.059 0.615 ± 0.042 0.462 ± 0.115
0.350 ± 0.050 0.533 ± 0.274 0.305 ± 0.028 0.912 ± 0.025 0.385 ± 0.061 0.943 ± 0.020 0.714 ± 0.260
0.080 ± 0.042 0.288 ± 0.141 0.108 ± 0.005 0.514 ± 0.060 0.118 ± 0.034 0.548 ± 0.059 0.365 ± 0.157
0.266 ± 0.140 0.520 ± 0.256 0.347 ± 0.048 0.830 ± 0.039 0.374 ± 0.032 0.869 ± 0.024 0.643 ± 0.284
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.374 ± 0.000 0.774 ± 0.136 0.839 ± 0.033
0.247 ± 0.043 0.504 ± 0.061 0.543 ± 0.045
0.330 ± 0.039 0.818 ± 0.199 0.871 ± 0.045
0.087 ± 0.046 0.422 ± 0.089 0.456 ± 0.082
0.295 ± 0.155 0.765 ± 0.144 0.800 ± 0.125
Model
Table 11: EC depth-2 — structure split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.100 ± 0.167 0.104 ± 0.122 0.080 ± 0.006 0.738 ± 0.037 0.088 ± 0.011 0.660 ± 0.211 0.621 ± 0.202
0.163 ± 0.036 0.158 ± 0.019 0.095 ± 0.005 0.406 ± 0.065 0.101 ± 0.016 0.351 ± 0.091 0.301 ± 0.063
0.111 ± 0.118 0.114 ± 0.044 0.074 ± 0.006 0.814 ± 0.037 0.081 ± 0.008 0.736 ± 0.229 0.670 ± 0.207
0.014 ± 0.023 0.024 ± 0.015 0.003 ± 0.002 0.316 ± 0.069 0.003 ± 0.005 0.245 ± 0.099 0.197 ± 0.071
0.018 ± 0.030 0.030 ± 0.020 0.013 ± 0.012 0.768 ± 0.040 0.007 ± 0.011 0.679 ± 0.244 0.605 ± 0.208
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.073 ± 0.015 0.116 ± 0.157 0.676 ± 0.061
0.113 ± 0.010 0.169 ± 0.058 0.337 ± 0.035
0.079 ± 0.007 0.141 ± 0.169 0.753 ± 0.053
0.001 ± 0.001 0.030 ± 0.056 0.205 ± 0.047
0.003 ± 0.007 0.077 ± 0.193 0.685 ± 0.060
Model
30
Table 12: EC depth-3 — structure split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.662 ± 0.136 0.170 ± 0.254 0.028 ± 0.007 0.724 ± 0.031 0.240 ± 0.282 0.718 ± 0.045 0.656 ± 0.062
0.259 ± 0.080 0.132 ± 0.110 0.046 ± 0.007 0.363 ± 0.038 0.132 ± 0.083 0.356 ± 0.034 0.297 ± 0.038
0.631 ± 0.189 0.205 ± 0.288 0.045 ± 0.005 0.779 ± 0.018 0.278 ± 0.301 0.779 ± 0.033 0.703 ± 0.060
0.102 ± 0.073 0.043 ± 0.090 0.000 ± 0.001 0.254 ± 0.043 0.049 ± 0.064 0.242 ± 0.064 0.169 ± 0.048
0.478 ± 0.358 0.141 ± 0.295 0.000 ± 0.001 0.743 ± 0.035 0.247 ± 0.319 0.750 ± 0.056 0.648 ± 0.066
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.250 ± 0.270 0.330 ± 0.354 0.678 ± 0.069
0.149 ± 0.106 0.207 ± 0.144 0.335 ± 0.037
0.236 ± 0.231 0.365 ± 0.373 0.745 ± 0.050
0.042 ± 0.053 0.110 ± 0.132 0.201 ± 0.044
0.174 ± 0.254 0.324 ± 0.398 0.699 ± 0.067
Model
sequence split Train and test proteins are sequence-disjoint (UniProt-clustered) but may share structural class. Table 13: EC depth-1 — sequence split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.688 ± 0.052 0.742 ± 0.031 0.607 ± 0.039 0.844 ± 0.037 0.715 ± 0.065 0.868 ± 0.020 0.754 ± 0.047
0.450 ± 0.102 0.604 ± 0.042 0.274 ± 0.007 0.741 ± 0.089 0.542 ± 0.063 0.762 ± 0.046 0.551 ± 0.092
0.778 ± 0.076 0.853 ± 0.020 0.559 ± 0.024 0.917 ± 0.017 0.832 ± 0.033 0.924 ± 0.015 0.860 ± 0.042
0.275 ± 0.066 0.403 ± 0.089 0.086 ± 0.044 0.578 ± 0.123 0.324 ± 0.092 0.627 ± 0.071 0.407 ± 0.099
0.699 ± 0.054 0.746 ± 0.030 0.334 ± 0.186 0.835 ± 0.028 0.726 ± 0.054 0.858 ± 0.017 0.775 ± 0.044
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.751 ± 0.039 0.736 ± 0.061 0.813 ± 0.054
0.612 ± 0.080 0.613 ± 0.065 0.631 ± 0.084
0.865 ± 0.028 0.848 ± 0.040 0.896 ± 0.027
0.460 ± 0.105 0.413 ± 0.118 0.482 ± 0.135
0.757 ± 0.033 0.742 ± 0.051 0.814 ± 0.042
Model
Table 14: EC depth-2 — sequence split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.556 ± 0.040 0.644 ± 0.074 0.398 ± 0.027 0.776 ± 0.038 0.583 ± 0.062 0.734 ± 0.070 0.673 ± 0.094
0.327 ± 0.061 0.476 ± 0.071 0.141 ± 0.011 0.601 ± 0.042 0.422 ± 0.066 0.569 ± 0.050 0.471 ± 0.102
0.745 ± 0.047 0.770 ± 0.044 0.341 ± 0.023 0.845 ± 0.020 0.720 ± 0.037 0.823 ± 0.033 0.792 ± 0.061
0.171 ± 0.040 0.274 ± 0.097 0.003 ± 0.006 0.419 ± 0.077 0.259 ± 0.067 0.375 ± 0.066 0.285 ± 0.114
0.675 ± 0.041 0.707 ± 0.056 0.021 ± 0.043 0.794 ± 0.025 0.676 ± 0.031 0.762 ± 0.041 0.736 ± 0.067
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.538 ± 0.135 0.667 ± 0.061 0.723 ± 0.078
0.330 ± 0.111 0.477 ± 0.045 0.537 ± 0.091
0.703 ± 0.087 0.784 ± 0.027 0.831 ± 0.042
0.198 ± 0.121 0.283 ± 0.094 0.355 ± 0.093
0.653 ± 0.070 0.723 ± 0.055 0.770 ± 0.051
Model
31
Table 15: EC depth-3 — sequence split. top-1
mAP
µAP
aa 3Di_tokens random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.453 ± 0.102 0.673 ± 0.052 0.085 ± 0.066 0.706 ± 0.044 0.512 ± 0.045 0.720 ± 0.049 0.641 ± 0.099
0.329 ± 0.062 0.490 ± 0.044 0.083 ± 0.006 0.584 ± 0.053 0.395 ± 0.039 0.544 ± 0.054 0.484 ± 0.060
0.757 ± 0.067 0.800 ± 0.024 0.138 ± 0.034 0.803 ± 0.030 0.689 ± 0.020 0.818 ± 0.039 0.785 ± 0.038
0.172 ± 0.040 0.299 ± 0.049 0.000 ± 0.000 0.394 ± 0.051 0.245 ± 0.061 0.385 ± 0.081 0.293 ± 0.073
0.727 ± 0.040 0.746 ± 0.027 0.000 ± 0.000 0.771 ± 0.027 0.692 ± 0.026 0.775 ± 0.034 0.729 ± 0.049
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.526 ± 0.118 0.642 ± 0.061 0.689 ± 0.078
0.400 ± 0.077 0.493 ± 0.033 0.525 ± 0.061
0.727 ± 0.060 0.791 ± 0.027 0.818 ± 0.031
0.214 ± 0.090 0.265 ± 0.051 0.350 ± 0.088
0.696 ± 0.058 0.735 ± 0.033 0.773 ± 0.053
Model
random split
Train and test proteins are drawn at random; the easiest setting and a sanity check. Table 16: EC depth-1 — random split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.929 ± 0.008 0.895 ± 0.014 0.511 ± 0.008 0.921 ± 0.016 0.901 ± 0.016 0.923 ± 0.008 0.905 ± 0.013
0.875 ± 0.027 0.794 ± 0.031 0.202 ± 0.005 0.847 ± 0.035 0.828 ± 0.037 0.840 ± 0.014 0.826 ± 0.031
0.970 ± 0.006 0.959 ± 0.006 0.485 ± 0.016 0.972 ± 0.006 0.958 ± 0.009 0.970 ± 0.004 0.963 ± 0.006
0.748 ± 0.121 0.736 ± 0.040 0.109 ± 0.008 0.788 ± 0.052 0.775 ± 0.052 0.795 ± 0.023 0.772 ± 0.039
0.932 ± 0.008 0.898 ± 0.013 0.460 ± 0.049 0.921 ± 0.014 0.903 ± 0.014 0.922 ± 0.006 0.907 ± 0.011
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.874 ± 0.023 0.902 ± 0.011 0.918 ± 0.014
0.772 ± 0.028 0.814 ± 0.016 0.832 ± 0.023
0.945 ± 0.012 0.960 ± 0.004 0.970 ± 0.004
0.703 ± 0.055 0.765 ± 0.036 0.803 ± 0.023
0.877 ± 0.021 0.907 ± 0.010 0.922 ± 0.010
Model
Table 17: EC depth-2 — random split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.903 ± 0.013 0.832 ± 0.015 0.281 ± 0.010 0.884 ± 0.012 0.838 ± 0.014 0.884 ± 0.009 0.857 ± 0.011
0.768 ± 0.047 0.638 ± 0.024 0.084 ± 0.010 0.697 ± 0.014 0.656 ± 0.016 0.698 ± 0.024 0.651 ± 0.019
0.940 ± 0.012 0.901 ± 0.008 0.242 ± 0.012 0.940 ± 0.007 0.908 ± 0.008 0.937 ± 0.006 0.914 ± 0.008
0.666 ± 0.082 0.540 ± 0.045 0.004 ± 0.003 0.615 ± 0.029 0.566 ± 0.037 0.604 ± 0.045 0.551 ± 0.035
0.914 ± 0.014 0.859 ± 0.010 0.030 ± 0.029 0.901 ± 0.009 0.868 ± 0.011 0.905 ± 0.008 0.874 ± 0.015
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.807 ± 0.013 0.862 ± 0.014 0.878 ± 0.012
0.618 ± 0.013 0.667 ± 0.024 0.676 ± 0.022
0.892 ± 0.009 0.922 ± 0.008 0.931 ± 0.007
0.494 ± 0.032 0.574 ± 0.026 0.594 ± 0.023
0.840 ± 0.011 0.883 ± 0.011 0.899 ± 0.005
Model
32
Table 18: EC depth-3 — random split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.887 ± 0.014 0.830 ± 0.016 0.156 ± 0.020 0.872 ± 0.019 0.802 ± 0.015 0.871 ± 0.010 0.857 ± 0.021
0.827 ± 0.019 0.730 ± 0.027 0.063 ± 0.010 0.782 ± 0.018 0.718 ± 0.025 0.769 ± 0.022 0.773 ± 0.032
0.932 ± 0.007 0.895 ± 0.010 0.138 ± 0.019 0.929 ± 0.009 0.888 ± 0.006 0.924 ± 0.004 0.920 ± 0.011
0.734 ± 0.042 0.576 ± 0.031 0.000 ± 0.001 0.655 ± 0.044 0.596 ± 0.020 0.668 ± 0.042 0.623 ± 0.041
0.913 ± 0.006 0.868 ± 0.012 0.001 ± 0.003 0.901 ± 0.010 0.860 ± 0.007 0.900 ± 0.009 0.884 ± 0.013
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.788 ± 0.019 0.850 ± 0.015 0.857 ± 0.015
0.708 ± 0.023 0.781 ± 0.021 0.763 ± 0.026
0.890 ± 0.008 0.915 ± 0.008 0.921 ± 0.008
0.542 ± 0.032 0.639 ± 0.043 0.625 ± 0.048
0.842 ± 0.013 0.885 ± 0.008 0.891 ± 0.010
Model
G.2
GO Label Prediction
We report top-50 Gene Ontology term prediction under all three splits. structure split
Table 19: GO top-50 — structure split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.354 ± 0.348 0.641 ± 0.052 0.294 ± 0.163 0.704 ± 0.048 0.468 ± 0.224 0.665 ± 0.073 0.619 ± 0.057
0.300 ± 0.141 0.400 ± 0.021 0.114 ± 0.005 0.444 ± 0.021 0.208 ± 0.108 0.428 ± 0.040 0.397 ± 0.033
0.500 ± 0.301 0.837 ± 0.026 0.247 ± 0.058 0.859 ± 0.018 0.492 ± 0.243 0.835 ± 0.101 0.781 ± 0.071
0.116 ± 0.186 0.304 ± 0.040 0.000 ± 0.000 0.374 ± 0.028 0.084 ± 0.127 0.342 ± 0.083 0.283 ± 0.073
0.299 ± 0.411 0.777 ± 0.072 0.000 ± 0.001 0.814 ± 0.022 0.257 ± 0.359 0.747 ± 0.220 0.681 ± 0.127
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.610 ± 0.073 0.676 ± 0.073 0.692 ± 0.073
0.199 ± 0.074 0.423 ± 0.020 0.428 ± 0.033
0.386 ± 0.166 0.849 ± 0.030 0.846 ± 0.047
0.040 ± 0.108 0.335 ± 0.038 0.347 ± 0.052
0.087 ± 0.252 0.790 ± 0.043 0.798 ± 0.065
Model
sequence split
Table 20: GO top-50 — sequence split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.511 ± 0.079 0.599 ± 0.046 0.250 ± 0.051 0.633 ± 0.070 0.574 ± 0.079 0.590 ± 0.073 0.470 ± 0.076
0.336 ± 0.078 0.480 ± 0.021 0.099 ± 0.007 0.504 ± 0.041 0.436 ± 0.033 0.470 ± 0.046 0.368 ± 0.074
0.527 ± 0.081 0.610 ± 0.024 0.175 ± 0.011 0.624 ± 0.056 0.589 ± 0.046 0.608 ± 0.048 0.512 ± 0.066
0.180 ± 0.080 0.365 ± 0.042 0.001 ± 0.004 0.356 ± 0.064 0.324 ± 0.048 0.344 ± 0.070 0.228 ± 0.066
0.451 ± 0.072 0.536 ± 0.032 0.004 ± 0.011 0.525 ± 0.060 0.499 ± 0.046 0.516 ± 0.064 0.434 ± 0.063
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.463 ± 0.069 0.657 ± 0.072 0.513 ± 0.073
0.373 ± 0.062 0.486 ± 0.046 0.406 ± 0.079
0.498 ± 0.062 0.622 ± 0.043 0.548 ± 0.080
0.193 ± 0.081 0.362 ± 0.072 0.264 ± 0.102
0.394 ± 0.065 0.550 ± 0.052 0.462 ± 0.081
Model
random split
33
Table 21: GO top-50 — random split. top-1
mAP
µAP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.900 ± 0.012 0.865 ± 0.008 0.302 ± 0.015 0.888 ± 0.010 0.855 ± 0.017 0.891 ± 0.015 0.881 ± 0.009
0.949 ± 0.003 0.925 ± 0.006 0.110 ± 0.004 0.934 ± 0.006 0.922 ± 0.008 0.941 ± 0.006 0.936 ± 0.004
0.972 ± 0.002 0.944 ± 0.004 0.196 ± 0.006 0.955 ± 0.004 0.946 ± 0.005 0.958 ± 0.004 0.956 ± 0.004
0.924 ± 0.006 0.870 ± 0.007 0.000 ± 0.000 0.887 ± 0.006 0.878 ± 0.014 0.893 ± 0.005 0.894 ± 0.008
0.940 ± 0.005 0.875 ± 0.006 0.000 ± 0.000 0.892 ± 0.006 0.880 ± 0.011 0.898 ± 0.005 0.897 ± 0.008
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.831 ± 0.016 0.887 ± 0.009 0.896 ± 0.015
0.907 ± 0.008 0.936 ± 0.005 0.944 ± 0.006
0.931 ± 0.006 0.955 ± 0.004 0.963 ± 0.004
0.846 ± 0.011 0.891 ± 0.008 0.908 ± 0.009
0.858 ± 0.009 0.893 ± 0.008 0.909 ± 0.007
Model
G.3
Binding Site
We report residue-level binding-site classification under all three splits. structure split
Table 22: Binding-site — structure split. Model
AUROC
AP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.649 ± 0.054 0.594 ± 0.012 0.500 ± 0.005 0.694 ± 0.032 0.626 ± 0.007 0.678 ± 0.043 0.719 ± 0.019
0.300 ± 0.095 0.184 ± 0.020 0.124 ± 0.002 0.318 ± 0.067 0.208 ± 0.013 0.291 ± 0.110 0.418 ± 0.032
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.583 ± 0.004 0.594 ± 0.015 0.750 ± 0.011
0.169 ± 0.005 0.191 ± 0.025 0.471 ± 0.010
sequence split
Table 23: Binding-site — sequence split. Model
AUROC
AP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.660 ± 0.010 0.650 ± 0.005 0.500 ± 0.007 0.684 ± 0.005 0.678 ± 0.005 0.685 ± 0.009 0.675 ± 0.013
0.315 ± 0.011 0.298 ± 0.007 0.104 ± 0.003 0.323 ± 0.010 0.307 ± 0.004 0.317 ± 0.014 0.274 ± 0.011
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.624 ± 0.033 0.648 ± 0.007 0.675 ± 0.012
0.204 ± 0.081 0.298 ± 0.006 0.279 ± 0.009
random split
34
Table 24: Binding-site — random split. Model
G.4
AUROC
AP
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.861 ± 0.002 0.727 ± 0.002 0.504 ± 0.002 0.789 ± 0.004 0.780 ± 0.003 0.794 ± 0.004 0.749 ± 0.003
0.612 ± 0.004 0.386 ± 0.004 0.107 ± 0.001 0.461 ± 0.008 0.466 ± 0.005 0.474 ± 0.006 0.379 ± 0.004
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.722 ± 0.017 0.754 ± 0.002 0.787 ± 0.004
0.381 ± 0.023 0.429 ± 0.003 0.433 ± 0.005
Binding Affinity
We report binding-affinity regression on misato (label: − log Kd /Ki ) under all three splits. Each prediction conditions on a 167-bit MACCS structural-key fingerprint of the bound ligand (resolved from the misato ligand_id via the PDB Chemical Component Dictionary), in addition to the perresidue protein representation. All numbers below are in the format of mean ± std over 10 seeds. Bold = best, underline = second-best, computed within each half per column. Higher is better for R2 and Spearman; lower is better for MSE. structure split
Table 25: Binding affinity (− log Kd /Ki ) — structure split, ligand-aware head (MACCS-167), mean ± std over 10 seeds. R2
Spearman
MSE
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.113 ± 0.082 0.080 ± 0.063 0.103 ± 0.053 0.039 ± 0.126 0.091 ± 0.052 0.077 ± 0.074 0.092 ± 0.082
0.411 ± 0.034 0.378 ± 0.041 0.389 ± 0.032 0.465 ± 0.034 0.376 ± 0.034 0.432 ± 0.060 0.412 ± 0.045
2.862 ± 0.265 3.000 ± 0.207 2.924 ± 0.173 3.133 ± 0.411 2.965 ± 0.169 3.009 ± 0.240 2.959 ± 0.267
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.058 ± 0.056 0.075 ± 0.060 0.090 ± 0.082
0.364 ± 0.042 0.400 ± 0.045 0.415 ± 0.052
3.070 ± 0.183 3.016 ± 0.194 2.966 ± 0.267
Model
35
sequence split
Table 26: Binding affinity (− log Kd /Ki ) — sequence split, ligand-aware head (MACCS-167), mean ± std over 10 seeds. R2
Spearman
MSE
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.424 ± 0.049 0.437 ± 0.040 0.366 ± 0.023 0.466 ± 0.024 0.386 ± 0.034 0.464 ± 0.022 0.445 ± 0.026
0.659 ± 0.031 0.644 ± 0.025 0.595 ± 0.015 0.671 ± 0.017 0.611 ± 0.024 0.653 ± 0.013 0.654 ± 0.019
2.401 ± 0.204 2.236 ± 0.160 2.520 ± 0.093 2.121 ± 0.096 2.438 ± 0.134 2.129 ± 0.088 2.207 ± 0.104
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.392 ± 0.025 0.455 ± 0.032 0.451 ± 0.023
0.619 ± 0.021 0.656 ± 0.021 0.650 ± 0.018
2.418 ± 0.100 2.168 ± 0.126 2.183 ± 0.092
Model
random split
Table 27: Binding affinity (− log Kd /Ki ) — random split, ligand-aware head (MACCS-167), mean ± std over 10 seeds. R2
Spearman
MSE
aa 3Di_tokens Random AminoAseed ProToken ESM3struct E NSEMBITS, P=1
0.557 ± 0.038 0.481 ± 0.020 0.411 ± 0.018 0.470 ± 0.028 0.484 ± 0.010 0.496 ± 0.017 0.473 ± 0.035
0.750 ± 0.022 0.696 ± 0.009 0.642 ± 0.012 0.701 ± 0.009 0.697 ± 0.011 0.706 ± 0.010 0.691 ± 0.013
1.886 ± 0.161 2.072 ± 0.078 2.352 ± 0.071 2.115 ± 0.110 2.062 ± 0.040 2.011 ± 0.069 2.105 ± 0.138
ProtProfileMD_K8 Vote_3Di E NSEMBITS, P=full
0.448 ± 0.024 0.504 ± 0.015 0.500 ± 0.024
0.677 ± 0.010 0.712 ± 0.015 0.711 ± 0.014
2.204 ± 0.097 1.979 ± 0.061 1.996 ± 0.094
Model
36
G.5
Compute Requirements
All experiments were run on H100/H200 GPUs (one training/eval job per GPU). Numbers below are wall-clock times we measured during the runs that produced this paper; per-seed-per-task figures generalize to a single H100/H200, and “N GPUs” figures assume embarrassingly-parallel scheduling across N devices. One-time caching.
Built once, reused across all experiments.
Step
Wall-clock
Real-backbone cache, mdCATH-div (308k residues, K=10) Real-backbone cache, MISATO (16,972 entries, K=10) ESM3 K=16 descriptor cache, combined (7.3 M, P =10) Per-tokenizer baseline token cache (mdCATH or MISATO) ESMFold atom14 cache for ProteinGym (∼400k variants)
∼8 min (32 CPU) ∼11 min (32 CPU) ∼1 h (1 GPU) 3–5 min (1 GPU) ∼2 GPU-days (4 GPU)
The ESMFold cache is the dominant cold-start cost. Once it exists, all downstream PG scoring is fast (final paragraph). Tokenizer training. The production E NSEMBITS checkpoint (RVQVAET OKENIZER, 2048×128×128 codebook, ∼3.4 M parameters, 192-D ESM3 K=16 descriptor input) trains on combined mdCATH-div + MISATO (6.6 M training residues, P =10, varP 1–10) and converges via patience-40 early stopping at epoch 195 in ∼7.2 h on a single H200 (batch size 4096, AdamW 10−3 , k-means init, 1000-step warm-up + cosine decay). The 12-cell ablation grid (3 corpora × 2 descriptors × 2 projection modes) costs ∼50–80 GPU-hours total; runs are independent and parallelize across GPUs.
Task
Per-seed wall-clock
EC depth-{1,2,3} (one split, 40 ep, bs 32) GO top-50 (one split, same trunk as EC) Binding-site (one split, per-residue head) Binding-affinity (all 3 splits, MACCS-167 ligand) RMSF probe (one split, per-residue MLP, 40 ep) ProteinGym tokenize per assay (cache warm) ProteinGym scoring (all 96 assays, one tokenizer) First-time variant tokenization (16,972 MISATO entries)
30–90 s 30–60 s 30–60 s ∼40 s 20–40 s 30 s – 2 min ∼40 min CPU 5–10 min
Downstream evaluation (per seed, per split, per task). Full reproduction budget. Compiling every cell in this paper from a fresh checkout (assuming the ESMFold and Real-backbone caches are already built): • Downstream sweep (10 seeds, 5 tasks, 3 splits, ∼9 tokenizer baselines + E NSEMBITS at P =full and P =1). ∼1,500 jobs averaging ∼1 min each. Sequential ∼25 GPU-h; ∼6.5 h wall on 4 GPUs. • Ligand-aware affinity sweep (10 seeds, 3 splits, 8 baselines + E NSEMBITS at P =full/P =1, all splits in one call). ∼100 jobs at ∼40 s; ∼15 min on 4 GPUs. • ProteinGym scoring (5 tokenizers, 96 assays each). Tokenize: ∼2.5 h per tokenizer (∼2.5 h on 4 GPUs). Scoring: ∼40 min CPU per tokenizer. • ANOVA suite (η 2 + 1000-permutation null). ∼4 CPU-minutes per tokenizer; negligible. • Total without ESMFold cache build: ∼80 GPU-hours, achievable in ∼10–12 h wall on a single 8×H100 node. • Including ESMFold cache build: add ∼2 GPU-days one-time. 37
Memory footprint. Peak: tokenizer training ∼55 GB on H200 (varP set encoder + chunked eval of ∼720k val residues); conv1d/MLP downstream head ∼2–5 GB; ESM3 descriptor generation ∼1–2 GB; the largest model load (AminoAseed, ∼140 M params, full encoder + decoder + quantizer for PG scoring) ∼30 GB. All experiments fit on a single 80 GB device; the production tokenizer training does not fit on 40 GB at batch size 4096 with full-P chunked eval.
38