自注意力机制python

### 实现自注意力机制 为了理解并实现自注意力机制，在Python中通常会借助于深度学习框架如PyTorch来简化操作。下面提供了一个基础版本的自注意力层实现，这有助于加深对于此概念的理解。 #### 自定义Self-Attention Layer类 通过创建一个继承自`nn.Module`的新类，可以封装自注意力计算逻辑： ```python import torch from torch import nn import math class SelfAttention(nn.Module): def __init__(self, embed_size, heads): super(SelfAttention, self).__init__() self.embed_size = embed_size self.heads = heads self.head_dim = embed_size // heads assert (self.head_dim * heads == embed_size), "Embed size needs to be divisible by heads" self.values = nn.Linear(self.head_dim, embed_size, bias=False) self.keys = nn.Linear(self.head_dim, embed_size, bias=False) self.queries = nn.Linear(self.head_dim, embed_size, bias=False) self.fc_out = nn.Linear(embed_size, embed_size) def forward(self, values, keys, query, mask): N = query.shape[0] value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # Split embedding into self.heads pieces values = values.reshape(N, value_len, self.heads, self.head_dim) keys = keys.reshape(N, key_len, self.heads, self.head_dim) queries = query.reshape(N, query_len, self.heads, self.head_dim) energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) if mask is not None: energy = energy.masked_fill(mask == 0, float("-1e20")) attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3) out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape( N, query_len, self.embed_size ) out = self.fc_out(out) return out ``` 这段代码展示了如何构建一个多头自注意模块[^1]。这里的关键在于将输入向量分割成多个部分（即多头），分别计算查询、键和值矩阵之间的相似度得分，并应用softmax函数得到权重分布；最后加权求和这些值作为输出的一部分再经过线性变换返回最终的结果。 此外，如果希望尝试更加高效的变体比如Nyström Attention，则可以通过安装对应的库来进行实验[^2]。