| 11 | |
| 12 | # Update MSA with biased self-attention. bias from Pair & Str |
| 13 | class MSAPairStr2MSA(nn.Module): |
| 14 | def __init__(self, d_msa=256, d_pair=128, n_head=8, d_state=16, |
| 15 | d_hidden=32, p_drop=0.15, use_global_attn=False): |
| 16 | super(MSAPairStr2MSA, self).__init__() |
| 17 | self.norm_pair = nn.LayerNorm(d_pair) |
| 18 | self.proj_pair = nn.Linear(d_pair+36, d_pair) |
| 19 | self.norm_state = nn.LayerNorm(d_state) |
| 20 | self.proj_state = nn.Linear(d_state, d_msa) |
| 21 | self.drop_row = Dropout(broadcast_dim=1, p_drop=p_drop) |
| 22 | self.row_attn = MSARowAttentionWithBias(d_msa=d_msa, d_pair=d_pair, |
| 23 | n_head=n_head, d_hidden=d_hidden) |
| 24 | if use_global_attn: |
| 25 | self.col_attn = MSAColGlobalAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) |
| 26 | else: |
| 27 | self.col_attn = MSAColAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) |
| 28 | self.ff = FeedForwardLayer(d_msa, 4, p_drop=p_drop) |
| 29 | |
| 30 | # Do proper initialization |
| 31 | self.reset_parameter() |
| 32 | |
| 33 | def reset_parameter(self): |
| 34 | # initialize weights to normal distrib |
| 35 | self.proj_pair = init_lecun_normal(self.proj_pair) |
| 36 | self.proj_state = init_lecun_normal(self.proj_state) |
| 37 | |
| 38 | # initialize bias to zeros |
| 39 | nn.init.zeros_(self.proj_pair.bias) |
| 40 | nn.init.zeros_(self.proj_state.bias) |
| 41 | |
| 42 | def forward(self, msa, pair, rbf_feat, state): |
| 43 | ''' |
| 44 | Inputs: |
| 45 | - msa: MSA feature (B, N, L, d_msa) |
| 46 | - pair: Pair feature (B, L, L, d_pair) |
| 47 | - rbf_feat: Ca-Ca distance feature calculated from xyz coordinates (B, L, L, 36) |
| 48 | - xyz: xyz coordinates (B, L, n_atom, 3) |
| 49 | - state: updated node features after SE(3)-Transformer layer (B, L, d_state) |
| 50 | Output: |
| 51 | - msa: Updated MSA feature (B, N, L, d_msa) |
| 52 | ''' |
| 53 | B, N, L = msa.shape[:3] |
| 54 | |
| 55 | # prepare input bias feature by combining pair & coordinate info |
| 56 | pair = self.norm_pair(pair) |
| 57 | pair = torch.cat((pair, rbf_feat), dim=-1) |
| 58 | pair = self.proj_pair(pair) # (B, L, L, d_pair) |
| 59 | # |
| 60 | # update query sequence feature (first sequence in the MSA) with feedbacks (state) from SE3 |
| 61 | state = self.norm_state(state) |
| 62 | state = self.proj_state(state).reshape(B, 1, L, -1) |
| 63 | msa = msa.index_add(1, torch.tensor([0,], device=state.device), state) |
| 64 | # |
| 65 | # Apply row/column attention to msa & transform |
| 66 | msa = msa + self.drop_row(self.row_attn(msa, pair)) |
| 67 | msa = msa + self.col_attn(msa) |
| 68 | msa = msa + self.ff(msa) |
| 69 | |
| 70 | return msa |