Towards All-Atom Protein World Models: Learning and Diagnosing Predictive Latent Representations
Abstract
All-atom protein models compress atomic structure into latent representations, and reconstruction and generative objectives require these latents to preserve precise coordinate information. Yet, residue properties depend on both local atomic environments and longer-range 3D interactions. Joint-Embedding Predictive Architectures (JEPAs) offer an alternative: predicting representations in latent space rather than reconstructing the raw input. In vision, this shifts learning away from pixel-level reconstruction toward higher-level features that transfer across tasks. We ask whether the same approach can help protein representations capture both residue-level information and broader structural context useful for protein structure and function. Using the La-Proteina encoder, we train JEPA-style models to predict representations of masked residues from the remaining protein structure, and evaluate how well these representations transfer to biological tasks using StructTokenBench. EMA-JEPA with 70\% random masking scores higher than La-Proteina's reconstruction baseline on 16 of 20 functional-site and flexibility scores. We also identify a strong residue-index shortcut caused by I-JEPA-style absolute positional embeddings. Relative positioning removes this shortcut and scores higher on 18 of 20 comparisons. Relative-position LeJEPA adaptation performs better on functional-site tasks, while EMA-JEPA remains stronger on flexibility and remote homology. These results show the potential of latent prediction for all-atom protein representations while highlighting the importance of protein-specific representation design.