我觉得源码写的很好懂，我就不加注释了，直接上计算流程图。
AFTFull
class AFTFull(nn.Module): def __init__(self, max_seqlen, dim, hidden_dim=64): super().__init__() ''' max_seqlen: the maximum number of timesteps (sequence length) to be fed in dim: the embedding dimension of the tokens hidden_dim: the hidden dimension used inside AFT Full Number of heads is 1 as done in the paper ''' self.dim = dim self.hidden_dim = hidden_dim self.to_q = nn.Linear(dim, hidden_dim) self.to_k = nn.Linear(dim, hidden_dim) self.to_v = nn.Linear(dim, hidden_dim) self.project = nn.Linear(hidden_dim, dim) self.wbias = nn.Parameter(torch.Tensor(max_seqlen, max_seqlen)) nn.init.xavier_uniform_(self.wbias) def forward(self, x): B, T, _ = x.shape Q = self.to_q(x).view(B, T, self.hidden_dim) K = self.to_k(x).view(B, T, self.hidden_dim) V = self.to_v(x).view(B, T, self.hidden_dim) temp_wbias = self.wbias[:T, :T].unsqueeze(0) # sequences can still be variable length ''' From the paper ''' Q_sig = torch.sigmoid(Q) temp = torch.exp(temp_wbias) @ torch.mul(torch.exp(K), V) weighted = temp / (torch.exp(temp_wbias) @ torch.exp(K)) Yt = torch.mul(Q_sig, weighted) Yt = Yt.view(B, T, self.hidden_dim) Yt = self.project(Yt) return Yt
AFTSimple
class AFTSimple(nn.Module): def __init__(self, max_seqlen, dim, hidden_dim=64): super().__init__() ''' max_seqlen: the maximum number of timesteps (sequence length) to be fed in dim: the embedding dimension of the tokens hidden_dim: the hidden dimension used inside AFT Full Number of Heads is 1 as done in the paper. ''' self.dim = dim self.hidden_dim = hidden_dim self.to_q = nn.Linear(dim, hidden_dim) self.to_k = nn.Linear(dim, hidden_dim) self.to_v = nn.Linear(dim, hidden_dim) self.project = nn.Linear(hidden_dim, dim) def forward(self, x): B, T, _ = x.shape Q = self.to_q(x).view(B, T, self.hidden_dim) K = self.to_k(x).view(B, T, self.hidden_dim) V = self.to_v(x).view(B, T, self.hidden_dim) ''' From the paper ''' weights = torch.mul(torch.softmax(K, 1), V).sum(dim=1, keepdim=True) Q_sig = torch.sigmoid(Q) Yt = torch.mul(Q_sig, weights) Yt = Yt.view(B, T, self.hidden_dim) Yt = self.project(Yt) return Yt
AFTLocal
class AFTLocal(nn.Module): def __init__(self, max_seqlen, dim, hidden_dim=64, s=256): super().__init__() ''' max_seqlen: the maximum number of timesteps (sequence length) to be fed in dim: the embedding dimension of the tokens hidden_dim: the hidden dimension used inside AFT Full s: the window size used for AFT-Local in the paper Number of heads is 1 as done in the paper ''' self.dim = dim self.hidden_dim = hidden_dim self.to_q = nn.Linear(dim, hidden_dim) self.to_k = nn.Linear(dim, hidden_dim) self.to_v = nn.Linear(dim, hidden_dim) self.project = nn.Linear(hidden_dim, dim) self.wbias = nn.Parameter(torch.Tensor(max_seqlen, max_seqlen)) self.max_seqlen = max_seqlen self.s = s nn.init.xavier_uniform_(self.wbias) def forward(self, x): B, T, _ = x.shape Q = self.to_q(x).view(B, T, self.hidden_dim) K = self.to_k(x).view(B, T, self.hidden_dim) V = self.to_v(x).view(B, T, self.hidden_dim) self.wbias = nn.Parameter(torch.Tensor([ [self.wbias[i][j] if math.fabs(i-j) < self.s else 0 for j in range(self.max_seqlen)] for i in range(self.max_seqlen) ])) temp_wbias = self.wbias[:T, :T].unsqueeze(0) # sequences can still be variable length ''' From the paper ''' Q_sig = torch.sigmoid(Q) temp = torch.exp(temp_wbias) @ torch.mul(torch.exp(K), V) weighted = temp / (torch.exp(temp_wbias) @ torch.exp(K)) Yt = torch.mul(Q_sig, weighted) Yt = Yt.view(B, T, self.hidden_dim) Yt = self.project(Yt) return Yt
