Fine-Tuning Transformers for protein sequence-to-function modeling¶

In the previous notebook Learning sequence-to-function relationship using language models, we built a protein sequence-to-function model by training a regressor on embedding from transformers protein language model. We froze the pretrained transformers body's weights during training and used the hidden states as fixed features for the regressor.

In this notebook, we will employ transfer learning to fine-tune pre-trained protein language models for specific tasks of interest. The entire model will be trained end-to-end, avoiding the use of fixed hidden states as features. This necessitates a differentiable classification head, which is why the method incorporates a neural network for regression purposes.

Goals

• Fine-tune a large protein language model on deep mutational scanning (DMS) data and use it for inference.

Load the data¶

The following data originated from the "Data processing" section of the prior notebook titled "Learning protein sequence-function relationship from Deep Mutational Scanning data." It stems from the original Deep Mutational Scanning (DMS) data associated with the "Comprehensive mutagenesis maps the effect of all single codon mutations in the AAV2 rep gene on AAV production" paper, which was published in 2023 by the George M. Church lab DOI: 10.7554/eLife.87730.1. The published data underwent transformation to generate essential columns for sequence-to-function modeling.

Brief description of the columns in the dataset:

aa_mutations: Indicates amino acid mutations, where, for example, "M1A" signifies a change from amino acid "M" at position 1 to amino acid "A". A blank line at row index -1 denotes no amino acid mutation, representing the wild-type sequence.

fitness_values: Represents the fitness values corresponding to each protein variant.

sequence: Displays the wild-type protein sequence at index -1 and mutated protein sequences for all other indices.

num_fitval: Indicates the number of fitness values associated with each protein variant.

median_fitness: Represents the median of the num_fitval column.

std: Signifies the standard deviation of the num_fitval column.

log10_fitness: Reflects the log10 transformation of the median_fitness column.

Data normalization¶

The dataset exhibits substantial noise, inherent to high-throughput measurements. In the modeling phase, we streamline our approach by exclusively utilizing the median fitness value for amino acid substitutions spanning the entire protein length. We also perform transformation to normalize the median fitness value such that wildtype fitness corresponds to a value of 1.

Splitting the data¶

In machine learning, the train-test split is a vital practice involving the division of a dataset into two subsets: the training set, used to train the model, and the test set, employed to evaluate the model's performance on unseen data. This separation ensures that the model is not simply memorizing the training data but can generalize well to new examples. Achieving the right balance in data allocation is crucial for effective model training and evaluation, fostering reliable insights into its real-world applicability.

Tokenizing the data¶

Data tokenization in protein language models entails breaking down protein sequences into numerical indices, or tokens, to enable input to nueral network. These tokens, often corresponding to individual amino acids or short sequences, allow the model to capture intricate patterns in protein structures. Every model on transformers comes with an associated tokenizer that handles tokenization for it. This holds true for protein language models as well.

Our sequence has been converted into input_ids, which is the tokenized sequence, and an attention_mask. The purpose of the attention mask is to manage variable-length sequences. In instances of shorter sequences, padding is applied using blank "padding" tokens, and the attention mask is padded with 0s, signaling the model to disregard these tokens during processing.

Pytorch compatible Dataset creation¶

Our next step involves shaping this data into a dataset compatible with PyTorch. To achieve this, we leverage the HuggingFace Dataset class.

Fine-tuning Transformers¶

Fine-tuning a pretrained protein language model involves adjusting the parameters of an existing model that has already been trained on a large dataset. Initially pretrained on a diverse set of protein sequences, the model has learned general features and patterns. Fine-tuning adapts the model to a specific task or dataset by exposing it to a smaller, task-specific dataset. During this process, the model's weights are updated to better capture the nuances and specificities of the new data.

Loading a pretrained model¶

Use exactly the same model as you used when loading the tokenizer, or your model might not understand the tokenization scheme you're using!

The warnings are signaling that the model is shedding certain weights related to language modeling (specifically, the lm_head) while incorporating new weights for sequence classification (the classifier). This aligns perfectly with our intention to fine-tune a pretrained language model for a distinct task. The reminder is crucial—it underscores the need to fine-tune the model as it isn't directly usable for inference in its current state.

Defining the performance metrics¶

Training the model¶

Plot results¶

Saving and Loading the finetuned model¶

To check if the original and the reloaded models are identical we can compare weights.

Inference from the model¶

Conclusion¶

We have developed two complementary approaches for sequence-to-function modeling using deep mutational scanning data: one relying on feature extraction and the other on fine-tuning. Typically, with a dataset of substantial size, fine-tuning tends to yield superior model performance compared to feature extraction. Regardless of whether the model is fine-tuned or based on feature extraction, it holds the capability to contribute to prospective design by generating sequence proposals through in silico directed evolution. The detailed exploration of in silico directed evolution will be the focus of an upcoming notebook.

In the context of model inference, it is imperative to prioritize scientific rigor over statistical outcomes. A thorough comprehension of the data generation process, acknowledgment of model limitations, and an understanding of the distribution within the inference regime are essential. Consideration of specific experimental details is paramount - the study's design, assay fidelity (distinguishing between low and high throughput measurements), data quality (with attention to sequencing depth in next generation sequencing measurements), and the alignment of the ultimate test data distribution with the training data distribution. It is crucial to refrain from falling into blind love with model's elegance and power, as explained upon in the previous notebook. Instead, a judicious evaluation of the scientific context is indispensable for robust and meaningful model inference.

In this notebook, we fine-tuned a transformer protein language model using assay-labeled data, comprising protein sequence variants and their corresponding fitness values. The question arises: Can we fine-tune a transformer model in the absence of assay-labeled data? This intriguing possibility will be the focus of a future notebook.

Moreover, a more pressing challenge surfaces: Can we identify promising protein sequence variants for experimental validation when no initial experimental measurements are available? This will serve as the central theme for the upcoming notebook—zero-shot prediction of evolutionarily plausible and promising protein sequence variants.