Saved in:
Bibliographic Details
Main Authors: Walton, Thomas, Goel, Ayan, Aghazadeh, Amirali
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2605.16581
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911690478583808
author Walton, Thomas
Goel, Ayan
Aghazadeh, Amirali
author_facet Walton, Thomas
Goel, Ayan
Aghazadeh, Amirali
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_16581
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Structure-Aware Masking for Protein Representation Learning
Walton, Thomas
Goel, Ayan
Aghazadeh, Amirali
Machine Learning
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.
title Structure-Aware Masking for Protein Representation Learning
topic Machine Learning
url https://arxiv.org/abs/2605.16581