Attention Free Transformer (AFT)-2020论文笔记

---

名称：

Attention Free Transformer (AFT)

来源：

[2105.14103] An Attention Free Transformer

相关工作：

#Approximatingthedotproduct #Sparselocalattention #Contextcompression #Eliminatingdotproductattention #MLPsforvision

创新点：

贡献：

• 提出了一种全新的注意力机制替代方案，完全摒弃了点积注意力。

• AFT的计算复杂度与输入长度和特征维度呈线性关系，适用于大规模数据。

• AFT-local和AFT-conv变体通过引入局部性和空间权重共享，进一步提高了模型的效率和性能。

代码：

# ---------------------------------------  
# 论文:An Attention Free Transformer (arxiv2021)  
# ---------------------------------------  
import torch  
from torch import nn  
from torch.nn import init  class AFT_FULL(nn.Module):  def __init__(self, d_model, n=49, simple=False):  super(AFT_FULL, self).__init__()  self.fc_q = nn.Linear(d_model, d_model)  self.fc_k = nn.Linear(d_model, d_model)  self.fc_v = nn.Linear(d_model, d_model)  if (simple):  self.position_biases = torch.zeros((n, n))  else:  self.position_biases = nn.Parameter(torch.ones((n, n)))  self.d_model = d_model  self.n = n  self.sigmoid = nn.Sigmoid()  self.init_weights()  def init_weights(self):  for m in self.modules():  if isinstance(m, nn.Conv2d):  init.kaiming_normal_(m.weight, mode='fan_out')  if m.bias is not None:  init.constant_(m.bias, 0)  elif isinstance(m, nn.BatchNorm2d):  init.constant_(m.weight, 1)  init.constant_(m.bias, 0)  elif isinstance(m, nn.Linear):  init.normal_(m.weight, std=0.001)  if m.bias is not None:  init.constant_(m.bias, 0)  def forward(self, input):  bs, n, dim = input.shape  q = self.fc_q(input)  # bs,n,dim  k = self.fc_k(input).view(1, bs, n, dim)  # 1,bs,n,dim  v = self.fc_v(input).view(1, bs, n, dim)  # 1,bs,n,dim  numerator = torch.sum(torch.exp(k + self.position_biases.view(n, 1, -1, 1)) * v, dim=2)  # n,bs,dim  denominator = torch.sum(torch.exp(k + self.position_biases.view(n, 1, -1, 1)), dim=2)  # n,bs,dim  out = (numerator / denominator)  # n,bs,dim  out = self.sigmoid(q) * (out.permute(1, 0, 2))  # bs,n,dim  return out  # 输入 B C N,  输出 B C Nif __name__ == '__main__':  block = AFT_FULL(d_model=512, n=64).cuda()  input = torch.rand(64, 64, 512).cuda()  output = block( input)  print(input.size(), output.size())