Multi-Query Attention (MQA) PyTorch 实现

和多头注意力机制的唯一区别:K、V在不同的head之间实现了复用,而对于不同的头,Q依然不同。

因此这里的代码和标准多头注意力的实现也是几乎完全一样:

python 复制代码
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5

        # 查询、键、值投影
        self.q_proj = nn.Linear(embed_dim, embed_dim)  # 多头查询
        self.k_proj = nn.Linear(embed_dim, self.head_dim)  # 单头键
        self.v_proj = nn.Linear(embed_dim, self.head_dim)  # 单头值
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.shape

        # 投影
        q = self.q_proj(x)  # (batch, seq_len, embed_dim)
        k = self.k_proj(x)  # (batch, seq_len, head_dim)
        v = self.v_proj(x)  # (batch, seq_len, head_dim)

        # 重塑查询为多头
        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # (batch, num_heads, seq_len, head_dim)
        
        # 键和值保持单头,扩展到多头维度
        k = k.unsqueeze(1)  # (batch, 1, seq_len, head_dim)
        v = v.unsqueeze(1)  # (batch, 1, seq_len, head_dim)

        # 注意力计算
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        # (batch, num_heads, seq_len, seq_len)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)  # (batch, num_heads, seq_len, head_dim)

        # 合并多头
        out = out.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
        out = self.out_proj(out)  # (batch, seq_len, embed_dim)

        return out

# 示例用法
embed_dim = 64
num_heads = 8
model = MultiQueryAttention(embed_dim, num_heads)
x = torch.randn(2, 10, embed_dim)  # (batch, seq_len, embed_dim)
output = model(x)
print(output.shape)  # torch.Size([2, 10, 64])
相关推荐
点云SLAM8 分钟前
PyTorch 中.backward() 详解使用
人工智能·pytorch·python·深度学习·算法·机器学习·机器人
B1118521Y4637 分钟前
flask的使用
后端·python·flask
Learn Beyond Limits1 小时前
Transfer Learning|迁移学习
人工智能·python·深度学习·神经网络·机器学习·ai·吴恩达
love530love3 小时前
【保姆级教程】阿里 Wan2.1-T2V-14B 模型本地部署全流程:从环境配置到视频生成(附避坑指南)
人工智能·windows·python·开源·大模型·github·音视频
He1955013 小时前
Go初级之十:错误处理与程序健壮性
开发语言·python·golang
和鲸社区4 小时前
《斯坦福CS336》作业1开源,从0手搓大模型|代码复现+免环境配置
人工智能·python·深度学习·计算机视觉·语言模型·自然语言处理·nlp
豌豆花下猫5 小时前
Python 潮流周刊#118:Python 异步为何不够流行?(摘要)
后端·python·ai
THMAIL5 小时前
深度学习从入门到精通 - LSTM与GRU深度剖析:破解长序列记忆遗忘困境
人工智能·python·深度学习·算法·机器学习·逻辑回归·lstm
wheeldown5 小时前
【数学建模】数据预处理入门:从理论到动手操作
python·数学建模·matlab·python3.11
多打代码6 小时前
2025.09.05 用队列实现栈 & 有效的括号 & 删除字符串中的所有相邻重复项
python·算法