# Enhancing Protein Representation Learning via Manifold Restore Mixing

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 79 references that are in the map.
- First posted on arXiv 21 Jun 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Learning from Protein Structure with Geometric Vector Perceptrons](https://mlsanity.com/paper/2009.01411/) (2020)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [DeepDTA: Deep Drug-Target Binding Affinity Prediction](https://mlsanity.com/paper/1801.10193/) (2018)
- [Protein Representation Learning by Geometric Structure Pretraining](https://mlsanity.com/paper/2203.06125/) (2022)
- [G-Mixup: Graph Data Augmentation for Graph Classification](https://mlsanity.com/paper/2202.07179/) (2022)
- [DiffuseMix: Label-Preserving Data Augmentation with Diffusion Models](https://mlsanity.com/paper/2405.14881/) (2024)
- [On Feature Normalization and Data Augmentation](https://mlsanity.com/paper/2002.11102/) (2020)
- [Learning Hierarchical Protein Representations via Complete 3D Graph Networks](https://mlsanity.com/paper/2207.12600/) (2022)
- [SeqMix: Augmenting Active Sequence Labeling via Sequence Mixup](https://mlsanity.com/paper/2010.02322/) (2020)
- [DeepSF: deep convolutional neural network for mapping protein sequences to folds](https://mlsanity.com/paper/1706.01010/) (2017)
- [Knowledge Graph Embedding: A Survey from the Perspective of Representation Spaces](https://mlsanity.com/paper/2211.03536/) (2022)
- [A Survey of Mix-based Data Augmentation: Taxonomy, Methods, Applications, and Explainability](https://mlsanity.com/paper/2212.10888/) (2022)
- [Contrastive Representation Learning for 3D Protein Structures](https://mlsanity.com/paper/2205.15675/) (2022)
- [Is Transfer Learning Necessary for Protein Landscape Prediction?](https://mlsanity.com/paper/2011.03443/) (2020)
- [A Latent Diffusion Model for Protein Structure Generation](https://mlsanity.com/paper/2305.04120/) (2023)
- [A Survey on Mixup Augmentations and Beyond](https://mlsanity.com/paper/2409.05202/) (2024)
- [Graph Invariant Learning with Subgraph Co-mixup for Out-Of-Distribution Generalization](https://mlsanity.com/paper/2312.10988/) (2023)
- [Protein Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2502.17504/) (2025)
- [R-Mixup: Riemannian Mixup for Biological Networks](https://mlsanity.com/paper/2306.02532/) (2023)
- [Clustering for Protein Representation Learning](https://mlsanity.com/paper/2404.00254/) (2024)
- [ProtCLIP: Function-Informed Protein Multi-Modal Learning](https://mlsanity.com/paper/2412.20014/) (2024)
- [Deep Manifold Transformation for Protein Representation Learning](https://mlsanity.com/paper/2402.09416/) (2024)
- [Contrastive Multi-Task Learning with Solvent-Aware Augmentation for Drug Discovery](https://mlsanity.com/paper/2508.01799/) (2025)

## Abstract (arXiv)

Data augmentation (DA) has been proven to be an effective means for improving protein representation learning (PRL) by generating additional training samples. Although mainstream perturbation- and sampling-based augmentation methods can produce data containing sufficient variations, they carry the risk of disrupting the protein structure and function. Some crafted protein homology modeling tools can generate conformations, but reduce structural diversity. The above dilemmas lead us to a question: Can we restore the disrupted structure caused by DA operations, providing data with both the original structure and diverse variations? In this work, we first analyze and empirically reveal the structure defect and performance degradation issues of existing DA methods. Based on the findings, we propose a simple yet effective DA method, Manifold Restore Mixing (MRM), for protein representation learning. Specifically, inspired by manifold mixup, we mix the hidden representations of original and augmented protein data to generate new samples that restore structural information lost in DA while introducing diverse variations. Furthermore, we develop a sample difficulty scheduler that adjusts the beta distribution in mixup to provide models with progressively challenging mixed samples during training, which improves the final performance. Comprehensive experiments on various PRL backbones and downstream tasks demonstrate the effectiveness and generalization of our method. The complete code and weights will be released upon acceptance. We provide a implementation at https://github.com/KingGugu/MRM.