En-transformer

Project Introduction: E(n)-Equivariant Transformer

The E(n)-Equivariant Transformer is an innovative computational model that enhances the Transformer architecture by integrating it with principles from the E(n)-Equivariant Graph Neural Network developed by Max Welling and his team. This model is particularly noted for its robust application in the field of protein and antibody design, specifically in modeling CDR loops in antibodies, as highlighted in recent research.

Installation

To use the En-transformer in your projects, you can easily install it using pip:

$ pip install En-transformer

Usage Overview

The E(n)-Equivariant Transformer is flexible and can be tailored to a variety of needs. Here’s a glimpse of how it can be used:

1. Basic Transformer Setup:

   The model configuration allows customization of dimensionality, depth, number of heads, and edge features, among others. Attention mechanisms can be configured to focus on a select number of nearest neighbors in the data, and various computational strategies like “talking heads” and checkpointing can be applied for optimized performance.

   import torch
   from en_transformer import EnTransformer
   
   model = EnTransformer(
       dim = 512,
       depth = 4,
       dim_head = 64,
       heads = 8,
       edge_dim = 4,
       neighbors = 64,
       talking_heads = True,
       checkpoint = True,
       use_cross_product = True,
       num_global_linear_attn_heads = 4
   )
   
   feats = torch.randn(1, 1024, 512)
   coors = torch.randn(1, 1024, 3)
   edges = torch.randn(1, 1024, 1024, 4)
   mask = torch.ones(1, 1024).bool()
   feats, coors = model(feats, coors, edges, mask = mask)
   

2. Atomic and Bond Embedding:

   The transformer can handle embeddings for atomic structures and types of bonds, which is particularly useful in molecular modeling.

   atoms = torch.randint(0, 10, (1, 16))
   bonds = torch.randint(0, 5, (1, 16, 16))
   coors = torch.randn(1, 16, 3)
   feats_out, coors_out = model(atoms, coors, edges = bonds)
   

3. Sparse Neighbor Attention:

   For situations where sparse neighbor attention is needed, adjacency matrices can be utilized.

   atoms = torch.randint(0, 10, (1, 16))
   coors = torch.randn(1, 16, 3)
   adj_mat = (torch.arange(atoms.shape[1])[:, None] <= (torch.arange(atoms.shape[1])[None, :] + 1)) & (torch.arange(atoms.shape[1])[:, None] >= (torch.arange(atoms.shape[1])[None, :] - 1))
   feats_out, coors_out = model(atoms, coors, adj_mat = adj_mat)
   

Advanced Features

• Continuous Edges and Edge Embeddings:

  The model also accommodates the use of continuous edges and their embeddings, adding another layer of detail to network configuration.

  feats = torch.randn(1, 16, 512)
  coors = torch.randn(1, 16, 3)
  edges = torch.randn(1, 16, 16, 4)
  feats1, coors1 = model(feats, coors, adj_mat = adj_mat, edges = edges)
  

Practical Applications

An example of the E(n)-Equivariant Transformer in action is the toy task of protein backbone coordinate denoising, which highlights its potential in structural biology. Preparing your environment with packages like sidechainnet allows you to experiment with such applications.

Future Developments

The project is continually evolving. There are plans to incorporate further advancements and broaden its applicability, especially towards modeling larger molecules.

References and Citations

The development of this model draws on a rich body of research, including foundational works on graph neural networks and attention mechanisms. Researchers are encouraged to consult the cited works for a deep dive into the underlying theories and algorithms.

The E(n)-Equivariant Transformer is a cutting-edge tool in computational biology and machine learning, offering vast customization possibilities and applications in structural predictions and molecular design. Whether you are a researcher looking to model complex systems or a developer interested in leveraging neural networks for spatial attention tasks, this transformer offers a powerful and flexible platform.