Structure-Aware Masking for Protein Representation Learning

2026-05-15

AI Summary

This work addresses a key limitation of conventional protein language models, which employ random masking and thereby overlook long-range residue interactions inherent in three-dimensional protein structures, hindering their ability to capture spatial couplings critical for function. To overcome this, the authors propose Bucket Masking—a novel strategy that integrates 3D structural information directly into the masking scheme. By grouping spatially proximal residues based on contact maps and jointly masking them, the method aligns the masking distribution with structural contacts, effectively introducing a structure-aware positional inductive bias. Evaluated across four protein fitness prediction tasks, Bucket Masking achieves performance gains of up to 14% over random masking, with particularly pronounced improvements in modeling higher-order mutational epistasis.

Abstract

Masked language modeling (MLM) is the standard objective for training protein language models, typically implemented by randomly masking individual residues at a fixed rate (e.g., 15%). This practice implicitly assumes that all sequence positions contribute equally to representation learning. In downstream fitness prediction tasks, however, protein sequences are governed by three-dimensional structural dependencies and long-range residue contacts that induce strong nonlocal couplings between residues. We introduce Bucket Masking, a structure-aware masking strategy that selects groups of residues based on their proximity in three-dimensional space, preferentially masking structurally coupled regions during training. By conditioning the masking distribution on residue contacts, Bucket Masking shifts the learning objective toward modeling long-range interactions that are critical for protein function. Across four downstream protein fitness prediction tasks, Bucket Masking enables up to a 14% improvement over standard random masking, excelling at predicting higher-order mutational interactions. Through controlled ablations, we show that these improvements arise from mask placement rather than span size, establishing masking as a positional inductive bias.

Problem

• protein representation learning
• masked language modeling
• structural dependencies
• long-range interactions
• fitness prediction

Innovation

• structure-aware masking
• protein language models
• masked language modeling
• residue contacts
• long-range interactions