| from typing import Optional
|
|
|
| import torch
|
| from einops.layers.torch import Rearrange
|
| from torch import Tensor, nn
|
|
|
| from . import vb_layers_initialize as init
|
|
|
|
|
| class AttentionPairBias(nn.Module):
|
| """Attention pair bias layer."""
|
|
|
| def __init__(
|
| self,
|
| c_s: int,
|
| c_z: Optional[int] = None,
|
| num_heads: Optional[int] = None,
|
| inf: float = 1e6,
|
| compute_pair_bias: bool = True,
|
| ) -> None:
|
| """Initialize the attention pair bias layer.
|
|
|
| Parameters
|
| ----------
|
| c_s : int
|
| The input sequence dimension.
|
| c_z : int
|
| The input pairwise dimension.
|
| num_heads : int
|
| The number of heads.
|
| inf : float, optional
|
| The inf value, by default 1e6
|
|
|
| """
|
| super().__init__()
|
|
|
| assert c_s % num_heads == 0
|
|
|
| self.c_s = c_s
|
| self.num_heads = num_heads
|
| self.head_dim = c_s // num_heads
|
| self.inf = inf
|
|
|
| self.proj_q = nn.Linear(c_s, c_s)
|
| self.proj_k = nn.Linear(c_s, c_s, bias=False)
|
| self.proj_v = nn.Linear(c_s, c_s, bias=False)
|
| self.proj_g = nn.Linear(c_s, c_s, bias=False)
|
|
|
| self.compute_pair_bias = compute_pair_bias
|
| if compute_pair_bias:
|
| self.proj_z = nn.Sequential(
|
| nn.LayerNorm(c_z),
|
| nn.Linear(c_z, num_heads, bias=False),
|
| Rearrange("b ... h -> b h ..."),
|
| )
|
| else:
|
| self.proj_z = Rearrange("b ... h -> b h ...")
|
|
|
| self.proj_o = nn.Linear(c_s, c_s, bias=False)
|
| init.final_init_(self.proj_o.weight)
|
|
|
| def forward(
|
| self,
|
| s: Tensor,
|
| z: Tensor,
|
| mask: Tensor,
|
| k_in: Tensor,
|
| multiplicity: int = 1,
|
| ) -> Tensor:
|
| """Forward pass.
|
|
|
| Parameters
|
| ----------
|
| s : torch.Tensor
|
| The input sequence tensor (B, S, D)
|
| z : torch.Tensor
|
| The input pairwise tensor or bias (B, N, N, D)
|
| mask : torch.Tensor
|
| The pairwise mask tensor (B, N, N)
|
|
|
| Returns
|
| -------
|
| torch.Tensor
|
| The output sequence tensor.
|
|
|
| """
|
| B = s.shape[0]
|
|
|
|
|
| q = self.proj_q(s).view(B, -1, self.num_heads, self.head_dim)
|
| k = self.proj_k(k_in).view(B, -1, self.num_heads, self.head_dim)
|
| v = self.proj_v(k_in).view(B, -1, self.num_heads, self.head_dim)
|
|
|
| bias = self.proj_z(z)
|
| bias = bias.repeat_interleave(multiplicity, 0)
|
|
|
| g = self.proj_g(s).sigmoid()
|
|
|
| with torch.autocast("cuda", enabled=False):
|
|
|
| attn = torch.einsum("bihd,bjhd->bhij", q.float(), k.float())
|
| attn = attn / (self.head_dim**0.5) + bias.float()
|
| attn = attn + (1 - mask[:, None, None].float()) * -self.inf
|
| attn = attn.softmax(dim=-1)
|
|
|
|
|
| o = torch.einsum("bhij,bjhd->bihd", attn, v.float()).to(v.dtype)
|
| o = o.reshape(B, -1, self.c_s)
|
| o = self.proj_o(g * o)
|
|
|
| return o
|
|
|
