2、维度分配原理：总维度d_model与头数h的关系

好，咱们直接进入正题。

多头注意力机制里，维度分配是最基础、也最容易踩坑的地方。我刚开始接触Transformer时，就被这个d_model、h、d_k的关系绕晕过。今天咱们把它彻底讲明白。

2.1 核心公式：d_k = d_model / h

说白了，就是一句话：

举个例子：

• 如果 d_model = 512，h = 8，那么 d_k = 512 / 8 = 64
• 如果 d_model = 768，h = 12，那么 d_k = 768 / 12 = 64
• 如果 d_model = 1024，h = 16，那么 d_k = 1024 / 16 = 64

你发现没？很多经典模型里d_k都等于64。这不是巧合，是设计上的经验值。

我个人习惯把d_k看作每个头的“视野宽度”。64维的向量，刚好能捕捉到一些细粒度特征，又不会太稀疏。

2.2 为什么需要整除？

这个问题，我曾经在项目里吃过亏。

有一次我图省事，设了d_model=512，h=7。你算算，512/7≈73.14，不是整数。结果模型跑起来，报错信息看得我头皮发麻。

为什么会这样？原因有三：

1. 张量形状必须对齐：每个头的输出要拼接回d_model维度。如果d_k不是整数，拼接后的形状就乱套了。
2. 计算效率问题：GPU对规整的张量形状有优化。非整数维度会触发大量padding操作，慢得你想哭。
3. 语义完整性：每个头应该分配到等量的维度。不均匀分配，有的头“看得多”，有的头“看得少”，这就不公平了。

assert d_model % h == 0, f"d_model({d_model}) 必须能被 h({h}) 整除"

2.3 维度分配的直观理解

你想想看，多头注意力就像是一个团队在开会。

• d_model：整个团队的知识总量
• h：团队里的人数
• d_k：每个人分到的知识量

如果总知识量是512，你非要分给7个人。那每个人分到73.14份知识？这怎么分？有人多拿一点，有人少拿一点，团队协作就出问题了。

所以，整除是硬性要求。不是“最好”，是“必须”。

2.4 常见配置速查表

我整理了一份常用配置，方便你参考：

嗯，这里要注意：d_k=64不是铁律。像GPT-2 Medium就用128。关键是d_model和h要匹配好。

2.5 代码实现中的维度分配

咱们看看实际代码怎么写。我习惯用PyTorch，直接上代码：

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h):
        super().__init__()
        # 这里必须检查整除
        assert d_model % h == 0, f"d_model({d_model}) % h({h}) != 0"
        
        self.d_model = d_model
        self.h = h
        self.d_k = d_model // h  # 每个头的维度
        
        # 定义Q、K、V的线性变换
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        # 线性变换后，拆分成多头
        Q = self.W_q(x).view(batch_size, seq_len, self.h, self.d_k)
        K = self.W_k(x).view(batch_size, seq_len, self.h, self.d_k)
        V = self.W_v(x).view(batch_size, seq_len, self.h, self.d_k)
        
        # 转置：把头的维度放到第1维
        Q = Q.transpose(1, 2)  # [batch, h, seq_len, d_k]
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        
        # 计算注意力（简化版）
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        
        # 拼接多头输出
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        # 最后的线性变换
        return self.W_o(out)

2.6 维度分配的深层思考

到这里，你可能想问：为什么非得是d_k = d_model / h？能不能让每个头维度不一样？

嗯，理论上可以。但实践中没人这么干。

原因很简单：对称性。每个头在模型中的地位是平等的。你给头A分配128维，给头B分配64维，那模型就会偏向头A。这违背了多头注意力的设计初衷——让不同头关注不同模式，而不是让某些头“权力更大”。

我记得有一次，团队里有人想试试非均匀分配。结果调了一周，效果还不如均匀分配。从那以后，我们就老老实实用d_k = d_model / h了。

2.7 总结一下

• 核心公式：d_k = d_model / h，必须整除
• 常见配置：d_k=64或128，d_model和h按需搭配
• 代码检查：加断言，避免运行时崩溃
• 硬件优化：d_k尽量选2的幂次方

下一章，咱们聊聊每个头到底在学什么。你会发现，不同头关注的东西，差别还挺大的。