Models and Methodology

Data Sourcing from ProteinGym

To ensure the breadth and reproducibility of our work, we chose ProteinGym as our primary data source. ProteinGym is a public database that provides standardized benchmark tests for protein language models and machine learning models. It aggregates a large amount of published Deep Mutational Scanning (DMS) experimental data. DMS experiments can measure the fitness scores of tens of thousands of protein variants (usually single-point mutations) at once, providing us with extremely valuable large-scale, high-quality sequence-function labeled data. Our team systematically selected and processed 50 different protein datasets from the ProteinGym database, covering a wide range from viral proteins[5], fluorescent proteins to various metabolic enzymes. This ensures that our model is exposed to a sufficiently diverse range of protein families, structures, and functions, laying the foundation for it to learn more generalizable principles.

Functional Classification of Datasets

The "performance" or "fitness" of a protein is a highly context-dependent concept. For an enzyme, performance might mean higher catalytic efficiency (kcat/Km); for an antibody, it could be stronger antigen-binding affinity; and for an industrial protein, it might be superior thermal stability. Directly predicting all these vastly different physicochemical properties with a single model is extremely difficult and imprecise. Therefore, our primary strategy is "divide and conquer." With 50 datasets in hand, we designed the following two-level classification standard:

• Primary Standard - Experimental Method (selection_assay): We used the selection_assay field recorded in each dataset as the primary classification criterion. This field describes the biological method used in the DMS experiment to screen for high-activity variants. For example, variants that remain active after heat treatment are classified as "Thermostability"[3]; variants that promote cell growth on a specific substrate are classified as "Growth-based Selection."
• Secondary Standard - Literature Source: For some datasets where the selection_assay description was vague or unique, we grouped them with other datasets from the same published study (reference). We assumed that the protein functions and evaluation systems focused on within the same study have inherent consistency.

Through this fine-grained classification process, we successfully organized the 50 datasets into several functional categories, including Stability, Growth, Viral Replication, Fluorescence, and Binding Affinity.

The Scoring Function: A "Virtual Experimenter"

The role of the scoring function in our workflow is that of a "computational biologist" or "virtual experimenter." After the model generates a new protein sequence, we need a reliable tool to quickly predict its performance. The development of our scoring function went through a significant iterative process.

Phase One: Separated, Specific Scoring Functions

In the initial phase, we followed the "divide and conquer" strategy, independently training a dedicated scoring function for each functional category (e.g., "Thermostability"). We used a classic and highly efficient machine learning model, specifically Ridge Regression, which is a form of regularized linear regression. This model was chosen for its fast training speed, strong interpretability, and its effectiveness in handling high-dimensional feature data (like sequence embeddings) by using regularization to prevent overfitting. We used sequence embeddings from a pre-trained ESM-2 model as input features and DMS fitness scores as the training target. The advantage of this method was its high specificity. However, it resulted in high management costs and data silos, as a model trained for "Thermostability" could not learn from "Binding Affinity" data, despite potential shared underlying principles.

Phase Two: Merged, Generalized Scoring Function

To overcome these limitations, we entered the second phase: merged training. We hypothesized that different biological functions may share more fundamental, universal sequence patterns. For example, good "enzymatic activity" likely requires a degree of "structural stability." We merged datasets from functionally similar categories to form a larger training set, which was then used to train a single, more powerful scoring function.

This is the relevant code implementation part:

def train_and_evaluate(X, y, test_size=0.2, random_state=42):
    """在嵌入上训练岭回归模型并评估其性能。"""
    print("\n--- 模型训练与评估 ---")
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state
    )
    print(f"训练集大小: {len(X_train)}")
    print(f"测试集大小: {len(X_test)}")

    model = Ridge(alpha=0.1)
    model.fit(X_train, y_train)

    y_test_pred = model.predict(X_test)
    test_r2 = r2_score(y_test, y_test_pred)
    test_rmse = np.sqrt(mean_squared_error(y_test, y_test_pred))

    print(f"\n模型在测试集上的性能:")
    print(f"  R² (决定系数): {test_r2:.4f}")
    print(f"  RMSE (均方根误差): {test_rmse:.4f}")

    return model, (y_test, y_test_pred)

Below are the results of our similarity analysis on the stability-related category datasets, which we used as a basis for clustering and merging.

This merged scoring function became the primary evaluation tool in our workflow. By training on richer and more diverse data, its generalization ability was significantly enhanced, providing a reliable basis for screening candidates.

The Core Engine: ESM-2 Protein Language Model

The core engine of all our computational work is the ESM-2 model, developed by the Meta AI team [6]. ESM-2 is a language model based on the Transformer architecture, specifically designed for protein sequences. Through self-supervised pre-training on a massive database of hundreds of millions of real protein sequences (like UniRef), it has learned the complex patterns of evolution, structure, and function embedded in amino acid sequences [1].

It comprehends both local grammar (adjacent amino acids) and long-range co-evolutionary relationships (global semantics). This profound "contextual understanding" enables ESM-2 to excel at its core task—Masked Language Modeling (MLM). The MLM task involves predicting the most likely original amino acid with high accuracy when one or more residues are masked. This is precisely the capability we need for protein sequence modification. By providing the model with a masked sequence, we are essentially asking: "Which amino acid substitution at this position would best maintain or improve the protein's overall fitness?" In this project, to balance efficiency and performance, we used two versions of the model: ESM2-650M and ESM2-35M.

Initial Approach: Fine-tuning on Positive Samples

• Goal: Our initial task was to verify the feasibility of the fine-tuning route—to prove that secondary training could make the model's performance surpass the original pre-trained model.
• Dataset & Method: We selected about 300 high-activity sequences from the A4GRB6_PSEAI_Chen_2020.csv dataset to form a pure "positive example" training set. We used the esm2_t33_650M_UR50D model and trained it for 5 epochs with a learning rate of 1e-5.
• Results and Reflection: A small-scale test showed that 7 out of 10 modified sequences had improved scores, proving the potential of fine-tuning. However, a serious problem emerged: the model's output showed high homogeneity and lacked diversity. This indicated that the model had overfit to the specific high-activity patterns, losing its ability to explore a broader sequence space, which is fatal for protein engineering.

Advanced Strategy: Integrated Contrastive Learning

After reviewing the problems from the first phase, we realized that a true expert knows both what is right and what is wrong. We fundamentally optimized our fine-tuning strategy.

Innovation 1: Generalized Integrated Contrastive Learning Dataset

We expanded our scope to all 50 ProteinGym datasets, creating a large-scale database of over one hundred thousand sequences. We defined high-activity sequences as "Positive Samples" and low-activity sequences as "Negative Samples." By training the model on both, it could not only learn the beneficial patterns from positive samples but also learn to avoid the harmful "non-target patterns" from negative samples. This contrastive learning approach enhanced the model's discriminative power and generalizability.

Innovation 2: Optimizing Training Efficiency and Model Selection

To accelerate iteration, we switched to the more lightweight ESM2-35M model. To prevent overfitting, we adopted an "Indirect Early Stopping" strategy. We set EPOCHS to 1, but after every training step, we evaluated the model's loss on a validation set. We only kept the model checkpoint from the moment the validation loss reached its minimum. This is a more intelligent and efficient training paradigm than a fixed number of epochs.

Sequence Generation: Point-by-Point Scanning Mask Prediction

We designed a systematic and interpretable strategy: Point-by-point Scanning Mask Prediction.

Core Algorithm: For a given sequence of length L, we loop from the first to the last position. In each iteration i, we replace the amino acid at position i with a <mask> token. This masked sequence is fed into our fine-tuned ESM-2 35M model. The model predicts the most likely amino acid for that position. If the predicted amino acid is different from the original and its confidence score is above a threshold, we log it as a potential beneficial mutation.

Output and Final Sequence Generation: This process generates a complete "mutation suggestion map." Based on this map, we can generate candidate sequences. In this project, we mainly focused on generating and validating high-confidence single-point mutation sequences.

The code implementation:

for i, single_pos_predictions in enumerate(pipeline_output):
        if not single_pos_predictions:
            predicted_seq += original_aas[i]
            continue
            
        top_pred = single_pos_predictions[0]
        
        if top_pred['score'] < confidence_threshold:
            predicted_seq += original_aas[i]
        else:
            predicted_seq += top_pred['token_str']
            if top_pred['token_str'] != original_aas[i]:
                analysis += f"\tposition {i}: predict {top_pred['token_str']} while original {original_aas[i]} (score: {top_pred['score']:.4f})\n"
                different_cnt += 1
                
    analysis_header = f'\ttotal different aa count: {different_cnt}\n'
    return predicted_seq, analysis_header + analysis

Strategy Advantages: The advantage of this method is its systematic and comprehensive nature. It ensures that the optimization potential of every single residue is examined by the model, providing valuable insights into the sequence-function relationship.

The figure below shows the sequence obtained after performing a point-by-point scanning mask prediction on the M1C mutant using the model fine-tuned on the A4GRB6_PSEAI_Chen_2020 dataset. The residues that differ from the original sequence, which are the potential beneficial mutations revealed by our fine-tuned model, are highlighted in yellow. The blue background indicates residues that are the same as the original sequence.

Scientific Validation: A Three-Way Comparative Framework

To scientifically evaluate our workflow, we established a three-sequence comparative framework. For every low-activity sequence, we compared the following three scores evaluated by our merged scoring function:

• s1 (Baseline Score): The score of the original low-activity sequence.
• s2 (Control Score): The score of the optimal sequence generated by the untuned, original ESM-2 model.
• s3 (Experimental Score): The score of the optimal sequence generated by our fine-tuned ESM2-35M model.

Our golden standard for a "successful modification" is: s3 > s2 > s1. This strict standard requires our method to not only improve upon the original but also significantly outperform a strong control based on the same base model. This proves the added value of our contrastive learning fine-tuning strategy. The figure below is a flowchart of our simulated experimental evaluation process.

Below are the visualized evaluation results from the model fine-tuned on the CSN4_MOUSE_Tsuboyama_2023_1UFM dataset, shown as box plots and violin plots. It can be intuitively seen that the quantiles, mean, and distribution of the modification effects of the model fine-tuned on this dataset have achieved better results compared to the original sequence and the original model.

Below are the specific score plots of the modification results on this dataset. Each red scatter point indicates that the sequence score represented by the vertical axis is better than the sequence score represented by the horizontal axis, while yellow scatter points indicate the opposite.