msa多头注意力机制

### 多头注意力机制在多重序列比对中的工作原理 多头注意力机制（Multi-Head Attention, MHA）允许模型在同一位置关注来自不同表示子空间的信息。对于多重序列比对（Multiple Sequence Alignment, MSA），这种机制能够捕捉到更复杂的模式和关系。 #### 输入处理 当应用于MSA时，每个多重序列会被视为一个查询（Query）、键（Key）以及值（Value）。这些向量通过线性投影映射至较低维度的空间中[^1]： ```python import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads == 0 self.d_k = d_model // num_heads self.num_heads = num_heads self.linears = clones(nn.Linear(d_model, d_model), 4) def forward(self, query, key, value, mask=None): batch_size = query.size(0) # Do all the linear projections at once query, key, value = [ lin(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) for lin, x in zip(self.linears, (query, key, value)) ] ``` #### 计算注意力分数 接着计算各个头部之间的相似度得分，并应用softmax函数得到权重分布。此过程重复多次以形成多个独立的关注视角: ```python def attention(query, key, value, mask=None, dropout=None): scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1)) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) p_attn = F.softmax(scores, dim=-1) if dropout is not None: p_attn = dropout(p_attn) return torch.matmul(p_attn, value), p_attn ``` #### 结果拼接 最后将各头部的结果重新组合起来并传递给下一层神经网络组件。这种方式使得模型能够在不同抽象层次上理解输入数据间的联系: ```python output, _ = attention(query, key, value) output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) return self.linears[-1](output) ```