Predicting Protein Secondary Structure with ProtT5

Introduction

While AlphaFold predicts full 3D atomic coordinates, there are many scenarios (like filtering massive metagenomic datasets or running on low-resource edge devices) where we only need to predict the secondary structure (alpha helices, beta sheets, coils) of a protein. ProtT5 is a massive transformer model trained on millions of protein sequences. By extracting its embeddings, we can predict secondary structure with incredible accuracy using a simple linear classifier on top.

Architecture Overview

graph TD
    FASTA(Raw FASTA Sequence) --> Tokenizer(T5 Tokenizer)
    Tokenizer --> IDs(Token IDs & Attention Mask)
    
    IDs --> ProtT5(ProtT5-XL-U50 Encoder)
    ProtT5 --> Embeddings(Per-Residue Embeddings: L x 1024)
    
    Embeddings --> Conv1(1D Convolutional Layer)
    Conv1 --> ReLU(ReLU Activation)
    ReLU --> Linear(Linear Classification Layer)
    Linear --> Softmax(Softmax over 3 states: H, E, C)
    
    Softmax --> Pred(Secondary Structure Sequence)
        

Prerequisites

• Python 3.8+
• PyTorch
• Transformers (Hugging Face)
• A basic understanding of FASTA sequences

Step 1: Install Dependencies

pip install torch transformers sentencepiece

Step 2: Load the Pre-trained ProtT5 Model

We'll load ProtT5-XL-U50 from the Hugging Face Hub. This model was trained on the UniRef50 dataset.

import torch
from transformers import T5EncoderModel, T5Tokenizer
import re

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Load the tokenizer and the encoder model (we only need the encoder for embeddings)
tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
model = T5EncoderModel.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc').to(device)
model.eval()

Step 3: Extract Sequence Embeddings

Protein language models treat amino acids as "words". We must separate the sequence into individual characters separated by spaces before tokenization.

# Example FASTA sequence
sequence = "MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTAGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVREIRQHKLRKLNPPDESGPGCMSCKCVLS"

# Preprocess: space out amino acids and replace rare/unknown amino acids with 'X'
processed_seq = " ".join(list(re.sub(r"[UZOB]", "X", sequence)))

# Tokenize
ids = tokenizer([processed_seq], add_special_tokens=True, padding="longest")
input_ids = torch.tensor(ids['input_ids']).to(device)
attention_mask = torch.tensor(ids['attention_mask']).to(device)

# Forward pass to get embeddings
with torch.no_grad():
    embedding_repr = model(input_ids=input_ids, attention_mask=attention_mask)

# The output has shape (batch_size, sequence_length + 1, embedding_dim)
# We ignore the last token (the  token) to get the per-residue embeddings
per_residue_embeddings = embedding_repr.last_hidden_state[0, :-1, :]
print(f"Embedding shape: {per_residue_embeddings.shape}")
</code></pre>

        
Step 4: Predicting Secondary Structure with a 1D CNN
        
To predict the 3-state secondary structure (Helix 'H', Strand 'E', Coil 'C'), we will train a small 1D Convolutional Neural Network (CNN) taking these 1024-dimensional embeddings as input. A CNN is excellent here because it looks at local neighborhoods of amino acids, which physically interact to form secondary structures.
        
import torch.nn as nn

class SecondaryStructureCNN(nn.Module):
    def __init__(self, embedding_dim=1024, num_classes=3):
        super().__init__()
        # 1D Conv layer to look at a neighborhood of 5 residues (kernel_size=5, padding=2)
        self.conv = nn.Conv1d(in_channels=embedding_dim, out_channels=256, kernel_size=5, padding=2)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(256, num_classes)
        
    def forward(self, x):
        # x is expected to be shape (batch, sequence_length, embedding_dim)
        # Conv1d expects (batch, channels, length), so we permute
        x = x.permute(0, 2, 1) 
        x = self.relu(self.conv(x))
        
        # Permute back to (batch, length, channels) for the Linear layer
        x = x.permute(0, 2, 1)
        return self.fc(x)

# Initialize the classifier
classifier = SecondaryStructureCNN().to(device)

# Example Forward Pass
# In a real scenario, you would train this classifier using CrossEntropyLoss against SS labels (e.g. from DSSP)
logits = classifier(per_residue_embeddings)
predictions = torch.argmax(logits, dim=-1)

# Map numeric predictions back to characters
ss_map = {0: 'H', 1: 'E', 2: 'C'}
predicted_ss = "".join([ss_map[p.item()] for p in predictions[0]])
print(f"Predicted SS: {predicted_ss}")


        
Conclusion
        
Protein Language Models like ProtT5 have revolutionized bioinformatics. By pre-training on millions of unlabeled sequences, they learn the fundamental physical and evolutionary grammar of biology. Extracting these embeddings allows us to solve complex downstream tasks—like secondary structure prediction, subcellular localization, and solubility prediction—using extremely lightweight models (like a single 1D Convolution layer) with very little labeled data, vastly accelerating computational biology pipelines.
    </section>
</article>