En-transformer

Use EnTransformer when you have pre-computed continuous node features and 3D coordinates.

Arguments:

• dim: Dimension of the node features.
• depth: Number of layers.
• dim_head: Dimension per attention head.
• heads: Number of attention heads.
• edge_dim: Dimension of the edge features.
• neighbors: Number of nearest neighbors to attend to (set to 0 to turn off).
• talking_heads: Whether to use Shazeer's talking heads mechanism.
• checkpoint: Whether to use gradient checkpointing.
• use_cross_product: Whether to use cross product vectors.
• num_global_linear_attn_heads: Number of heads to use for global linear attention if neighbors is low.

Input Shapes:

• feats: (batch, nodes, dim)
• coors: (batch, nodes, 3)
• edges: (batch, nodes, nodes, edge_dim)
• mask: (batch, nodes) (boolean mask)

Returns:

• feats, coors: Updated features and coordinates with the same shapes as inputs.

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)  # (1, 1024, 512), (1, 1024, 3)

If you want to attend only to sparse neighbors defined by an adjacency matrix, configure the model with only_sparse_neighbors=True.

Arguments:

• only_sparse_neighbors: Must be set to True to use the adj_mat argument.
• neighbors: Set to 0.
• num_adj_degrees: Number of degrees to derive from 1st degree neighbors passed in.
• adj_dim: Dimension of the adjacency degree information if passed as edge embedding.

Input Shapes:

• atoms: (batch, nodes) (integer tokens)
• coors: (batch, nodes, 3)
• adj_mat: (nodes, nodes) (boolean adjacency matrix)

Returns:

• feats_out, coors_out: Updated features and coordinates.

import torch
from en_transformer import EnTransformer

model = EnTransformer(
    num_tokens = 10,
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    neighbors = 0,
    only_sparse_neighbors = True,
    num_adj_degrees = 2,
    adj_dim = 8
)

atoms = torch.randint(0, 10, (1, 16))
coors = torch.randn(1, 16, 3)

# Example: naively assume a single chain of atoms
i = torch.arange(atoms.shape[1])
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

feats_out, coors_out = model(atoms, coors, adj_mat = adj_mat) # (1, 16, 512), (1, 16, 3)

Use EnTransformer to let the network handle embeddings for discrete node types (e.g., atoms) and edge types (e.g., bond types).

Arguments:

• num_tokens: Number of unique node types.
• rel_pos_emb: Set to True if the sequence is not an unordered set to accelerate convergence.
• num_edge_tokens: Number of unique edge types.
• dim: Node feature dimension.
• edge_dim: Edge feature dimension.

Input Shapes:

• atoms: (batch, nodes) (integer tokens)
• bonds: (batch, nodes, nodes) (integer tokens representing edge types)
• coors: (batch, nodes, 3)

Returns:

• feats_out, coors_out: Updated features and coordinates.

import torch
from en_transformer import EnTransformer

model = EnTransformer(
    num_tokens = 10,
    rel_pos_emb = True,
    num_edge_tokens = 5,
    dim = 128,
    edge_dim = 16,
    depth = 3,
    heads = 4,
    dim_head = 32,
    neighbors = 8
)

atoms = torch.randint(0, 10, (1, 16))    # 10 different types of atoms
bonds = torch.randint(0, 5, (1, 16, 16)) # 5 different types of bonds (n x n)
coors = torch.randn(1, 16, 3)            # atomic spatial coordinates

feats_out, coors_out = model(atoms, coors, edges = bonds) # (1, 16, 512), (1, 16, 3)

You can combine sparse neighbor attention (via an adjacency matrix) with continuous edge features.

Arguments:

• only_sparse_neighbors: Set to True.
• neighbors: Set to 0.

Input Shapes:

• feats: (batch, nodes, dim)
• coors: (batch, nodes, 3)
• adj_mat: (nodes, nodes) (boolean adjacency matrix)
• edges: (batch, nodes, nodes, edge_dim)

Returns:

• feats1, coors1: Updated features and coordinates.

import torch
from en_transformer import EnTransformer
from en_transformer.utils import rot

model = EnTransformer(
    dim = 512,
    depth = 1,
    heads = 4,
    dim_head = 32,
    edge_dim = 4,
    num_nearest_neighbors = 0,
    only_sparse_neighbors = True
)

feats = torch.randn(1, 16, 512)
coors = torch.randn(1, 16, 3)
edges = torch.randn(1, 16, 16, 4)

i = torch.arange(feats.shape[1])
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

feats1, coors1 = model(feats, coors, adj_mat = adj_mat, edges = edges)