Equiformer-Pytorch: 革命性的SE3/E3等变注意力网络

Equiformer-Pytorch:革命性的SE3/E3等变注意力网络

Equiformer architecture diagram

Equiformer-Pytorch是一个突破性的深度学习模型,它实现了SE3/E3等变注意力网络,在分子和蛋白质建模等任务中达到了新的性能水平。该模型由Phil Wang开发并开源,目前正在被广泛关注和应用。

模型简介

Equiformer是基于SE3 Transformers的设计思路,但做了一些关键的改进:

将点积注意力机制替换为MLP注意力
引入了来自GATv2的非线性消息传递机制
使用深度张量积来提高效率

这些创新使得Equiformer在保持SE3等变性的同时,能更好地捕捉复杂的3D结构信息。

安装与使用

Equiformer-Pytorch可以通过pip轻松安装:

pip install equiformer-pytorch

下面是一个简单的使用示例:

import torch
from equiformer_pytorch import Equiformer

model = Equiformer(
    num_tokens = 24,
    dim = (4, 4, 2),               # 每种类型的维度,升序,长度必须与degree数量匹配
    dim_head = (4, 4, 4),          # 每个注意力头的维度
    heads = (2, 2, 2),             # 注意力头数量
    num_linear_attn_heads = 0,     # 全局线性注意力头数量
    num_degrees = 3,               # degree数量
    depth = 4,                     # 等变Transformer的深度
    attend_self = True,            # 是否自注意
    reduce_dim_out = True,         # 是否将输出维度降为1
    l2_dist_attention = False      # 设为False以尝试MLP注意力
).cuda()

feats = torch.randint(0, 24, (1, 128)).cuda()
coors = torch.randn(1, 128, 3).cuda()
mask  = torch.ones(1, 128).bool().cuda()

out = model(feats, coors, mask) # (1, 128)

out.type0 # 不变量 type 0    - (1, 128)
out.type1 # 等变量 type 1  - (1, 128, 3)

内存优化

Equiformer-Pytorch还提供了一种基于可逆网络的内存优化方法。这使得即使增加网络深度,内存使用也只会保持在一个Equiformer Transformer块(注意力和前馈)的水平:

model = Equiformer(
    num_tokens = 24,
    dim = (4, 4, 2),
    dim_head = (4, 4, 4),
    heads = (2, 2, 2),
    num_degrees = 3,
    depth = 48,          # 深度增加到48
    reversible = True,   # 启用可逆网络
).cuda()

边缘和邻接矩阵支持

Equiformer-Pytorch还支持处理边缘信息(如原子键)和邻接矩阵:

# 带边缘信息的示例
model = Equiformer(
    num_tokens = 28,
    dim = 64,
    num_edge_tokens = 4,       # 边缘类型数量,如4种键类型
    edge_dim = 16,             # 边缘嵌入维度
    depth = 2,
    input_degrees = 1,
    num_degrees = 3,
    reduce_dim_out = True
)

atoms = torch.randint(0, 28, (2, 32))
bonds = torch.randint(0, 4, (2, 32, 32))
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

out = model(atoms, coors, mask, edges = bonds)

# 带邻接矩阵的示例
model = Equiformer(
    dim = 32,
    heads = 8,
    depth = 1,
    dim_head = 64,
    num_degrees = 2,
    valid_radius = 10,
    reduce_dim_out = True,
    attend_sparse_neighbors = True,
    num_neighbors = 0,               
    num_adj_degrees_embed = 2,       
    max_sparse_neighbors = 8         
)

feats = torch.randn(1, 128, 32)
coors = torch.randn(1, 128, 3)
mask  = torch.ones(1, 128).bool()

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

out = model(feats, coors, mask, adj_mat = adj_mat)