| 229 | return out |
| 230 | |
| 231 | class MSAColAttention(nn.Module): |
| 232 | def __init__(self, d_msa=256, n_head=8, d_hidden=32): |
| 233 | super(MSAColAttention, self).__init__() |
| 234 | self.norm_msa = nn.LayerNorm(d_msa) |
| 235 | # |
| 236 | self.to_q = nn.Linear(d_msa, n_head*d_hidden, bias=False) |
| 237 | self.to_k = nn.Linear(d_msa, n_head*d_hidden, bias=False) |
| 238 | self.to_v = nn.Linear(d_msa, n_head*d_hidden, bias=False) |
| 239 | self.to_g = nn.Linear(d_msa, n_head*d_hidden) |
| 240 | self.to_out = nn.Linear(n_head*d_hidden, d_msa) |
| 241 | |
| 242 | self.scaling = 1/math.sqrt(d_hidden) |
| 243 | self.h = n_head |
| 244 | self.dim = d_hidden |
| 245 | |
| 246 | self.reset_parameter() |
| 247 | |
| 248 | def reset_parameter(self): |
| 249 | # query/key/value projection: Glorot uniform / Xavier uniform |
| 250 | nn.init.xavier_uniform_(self.to_q.weight) |
| 251 | nn.init.xavier_uniform_(self.to_k.weight) |
| 252 | nn.init.xavier_uniform_(self.to_v.weight) |
| 253 | |
| 254 | # gating: zero weights, one biases (mostly open gate at the begining) |
| 255 | nn.init.zeros_(self.to_g.weight) |
| 256 | nn.init.ones_(self.to_g.bias) |
| 257 | |
| 258 | # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining |
| 259 | nn.init.zeros_(self.to_out.weight) |
| 260 | nn.init.zeros_(self.to_out.bias) |
| 261 | |
| 262 | def forward(self, msa): |
| 263 | B, N, L = msa.shape[:3] |
| 264 | # |
| 265 | msa = self.norm_msa(msa) |
| 266 | # |
| 267 | query = self.to_q(msa).reshape(B, N, L, self.h, self.dim) |
| 268 | key = self.to_k(msa).reshape(B, N, L, self.h, self.dim) |
| 269 | value = self.to_v(msa).reshape(B, N, L, self.h, self.dim) |
| 270 | gate = torch.sigmoid(self.to_g(msa)) |
| 271 | # |
| 272 | query = query * self.scaling |
| 273 | attn = einsum('bqihd,bkihd->bihqk', query, key) |
| 274 | attn = F.softmax(attn, dim=-1) |
| 275 | # |
| 276 | out = einsum('bihqk,bkihd->bqihd', attn, value).reshape(B, N, L, -1) |
| 277 | out = gate * out |
| 278 | # |
| 279 | out = self.to_out(out) |
| 280 | return out |
| 281 | |
| 282 | class MSAColGlobalAttention(nn.Module): |
| 283 | def __init__(self, d_msa=64, n_head=8, d_hidden=8): |