Technical Summary: Enhancing Protein Representation Learning via Manifold Restore Mixing

1. Problem Statement

Protein Representation Learning (PRL) relies heavily on Data Augmentation (DA) to overcome the scarcity of experimental 3D structural data. However, mainstream DA methods—such as residue substitution, coordinate noise injection, and substructure sampling—often disrupt the critical correspondence between a protein's 1D amino acid sequence and its 3D conformation.

The authors identify a fundamental dilemma: while perturbation-based methods generate diverse variations, they risk destroying the structural integrity of core functional domains, rendering the augmented data biologically implausible or functionally inconsistent. Conversely, homology modeling or diffusion-based generation tools often lack sufficient structural diversity. Empirical analysis in the paper demonstrates that applying these standard augmentation techniques to downstream tasks frequently leads to significant performance degradation rather than improvement, as models struggle to learn accurate representations from structurally compromised samples.

2. Methodology: Manifold Restore Mixing (MRM)

To address the structural defects of existing DA, the authors propose Manifold Restore Mixing (MRM), a method that operates in the latent representation space rather than the raw input space. The methodology consists of three core components:

A. Manifold Restore Mixing

Inspired by Manifold Mixup, MRM fuses the hidden representations of original protein data ($GP$) and augmented protein data ($GP_a$) at a specific intermediate layer $k$ of the neural network.

• Process: The network processes both $GP$ and $GP_a$ up to layer $k$, yielding representations $R_k(GP)$ and $R_k(GP_a)$.
• Mixing: A mixed representation $\tilde{R}_k$ is generated via linear interpolation:
  $$\tilde{R}_k = \lambda \cdot R_k(GP_a) + (1 - \lambda) \cdot R_k(GP)$$
  where $\lambda \sim \text{Beta}(\alpha_1, \alpha_2)$.
• Restoration Mechanism: By mixing at the representation level, the method aims to "restore" the structural information lost during the augmentation of $GP_a$ by re-introducing the intact structural features of $GP$, while retaining the diverse variations introduced by the augmentation. The mixed representation is then passed through the remaining layers to compute the loss.

B. Sample Difficulty Scheduler

Standard Mixup uses a fixed Beta distribution (typically $\alpha_1 = \alpha_2 < 1$), which can lead to a "monopoly effect" where samples are dominated entirely by either the original or augmented data. MRM introduces a curriculum learning approach:

• Initial State: Training begins with $\alpha_1 > \alpha_2 > 1$, biasing the mixture toward the original data (easier samples).
• Progression: Over training epochs, the parameters are adjusted such that $\alpha_1$ decreases and $\alpha_2$ increases, shifting the distribution to include more augmented data (harder samples).
• Goal: This provides the model with a progressive learning curve, moving from easy (structure-preserving) to hard (variation-rich) samples.

C. Two-Stage Regularized Training

To prevent convergence instability caused by mixing representations when model parameters are randomly initialized, MRM employs a two-stage strategy:

1. Stage 1: The model is trained exclusively on original data to establish a stable baseline representation.
2. Stage 2: The MRM module is activated. The total loss combines Empirical Risk Minimization (ERM) on the original data and Vicinal Risk Minimization (VRM) on the mixed samples:
   $$L_{second} = L_{CLS}(Net(GP), y) + \gamma L_{CLS}(\tilde{R}_f, y)$$
   This ensures the model maintains in-distribution fitting accuracy while exploring the vicinal space of restored representations.

3. Key Contributions

1. Empirical Analysis of DA Defects: The paper provides a systematic analysis and empirical evidence showing that standard protein DA operations disrupt 1D-3D structural correspondence, leading to performance degradation in downstream tasks.
2. MRM Framework: The proposal of Manifold Restore Mixing, which restores structural information lost during augmentation by fusing original and augmented representations in the latent space.
3. Training Enhancements: The development of a sample difficulty scheduler and a two-stage regularized training strategy to stabilize training and optimize the balance between structural fidelity and variation.
4. Comprehensive Validation: Extensive experiments across multiple PRL backbones (e.g., GearNet, CDConv) and diverse downstream tasks (Fold Classification, Enzyme Reaction, Gene Ontology, EC Number Prediction) demonstrating superior performance over existing DA methods and mixup variants.

4. Experimental Results

The authors evaluated MRM on four major tasks: Protein Fold Classification, Enzyme Reaction Classification, Gene Ontology (GO) Term Prediction, and Enzyme Commission (EC) Number Prediction.

• Performance Gains: MRM consistently outperformed baseline models and other augmentation strategies. For instance, on the Fold Classification task, MRM improved the accuracy of GearNet-Edge by 3.1% and CDConv by 4.1%. In Gene Ontology prediction, it achieved improvements across all sub-tasks (BP, MF, CC) where traditional augmentation methods failed or degraded performance.
• Comparison with Mixup Variants: Direct application of standard Manifold Mixup or NFM to protein data resulted in performance drops. The "Two-Stage" (TS) variants of these methods showed slight improvements but still lagged behind MRM, confirming that the specific "restore" mechanism of MRM is critical.
• Comparison with Protein LLMs: Despite not using large-scale pre-training, MRM-enhanced models achieved state-of-the-art or competitive results compared to pre-trained Protein Large Language Models (e.g., ESM-2, ProtBERT) on Fold and Enzyme Reaction tasks, achieving the best average ranking across all evaluated methods.
• Generalization: The method demonstrated robustness across different sequence similarity cutoffs (30% to 95%), indicating strong generalization capabilities even for proteins with low homology to the training set.
• Ablation Studies: Removing any component (1D/3D augmentation, two-stage training, or the difficulty scheduler) resulted in performance drops, validating the necessity of the full framework.

5. Significance and Claims

The paper claims that MRM offers a novel paradigm for data augmentation in computational biology by addressing the structural fragility of protein data. Its significance lies in:

• Restoring Structural Integrity: It enables the use of diverse augmented data without sacrificing the essential structural information required for accurate protein function prediction.
• Bridging the Gap: It allows traditional augmentation operations (which are often too destructive for PRL) to be effectively utilized in downstream tasks.
• Potential for Pre-training: The authors suggest that the "restore mixing" paradigm could offer insights for constructing self-supervised signals in the pre-training of Protein LLMs, potentially improving their ability to learn from augmented views.

The authors remain modest regarding biological interpretation, noting that the "restoration" occurs in the latent space and does not reconstruct a fully physically realizable 3D conformation in coordinate space. They acknowledge that future work could integrate Bayesian optimization or reinforcement learning to dynamically adjust augmentation operations.