# Summary
# Cues
# Notes
这个操作是**给mask添加维度**,让它能与多头注意力的Q、K、V张量进行**广播(broadcasting)**!
## 🎯 **背景:需要匹配的张量形状**
```python
# 多头注意力中的张量形状:
Q.shape = [batch_size, num_heads, seq_len, d_k] # [2, 4, 10, 16]
K.shape = [batch_size, num_heads, seq_len, d_k] # [2, 4, 10, 16]
V.shape = [batch_size, num_heads, seq_len, d_k] # [2, 4, 10, 16]
# 注意力分数的形状:
scores.shape = [batch_size, num_heads, seq_len, seq_len] # [2, 4, 10, 10]
```
## 📊 **mask的维度扩展过程**
### **原始mask(操作前):**
```python
# 因果mask(下三角矩阵)
mask = torch.triu(torch.ones(seq_len, seq_len) * float("-inf"), diagonal=1)
原始形状: [seq_len, seq_len] = [10, 10]
具体数值:
mask = [[ 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, 0, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, 0, 0, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, 0, 0, 0, -inf, -inf, -inf, -inf],
[ 0, 0, 0, 0, 0, 0, 0, -inf, -inf, -inf],
[ 0, 0, 0, 0, 0, 0, 0, 0, -inf, -inf],
[ 0, 0, 0, 0, 0, 0, 0, 0, 0, -inf],
[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
```
### **第1次 unsqueeze(0):添加batch维度**
```python
mask = mask.unsqueeze(0)
新形状: [1, seq_len, seq_len] = [1, 10, 10]
# 现在mask变成:
mask = [[[ 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
...
[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]]
```
### **第2次 unsqueeze(0):添加head维度**
```python
mask = mask.unsqueeze(0)
最终形状: [1, 1, seq_len, seq_len] = [1, 1, 10, 10]
# 现在mask变成4维张量:
mask = [[[[ 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[ 0, 0, 0, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
...
[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]]]
```
## 🔄 **广播机制的工作原理**
### **张量形状对比:**
```python
scores.shape = [2, 4, 10, 10] # [batch, heads, seq, seq]
mask.shape = [1, 1, 10, 10] # [1, 1, seq, seq]
```
### **广播规则:**
```python
# PyTorch会自动扩展较小的维度
scores + mask 时:
mask[1, 1, 10, 10] 自动广播为 → mask[2, 4, 10, 10]
# 相当于:
expanded_mask = mask.expand(2, 4, 10, 10)
```
## 🎭 **具体的广播过程**
### **广播前:**
```python
# 每个batch、每个head都需要相同的因果mask
batch_0_head_0: 需要 [10, 10] 的因果mask
batch_0_head_1: 需要 [10, 10] 的因果mask
batch_0_head_2: 需要 [10, 10] 的因果mask
batch_0_head_3: 需要 [10, 10] 的因果mask
batch_1_head_0: 需要 [10, 10] 的因果mask
...
```
### **广播后:**
```python
# 一个 [1, 1, 10, 10] 的mask被复制到所有位置
final_mask.shape = [2, 4, 10, 10]
# 实际上是:
final_mask[0, 0, :, :] = 原始mask # batch0_head0
final_mask[0, 1, :, :] = 原始mask # batch0_head1
final_mask[0, 2, :, :] = 原始mask # batch0_head2
final_mask[0, 3, :, :] = 原始mask # batch0_head3
final_mask[1, 0, :, :] = 原始mask # batch1_head0
final_mask[1, 1, :, :] = 原始mask # batch1_head1
final_mask[1, 2, :, :] = 原始mask # batch1_head2
final_mask[1, 3, :, :] = 原始mask # batch1_head3
```
## 💡 **为什么这样做?**
### **内存效率:**
```python
# 不需要实际创建 [2, 4, 10, 10] 的大tensor
# 只需 [1, 1, 10, 10],让PyTorch自动广播
```
### **计算便利:**
```python
# 现在可以直接相加:
masked_scores = scores + mask # 形状完全匹配!
```
## 🎯 **总结**
这两次`unsqueeze(0)`操作就是:
**把2D的因果mask → 变成4D的广播兼容mask**
让同一个mask pattern能够应用到所有batch和所有head上,实现"一次定义,到处使用"!🎯