Enhancing Protein Representation Learning via Manifold Restore Mixing

Protein representation learning (PRL) often relies on data augmentation (DA) to generate diverse training samples. However, standard DA techniques—such as residue substitution, coordinate noise, or substructure sampling—frequently disrupt the delicate 3D structural integrity of proteins. This paper investigates whether one can restore this disrupted structural information while retaining the diversity introduced by augmentation, ultimately aiming to improve downstream task performance.

The authors first empirically demonstrate that common DA methods often lead to performance degradation in representative PRL backbones like GearNet and CDConv. To address this, they propose Manifold Restore Mixing (MRM). Instead of augmenting raw data, MRM performs mixup operations in the latent space, blending the hidden representations of original and augmented protein structures. This approach allows the model to leverage the structural information of the original protein while incorporating the variations from the augmented version.

To further optimize learning, the authors introduce a difficulty scheduler that adjusts the beta distribution of the mixup process during training. By gradually shifting the weight of the augmented representation, the model is exposed to progressively more challenging samples. A two-stage training strategy is also employed to ensure convergence stability, preventing the noise introduced by hybrid representations from hindering the initial learning phase.

Experimental results across multiple PRL backbones and diverse downstream tasks (including fold classification and enzyme function prediction) show that MRM consistently outperforms traditional DA methods. By repairing the structural defects inherent in standard augmentation, MRM enables models to learn more robust representations. The inclusion of the difficulty scheduler further enhances performance by ensuring the model transitions from easier, more stable samples to more complex, informative ones throughout the training process.