| from typing import Optional, Tuple |
| import torch |
| from torch import Tensor, nn |
| from torch_geometric.nn import MessagePassing |
| from torch_geometric.utils import softmax |
| from torch_scatter import scatter |
| from torch_sparse import SparseTensor |
| import loralib as lora |
| from esm.multihead_attention import MultiheadAttention |
| import math |
| from torch import _dynamo |
| _dynamo.config.suppress_errors = True |
| from ..module.utils import ( |
| CosineCutoff, |
| act_class_mapping, |
| get_template_fn, |
| gelu |
| ) |
|
|
|
|
| |
| class EquivariantMultiHeadAttention(MessagePassing): |
| """Equivariant multi-head attention layer.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantMultiHeadAttention, self).__init__( |
| aggr="mean", node_dim=0) |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
| self.share_kv = share_kv |
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
| self.cutoff = CosineCutoff(cutoff_lower, cutoff_upper) |
|
|
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.k_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) if not share_kv else None |
| self.v_proj = lora.Linear( |
| x_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| self.o_proj = lora.Linear( |
| x_hidden_channels, x_channels * 2 + vec_channels, r=use_lora) |
| self.vec_proj = lora.Linear( |
| vec_channels, vec_hidden_channels * 2 + vec_channels, bias=False, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.k_proj = nn.Linear(x_channels, x_hidden_channels) if not share_kv else None |
| self.v_proj = nn.Linear( |
| x_channels, x_hidden_channels + vec_channels * 2) |
| self.o_proj = nn.Linear( |
| x_hidden_channels, x_channels * 2 + vec_channels) |
| self.vec_proj = nn.Linear( |
| vec_channels, vec_hidden_channels * 2 + vec_channels, bias=False) |
|
|
| self.dk_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
|
|
| self.dv_proj = None |
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_proj = lora.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.dv_proj = nn.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.k_proj.weight) |
| self.k_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.v_proj.weight) |
| self.v_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.vec_proj.weight) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, r_ij, f_ij, d_ij, return_attn=False): |
| x = self.layernorm(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| v = self.v_proj(x).reshape(-1, self.num_heads, |
| self.x_head_dim + self.vec_head_dim * 2) |
| if self.share_kv: |
| k = v[:, :, :self.x_head_dim] |
| else: |
| k = self.k_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
|
|
| vec1, vec2, vec3 = torch.split(self.vec_proj(vec), |
| [self.vec_hidden_channels, self.vec_hidden_channels, self.vec_channels], dim=-1) |
| vec = vec.reshape(-1, 3, self.num_heads, self.vec_head_dim) |
| vec_dot = (vec1 * vec2).sum(dim=1) |
|
|
| dk = ( |
| self.act(self.dk_proj(f_ij)).reshape(-1, |
| self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None |
| else None |
| ) |
| dv = ( |
| self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, |
| self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_proj is not None |
| else None |
| ) |
|
|
| |
| |
| x, vec, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| dv=dv, |
| r_ij=r_ij, |
| d_ij=d_ij, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| vec = vec.reshape(-1, 3, self.vec_channels) |
|
|
| o1, o2, o3 = torch.split(self.o_proj( |
| x), [self.vec_channels, self.x_channels, self.x_channels], dim=1) |
| dx = vec_dot * o2 + o3 |
| dvec = vec3 * o1.unsqueeze(1) + vec |
| if return_attn: |
| return dx, dvec, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return dx, dvec, None |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dv, r_ij, d_ij): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j * dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) * self.cutoff(r_ij).unsqueeze(1) |
|
|
| |
| if dv is not None: |
| v_j = v_j * dv |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = vec_j * vec1.unsqueeze(1) + vec2.unsqueeze(1) * \ |
| d_ij.unsqueeze(2).unsqueeze(3) |
| return x, vec, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, vec, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size) |
| vec = scatter(vec, index, dim=self.node_dim, dim_size=dim_size) |
| return x, vec, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class ESMMultiheadAttention(MultiheadAttention): |
| """Multi-headed attention. |
| |
| See "Attention Is All You Need" for more details. |
| """ |
|
|
| def __init__( |
| self, |
| embed_dim, |
| num_heads, |
| kdim=None, |
| vdim=None, |
| dropout=0.0, |
| bias=True, |
| add_bias_kv: bool = False, |
| add_zero_attn: bool = False, |
| self_attention: bool = False, |
| encoder_decoder_attention: bool = False, |
| use_rotary_embeddings: bool = False, |
| ): |
| super().__init__(embed_dim, num_heads, kdim, vdim, dropout, bias, add_bias_kv, add_zero_attn, self_attention, |
| encoder_decoder_attention, use_rotary_embeddings) |
| |
| self.k_proj = lora.Linear(self.kdim, embed_dim, bias=bias, r=16) |
| self.v_proj = lora.Linear(self.vdim, embed_dim, bias=bias, r=16) |
| self.q_proj = lora.Linear(embed_dim, embed_dim, bias=bias, r=16) |
| self.out_proj = lora.Linear(embed_dim, embed_dim, bias=bias, r=16) |
|
|
|
|
| |
| class EquivariantPAEMultiHeadAttention(EquivariantMultiHeadAttention): |
| """Equivariant multi-head attention layer.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantPAEMultiHeadAttention, self).__init__( |
| x_channels=x_channels, |
| x_hidden_channels=x_hidden_channels, |
| vec_channels=vec_channels, |
| vec_hidden_channels=vec_hidden_channels, |
| share_kv=share_kv, |
| edge_attr_channels=edge_attr_channels, |
| distance_influence=distance_influence, |
| num_heads=num_heads, |
| activation=activation, |
| attn_activation=attn_activation, |
| cutoff_lower=cutoff_lower, |
| cutoff_upper=cutoff_upper, |
| use_lora=use_lora) |
| |
| self.cutoff = None |
| |
| self.dk_dist_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.dk_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels) |
| self.dv_dist_proj = None |
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.dv_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2) |
| if self.dk_dist_proj: |
| nn.init.xavier_uniform_(self.dk_dist_proj.weight) |
| self.dk_dist_proj.bias.data.fill_(0) |
| if self.dv_dist_proj: |
| nn.init.xavier_uniform_(self.dv_dist_proj.weight) |
| self.dv_dist_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, w_ij, f_dist_ij, f_ij, d_ij, plddt, return_attn=False): |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| v = self.v_proj(x).reshape(-1, self.num_heads, |
| self.x_head_dim + self.vec_head_dim * 2) |
| if self.share_kv: |
| k = v[:, :, :self.x_head_dim] |
| else: |
| k = self.k_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
|
|
| vec1, vec2, vec3 = torch.split(self.vec_proj(vec), |
| [self.vec_hidden_channels, self.vec_hidden_channels, self.vec_channels], dim=-1) |
| vec = vec.reshape(-1, 3, self.num_heads, self.vec_head_dim) |
| vec_dot = (vec1 * vec2).sum(dim=1) |
|
|
| dk = ( |
| self.act(self.dk_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None |
| else None |
| ) |
| dk_dist = ( |
| self.act(self.dk_dist_proj(f_dist_ij)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_dist_proj is not None |
| else None |
| ) |
| dv = ( |
| self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_proj is not None |
| else None |
| ) |
| dv_dist = ( |
| self.act(self.dv_dist_proj(f_dist_ij)).reshape(-1, self.num_heads, self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_dist_proj is not None |
| else None |
| ) |
|
|
| |
| |
| x, vec, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| dk_dist=dk_dist, |
| dv=dv, |
| dv_dist=dv_dist, |
| d_ij=d_ij, |
| w_ij=w_ij, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| vec = vec.reshape(-1, 3, self.vec_channels) |
|
|
| o1, o2, o3 = torch.split(self.o_proj( |
| x), [self.vec_channels, self.x_channels, self.x_channels], dim=1) |
| dx = vec_dot * o2 * plddt.unsqueeze(1) + o3 |
| dvec = vec3 * o1.unsqueeze(1) * plddt.unsqueeze(1).unsqueeze(2) + vec |
| if return_attn: |
| return dx, dvec, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return dx, dvec, None |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dk_dist, dv, dv_dist, d_ij, w_ij): |
| |
| attn = (q_i * k_j) |
| if dk is not None: |
| attn += dk |
| if dk_dist is not None: |
| attn += dk_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| attn = attn.sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) |
|
|
| |
| if dv is not None: |
| v_j += dv |
| if dv_dist is not None: |
| v_j += dv_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = vec_j * vec1.unsqueeze(1) + vec2.unsqueeze(1) * \ |
| d_ij.unsqueeze(2).unsqueeze(3) |
| return x, vec, attn |
|
|
|
|
| |
| class EquivariantWeightedPAEMultiHeadAttention(EquivariantMultiHeadAttention): |
| """Equivariant multi-head attention layer.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantWeightedPAEMultiHeadAttention, self).__init__( |
| x_channels=x_channels, |
| x_hidden_channels=x_hidden_channels, |
| vec_channels=vec_channels, |
| vec_hidden_channels=vec_hidden_channels, |
| share_kv=share_kv, |
| edge_attr_channels=edge_attr_channels, |
| distance_influence=distance_influence, |
| num_heads=num_heads, |
| activation=activation, |
| attn_activation=attn_activation, |
| cutoff_lower=cutoff_lower, |
| cutoff_upper=cutoff_upper, |
| use_lora=use_lora) |
| |
| self.cutoff = None |
| |
| self.pae_weight = nn.Linear(1, 1, bias=True) |
| self.pae_weight.weight.data.fill_(-0.5) |
| self.pae_weight.bias.data.fill_(7.5) |
| |
| self.dk_dist_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.dk_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels) |
| self.dv_dist_proj = None |
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.dv_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2) |
| if self.dk_dist_proj: |
| nn.init.xavier_uniform_(self.dk_dist_proj.weight) |
| self.dk_dist_proj.bias.data.fill_(0) |
| if self.dv_dist_proj: |
| nn.init.xavier_uniform_(self.dv_dist_proj.weight) |
| self.dv_dist_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, w_ij, f_dist_ij, f_ij, d_ij, plddt, return_attn=False): |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| v = self.v_proj(x).reshape(-1, self.num_heads, |
| self.x_head_dim + self.vec_head_dim * 2) |
| if self.share_kv: |
| k = v[:, :, :self.x_head_dim] |
| else: |
| k = self.k_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
|
|
| vec1, vec2, vec3 = torch.split(self.vec_proj(vec), |
| [self.vec_hidden_channels, self.vec_hidden_channels, self.vec_channels], dim=-1) |
| vec = vec.reshape(-1, 3, self.num_heads, self.vec_head_dim) |
| vec_dot = (vec1 * vec2).sum(dim=1) |
|
|
| dk = ( |
| self.act(self.dk_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None |
| else None |
| ) |
| dk_dist = ( |
| self.act(self.dk_dist_proj(f_dist_ij)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_dist_proj is not None |
| else None |
| ) |
| dv = ( |
| self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_proj is not None |
| else None |
| ) |
| dv_dist = ( |
| self.act(self.dv_dist_proj(f_dist_ij)).reshape(-1, self.num_heads, self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_dist_proj is not None |
| else None |
| ) |
|
|
| |
| |
| x, vec, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| dk_dist=dk_dist, |
| dv=dv, |
| dv_dist=dv_dist, |
| d_ij=d_ij, |
| w_ij=nn.functional.sigmoid(self.pae_weight(w_ij.unsqueeze(-1)).squeeze(-1)), |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| vec = vec.reshape(-1, 3, self.vec_channels) |
|
|
| o1, o2, o3 = torch.split(self.o_proj( |
| x), [self.vec_channels, self.x_channels, self.x_channels], dim=1) |
| dx = vec_dot * o2 * plddt.unsqueeze(1) + o3 |
| dvec = vec3 * o1.unsqueeze(1) * plddt.unsqueeze(1).unsqueeze(2) + vec |
| if return_attn: |
| return dx, dvec, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return dx, dvec, None |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dk_dist, dv, dv_dist, d_ij, w_ij): |
| |
| attn = (q_i * k_j) |
| if dk_dist is not None: |
| if dk is not None: |
| attn *= (dk + dk_dist * w_ij.unsqueeze(1).unsqueeze(2)) |
| else: |
| attn *= dk_dist * w_ij |
| else: |
| if dk is not None: |
| attn *= dk |
| attn = attn.sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) |
|
|
| |
| if dv is not None: |
| v_j += dv |
| if dv_dist is not None: |
| v_j += dv_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = vec_j * vec1.unsqueeze(1) + vec2.unsqueeze(1) * \ |
| d_ij.unsqueeze(2).unsqueeze(3) |
| return x, vec, attn |
|
|
|
|
| class EquivariantPAEMultiHeadAttentionSoftMaxFullGraph(nn.Module): |
| """Equivariant multi-head attention layer with softmax, apply attention on full graph by default""" |
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| |
| super(EquivariantPAEMultiHeadAttentionSoftMaxFullGraph, self).__init__() |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
| self.share_kv = share_kv |
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.cutoff = None |
| self.scaling = self.x_head_dim**-0.5 |
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.k_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) if not share_kv else None |
| self.v_proj = lora.Linear(x_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| self.o_proj = lora.Linear(x_hidden_channels, x_channels * 2 + vec_channels, r=use_lora) |
| self.vec_proj = lora.Linear(vec_channels, vec_hidden_channels * 2 + vec_channels, bias=False, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.k_proj = nn.Linear(x_channels, x_hidden_channels) if not share_kv else None |
| self.v_proj = nn.Linear(x_channels, x_hidden_channels + vec_channels * 2) |
| self.o_proj = nn.Linear(x_hidden_channels, x_channels * 2 + vec_channels) |
| self.vec_proj = nn.Linear(vec_channels, vec_hidden_channels * 2 + vec_channels, bias=False) |
|
|
| self.dk_proj = None |
| self.dk_dist_proj = None |
| self.dv_proj = None |
| self.dv_dist_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| self.dk_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| self.dk_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels) |
|
|
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_proj = lora.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| self.dv_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.dv_proj = nn.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2) |
| self.dv_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2) |
| |
| self.pae_weight = nn.Linear(1, 1, bias=True) |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.k_proj.weight) |
| self.k_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.v_proj.weight) |
| self.v_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.vec_proj.weight) |
| self.pae_weight.weight.data.fill_(-0.5) |
| self.pae_weight.bias.data.fill_(7.5) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
| if self.dk_dist_proj: |
| nn.init.xavier_uniform_(self.dk_dist_proj.weight) |
| self.dk_dist_proj.bias.data.fill_(0) |
| if self.dv_dist_proj: |
| nn.init.xavier_uniform_(self.dv_dist_proj.weight) |
| self.dv_dist_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, w_ij, f_dist_ij, f_ij, d_ij, plddt, key_padding_mask, return_attn=False): |
| |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x) * self.scaling |
| v = self.v_proj(x) |
| |
| |
| |
| k = self.k_proj(x) |
|
|
| vec1, vec2, vec3 = torch.split(self.vec_proj(vec), |
| [self.vec_hidden_channels, self.vec_hidden_channels, self.vec_channels], dim=-1) |
| vec_dot = (vec1 * vec2).sum(dim=-2) |
|
|
| dk = self.act(self.dk_proj(f_ij)) |
| dk_dist = self.act(self.dk_dist_proj(f_dist_ij)) |
| dv = self.act(self.dv_proj(f_ij)) |
| dv_dist = self.act(self.dv_dist_proj(f_dist_ij)) |
|
|
| |
| x, vec, attn = self.attention( |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| dk_dist=dk_dist, |
| dv=dv, |
| dv_dist=dv_dist, |
| d_ij=d_ij, |
| w_ij=nn.functional.sigmoid(self.pae_weight(w_ij.unsqueeze(-1)).squeeze(-1)), |
| key_padding_mask=key_padding_mask, |
| ) |
| o1, o2, o3 = torch.split(self.o_proj(x), [self.vec_channels, self.x_channels, self.x_channels], dim=-1) |
| dx = vec_dot * o2 * plddt.unsqueeze(-1) + o3 |
| dvec = vec3 * o1.unsqueeze(-2) * plddt.unsqueeze(-1).unsqueeze(-2) + vec |
| |
| dx = dx.masked_fill(key_padding_mask.unsqueeze(-1), 0) |
| dvec = dvec.masked_fill(key_padding_mask.unsqueeze(-1).unsqueeze(-1), 0) |
| if return_attn: |
| return dx, dvec, attn |
| else: |
| return dx, dvec, None |
|
|
| def attention(self, q, k, v, vec, dk, dk_dist, dv, dv_dist, d_ij, w_ij, key_padding_mask=None, need_head_weights=False): |
| |
| |
| |
| |
| |
| |
| |
| bsz, tgt_len, _ = q.size() |
| src_len = k.size(1) |
| |
| |
| q = q.transpose(0, 1).reshape(tgt_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| k = k.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| v = v.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim + 2 * self.vec_head_dim).transpose(0, 1).contiguous() |
| |
| vec = vec.permute(1, 2, 0, 3).reshape(src_len, 3, bsz * self.num_heads, self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| dk = dk.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| dk_dist = dk_dist.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| dv = dv.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim + 2 * self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| dv_dist = dv_dist.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim + 2 * self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| assert key_padding_mask.size(0) == bsz |
| assert key_padding_mask.size(1) == src_len |
| |
| attn_weights = torch.multiply(q[:, :, None, :], k[:, None, :, :]) |
| |
| |
| |
| |
| assert w_ij is not None |
| |
| attn_weights *= (dk + dk_dist * w_ij[:, :, :, None].repeat(self.num_heads, 1, 1, self.x_head_dim)) |
| |
| v = v.unsqueeze(1) + dv + dv_dist * w_ij[:, :, :, None].repeat(self.num_heads, 1, 1, self.x_head_dim + 2 * self.vec_head_dim) |
| |
| |
| |
| |
| |
| |
| attn_weights = attn_weights.sum(dim=-1) |
| |
| |
| |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).contiguous() |
| attn_weights = attn_weights.masked_fill( |
| key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf") |
| ) |
| attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len).contiguous() |
| |
| attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32) |
| |
| x, vec1, vec2 = torch.split(v, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=-1) |
| |
| x_out = torch.einsum('bts,btsh->bth', attn_weights, x) |
| |
| vec_out_1 = torch.einsum('bsih,btsh->btih', vec, vec1) |
| |
| vec_out_2 = torch.einsum('btsi,btsh->btih', d_ij, vec2) |
| |
| vec_out = vec_out_1 + vec_out_2 |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) |
| |
| |
| attn_weights = attn_weights.mean(dim=0) |
| |
| x_out = x_out.transpose(1, 0).reshape(tgt_len, bsz, self.num_heads * self.x_head_dim).transpose(1, 0).contiguous() |
| |
| vec_out = vec_out.permute(1, 2, 0, 3).reshape(tgt_len, 3, bsz, self.num_heads * self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| return x_out, vec_out, attn_weights |
|
|
|
|
| class MultiHeadAttentionSoftMaxFullGraph(nn.Module): |
| """ |
| Multi-head attention layer with softmax, apply attention on full graph by default |
| No equivariant property, but can take structure information as input, just didn't use it |
| """ |
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| |
| super(MultiHeadAttentionSoftMaxFullGraph, self).__init__() |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
| self.share_kv = share_kv |
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.cutoff = None |
| self.scaling = self.x_head_dim**-0.5 |
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.k_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) if not share_kv else None |
| self.v_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.o_proj = lora.Linear(x_hidden_channels, x_channels, r=use_lora) |
| |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.k_proj = nn.Linear(x_channels, x_hidden_channels) if not share_kv else None |
| self.v_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.o_proj = nn.Linear(x_hidden_channels, x_channels) |
| |
|
|
| self.dk_proj = None |
| self.dk_dist_proj = None |
| self.dv_proj = None |
| self.dv_dist_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| |
| else: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| |
|
|
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| |
| else: |
| self.dv_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| |
| |
| |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.k_proj.weight) |
| self.k_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.v_proj.weight) |
| self.v_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| |
| |
| |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, w_ij, f_dist_ij, f_ij, d_ij, plddt, key_padding_mask, return_attn=False): |
| |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x) * self.scaling |
| v = self.v_proj(x) |
| |
| |
| |
| k = self.k_proj(x) |
|
|
| |
| |
| |
|
|
| dk = self.act(self.dk_proj(f_ij)) |
| |
| dv = self.act(self.dv_proj(f_ij)) |
| |
|
|
| |
| x, vec, attn = self.attention( |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| |
| dv=dv, |
| |
| |
| |
| key_padding_mask=key_padding_mask, |
| ) |
| |
| |
| dx = self.o_proj(x) |
| |
| |
| dx = dx.masked_fill(key_padding_mask.unsqueeze(-1), 0) |
| |
| if return_attn: |
| return dx, vec, attn |
| else: |
| return dx, vec, None |
|
|
| def attention(self, q, k, v, vec, dk, dv, key_padding_mask=None, need_head_weights=False): |
| |
| |
| |
| |
| |
| |
| |
| bsz, tgt_len, _ = q.size() |
| src_len = k.size(1) |
| |
| |
| q = q.transpose(0, 1).reshape(tgt_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| k = k.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| v = v.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| |
| |
| |
| |
| |
| dk = dk.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| |
| |
| |
| dv = dv.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| |
| |
| assert key_padding_mask.size(0) == bsz |
| assert key_padding_mask.size(1) == src_len |
| |
| attn_weights = torch.multiply(q[:, :, None, :], k[:, None, :, :]) |
| |
| |
| |
| |
| |
| |
| attn_weights *= dk |
| |
| v = v.unsqueeze(1) + dv |
| |
| |
| |
| |
| |
| |
| attn_weights = attn_weights.sum(dim=-1) |
| |
| |
| |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).contiguous() |
| attn_weights = attn_weights.masked_fill( |
| key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf") |
| ) |
| attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len).contiguous() |
| |
| attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32) |
| |
| |
| |
| x_out = torch.einsum('bts,btsh->bth', attn_weights, v) |
| |
| |
| |
| |
| |
| |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) |
| |
| |
| attn_weights = attn_weights.mean(dim=0) |
| |
| x_out = x_out.transpose(1, 0).reshape(tgt_len, bsz, self.num_heads * self.x_head_dim).transpose(1, 0).contiguous() |
| |
| |
| return x_out, vec, attn_weights |
|
|
|
|
| class PAEMultiHeadAttentionSoftMaxStarGraph(nn.Module): |
| """Equivariant multi-head attention layer with softmax, apply attention on full graph by default""" |
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| |
| super(PAEMultiHeadAttentionSoftMaxStarGraph, self).__init__() |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
| self.share_kv = share_kv |
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.cutoff = None |
| self.scaling = self.x_head_dim**-0.5 |
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.k_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) if not share_kv else None |
| self.v_proj = lora.Linear(x_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.k_proj = nn.Linear(x_channels, x_hidden_channels) if not share_kv else None |
| self.v_proj = nn.Linear(x_channels, x_hidden_channels) |
|
|
| self.dk_proj = None |
| self.dk_dist_proj = None |
| self.dv_proj = None |
| self.dv_dist_proj = None |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| self.dk_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| self.dk_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels) |
|
|
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_proj = lora.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| self.dv_dist_proj = lora.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2, r=use_lora) |
| else: |
| self.dv_proj = nn.Linear(edge_attr_channels, x_hidden_channels + vec_channels * 2) |
| self.dv_dist_proj = nn.Linear(edge_attr_dist_channels, x_hidden_channels + vec_channels * 2) |
| |
| self.pae_weight = nn.Linear(1, 1, bias=True) |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.k_proj.weight) |
| self.k_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.v_proj.weight) |
| self.v_proj.bias.data.fill_(0) |
| self.pae_weight.weight.data.fill_(-0.5) |
| self.pae_weight.bias.data.fill_(7.5) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
| if self.dk_dist_proj: |
| nn.init.xavier_uniform_(self.dk_dist_proj.weight) |
| self.dk_dist_proj.bias.data.fill_(0) |
| if self.dv_dist_proj: |
| nn.init.xavier_uniform_(self.dv_dist_proj.weight) |
| self.dv_dist_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, x_center_index, w_ij, f_dist_ij, f_ij, key_padding_mask, return_attn=False): |
| |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x[x_center_index].unsqueeze(1)) * self.scaling |
| v = self.v_proj(x) |
| |
| |
| |
| k = self.k_proj(x) |
|
|
| dk = self.act(self.dk_proj(f_ij)) |
| dk_dist = self.act(self.dk_dist_proj(f_dist_ij)) |
| dv = self.act(self.dv_proj(f_ij)) |
| dv_dist = self.act(self.dv_dist_proj(f_dist_ij)) |
|
|
| |
| x, attn = self.attention( |
| q=q, |
| k=k, |
| v=v, |
| dk=dk, |
| dk_dist=dk_dist, |
| dv=dv, |
| dv_dist=dv_dist, |
| w_ij=nn.functional.sigmoid(self.pae_weight(w_ij.unsqueeze(-1)).squeeze(-1)), |
| key_padding_mask=key_padding_mask, |
| ) |
| if return_attn: |
| return x, attn |
| else: |
| return x, None |
|
|
| def attention(self, q, k, v, dk, dk_dist, dv, dv_dist, w_ij, key_padding_mask=None, need_head_weights=False): |
| |
| |
| |
| |
| |
| |
| |
| bsz, tgt_len, _ = q.size() |
| src_len = k.size(1) |
| |
| |
| q = q.transpose(0, 1).reshape(tgt_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| k = k.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| v = v.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| |
| |
| |
| dk = dk.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| dk_dist = dk_dist.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| dv = dv.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim + 2 * self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| dv_dist = dv_dist.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim + 2 * self.vec_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| assert key_padding_mask.size(0) == bsz |
| assert key_padding_mask.size(1) == src_len |
| |
| attn_weights = torch.multiply(q[:, :, None, :], k[:, None, :, :]) |
| |
| |
| |
| |
| assert w_ij is not None |
| |
| attn_weights *= (dk + dk_dist * w_ij[:, :, :, None].repeat(self.num_heads, 1, 1, self.x_head_dim)) |
| |
| v = v.unsqueeze(1) + dv + dv_dist * w_ij[:, :, :, None].repeat(self.num_heads, 1, 1, self.x_head_dim + 2 * self.vec_head_dim) |
| |
| |
| |
| |
| |
| |
| attn_weights = attn_weights.sum(dim=-1) |
| |
| |
| |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).contiguous() |
| attn_weights = attn_weights.masked_fill( |
| key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf") |
| ) |
| attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len).contiguous() |
| |
| attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32) |
| |
| x_out = torch.einsum('bts,btsh->bth', attn_weights, v) |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) |
| |
| |
| attn_weights = attn_weights.mean(dim=0) |
| |
| x_out = x_out.transpose(1, 0).reshape(tgt_len, bsz, self.num_heads * self.x_head_dim).transpose(1, 0).contiguous() |
| return x_out, attn_weights |
|
|
|
|
| class MultiHeadAttentionSoftMaxStarGraph(nn.Module): |
| """Equivariant multi-head attention layer with softmax, apply attention on full graph by default""" |
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| |
| super(MultiHeadAttentionSoftMaxStarGraph, self).__init__() |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
| self.share_kv = share_kv |
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.cutoff = None |
| self.scaling = self.x_head_dim**-0.5 |
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.k_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) if not share_kv else None |
| self.v_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.k_proj = nn.Linear(x_channels, x_hidden_channels) if not share_kv else None |
| self.v_proj = nn.Linear(x_channels, x_hidden_channels) |
|
|
| self.dk_proj = None |
| |
| self.dv_proj = None |
| |
| if distance_influence in ["keys", "both"]: |
| if use_lora is not None: |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| |
| else: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| |
|
|
| if distance_influence in ["values", "both"]: |
| if use_lora is not None: |
| self.dv_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| |
| else: |
| self.dv_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| |
| |
| |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.k_proj.weight) |
| self.k_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.v_proj.weight) |
| self.v_proj.bias.data.fill_(0) |
| |
| |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
| |
| |
| |
| |
| |
| |
|
|
| def forward(self, x, x_center_index, w_ij, f_dist_ij, f_ij, key_padding_mask, return_attn=False): |
| |
| |
| |
| x = self.layernorm(x) |
| q = self.q_proj(x[x_center_index].unsqueeze(1)) * self.scaling |
| v = self.v_proj(x) |
| |
| |
| |
| k = self.k_proj(x) |
|
|
| dk = self.act(self.dk_proj(f_ij)) |
| |
| dv = self.act(self.dv_proj(f_ij)) |
| |
|
|
| |
| x, attn = self.attention( |
| q=q, |
| k=k, |
| v=v, |
| dk=dk, |
| |
| dv=dv, |
| |
| |
| key_padding_mask=key_padding_mask, |
| ) |
| if return_attn: |
| return x, attn |
| else: |
| return x, None |
|
|
| def attention(self, q, k, v, dk, dv, key_padding_mask=None, need_head_weights=False): |
| |
| |
| |
| |
| |
| |
| |
| bsz, tgt_len, _ = q.size() |
| src_len = k.size(1) |
| |
| |
| q = q.transpose(0, 1).reshape(tgt_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| k = k.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| v = v.transpose(0, 1).reshape(src_len, bsz * self.num_heads, self.x_head_dim).transpose(0, 1).contiguous() |
| |
| |
| |
| dk = dk.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| |
| |
| |
| dv = dv.permute(1, 2, 0, 3).reshape(tgt_len, src_len, bsz * self.num_heads, self.x_head_dim).permute(2, 0, 1, 3).contiguous() |
| |
| |
| |
| |
| |
| assert key_padding_mask.size(0) == bsz |
| assert key_padding_mask.size(1) == src_len |
| |
| attn_weights = torch.multiply(q[:, :, None, :], k[:, None, :, :]) |
| |
| |
| |
| |
| |
| |
| attn_weights *= dk |
| |
| v = v.unsqueeze(1) + dv |
| |
| |
| |
| |
| |
| |
| attn_weights = attn_weights.sum(dim=-1) |
| |
| |
| |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).contiguous() |
| attn_weights = attn_weights.masked_fill( |
| key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf") |
| ) |
| attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len).contiguous() |
| |
| attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32) |
| |
| x_out = torch.einsum('bts,btsh->bth', attn_weights, v) |
| attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) |
| |
| |
| attn_weights = attn_weights.mean(dim=0) |
| |
| x_out = x_out.transpose(1, 0).reshape(tgt_len, bsz, self.num_heads * self.x_head_dim).transpose(1, 0).contiguous() |
| return x_out, attn_weights |
|
|
|
|
| |
| class EquivariantProMultiHeadAttention(MessagePassing): |
| """Equivariant multi-head attention layer.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| ): |
| super(EquivariantMultiHeadAttention, self).__init__( |
| aggr="mean", node_dim=0) |
| assert x_hidden_channels % num_heads == 0 \ |
| and vec_channels % num_heads == 0, ( |
| f"The number of hidden channels x_hidden_channels ({x_hidden_channels}) " |
| f"and vec_channels ({vec_channels}) " |
| f"must be evenly divisible by the number of " |
| f"attention heads ({num_heads})" |
| ) |
| assert vec_hidden_channels == x_channels, ( |
| f"The number of hidden channels x_channels ({x_channels}) " |
| f"and vec_hidden_channels ({vec_hidden_channels}) " |
| f"must be equal" |
| ) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| |
| self.vec_head_dim = vec_channels // num_heads |
|
|
| self.layernorm = nn.LayerNorm(x_channels) |
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
| self.cutoff = CosineCutoff(cutoff_lower, cutoff_upper) |
|
|
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| |
| self.kv_proj = nn.Linear( |
| x_channels, x_hidden_channels + vec_channels * 2) |
| self.o_proj = nn.Linear( |
| x_hidden_channels, x_channels * 2 + vec_channels) |
|
|
| self.vec_proj = nn.Linear( |
| vec_channels, vec_hidden_channels * 2 + vec_channels, bias=False) |
|
|
| self.dk_proj = None |
| if distance_influence in ["keys", "both"]: |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
|
|
| self.dv_proj = None |
| if distance_influence in ["values", "both"]: |
| self.dv_proj = nn.Linear( |
| edge_attr_channels, x_hidden_channels + vec_channels * 2) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| |
| |
| nn.init.xavier_uniform_(self.kv_proj.weight) |
| self.kv_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.vec_proj.weight) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
| if self.dv_proj: |
| nn.init.xavier_uniform_(self.dv_proj.weight) |
| self.dv_proj.bias.data.fill_(0) |
|
|
| def forward(self, x, vec, edge_index, r_ij, f_ij, d_ij, return_attn=False): |
| x = self.layernorm(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| |
| v = self.kv_proj(x).reshape(-1, self.num_heads, |
| self.x_head_dim + self.vec_head_dim * 2) |
| k = v[:, :, :self.x_head_dim] |
|
|
| vec1, vec2, vec3 = torch.split(self.vec_proj(vec), |
| [self.vec_hidden_channels, self.vec_hidden_channels, self.vec_channels], dim=-1) |
| vec = vec.reshape(-1, 3, self.num_heads, self.vec_head_dim) |
| vec_dot = (vec1 * vec2).sum(dim=1) |
|
|
| dk = ( |
| self.act(self.dk_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None |
| else None |
| ) |
| dv = ( |
| self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, self.x_head_dim + self.vec_head_dim * 2) |
| if self.dv_proj is not None |
| else None |
| ) |
|
|
| |
| |
| x, vec, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| vec=vec, |
| dk=dk, |
| dv=dv, |
| r_ij=r_ij, |
| d_ij=d_ij, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| vec = vec.reshape(-1, 3, self.vec_channels) |
|
|
| o1, o2, o3 = torch.split(self.o_proj( |
| x), [self.vec_channels, self.x_channels, self.x_channels], dim=1) |
| dx = vec_dot * o2 + o3 |
| dvec = vec3 * o1.unsqueeze(1) + vec |
| if return_attn: |
| return dx, dvec, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return dx, dvec, None |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dv, r_ij, d_ij): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j * dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) * self.cutoff(r_ij).unsqueeze(1) |
|
|
| |
| if dv is not None: |
| v_j = v_j * dv |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = vec_j * vec1.unsqueeze(1) + vec2.unsqueeze(1) * \ |
| d_ij.unsqueeze(2).unsqueeze(3) |
| return x, vec, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, vec, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size) |
| vec = scatter(vec, index, dim=self.node_dim, dim_size=dim_size) |
| return x, vec, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class EquivariantMultiHeadAttentionSoftMax(EquivariantMultiHeadAttention): |
| """Equivariant multi-head attention layer with softmax""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantMultiHeadAttentionSoftMax, self).__init__(x_channels=x_channels, |
| x_hidden_channels=x_hidden_channels, |
| vec_channels=vec_channels, |
| vec_hidden_channels=vec_hidden_channels, |
| share_kv=share_kv, |
| edge_attr_channels=edge_attr_channels, |
| distance_influence=distance_influence, |
| num_heads=num_heads, |
| activation=activation, |
| attn_activation=attn_activation, |
| cutoff_lower=cutoff_lower, |
| cutoff_upper=cutoff_upper, |
| use_lora=use_lora) |
| self.attn_activation = nn.LeakyReLU(0.2) |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dv, r_ij, d_ij, |
| index: Tensor, |
| ptr: Optional[Tensor], |
| size_i: Optional[int]): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j * dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) * self.cutoff(r_ij).unsqueeze(1) |
| attn = softmax(attn, index, ptr, size_i) |
| |
| |
| |
| if dv is not None: |
| v_j = v_j * dv |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = (vec1.unsqueeze(1) * vec_j + vec2.unsqueeze(1) * d_ij.unsqueeze(2).unsqueeze(3)) \ |
| * attn.unsqueeze(1).unsqueeze(3) |
| return x, vec, attn |
|
|
|
|
| |
| class EquivariantPAEMultiHeadAttentionSoftMax(EquivariantPAEMultiHeadAttention): |
| """Equivariant multi-head attention layer with softmax""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantPAEMultiHeadAttentionSoftMax, self).__init__( |
| x_channels=x_channels, |
| x_hidden_channels=x_hidden_channels, |
| vec_channels=vec_channels, |
| vec_hidden_channels=vec_hidden_channels, |
| share_kv=share_kv, |
| edge_attr_channels=edge_attr_channels, |
| edge_attr_dist_channels=edge_attr_dist_channels, |
| distance_influence=distance_influence, |
| num_heads=num_heads, |
| activation=activation, |
| attn_activation=attn_activation, |
| cutoff_lower=cutoff_lower, |
| cutoff_upper=cutoff_upper, |
| use_lora=use_lora) |
| self.attn_activation = nn.LeakyReLU(0.2) |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dk_dist, dv, dv_dist, d_ij, w_ij, |
| index: Tensor, |
| ptr: Optional[Tensor], |
| size_i: Optional[int]): |
| |
| attn = (q_i * k_j) |
| if dk is not None: |
| attn += dk |
| if dk_dist is not None: |
| attn += dk_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| attn = attn.sum(dim=-1) |
| |
| attn = self.attn_activation(attn) |
| attn = softmax(attn, index, ptr, size_i) |
| |
| |
| |
| if dv is not None: |
| v_j += dv |
| if dv_dist is not None: |
| v_j += dv_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = (vec1.unsqueeze(1) * vec_j + vec2.unsqueeze(1) * d_ij.unsqueeze(2).unsqueeze(3)) \ |
| * attn.unsqueeze(1).unsqueeze(3) |
| return x, vec, attn |
|
|
| |
| class EquivariantWeightedPAEMultiHeadAttentionSoftMax(EquivariantWeightedPAEMultiHeadAttention): |
| """Equivariant multi-head attention layer with softmax""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| share_kv, |
| edge_attr_channels, |
| edge_attr_dist_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| use_lora=None, |
| ): |
| super(EquivariantWeightedPAEMultiHeadAttentionSoftMax, self).__init__( |
| x_channels=x_channels, |
| x_hidden_channels=x_hidden_channels, |
| vec_channels=vec_channels, |
| vec_hidden_channels=vec_hidden_channels, |
| share_kv=share_kv, |
| edge_attr_channels=edge_attr_channels, |
| edge_attr_dist_channels=edge_attr_dist_channels, |
| distance_influence=distance_influence, |
| num_heads=num_heads, |
| activation=activation, |
| attn_activation=attn_activation, |
| cutoff_lower=cutoff_lower, |
| cutoff_upper=cutoff_upper, |
| use_lora=use_lora) |
| self.attn_activation = nn.LeakyReLU(0.2) |
|
|
| def message(self, q_i, k_j, v_j, vec_j, dk, dk_dist, dv, dv_dist, d_ij, w_ij, |
| index: Tensor, |
| ptr: Optional[Tensor], |
| size_i: Optional[int]): |
| |
| attn = (q_i * k_j) |
| if dk_dist is not None: |
| if dk is not None: |
| attn *= (dk + dk_dist * w_ij.unsqueeze(1).unsqueeze(2)) |
| else: |
| attn *= dk_dist * w_ij |
| else: |
| if dk is not None: |
| attn *= dk |
| attn = attn.sum(dim=-1) |
| |
| attn = self.attn_activation(attn) |
| attn = softmax(attn, index, ptr, size_i) |
| |
| |
| |
| if dv is not None: |
| v_j += dv |
| if dv_dist is not None: |
| v_j += dv_dist * w_ij.unsqueeze(1).unsqueeze(2) |
| x, vec1, vec2 = torch.split( |
| v_j, [self.x_head_dim, self.vec_head_dim, self.vec_head_dim], dim=2) |
|
|
| |
| x = x * attn.unsqueeze(2) |
| |
| vec = (vec1.unsqueeze(1) * vec_j + vec2.unsqueeze(1) * d_ij.unsqueeze(2).unsqueeze(3)) \ |
| * attn.unsqueeze(1).unsqueeze(3) |
| return x, vec, attn |
|
|
|
|
| |
| class MSAEncoder(nn.Module): |
| def __init__(self, num_species, pairwise_type, weighting_schema): |
| """[summary] |
| |
| Args: |
| num_species (int): Number of species to use from MSA. [1,200] // 200 used in default gMVP |
| pairwise_type ([str]): method for calculating pairwise coevolution. only "cov" supported |
| weighting_schema ([str]): species weighting type; "spe" -> use dense layer to weight speices |
| "none" -> uniformly weight species |
| |
| Raises: |
| NotImplementedError: [description] |
| """ |
| super(MSAEncoder, self).__init__() |
| self.num_species = num_species |
| self.pairwise_type = pairwise_type |
| self.weighting_schema = weighting_schema |
| if self.weighting_schema == 'spe': |
| self.W = nn.parameter.Parameter( |
| torch.zeros((1, num_species)), |
| requires_grad=True) |
|
|
| elif self.weighting_schema == 'none': |
| self.W = torch.tensor(1.0 / self.num_species).repeat(self.num_species) |
| else: |
| raise NotImplementedError |
| |
| def forward(self, x, edge_index): |
| |
| shape = x.shape |
| L, N = shape[0], shape[1] |
| E = edge_index.shape[1] |
| |
| A = 21 |
| x = x[:, :self.num_species] |
| if self.weighting_schema == 'spe': |
| sm = torch.nn.Softmax(dim=-1) |
| W = sm(self.W) |
| else: |
| W = self.W |
| x = nn.functional.one_hot(x.type(torch.int64), A).type(torch.float32) |
| x1 = torch.matmul(W[:, None], x) |
|
|
| if self.pairwise_type == 'fre': |
| x2 = torch.matmul(x[edge_index[0], :, :, None], x[edge_index[1], :, None, :]) |
| x2 = x2.reshape((E, N, A * A)) |
| x2 = (W[:, :, None] * x2).sum(dim=1) |
| elif self.pairwise_type == 'cov': |
| |
| x2 = torch.matmul(x[edge_index[0], :, :, None], x[edge_index[1], :, None, :]) |
| x2 = (W[:, :, None, None] * x2).sum(dim=1) |
| x2_t = x1[edge_index[0], 0, :, None] * x1[edge_index[1], 0, None, :] |
| x2 = (x2 - x2_t).reshape(E, A * A) |
| x2 = x2.reshape(E, A * A) |
| norm = torch.sqrt(torch.sum(torch.square(x2), dim=-1, keepdim=True) + 1e-12) |
| x2 = torch.cat([x2, norm], dim=-1) |
| elif self.pairwise_type == 'cov_all': |
| print('cov_all not implemented in EvolEncoder2') |
| raise NotImplementedError |
| elif self.pairwise_type == 'inv_cov': |
| print('in_cov not implemented in EvolEncoder2') |
| raise NotImplementedError |
| elif self.pairwise_type == 'none': |
| x2 = None |
| else: |
| raise NotImplementedError( |
| f'pairwise_type {self.pairwise_type} not implemented') |
|
|
| x1 = torch.squeeze(x1, dim=1) |
|
|
| return x1, x2 |
|
|
|
|
| |
| class MSAEncoderFullGraph(nn.Module): |
| def __init__(self, num_species, pairwise_type, weighting_schema): |
| """[summary] |
| |
| Args: |
| num_species (int): Number of species to use from MSA. [1,200] // 200 used in default gMVP |
| pairwise_type ([str]): method for calculating pairwise coevolution. only "cov" supported |
| weighting_schema ([str]): species weighting type; "spe" -> use dense layer to weight speices |
| "none" -> uniformly weight species |
| |
| Raises: |
| NotImplementedError: [description] |
| """ |
| super(MSAEncoderFullGraph, self).__init__() |
| self.num_species = num_species |
| self.pairwise_type = pairwise_type |
| self.weighting_schema = weighting_schema |
| if self.weighting_schema == 'spe': |
| self.W = nn.parameter.Parameter( |
| torch.zeros((num_species)), |
| requires_grad=True) |
|
|
| elif self.weighting_schema == 'none': |
| self.W = torch.tensor(1.0 / self.num_species).repeat(self.num_species) |
| else: |
| raise NotImplementedError |
| |
| def forward(self, x): |
| |
| shape = x.shape |
| B, L, N = shape[0], shape[1], shape[2] |
| A = 21 |
| x = x[:, :, :self.num_species] |
| if self.weighting_schema == 'spe': |
| W = torch.nn.functional.softmax(self.W, dim=-1) |
| else: |
| W = self.W |
| x = nn.functional.one_hot(x.type(torch.int64), A).type(torch.float32) |
| x1 = torch.einsum('blna,n->bla', x, W) |
|
|
| if self.pairwise_type == 'cov': |
| |
| |
| |
| |
| |
| x2 = (torch.einsum('bLnA,blna,n->bLlAa', x, x, W) - x1[:, :, None, :, None] * x1[:, None, :, None, :]).reshape(B, L, L, A * A) |
| norm = torch.sqrt(torch.sum(torch.square(x2), dim=-1, keepdim=True) + 1e-12) |
| x2 = torch.cat([x2, norm], dim=-1) |
| elif self.pairwise_type == 'cov_all': |
| print('cov_all not implemented in EvolEncoder2') |
| raise NotImplementedError |
| elif self.pairwise_type == 'inv_cov': |
| print('in_cov not implemented in EvolEncoder2') |
| raise NotImplementedError |
| elif self.pairwise_type == 'none': |
| x2 = None |
| else: |
| raise NotImplementedError( |
| f'pairwise_type {self.pairwise_type} not implemented') |
| return x1, x2 |
|
|
|
|
| class NodeToEdgeAttr(nn.Module): |
| def __init__(self, node_channel, hidden_channel, edge_attr_channel, use_lora=None, layer_norm=False): |
| super().__init__() |
| self.layer_norm = layer_norm |
| if layer_norm: |
| self.layernorm = nn.LayerNorm(node_channel) |
| if use_lora is not None: |
| self.proj = lora.Linear(node_channel, hidden_channel * 2, bias=True, r=use_lora) |
| self.o_proj = lora.Linear(2 * hidden_channel, edge_attr_channel, r=use_lora) |
| else: |
| self.proj = nn.Linear(node_channel, hidden_channel * 2, bias=True) |
| self.o_proj = nn.Linear(2 * hidden_channel, edge_attr_channel, bias=True) |
|
|
| torch.nn.init.zeros_(self.proj.bias) |
| torch.nn.init.zeros_(self.o_proj.bias) |
|
|
| def forward(self, x, edge_index): |
| """ |
| Inputs: |
| x: N x sequence_state_dim |
| |
| Output: |
| edge_attr: edge_index.shape[0] x pairwise_state_dim |
| |
| Intermediate state: |
| B x L x L x 2*inner_dim |
| """ |
| x = self.layernorm(x) if self.layer_norm else x |
| q, k = self.proj(x).chunk(2, dim=-1) |
|
|
| prod = q[edge_index[0], :] * k[edge_index[1], :] |
| diff = q[edge_index[0], :] - k[edge_index[1], :] |
|
|
| edge_attr = torch.cat([prod, diff], dim=-1) |
| edge_attr = self.o_proj(edge_attr) |
|
|
| return edge_attr |
|
|
|
|
| class MultiplicativeUpdate(MessagePassing): |
| def __init__(self, vec_in_channel, hidden_channel, hidden_vec_channel, ee_channels=None, use_lora=None, layer_norm=True) -> None: |
| super(MultiplicativeUpdate, self).__init__(aggr="mean") |
| self.vec_in_channel = vec_in_channel |
| self.hidden_channel = hidden_channel |
| self.hidden_vec_channel = hidden_vec_channel |
|
|
| if use_lora is not None: |
| self.linear_a_p = lora.Linear(self.vec_in_channel, self.hidden_vec_channel, bias=False, r=use_lora) |
| self.linear_b_p = lora.Linear(self.vec_in_channel, self.hidden_vec_channel, bias=False, r=use_lora) |
| self.linear_g = lora.Linear(self.hidden_vec_channel, self.hidden_channel, r=use_lora) |
| else: |
| self.linear_a_p = nn.Linear(self.vec_in_channel, self.hidden_vec_channel, bias=False) |
| self.linear_b_p = nn.Linear(self.vec_in_channel, self.hidden_vec_channel, bias=False) |
| self.linear_g = nn.Linear(self.hidden_vec_channel, self.hidden_channel) |
| if ee_channels is not None: |
| if use_lora is not None: |
| self.linear_ee = lora.Linear(ee_channels, self.hidden_channel, r=use_lora) |
| else: |
| self.linear_ee = nn.Linear(ee_channels, self.hidden_channel) |
| else: |
| self.linear_ee = None |
| self.layer_norm = layer_norm |
| if layer_norm: |
| self.layer_norm_in = nn.LayerNorm(self.hidden_channel) |
| self.layer_norm_out = nn.LayerNorm(self.hidden_channel) |
|
|
| self.sigmoid = nn.Sigmoid() |
|
|
| def forward(self, |
| edge_attr: torch.Tensor, |
| edge_vec: torch.Tensor, |
| edge_edge_index: torch.Tensor, |
| edge_edge_attr: torch.Tensor, |
| ) -> torch.Tensor: |
| """ |
| Args: |
| edge_vec: |
| [*, 3, in_channel] input tensor |
| edge_attr: |
| [*, hidden_channel] input mask |
| Returns: |
| [*, hidden_channel] output tensor |
| """ |
| if self.layer_norm: |
| x = self.layer_norm_in(edge_attr) |
| x = self.propagate(edge_index=edge_edge_index, |
| a=self.linear_a_p(edge_vec).reshape(edge_attr.shape[0], -1), |
| b=self.linear_b_p(edge_vec).reshape(edge_attr.shape[0], -1), |
| edge_attr=x, |
| ee_ij=edge_edge_attr, ) |
| if self.layer_norm: |
| x = self.layer_norm_out(x) |
| edge_attr = edge_attr + x |
| return edge_attr |
|
|
| def message(self, a_i: Tensor, b_j: Tensor, edge_attr_j: Tensor, ee_ij: Tensor,) -> Tensor: |
| |
| |
| s = (a_i.reshape(-1, 3, self.hidden_vec_channel).permute(0, 2, 1) \ |
| * b_j.reshape(-1, 3, self.hidden_vec_channel).permute(0, 2, 1)).sum(dim=-1) |
| if ee_ij is not None and self.linear_ee is not None: |
| s = self.sigmoid(self.linear_ee(ee_ij) + self.linear_g(s)) |
| else: |
| s = self.sigmoid(self.linear_g(s)) |
| return s * edge_attr_j |
|
|
|
|
| |
| class EquivariantTriAngularMultiHeadAttention(MessagePassing): |
| """Equivariant multi-head attention layer. Add Triangular update between edges.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| triangular_update=False, |
| ee_channels=None, |
| ): |
| super(EquivariantTriAngularMultiHeadAttention, self).__init__(aggr="mean", node_dim=0) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| self.ee_channels = ee_channels |
| |
|
|
| self.layernorm_in = nn.LayerNorm(x_channels) |
| self.layernorm_out = nn.LayerNorm(x_hidden_channels) |
|
|
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
|
|
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.kv_proj = nn.Linear(x_channels, x_hidden_channels) |
| |
| self.o_proj = nn.Linear(x_hidden_channels, x_hidden_channels) |
| self.out = nn.Linear(x_hidden_channels, x_channels) |
| |
| |
|
|
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| self.triangular_update = triangular_update |
| if self.triangular_update: |
| self.edge_triangle_start_update = MultiplicativeUpdate(vec_in_channel=vec_channels, |
| hidden_channel=edge_attr_channels, |
| hidden_vec_channel=vec_hidden_channels, |
| ee_channels=ee_channels, ) |
| self.edge_triangle_end_update = MultiplicativeUpdate(vec_in_channel=vec_channels, |
| hidden_channel=edge_attr_channels, |
| hidden_vec_channel=vec_hidden_channels, |
| ee_channels=ee_channels, ) |
| self.node_to_edge_attr = NodeToEdgeAttr(node_channel=x_channels, |
| hidden_channel=x_hidden_channels, |
| edge_attr_channel=edge_attr_channels) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| self.layernorm_in.reset_parameters() |
| self.layernorm_out.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.kv_proj.weight) |
| self.kv_proj.bias.data.fill_(0) |
| |
| |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
|
|
| def get_start_index(self, edge_index): |
| edge_start_index = [] |
| start_node_count = edge_index[0].unique(return_counts=True) |
| start_nodes = start_node_count[0][start_node_count[1] > 1] |
| for i in start_nodes: |
| node_start_index = torch.where(edge_index[0] == i)[0] |
| candidates = torch.combinations(node_start_index, r=2).T |
| edge_start_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_start_index = torch.concat(edge_start_index, dim=1) |
| edge_start_index = edge_start_index[:, edge_start_index[0] != edge_start_index[1]] |
| return edge_start_index |
|
|
| def get_end_index(self, edge_index): |
| edge_end_index = [] |
| end_node_count = edge_index[1].unique(return_counts=True) |
| end_nodes = end_node_count[0][end_node_count[1] > 1] |
| for i in end_nodes: |
| node_end_index = torch.where(edge_index[1] == i)[0] |
| candidates = torch.combinations(node_end_index, r=2).T |
| edge_end_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_end_index = torch.concat(edge_end_index, dim=1) |
| edge_end_index = edge_end_index[:, edge_end_index[0] != edge_end_index[1]] |
| return edge_end_index |
|
|
| def forward(self, x, coords, edge_index, edge_attr, edge_vec, return_attn=False): |
| residue = x |
| x = self.layernorm_in(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| k = self.kv_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| v = k |
|
|
| |
| if self.triangular_update: |
| edge_attr += self.node_to_edge_attr(x, edge_index) |
|
|
| |
| |
| edge_edge_index = self.get_start_index(edge_index) |
| if self.ee_channels is not None: |
| edge_edge_attr = coords[edge_index[1][edge_edge_index[0]], :, [0]] - coords[edge_index[1][edge_edge_index[1]], :, [0]] |
| edge_edge_attr = torch.norm(edge_edge_attr, dim=-1, keepdim=True) |
| else: |
| edge_edge_attr = None |
| edge_attr = self.edge_triangle_start_update( |
| edge_attr, edge_vec, |
| edge_edge_index, |
| edge_edge_attr |
| ) |
| edge_edge_index = self.get_end_index(edge_index) |
| if self.ee_channels is not None: |
| edge_edge_attr = coords[edge_index[0][edge_edge_index[0]], :, [0]] - coords[edge_index[0][edge_edge_index[1]], :, [0]] |
| edge_edge_attr = torch.norm(edge_edge_attr, dim=-1, keepdim=True) |
| else: |
| edge_edge_attr = None |
| edge_attr = self.edge_triangle_end_update( |
| edge_attr, edge_vec, |
| edge_edge_index, |
| edge_edge_attr |
| ) |
| del edge_edge_attr, edge_edge_index |
|
|
| dk = ( |
| self.act(self.dk_proj(edge_attr)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None else None |
| ) |
|
|
| |
| |
| x, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| dk=dk, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| x = residue + x |
| x = self.layernorm_out(x) |
| x = gelu(self.o_proj(x)) |
| x = self.out(x) |
| del residue, q, k, v, dk |
| if return_attn: |
| return x, edge_attr, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return x, edge_attr, None |
|
|
| def message(self, q_i, k_j, v_j, dk): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j * dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) |
|
|
| |
| x = v_j * attn.unsqueeze(2) |
| return x, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr) |
| return x, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class EquivariantTriAngularDropMultiHeadAttention(MessagePassing): |
| """Equivariant multi-head attention layer. Add Triangular update between edges.""" |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| rbf_channels, |
| triangular_update=False, |
| ee_channels=None, |
| drop_out_rate=0.0, |
| use_lora=None, |
| layer_norm=True, |
| ): |
| super(EquivariantTriAngularDropMultiHeadAttention, self).__init__(aggr="mean", node_dim=0) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| self.ee_channels = ee_channels |
| self.rbf_channels = rbf_channels |
| self.layer_norm = layer_norm |
| |
| if layer_norm: |
| self.layernorm_in = nn.LayerNorm(x_channels) |
| self.layernorm_out = nn.LayerNorm(x_hidden_channels) |
|
|
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
|
|
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.kv_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| self.o_proj = lora.Linear(x_hidden_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.kv_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| self.o_proj = nn.Linear(x_hidden_channels, x_hidden_channels) |
|
|
| self.triangular_drop = nn.Dropout(drop_out_rate) |
| self.rbf_drop = nn.Dropout(drop_out_rate) |
| self.dense_drop = nn.Dropout(drop_out_rate) |
| self.dropout = nn.Dropout(drop_out_rate) |
| self.triangular_update = triangular_update |
| if self.triangular_update: |
| self.edge_triangle_end_update = MultiplicativeUpdate(vec_in_channel=vec_channels, |
| hidden_channel=edge_attr_channels, |
| hidden_vec_channel=vec_hidden_channels, |
| ee_channels=ee_channels, |
| layer_norm=layer_norm, |
| use_lora=use_lora) |
| self.node_to_edge_attr = NodeToEdgeAttr(node_channel=x_channels, |
| hidden_channel=x_hidden_channels, |
| edge_attr_channel=edge_attr_channels, |
| use_lora=use_lora) |
| self.triangle_update_dropout = nn.Dropout(0.5) |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| if self.layer_norm: |
| self.layernorm_in.reset_parameters() |
| self.layernorm_out.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.kv_proj.weight) |
| self.kv_proj.bias.data.fill_(0) |
| |
| |
| nn.init.xavier_uniform_(self.o_proj.weight) |
| self.o_proj.bias.data.fill_(0) |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
|
|
| def get_start_index(self, edge_index): |
| edge_start_index = [] |
| start_node_count = edge_index[0].unique(return_counts=True) |
| start_nodes = start_node_count[0][start_node_count[1] > 1] |
| for i in start_nodes: |
| node_start_index = torch.where(edge_index[0] == i)[0] |
| candidates = torch.combinations(node_start_index, r=2).T |
| edge_start_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_start_index = torch.concat(edge_start_index, dim=1) |
| edge_start_index = edge_start_index[:, edge_start_index[0] != edge_start_index[1]] |
| return edge_start_index |
|
|
| def get_end_index(self, edge_index): |
| edge_end_index = [] |
| end_node_count = edge_index[1].unique(return_counts=True) |
| end_nodes = end_node_count[0][end_node_count[1] > 1] |
| for i in end_nodes: |
| node_end_index = torch.where(edge_index[1] == i)[0] |
| candidates = torch.combinations(node_end_index, r=2).T |
| edge_end_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_end_index = torch.concat(edge_end_index, dim=1) |
| edge_end_index = edge_end_index[:, edge_end_index[0] != edge_end_index[1]] |
| return edge_end_index |
|
|
| def forward(self, x, coords, edge_index, edge_attr, edge_vec, return_attn=False): |
| residue = x |
| if self.layer_norm: |
| x = self.layernorm_in(x) |
| q = self.q_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| k = self.kv_proj(x).reshape(-1, self.num_heads, self.x_head_dim) |
| v = k |
|
|
| |
| if self.triangular_update: |
| edge_attr += self.node_to_edge_attr(x, edge_index) |
| |
| |
| edge_edge_index = self.get_end_index(edge_index) |
| edge_edge_index = edge_edge_index[:, self.triangular_drop( |
| torch.ones(edge_edge_index.shape[1], device=edge_edge_index.device) |
| ).to(torch.bool)] |
| if self.ee_channels is not None: |
| edge_edge_attr = coords[edge_index[0][edge_edge_index[0]], :, [0]] - coords[edge_index[0][edge_edge_index[1]], :, [0]] |
| edge_edge_attr = torch.norm(edge_edge_attr, dim=-1, keepdim=True) |
| else: |
| edge_edge_attr = None |
| edge_attr = self.edge_triangle_end_update( |
| edge_attr, edge_vec, |
| edge_edge_index, |
| edge_edge_attr |
| ) |
| del edge_edge_attr, edge_edge_index |
|
|
| |
| edge_attr = torch.cat((edge_attr[:, :-self.rbf_channels], |
| self.rbf_drop(edge_attr[:, -self.rbf_channels:])), |
| dim=-1) |
|
|
| dk = ( |
| self.act(self.dk_proj(edge_attr)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None else None |
| ) |
|
|
| |
| |
| x, attn = self.propagate( |
| edge_index, |
| q=q, |
| k=k, |
| v=v, |
| dk=dk, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| if self.layer_norm: |
| x = self.layernorm_out(x) |
| x = self.dense_drop(x) |
| x = residue + gelu(x) |
| x = self.o_proj(x) |
| x = self.dropout(x) |
| del residue, q, k, v, dk |
| if return_attn: |
| return x, edge_attr, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return x, edge_attr, None |
|
|
| def message(self, q_i, k_j, v_j, dk): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j * dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) |
|
|
| |
| x = v_j * attn.unsqueeze(2) |
| return x, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr) |
| return x, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class EquivariantTriAngularStarMultiHeadAttention(MessagePassing): |
| """ |
| Equivariant multi-head attention layer. Add Triangular update between edges. Only update the center node. |
| """ |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| cutoff_lower, |
| cutoff_upper, |
| triangular_update=False, |
| ee_channels=None, |
| ): |
| super(EquivariantTriAngularStarMultiHeadAttention, self).__init__(aggr="mean", node_dim=0) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| self.ee_channels = ee_channels |
| |
|
|
| |
| self.layernorm_out = nn.LayerNorm(x_hidden_channels) |
|
|
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
|
|
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.kv_proj = nn.Linear(x_channels, x_hidden_channels) |
| |
| |
| |
| |
| |
| self.gru = nn.GRUCell(x_channels, x_channels) |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| self.triangular_update = triangular_update |
| if self.triangular_update: |
| |
| |
| |
| |
| self.edge_triangle_end_update = MultiplicativeUpdate(vec_in_channel=vec_channels, |
| hidden_channel=edge_attr_channels, |
| hidden_vec_channel=vec_hidden_channels, |
| ee_channels=ee_channels, ) |
| self.node_to_edge_attr = NodeToEdgeAttr(node_channel=x_channels, |
| hidden_channel=x_hidden_channels, |
| edge_attr_channel=edge_attr_channels) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| |
| self.layernorm_out.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.kv_proj.weight) |
| self.kv_proj.bias.data.fill_(0) |
| |
| |
| |
| |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
|
|
| def get_start_index(self, edge_index): |
| edge_start_index = [] |
| start_node_count = edge_index[0].unique(return_counts=True) |
| start_nodes = start_node_count[0][start_node_count[1] > 1] |
| for i in start_nodes: |
| node_start_index = torch.where(edge_index[0] == i)[0] |
| candidates = torch.combinations(node_start_index, r=2).T |
| edge_start_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_start_index = torch.concat(edge_start_index, dim=1) |
| edge_start_index = edge_start_index[:, edge_start_index[0] != edge_start_index[1]] |
| return edge_start_index |
|
|
| def get_end_index(self, edge_index): |
| edge_end_index = [] |
| end_node_count = edge_index[1].unique(return_counts=True) |
| end_nodes = end_node_count[0][end_node_count[1] > 1] |
| for i in end_nodes: |
| node_end_index = torch.where(edge_index[1] == i)[0] |
| candidates = torch.combinations(node_end_index, r=2).T |
| edge_end_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_end_index = torch.concat(edge_end_index, dim=1) |
| edge_end_index = edge_end_index[:, edge_end_index[0] != edge_end_index[1]] |
| return edge_end_index |
|
|
| def forward(self, x, coords, edge_index, edge_attr, edge_vec, return_attn=False): |
| |
| end_node_count = edge_index[1].unique(return_counts=True) |
| center_nodes = end_node_count[0][end_node_count[1] > 1] |
| other_nodes = end_node_count[0][end_node_count[1] <= 1] |
| residue = x[center_nodes] |
| |
| edge_attr = edge_attr[torch.isin(edge_index[1], center_nodes), :] |
| edge_vec = edge_vec[torch.isin(edge_index[1], center_nodes), :] |
| edge_index = edge_index[:, torch.isin(edge_index[1], center_nodes)] |
| |
| q = self.q_proj(residue).reshape(-1, self.num_heads, self.x_head_dim) |
| kv = self.kv_proj(x[other_nodes]).reshape(-1, self.num_heads, self.x_head_dim) |
| qkv = torch.zeros(x.shape[0], self.num_heads, self.x_head_dim).to(x.device, non_blocking=True) |
| qkv[center_nodes] = q |
| qkv[other_nodes] = kv |
| |
| if self.triangular_update: |
| edge_attr += self.node_to_edge_attr(x, edge_index) |
| |
| |
| edge_edge_index = self.get_end_index(edge_index) |
| if self.ee_channels is not None: |
| edge_edge_attr = coords[edge_index[0][edge_edge_index[0]], :, [0]] - coords[edge_index[0][edge_edge_index[1]], :, [0]] |
| edge_edge_attr = torch.norm(edge_edge_attr, dim=-1, keepdim=True) |
| else: |
| edge_edge_attr = None |
| edge_attr = self.edge_triangle_end_update( |
| edge_attr, edge_vec, |
| edge_edge_index, |
| edge_edge_attr |
| ) |
| del edge_edge_attr, edge_edge_index |
|
|
| dk = ( |
| self.act(self.dk_proj(edge_attr)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None else None |
| ) |
|
|
| |
| |
| x, attn = self.propagate( |
| edge_index, |
| q=qkv, |
| k=qkv, |
| v=qkv, |
| dk=dk, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| |
| x = x[center_nodes] |
| x = self.layernorm_out(x) |
| x = self.gru(residue, x) |
| del residue, dk |
| if return_attn: |
| return x, edge_attr, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return x, edge_attr, None |
|
|
| def message(self, q_i, k_j, v_j, dk): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j + dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) / self.x_head_dim |
|
|
| |
| x = v_j * attn.unsqueeze(2) |
| return x, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr) |
| return x, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class EquivariantTriAngularStarDropMultiHeadAttention(MessagePassing): |
| """ |
| Equivariant multi-head attention layer. Add Triangular update between edges. Only update the center node. |
| """ |
|
|
| def __init__( |
| self, |
| x_channels, |
| x_hidden_channels, |
| vec_channels, |
| vec_hidden_channels, |
| edge_attr_channels, |
| distance_influence, |
| num_heads, |
| activation, |
| attn_activation, |
| rbf_channels, |
| triangular_update=False, |
| ee_channels=None, |
| drop_out_rate=0.0, |
| use_lora=None, |
| ): |
| super(EquivariantTriAngularStarDropMultiHeadAttention, self).__init__(aggr="mean", node_dim=0) |
|
|
| self.distance_influence = distance_influence |
| self.num_heads = num_heads |
| self.x_channels = x_channels |
| self.x_hidden_channels = x_hidden_channels |
| self.x_head_dim = x_hidden_channels // num_heads |
| self.vec_channels = vec_channels |
| self.vec_hidden_channels = vec_hidden_channels |
| self.ee_channels = ee_channels |
| self.rbf_channels = rbf_channels |
| |
|
|
| |
| self.layernorm_out = nn.LayerNorm(x_hidden_channels) |
|
|
| self.act = activation() |
| self.attn_activation = act_class_mapping[attn_activation]() |
|
|
| if use_lora is not None: |
| self.q_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.kv_proj = lora.Linear(x_channels, x_hidden_channels, r=use_lora) |
| self.dk_proj = lora.Linear(edge_attr_channels, x_hidden_channels, r=use_lora) |
| else: |
| self.q_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.kv_proj = nn.Linear(x_channels, x_hidden_channels) |
| self.dk_proj = nn.Linear(edge_attr_channels, x_hidden_channels) |
| |
| |
| |
| |
| |
| self.gru = nn.GRUCell(x_channels, x_channels) |
| |
| self.triangular_drop = nn.Dropout(drop_out_rate) |
| self.rbf_drop = nn.Dropout(drop_out_rate) |
| self.dense_drop = nn.Dropout(drop_out_rate) |
| self.dropout = nn.Dropout(drop_out_rate) |
| self.triangular_update = triangular_update |
| if self.triangular_update: |
| self.edge_triangle_end_update = MultiplicativeUpdate(vec_in_channel=vec_channels, |
| hidden_channel=edge_attr_channels, |
| hidden_vec_channel=vec_hidden_channels, |
| ee_channels=ee_channels, |
| use_lora=use_lora) |
| self.node_to_edge_attr = NodeToEdgeAttr(node_channel=x_channels, |
| hidden_channel=x_hidden_channels, |
| edge_attr_channel=edge_attr_channels, |
| use_lora=use_lora) |
| self.triangle_update_dropout = nn.Dropout(0.5) |
|
|
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| |
| self.layernorm_out.reset_parameters() |
| nn.init.xavier_uniform_(self.q_proj.weight) |
| self.q_proj.bias.data.fill_(0) |
| nn.init.xavier_uniform_(self.kv_proj.weight) |
| self.kv_proj.bias.data.fill_(0) |
| |
| |
| |
| |
| if self.dk_proj: |
| nn.init.xavier_uniform_(self.dk_proj.weight) |
| self.dk_proj.bias.data.fill_(0) |
|
|
| def get_start_index(self, edge_index): |
| edge_start_index = [] |
| start_node_count = edge_index[0].unique(return_counts=True) |
| start_nodes = start_node_count[0][start_node_count[1] > 1] |
| for i in start_nodes: |
| node_start_index = torch.where(edge_index[0] == i)[0] |
| candidates = torch.combinations(node_start_index, r=2).T |
| edge_start_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_start_index = torch.concat(edge_start_index, dim=1) |
| edge_start_index = edge_start_index[:, edge_start_index[0] != edge_start_index[1]] |
| return edge_start_index |
|
|
| def get_end_index(self, edge_index): |
| edge_end_index = [] |
| end_node_count = edge_index[1].unique(return_counts=True) |
| end_nodes = end_node_count[0][end_node_count[1] > 1] |
| for i in end_nodes: |
| node_end_index = torch.where(edge_index[1] == i)[0] |
| candidates = torch.combinations(node_end_index, r=2).T |
| edge_end_index.append(torch.cat([candidates, candidates.flip(0)], dim=1)) |
| edge_end_index = torch.concat(edge_end_index, dim=1) |
| edge_end_index = edge_end_index[:, edge_end_index[0] != edge_end_index[1]] |
| return edge_end_index |
|
|
| def forward(self, x, coords, edge_index, edge_attr, edge_vec, return_attn=False): |
| |
| end_node_count = edge_index[1].unique(return_counts=True) |
| center_nodes = end_node_count[0][end_node_count[1] > 1] |
| other_nodes = end_node_count[0][end_node_count[1] <= 1] |
| residue = x[center_nodes] |
| |
| edge_attr = edge_attr[torch.isin(edge_index[1], center_nodes), :] |
| edge_vec = edge_vec[torch.isin(edge_index[1], center_nodes), :] |
| edge_index = edge_index[:, torch.isin(edge_index[1], center_nodes)] |
| |
| q = self.q_proj(residue).reshape(-1, self.num_heads, self.x_head_dim) |
| kv = self.kv_proj(x[other_nodes]).reshape(-1, self.num_heads, self.x_head_dim) |
| qkv = torch.zeros(x.shape[0], self.num_heads, self.x_head_dim).to(x.device, non_blocking=True) |
| qkv[center_nodes] = q |
| qkv[other_nodes] = kv |
| |
| if self.triangular_update: |
| edge_attr += self.node_to_edge_attr(x, edge_index) |
| |
| |
| edge_edge_index = self.get_end_index(edge_index) |
| edge_edge_index = edge_edge_index[:, self.triangular_drop( |
| torch.ones(edge_edge_index.shape[1], device=edge_edge_index.device) |
| ).to(torch.bool)] |
| if self.ee_channels is not None: |
| edge_edge_attr = coords[edge_index[0][edge_edge_index[0]], :, [0]] - coords[edge_index[0][edge_edge_index[1]], :, [0]] |
| edge_edge_attr = torch.norm(edge_edge_attr, dim=-1, keepdim=True) |
| else: |
| edge_edge_attr = None |
| edge_attr = self.edge_triangle_end_update( |
| edge_attr, edge_vec, |
| edge_edge_index, |
| edge_edge_attr |
| ) |
| del edge_edge_attr, edge_edge_index |
| |
| |
| edge_attr = torch.cat((edge_attr[:, :-self.rbf_channels], |
| self.rbf_drop(edge_attr[:, -self.rbf_channels:])), |
| dim=-1) |
|
|
| dk = ( |
| self.act(self.dk_proj(edge_attr)).reshape(-1, self.num_heads, self.x_head_dim) |
| if self.dk_proj is not None else None |
| ) |
|
|
| |
| |
| x, attn = self.propagate( |
| edge_index, |
| q=qkv, |
| k=qkv, |
| v=qkv, |
| dk=dk, |
| size=None, |
| ) |
| x = x.reshape(-1, self.x_hidden_channels) |
| |
| x = x[center_nodes] |
| x = self.layernorm_out(x) |
| x = self.dense_drop(x) |
| x = self.gru(residue, x) |
| x = self.dropout(x) |
| del residue, dk |
| if return_attn: |
| return x, edge_attr, torch.concat((edge_index.T, attn), dim=1) |
| else: |
| return x, edge_attr, None |
|
|
| def message(self, q_i, k_j, v_j, dk): |
| |
| if dk is None: |
| attn = (q_i * k_j).sum(dim=-1) |
| else: |
| attn = (q_i * k_j + dk).sum(dim=-1) |
|
|
| |
| attn = self.attn_activation(attn) / self.x_head_dim |
|
|
| |
| x = v_j * attn.unsqueeze(2) |
| return x, attn |
|
|
| def aggregate( |
| self, |
| features: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], |
| index: torch.Tensor, |
| ptr: Optional[torch.Tensor], |
| dim_size: Optional[int], |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| x, attn = features |
| x = scatter(x, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr) |
| return x, attn |
|
|
| def update( |
| self, inputs: Tuple[torch.Tensor, torch.Tensor] |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| return inputs |
|
|
| def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: |
| pass |
|
|
| def edge_update(self) -> Tensor: |
| pass |
|
|
|
|
| |
| class PairFeatureNet(nn.Module): |
|
|
| def __init__(self, c_s, c_p, relpos_k=32, template_type="exp-normal-smearing-distance"): |
| super(PairFeatureNet, self).__init__() |
|
|
| self.c_s = c_s |
| self.c_p = c_p |
|
|
| self.linear_s_p_i = nn.Linear(c_s, c_p) |
| self.linear_s_p_j = nn.Linear(c_s, c_p) |
|
|
| self.relpos_k = relpos_k |
| self.n_bin = 2 * relpos_k + 1 |
| self.linear_relpos = nn.Linear(self.n_bin, c_p) |
|
|
| |
| self.template_fn, c_template = get_template_fn(template_type) |
| self.linear_template = nn.Linear(c_template, c_p) |
|
|
| def relpos(self, r): |
| |
| |
| |
|
|
| |
| d = r[:, :, None] - r[:, None, :] |
|
|
| |
| v = torch.arange(-self.relpos_k, self.relpos_k + 1).to(r.device, non_blocking=True) |
|
|
| |
| v_reshaped = v.view(*((1,) * len(d.shape) + (len(v),))) |
|
|
| |
| b = torch.argmin(torch.abs(d[:, :, :, None] - v_reshaped), dim=-1) |
|
|
| |
| oh = nn.functional.one_hot(b, num_classes=len(v)).float() |
|
|
| |
| p = self.linear_relpos(oh) |
|
|
| return p |
|
|
| def template(self, t): |
| return self.linear_template(self.template_fn(t)) |
|
|
| def forward(self, s, t, r, mask): |
| |
| p_mask = mask.unsqueeze(1) * mask.unsqueeze(2) |
| |
| p_i = self.linear_s_p_i(s) |
| p_j = self.linear_s_p_j(s) |
|
|
| |
| p = p_i[:, :, None, :] + p_j[:, None, :, :] |
|
|
| |
| p += self.relpos(r) |
| p += self.template(t) |
|
|
| |
| p *= p_mask.unsqueeze(-1) |
|
|
| return p |
|
|
|
|
| |
| |
| class TriangularSelfAttentionBlock(nn.Module): |
| def __init__( |
| self, |
| sequence_state_dim, |
| pairwise_state_dim, |
| sequence_head_width, |
| pairwise_head_width, |
| dropout=0, |
| **__kwargs, |
| ): |
| super().__init__() |
| from openfold.model.triangular_multiplicative_update import ( |
| TriangleMultiplicationIncoming, |
| TriangleMultiplicationOutgoing, |
| ) |
| from esm.esmfold.v1.misc import ( |
| Attention, |
| Dropout, |
| PairToSequence, |
| ResidueMLP, |
| SequenceToPair, |
| ) |
| assert sequence_state_dim % sequence_head_width == 0 |
| assert pairwise_state_dim % pairwise_head_width == 0 |
| sequence_num_heads = sequence_state_dim // sequence_head_width |
| pairwise_num_heads = pairwise_state_dim // pairwise_head_width |
| assert sequence_state_dim == sequence_num_heads * sequence_head_width |
| assert pairwise_state_dim == pairwise_num_heads * pairwise_head_width |
| assert pairwise_state_dim % 2 == 0 |
|
|
| self.sequence_state_dim = sequence_state_dim |
| self.pairwise_state_dim = pairwise_state_dim |
|
|
| self.layernorm_1 = nn.LayerNorm(sequence_state_dim) |
|
|
| self.sequence_to_pair = SequenceToPair( |
| sequence_state_dim, pairwise_state_dim // 2, pairwise_state_dim |
| ) |
| self.pair_to_sequence = PairToSequence( |
| pairwise_state_dim, sequence_num_heads) |
|
|
| self.seq_attention = Attention( |
| sequence_state_dim, sequence_num_heads, sequence_head_width, gated=True |
| ) |
| self.tri_mul_out = TriangleMultiplicationOutgoing( |
| pairwise_state_dim, |
| pairwise_state_dim, |
| ) |
| self.tri_mul_in = TriangleMultiplicationIncoming( |
| pairwise_state_dim, |
| pairwise_state_dim, |
| ) |
|
|
| self.mlp_seq = ResidueMLP( |
| sequence_state_dim, 4 * sequence_state_dim, dropout=dropout) |
| self.mlp_pair = ResidueMLP( |
| pairwise_state_dim, 4 * pairwise_state_dim, dropout=dropout) |
|
|
| assert dropout < 0.4 |
| self.drop = nn.Dropout(dropout) |
| self.row_drop = Dropout(dropout * 2, 2) |
| self.col_drop = Dropout(dropout * 2, 1) |
|
|
| torch.nn.init.zeros_(self.tri_mul_in.linear_z.weight) |
| torch.nn.init.zeros_(self.tri_mul_in.linear_z.bias) |
| torch.nn.init.zeros_(self.tri_mul_out.linear_z.weight) |
| torch.nn.init.zeros_(self.tri_mul_out.linear_z.bias) |
|
|
| torch.nn.init.zeros_(self.sequence_to_pair.o_proj.weight) |
| torch.nn.init.zeros_(self.sequence_to_pair.o_proj.bias) |
| torch.nn.init.zeros_(self.pair_to_sequence.linear.weight) |
| torch.nn.init.zeros_(self.seq_attention.o_proj.weight) |
| torch.nn.init.zeros_(self.seq_attention.o_proj.bias) |
| torch.nn.init.zeros_(self.mlp_seq.mlp[-2].weight) |
| torch.nn.init.zeros_(self.mlp_seq.mlp[-2].bias) |
| torch.nn.init.zeros_(self.mlp_pair.mlp[-2].weight) |
| torch.nn.init.zeros_(self.mlp_pair.mlp[-2].bias) |
|
|
| def forward(self, sequence_state, pairwise_state, mask=None, chunk_size=None, **__kwargs): |
| """ |
| Inputs: |
| sequence_state: B x L x sequence_state_dim |
| pairwise_state: B x L x L x pairwise_state_dim |
| mask: B x L boolean tensor of valid positions |
| |
| Output: |
| sequence_state: B x L x sequence_state_dim |
| pairwise_state: B x L x L x pairwise_state_dim |
| """ |
| assert len(sequence_state.shape) == 3 |
| assert len(pairwise_state.shape) == 4 |
| if mask is not None: |
| assert len(mask.shape) == 2 |
|
|
| batch_dim, seq_dim, sequence_state_dim = sequence_state.shape |
| pairwise_state_dim = pairwise_state.shape[3] |
| assert sequence_state_dim == self.sequence_state_dim |
| assert pairwise_state_dim == self.pairwise_state_dim |
| assert batch_dim == pairwise_state.shape[0] |
| assert seq_dim == pairwise_state.shape[1] |
| assert seq_dim == pairwise_state.shape[2] |
|
|
| |
| bias = self.pair_to_sequence(pairwise_state) |
|
|
| |
| y = self.layernorm_1(sequence_state) |
| y, _ = self.seq_attention(y, mask=mask, bias=bias) |
| sequence_state = sequence_state + self.drop(y) |
| sequence_state = self.mlp_seq(sequence_state) |
|
|
| |
| pairwise_state = pairwise_state + self.sequence_to_pair(sequence_state) |
|
|
| |
| tri_mask = mask.unsqueeze( |
| 2) * mask.unsqueeze(1) if mask is not None else None |
| pairwise_state = pairwise_state + self.row_drop( |
| self.tri_mul_out(pairwise_state, mask=tri_mask) |
| ) |
| pairwise_state = pairwise_state + self.col_drop( |
| self.tri_mul_in(pairwise_state, mask=tri_mask) |
| ) |
|
|
| |
| pairwise_state = self.mlp_pair(pairwise_state) |
|
|
| return sequence_state, pairwise_state |
|
|
|
|
| |
| class SeqPairAttentionOutput(nn.Module): |
| def __init__(self, seq_state_dim, pairwise_state_dim, num_heads, output_dim, dropout): |
| super(SeqPairAttentionOutput, self).__init__() |
| from esm.esmfold.v1.misc import ( |
| Attention, |
| PairToSequence, |
| ResidueMLP, |
| ) |
| self.seq_state_dim = seq_state_dim |
| self.pairwise_state_dim = pairwise_state_dim |
| self.output_dim = output_dim |
| seq_head_width = seq_state_dim // num_heads |
|
|
| self.layernorm = nn.LayerNorm(seq_state_dim) |
| self.seq_attention = Attention( |
| seq_state_dim, num_heads, seq_head_width, gated=True |
| ) |
| self.pair_to_sequence = PairToSequence(pairwise_state_dim, num_heads) |
| self.mlp_seq = ResidueMLP( |
| seq_state_dim, 4 * seq_state_dim, dropout=dropout) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, sequence_state, pairwise_state, mask=None): |
| |
| bias = self.pair_to_sequence(pairwise_state) |
|
|
| |
| y = self.layernorm(sequence_state) |
| y, _ = self.seq_attention(y, mask=mask, bias=bias) |
| sequence_state = sequence_state + self.drop(y) |
| sequence_state = self.mlp_seq(sequence_state) |
|
|
| return sequence_state |
|
|
