矩阵运算维度设计与Attention机制的底层逻辑

1. 矩阵运算维度设计

深度学习框架选择 $xW^T$并非随意之举，而是基于以下核心原则的精心设计：

1.1. 内存访问效率优先

现代硬件（CPU/GPU）的性能瓶颈主要在于内存访问速度，而非计算速度。$xW^T$能够实现连续内存访问：

对于 x的每一行（一个样本），程序可以顺序读取其所有特征，这与数据在内存中的物理存储方式一致。

这种模式极大提高了缓存命中率，从而显著提升计算速度。

相比之下，若采用 $Wx$ 模式，则需将 $x$ 转置或将特征维度前置，导致访问模式变为跳跃式（strided access），引发大量缓存缺失（cache miss），显著降低吞吐效率。

1.2. 批处理的天然适配

深度学习框架普遍采用 (B, L, D) 作为标准张量布局，其中：

• B：批次大小（Batch size）
• L：序列长度（Sequence length）
• D：特征维度（Feature dimension）

$xW^T$ 的运算天然适配该结构：

• 当 x.shape = [B, D] 时，$xW^T$ 直接输出 [B, d_out]，批次维度自动保留，无需额外维度变换。
• 若使用 $Wx$，则必须将 $x$ 转置为 $\text{[D, B]}$，破坏了批处理的直观性与编程习惯，增加代码复杂度与出错风险。

1.3. 框架的设计一致性

主流框架（如PyTorch、TensorFlow）中，nn.Linear(D_in, D_out) 的权重矩阵默认形状为 (D_out, D_in)，即：

• 每一行代表一个输出神经元的权重向量。
• $xW^T$ 实质上是将输入样本与每个输出神经元做内积。

这种设计还支持：

• SIMD指令集加速：连续内存访问允许CPU/GPU并行处理多个数据点。
• 自动微分友好：梯度传播路径清晰，无需复杂转置逻辑。

• 函数定义的时候，直接传参是输入维度和输出维度，便于代码开发的时候理解
• 我们实际计算的时候数据是行读取的，x的一行为(D_in)，我们希望W的一行也是(D_in)
• 因此将矩阵的维度定义成(D_out, D_in)，这样可以直接取出一行权重（维度D_in）
• 因此理论上X与W进行相乘的时候，W需要进行转置(D_in, D_out)

2. slef Attention矩阵运算推理

2.1 单头注意力机制

假如有B个输入，每个输入为一系列的token，序列长度为L，每个token转变成D维的embedding，则输入X.shape = [B, L, D]

1. 原始输入的文本，通过分词器得到tokenid
2. tokenid通过nn.embedding得到对应的词向量$\mathbf{x}$
3. 词向量与$W_Q$，$W_K$，$W_V$运算得到$\mathbf{q}$，$\mathbf{k}$，$\mathbf{v}$三个向量

graph BT
    a1[token1]-->|tokenizer|b1[token id1]
    b1-->|nn.embedding|c1[x1]
    c1-->|Wq|d1[q1]
    c1-->|Wk|e1[k1]
    c1-->|Wv|f1[v1]
    
    a2[token2]-->|tokenizer|b2[token id2]
    b2-->|nn.embedding|c2[x2]
    c2-->|Wq|d2[q2]
    c2-->|Wk|e2[k2]
    c2-->|Wv|f2[v2]
    
    a3[token3]-->|tokenizer|b3[token id3]
    b3-->|nn.embedding|c3[x3]
    c3-->|Wq|d3[q3]
    c3-->|Wk|e3[k3]
    c3-->|Wv|f3[v3]

上一个小节讲了，在做计算的时候$y=xW^T$，假如$x$的维度为(1, D), 则：

• $W_q$维度为[d_q, D],$W_q^T$的维度就是[D,d_q],$\mathbf{q}$的维度就是(1,D)x(D,d_q)->(1,d_q)
• $W_k$维度为[d_k, D],$W_k^T$的维度就是[D,d_k],$\mathbf{k}$的维度就是(1,D)x(D,d_k)->(1,d_k)
• $W_v$维度为[d_v, D],$W_v^T$的维度就是[D,d_v],$\mathbf{v}$的维度就是(1,D)x(D,d_v)->(1,d_v)

以token1为例，我们可以计算出

• $q_1 = x_1W_q^T$，q1.shape=[1,d_q]
• $k_1 = x_1W_k^T$，k1.shape=[1,d_k]
• $v_1 = x_1W_v^T$，v1.shape=[1,d_v]

每个序列有L个token，因此就有L个向量，可以发现Q，K，V都是一个二维矩阵，例如Q有L行，每行有d_q列

• $Q = [q_1,q_2,...q_L]$，Q.shape = [L, d_q]
• $K = [k_1,k_2,...k_L]$，K.shape = [L, d_k]
• $V = [v_1,v_2,...v_L]$，V.shape = [L, d_v]

我们要计算token之间的相关性，例如token0与其他所有token的相关性，

• $a_{11} = q_1 k_1^T$
• $a_{12} = q_1 k_2^T$
• $a_{13} = q_1 k_3^T$

q和k的维度要相同，否则无法进行点积d_q=d_k

token0和其他所有token的权重组织一个向量

所有的$k_i^T$组织起来的矩阵就是$K^T$，同理可以得到a2和其他所有

• $\mathbf{a_2} = \mathbf{q_2} * \mathbf{K^T}$
• $\mathbf{a_L} = \mathbf{q_L} * \mathbf{K^T}$

我们讲这个注意力权重拼接成一个矩阵

就得到了最终的注意力矩阵，也就是每个token对其他token的权重

• A.shape = [L, L]
• o = AV，o.shape = [L, o_v]

2.2 多头注意力机制

引入多头机制（Multi-Head Attention）的核心目的：让不同“子空间”学习不同的注意力模式。

1. 单个头内的运算 

假设我们有 $H=8$ 个头，总维度 $D=512$，那么每个头的维度 $d_{h}=64$。  对于其中某一个头（比如 Head 1）： 

• $Q_{1},K_{1}$ 的形状是 $[L,64]$。

• 注意力矩阵 $A_{1}$：

  $A_{1}=\text{softmax}\left(\frac{Q_{1}K_{1}^{T}}{\sqrt{64}}\right)$

注意：$A_{1}$ 的形状依然是 $[L,L]$。它表示的是这个头视角下，Token 之间的相关性。

• $V_{1}$ 的形状是 $[L,64]$。
• 相乘 $O_{1}=A_{1}V_{1}$：

$[L,L]\times [L,64]\rightarrow [L,64]$ 所以，Head 1 的输出是一个 $[L,64]$ 的矩阵。 

2. 为什么说“维度缩小了”？ 

在单头注意力中，$V$ 是 $[L,512]$，最后得到 $[L,512]$。 在多头中，每个头确实只产出了 $[L,64]$ 的结果。 但是有 8 个头： 

• Head 1 输出 $O_{1}:[L,64]$
• Head 2 输出 $O_{2}:[L,64]$
• ...
• Head 8 输出 $O_{8}:[L,64]$ 

3. 最终的“拼接” (Concatenation)

我们会把这 8 个头的输出在最后一个维度（特征维度）拼起来： $O_{concat}=[O_{1},O_{2},...,O_{8}]$ $[L,64]\text{\ 拼接\ }8\text{\ 次\ }\rightarrow [L,64\times 8]\rightarrow [L,512]$

4. 深度理解：为什么这样做？  可以把 $V$ 的 $512$ 个维度想象成 512 个不同的特征属性。 

• 单头： 用同一个注意力权重 $A$ 去加权所有的 512 个属性。
• 多头：
  • 前 64 个属性用 $A_{1}$ 来加权（可能关注语法）；
  • 中间 64 个属性用 $A_{2}$ 来加权（可能关注指代消解）；
  • ...
  • 最后 64 个属性用 $A_{8}$ 来加权。 

结论：

$V$ 的维度在每个头内部计算时确实是缩小的，但这种缩小是为了分工。每个头只负责处理一部分特征维度。最后拼在一起时，总的特征维度 $D$ 又回来了。