# Structure-Aware Masking for Protein Representation Learning

arXiv: https://arxiv.org/abs/2605.16581 · HTML: https://mlsanity.com/paper/2605.16581/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 67 references that are in the map.
- First posted on arXiv 15 May 2026.

## Builds on (most cited first, 19 of 19)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [SpanBERT: Improving Pre-training by Representing and Predicting Spans](https://mlsanity.com/paper/1907.10529/) (2019)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [LUKE: Deep Contextualized Entity Representations with Entity-aware Self-attention](https://mlsanity.com/paper/2010.01057/) (2020)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [Learning Better Masking for Better Language Model Pre-training](https://mlsanity.com/paper/2208.10806/) (2022)
- [Efficient Pre-training of Masked Language Model via Concept-based Curriculum Masking](https://mlsanity.com/paper/2212.07617/) (2022)
- [Training Compute-Optimal Protein Language Models](https://mlsanity.com/paper/2411.02142/) (2024)
- [Sparse Autoencoders for Low-N Protein Function Prediction and Design](https://mlsanity.com/paper/2508.18567/) (2025)
- [Fine-tuning Protein Language Models with Deep Mutational Scanning improves Variant Effect Prediction](https://mlsanity.com/paper/2405.06729/) (2024)
- [Understanding and Enhancing Mask-Based Pretraining towards Universal Representations](https://mlsanity.com/paper/2509.21650/) (2025)
- [GeneMask: Fast Pretraining of Gene Sequences to Enable Few-Shot Learning](https://mlsanity.com/paper/2307.15933/) (2023)
- [Unlocking Efficiency: Adaptive Masking for Gene Transformer Models](https://mlsanity.com/paper/2408.07180/) (2024)
- [SpecMER: Fast Protein Generation with K-mer Guided Speculative Decoding](https://mlsanity.com/paper/2509.21689/) (2025)

## Built on by (most cited first, 1 of 1)

- [Learning Task-Specific Antibody Representations via Function-Aware Masking](https://mlsanity.com/paper/2609.00518/) (2026)

## Abstract (arXiv)

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.