使用 pytorch 实现 MultiheadAttention

huggon1

894人浏览 · 2025-03-01 13:16:46

huggon1 · 2025-03-01 13:16:46 发布

概念

注意力机制的核心思想是通过引入 q (查询)、k (键) 和 v (值)，使得模型在进行预测时能够更加聚焦于序列中的某些部分。具体来说，q 代表当前的查询信息，k 和 v 是信息序列中的键值对，表示输入序列中不同部分的信息。

q 代表的是“查询”，即当前预测的目标值（例如，在语言模型中，查询可以是当前词或下一个词的预测）。
k 和 v 代表的是“键”与“值”，在信息处理过程中可以理解为对输入信息的编码，k 是用来与查询进行匹配的部分，v 则是根据匹配得到的加权信息。
在自注意力机制（Self-Attention）中，q = k = v，这意味着查询、键和值都来自相同的序列。一个典型的例子是句子续写任务，在这个任务中，模型需要通过已经生成的文本来预测下一个词。在这种情况下，查询、键和值都来自于当前的句子本身。

注意力机制还广泛应用于 Encoder-Decoder 模型中，例如机器翻译。在翻译任务中，查询（q）通常代表需要翻译的源语言（如中文）的词，键（k）和值（v）则来自于目标语言（如英文）的中间表示。具体来说，查询（中文）与键（英文的不同部分）通过注意力得分进行匹配，模型会将注意力集中在源语言中的相关部分，以便生成目标语言的翻译。

qkv 本身并不一定具有一个明确的抽象含义。它们更多的是计算框架中的一些具体计算单元，用于帮助模型通过加权的方式决定每个输入部分的重要性。实际应用中，q、k 和 v 的具体定义和来源可能因任务的不同而有所变化。

多头注意力是指将一个输入 token 的表示映射到多个子空间中，分别计算多个注意力权重。每个头（head）可以关注输入序列中的不同子任务或不同的特征维度。最终，这些头的输出会被合并，形成一个完整的输出。多头注意力的好处是，它允许模型在同一层中同时从不同的角度学习序列中不同部分的依赖关系，从而提高模型的表达能力。

实现

实现多头注意力需掌握以下核心计算公式：

$MultiHead(Q,K,V)=Concat(head1,…,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O$

$headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

$Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

以下是使用 pytorch 对多头注意力机制的实现，从最基础的单头自注意力，扩展到后面比较完善的多头注意力。

import torch
import torch.nn as nn
import math

'''
所有输入的张量的形状为 [seq_len,batch_size,dim]，
pytorch的中MultiheadAttention默认batch_first为false，即seq_len在第0维
'''

# 版本一：简单自注意力机制
class MultiheadAttention_v1(nn.Module):
    def __init__(self,embed_dim):
        super().__init__()
        self.embed_dim = embed_dim

        self.q_proj = nn.Linear(embed_dim,embed_dim)
        self.k_proj = nn.Linear(embed_dim,embed_dim)
        self.v_proj = nn.Linear(embed_dim,embed_dim)
        self.out_proj = nn.Linear(embed_dim,embed_dim)

    def forward(self,x):
        # [seq_len,batch_size,embed_dim]
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # attention_weight计算
        q,k,v = q.transpose(0,1),k.transpose(0,1),v.transpose(0,1) # [batch_size,seq_len,embed_dim]
        attention_weight = torch.matmul(  # [batch_size,seq_len,seq_len]
            q,k.transpose(-2,-1)
        )
        attention_weight = attention_weight / math.sqrt(self.embed_dim)
        attention_weight = nn.functional.softmax(attention_weight,dim=-1)
        attention_weight = torch.matmul(attention_weight,v) # [batch_size,seq_len,embed_dim]

        attention_weight = self.out_proj(attention_weight).transpose(0,1) # [seq_len,batch_size,embed_dim]
        return attention_weight

# 版本二：在版本一基础上添加多头、dropout、mask
# pytorch中的MultiheadAttention结合了key_padding_mask、attn_mask
# 这里只添加基础mask，mask逻辑在外部实现
class MultiheadAttention_v2(nn.Module):
    def __init__(self,embed_dim,head_num,dropout=0.0):
        super().__init__()
        self.head_num = head_num
        self.embed_dim = embed_dim
        self.dropout = dropout

        self.head_dim = embed_dim // head_num
        assert self.head_dim * head_num == embed_dim, "embed_dim must be divisible by head_num"

        self.q_proj = nn.Linear(embed_dim,embed_dim)
        self.k_proj = nn.Linear(embed_dim,embed_dim)
        self.v_proj = nn.Linear(embed_dim,embed_dim)
        self.out_proj = nn.Linear(embed_dim,embed_dim)

        if dropout > 0.0:
            self.dropout = nn.Dropout(dropout)

    def forward(self,x,mask=None):

        # [seq_len,batch_size,embed_dim]
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # 分头
        # [seq_len,batch_size,embed_dim] -> [head_num,batch_size,seq_len,head_dim]
        seq_len,batch_size,embed_dim = q.size()
        q = q.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        k = k.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        v = v.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)

        # attention_weight计算
        # (..,seq_len,head_dim)*(..,head_dim,seq_len) -> (..,seq_len,seq_len)
        # (..,seq_len,seq_len)*(..,seq_len,head_dim)->(head_num,batch_size,seq_len,head_dim)
        attention_weight = torch.matmul(
            q,k.transpose(-2,-1)
        )
        attention_weight = attention_weight / math.sqrt(self.embed_dim)
        if mask is not None:
            attention_weight = attention_weight.masked_fill(mask,-1e9)

        attention_weight = nn.functional.softmax(attention_weight,dim=-1)
        if self.dropout > 0.0:
            attention_weight = self.dropout(attention_weight,p=self.dropout)
        attention_weight = torch.matmul(attention_weight,v)

        # 合并多头
        '''
        ChatGPT:什么时候需要 .contiguous()：在对张量进行 transpose()、permute()、view() 等操作后，
        张量在内存中的布局可能变成非连续的。如果你需要继续对这些张量进行操作（如 view() 或 CUDA 操作），
        则需要调用 .contiguous() 来确保它们在内存中是连续的。
        '''
        attention_weight = attention_weight.transpose(0,2).contiguous().view(seq_len,batch_size,embed_dim)

        attention_weight = self.out_proj(attention_weight)
        return attention_weight


# 版本三：在版本二基础上修改为通用多头注意力机制。q,k,v维度可不同,tgt_len（q）和src_len(k,v)可不同
class MultiheadAttention_v3(nn.Module):
    def __init__(self,embed_dim,head_num,kdim=None,vdim=None,dropout=0.0):
        super().__init__()
        self.head_num = head_num
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self.head_dim = embed_dim // head_num
        self.dropout = dropout


        assert self.head_dim * head_num == embed_dim, "embed_dim must be divisible by head_num"

        self.q_proj = nn.Linear(embed_dim,embed_dim)
        self.k_proj = nn.Linear(embed_dim,self.kdim)
        self.v_proj = nn.Linear(embed_dim,self.vdim)

        self.out_proj = nn.Linear(embed_dim,embed_dim)

        if dropout > 0.0:
            self.dropout = nn.Dropout(dropout)

    def forward(self,q,k,v,mask=None):
        # 映射为相同维度
        # [seq_len,batch_size,embed_dim]
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)

        # 分头
        # [seq_len,batch_size,embed_dim] -> [head_num,batch_size,seq_len,head_dim]
        # q的seq_len为tgt_len,k,v的seq_len为src_len，可能不同
        seq_len,batch_size,embed_dim = q.size()
        q = q.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        k = k.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        v = v.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)

        # attention_weight计算
        # (..,tgt_seq,head_dim)*(..,head_dim,src_len) -> (..,tgt_seq,src_len)
        # (..,tgt_seq,src_len)*(..,src_len,head_dim)->(head_num,batch_size,tgt_len,head_dim)
        attention_weight = torch.matmul(
            q,k.transpose(-2,-1)
        )
        attention_weight = attention_weight / math.sqrt(self.embed_dim)
        if mask is not None:
            attention_weight = attention_weight.masked_fill(mask,-1e9)
        attention_weight = nn.functional.softmax(attention_weight,dim=-1)
        if self.dropout > 0.0:
            attention_weight = self.dropout(attention_weight,p=self.dropout)
        attention_weight = torch.matmul(attention_weight,v)

        # 合并多头
        attention_weight = attention_weight.transpose(0,2).contiguous().view(seq_len,batch_size,embed_dim)
        attention_weight = self.out_proj(attention_weight)
        return attention_weight

# 版本四：在版本二基础上，针对自注意力q=k=v的特性简化计算
class MultiheadAttention_v4(nn.Module):
    def __init__(self,embed_dim,head_num,kdim=None,vdim=None,dropout=0.0):
        super().__init__()
        self.head_num = head_num
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self.head_dim = embed_dim // head_num
        self.dropout = dropout


        assert self.head_dim * head_num == embed_dim, "embed_dim must be divisible by head_num"

        self.in_proj = nn.Linear(embed_dim,3*embed_dim)

        self.out_proj = nn.Linear(embed_dim,embed_dim)

        if dropout > 0.0:
            self.dropout = nn.Dropout(dropout)

    def forward(self,q,k,v,mask=None):
        # [seq_len,batch_size,embed_dim]
        q,k,v = self.in_proj(q).chunk(3,dim=-1)

        # 分头
        # [seq_len,batch_size,embed_dim] -> [head_num,batch_size,seq_len,head_dim]
        # q的seq_len为tgt_len,k,v的seq_len为src_len，可能不同
        seq_len,batch_size,embed_dim = q.size()
        q = q.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        k = k.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)
        v = v.view(seq_len,batch_size,self.head_num,self.head_dim).transpose(0,2)

        # attention_weight计算
        # (..,tgt_seq,head_dim)*(..,head_dim,src_len) -> (..,tgt_seq,src_len)
        # (..,tgt_seq,src_len)*(..,src_len,head_dim)->(head_num,batch_size,tgt_len,head_dim)
        attention_weight = torch.matmul(
            q,k.transpose(-2,-1)
        )
        attention_weight = attention_weight / math.sqrt(self.embed_dim)
        if mask is not None:
            attention_weight = attention_weight.masked_fill(mask,-1e9)
        attention_weight = nn.functional.softmax(attention_weight,dim=-1)
        if self.dropout > 0.0:
            attention_weight = self.dropout(attention_weight,p=self.dropout)
        attention_weight = torch.matmul(attention_weight,v)

        # 合并多头
        attention_weight = attention_weight.transpose(0,2).contiguous().view(seq_len,batch_size,embed_dim)
        attention_weight = self.out_proj(attention_weight)
        return attention_weight

pytorch 中的 MultiheadAttention 的特点

pytorch 中并没有直接使用线性层，而是通过初始化可训练参数作为权重，在 forward 中需要 linear 的时将权重和输入一同传给由动态链接库引入的 c 实现的 linear 方法。这种用法的好处是对于 qkv 输入维度相同时，pytorch 能构建一个大矩阵作为权重参数，在 forward 中再根据 q、k、v 之间的等价关系来划分权重进行 linear 计算。

torch>nn>functional.py>_in_projection_packed()
E = q.size(-1)
if k is v:
    if q is k:
        # self-attention
        proj = linear(q, w, b)
        # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
        proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
        return proj[0], proj[1], proj[2]
    else:
        # encoder-decoder attention
        w_q, w_kv = w.split([E, E * 2])
        if b is None:
            b_q = b_kv = None
        else:
            b_q, b_kv = b.split([E, E * 2])
        q_proj = linear(q, w_q, b_q)
        kv_proj = linear(k, w_kv, b_kv)
        # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
        kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
        return (q_proj, kv_proj[0], kv_proj[1])
else:
    w_q, w_k, w_v = w.chunk(3)
    if b is None:
        b_q = b_k = b_v = None
    else:
        b_q, b_k, b_v = b.chunk(3)
    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)

对于一些更特殊的情况，即满足以下条件，则会直接通过_native_multi_head_attention进行计算加速

self attention is being computed (i.e., query, key, and value are the same tensor).
inputs are batched (3D) with batch_first==True
Either autograd is disabled (using torch.inference_mode or torch.no_grad) or no tensor argument requires_grad
training is disabled (using .eval())
add_bias_kv is False
add_zero_attn is False
batch_first is True and the input is batched
kdim and vdim are equal to embed_dim
if a NestedTensor _ is passed, neither key_padding_mask
nor attn_mask is passed
autocast is disabled

PS:
python == 3.10
torch == 2.2.2

乐奇 Rokid 开放社区

更多推荐

智能眼镜新品扎堆：AR、VR、MR到底怎么分？

乐奇 Rokid 开放社区

FPC与机器人的双向赋能，解锁智能装备进化新势能

如果说AI算法、算力芯片是智能眼镜的“智慧大脑”，光学模组是智能眼镜的“视觉窗口”，那么FPC就是贯穿整机、连接所有功能模块的“柔性神经脉络”。方寸之间的柔性突破，破解了智能眼镜“轻薄形态”与“全能功能”的核心矛盾，是穿戴智能设备实现形态革新、体验升级的底层核心。未来，随着AR沉浸式技术、AI大模型、穿戴硬件技术的持续进化，智能眼镜将朝着更轻薄、更智能、更多场景、更高稳定性的方向持续迭代，对高端精