注意力机制中softmax函数

### 注意力机制中Softmax函数的作用

在神经网络的注意力机制里，Softmax函数扮演着至关重要的角色。通过Softmax变换，可以将任意实数值转换成概率分布形式，使得输出向量中的各个元素代表不同部分的重要性权重[^1]。

具体来说，在计算过程中，对于给定查询（Query）、键（Key）以及值（Value），首先会基于查询和键之间的相似度得分构建加权矩阵；接着利用Softmax对该得分进行标准化处理，从而获得一组介于0至1之间且总和等于1的概率值作为最终的关注力度分配方案[^2]。

这种做法不仅能够突出显示输入序列中最值得关注的部分，同时也保留了一定程度上的全局视角，允许模型兼顾次要特征的信息贡献。此外，由于经过了指数运算再归一化的过程，即使原始打分差异较小也能被放大出来形成明显的区分效果[^3]。

import torch
import torch.nn.functional as F

def attention(query, key, value):
    # 计算Q与K的点积得到scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(key.size(-1))
    
    # 对scores做softmax以获取attention weights
    attn_weights = F.softmax(scores, dim=-1)

    # 将weights应用于values上得出context vector
    context_vector = torch.matmul(attn_weights, value)
    
    return context_vector, attn_weights