Protein Distance Matrix Prediction

I had wanted to work on a biological related machine learning project so I had started to explore what public datasets were available. I came across PBD, the Protein Folding Databank. I thought this could be a compelling biological machine learning project as I know predicting protein folding is a complex task that is very important and applicable to the biomedical field. I also know that there has been a lot of work done on this particular task so I knew it was a viable solution.

The structure I initially thought of was a transformer type input to read the amino acid sequence, followed by a dense encoder to represent the latent structual featues, and then a tranformer type output to write the CIF file. As I started to look into how that architecture might perform, I found that while on the surface the proposed architecture had merit, but due to the complex local and non-local interactions and gemometric positioning generation it would likely fall flat. In my research, I found that in order to handle the complex spatial interactions and coordinates, Geometric models were used. I had never worked with gemetric models and so I was excited to learn more about them.

I came across a few different gemotric modeling concepts: Graph Neural Networks (GNNs), Geometric Deep Learning (GDLs), Geometric Transformers, Distance Matracies, and Point Clouds.

After researching these geometric modeling concepts, I realized that they would be a better fit for the protein folding task than the initial transformer-based architecture I had proposed. As I had not used these concepts before, I decided to work on some smaller projects to get familiar with them before diving into the protein folding project.

To start, I decided to work on a project that used GNNs to predict a distance matrix for small segments of a protien sequence. The goal was to work with segments that were 10-100 residues long and build a model that would input the amino acid sequence and output a distance matrix. I could use the PDB data to train the model and evaluate its performance. This project would help me get familiar with GNNs and distance matracies, which would be important for the protein folding project.

Project Goals

Download and preprocess PDB data using BioPython to create a dataset for training and evaluation.
Design and implement a model that utilized an Embedding Layer, Transformer Layer, and then GNN layers to predict distance matrices from amino acid sequences.
Train the model using PyTorch and evaluate its performance on the PDB dataset.
Use the model to predict distance matracies for new protein sequences and evaluate its performance.

Tech Stack

Protein Data Processing Biopython, Scipy
Machine Learning: PyTorch, DGL

Example

Input Image:

Challenges

Complex Dataset: The dataset was moderately complex as many PDB files contained multiple chains and ligands, which required careful processing to extract the relevant data and only operate on protein sequences that didn't require cofactors for proper folding. I had not recognized that the filtering method I initially used was not sufficient to remove all the cofactors and ligands from the protein sequences and thus a secondary filtering method was required when calculating the distance matracies to ensure there were no "HETATMS" in the file.
Padding and Attention: The model architecture utilized padding and attention. This was a challenge as the model had a large variance in the number of amino acids in the sequences. Initially, I had actually encoded the padding backwards, signalling the model to ignore the sequence and only focus on the padding. After that was solved, a custom loss function had to be written as when calculating MSE on the distance matrix, large portions of the matrix were 0 as they were padding. This led to an arbitrarily low loss on short sequences.
Trivial Solution Issue: I came across an issue where the model would learn a trivial solution to the problem, that is, it would learn to output the same distance matrix for all inputs. This was due to the model not being able to learn the complex spatial interactions and instead learning to output the most common distance matrix. I had to implement a custom loss function that penalized the model for outputting the same distance matrix for all inputs that minimized the loss function across the training set.
Novel Architecture: As the architecture was new to me, when I ran into issues, I wasn't sure if it was due to the model architecture, a lack of training, the dataset, or my implementation. I thought that this was going to be an easy first step into graph neural networks, but it turned out to be a much more complex task than I had anticipated.

Lessons Learned

Future Improvements

Implement a GNN model on an easier problem to get familiar with the architecture before diving into the protein folding project.
Return to this project when I have a better understanding of GNNs and how to prevent trivial solutions.