1. 说明
GPT 模型学习的一步
根据 Andrej Karpathy 大佬的课程视频实践得到
2. Jupyter Notebook
了解 Self-attention 的基本数学思路¶
1. 生成几批序列数据¶
In [4]:
import torch
# 生成一个3维张量作为输入
# batch, time, channels。批次，token位置，token向量各维度参数
B, T, C = 2, 8, 2
x = torch.rand(B,T,C)
# 查看x形状
x.shape
Out[4]:
torch.Size([2, 8, 2])
In [5]:
# 查看x内容，把这看做两批数据(B),每批有8个token(T)，每个token用2个参数描述(C)
x
Out[5]:
tensor([[[0.1720, 0.5959],
[0.7952, 0.6715],
[0.9896, 0.4965],
[0.3968, 0.4749],
[0.3561, 0.1178],
[0.4500, 0.8200],
[0.0763, 0.7931],
[0.8953, 0.0290]],
[[0.8279, 0.8667],
[0.8575, 0.1386],
[0.8771, 0.6202],
[0.0974, 0.0496],
[0.4497, 0.8949],
[0.7783, 0.1133],
[0.0800, 0.3141],
[0.0909, 0.8728]]])
2. 考虑使token注意到前面的token¶
- 我们希望这8个token能够相互交流（注意 attention 到）
- 同时我们希望每个 token 仅注意到它之前的 token (以更像在生成)，如第5个token仅能注意到前4个token
- 最简单的让 token 获取前方信息的方式就是把前面全部取平均，得到一个向量去描述前面信息（平均通常没什么用，且丢失了前方的token的位置信息，但用于示例很方便）。未来再考虑用其他的方式获取前面信息
2.1 先考虑最简单粗暴的取前面平均值的方式：循环¶
In [8]:
# 对于张量可以用这种方式抽取其中一部分
x[0,:2]
Out[8]:
tensor([[0.1720, 0.5959],
[0.7952, 0.6715]])
In [9]:
# bow means bag of words，表征抽取前面的信息
xbow = torch.zeros((B,T,C))
for b in range(B):
for t in range(T):
xprev = x[b,:t+1] # (t,C)
# 对到自己的位置的值取平均
xbow[b,t] = torch.mean(xprev,0)
In [10]:
# 原始输入
x[0]
Out[10]:
tensor([[0.1720, 0.5959],
[0.7952, 0.6715],
[0.9896, 0.4965],
[0.3968, 0.4749],
[0.3561, 0.1178],
[0.4500, 0.8200],
[0.0763, 0.7931],
[0.8953, 0.0290]])
In [11]:
# 对前面取平均后的结果
xbow[0]
Out[11]:
tensor([[0.1720, 0.5959],
[0.4836, 0.6337],
[0.6523, 0.5880],
[0.5884, 0.5597],
[0.5420, 0.4713],
[0.5266, 0.5294],
[0.4623, 0.5671],
[0.5164, 0.4998]])
2.2 换用矩阵乘法取平均的方式¶
循环的方式比较直观，但低效，且不方便推广。 而矩阵是更方便的
In [13]:
# 由于只对前面取平均，可以用 下三角矩阵。0位置自然就会在乘法后被忽略掉
a = torch.tril(torch.ones(4,4))
a
Out[13]:
tensor([[1., 0., 0., 0.],
[1., 1., 0., 0.],
[1., 1., 1., 0.],
[1., 1., 1., 1.]])
In [14]:
# 用这样的矩阵乘一个列向量等效在对列向量进行加权求和了（每行为一系列权重）。但由于我们是在取平均值，所以需要让每一行的和归一化
# 先按行求和
a_sum_by_1 = torch.sum(a,1,keepdim = True)
a_sum_by_1
Out[14]:
tensor([[1.],
[2.],
[3.],
[4.]])
In [15]:
# 按行除
a = a / a_sum_by_1
a
Out[15]:
tensor([[1.0000, 0.0000, 0.0000, 0.0000],
[0.5000, 0.5000, 0.0000, 0.0000],
[0.3333, 0.3333, 0.3333, 0.0000],
[0.2500, 0.2500, 0.2500, 0.2500]])
In [16]:
# 基于上面的思路再求前面平均值的 xbow2
# 先得到权重矩阵
wei = torch.tril(torch.ones(T,T))
wei = wei / wei.sum(1,keepdim = True)
wei
Out[16]:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],
[0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],
[0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],
[0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])
In [17]:
# 再用这个矩阵乘 x 得到 xbow2
# (B,T,T) @ (B,T,C) ---> (B,T,C)
xbow2 = wei @ x
xbow2
Out[17]:
tensor([[[0.1720, 0.5959],
[0.4836, 0.6337],
[0.6523, 0.5880],
[0.5884, 0.5597],
[0.5420, 0.4713],
[0.5266, 0.5294],
[0.4623, 0.5671],
[0.5164, 0.4998]],
[[0.8279, 0.8667],
[0.8427, 0.5026],
[0.8541, 0.5418],
[0.6650, 0.4188],
[0.6219, 0.5140],
[0.6480, 0.4472],
[0.5668, 0.4282],
[0.5073, 0.4838]]])
In [18]:
# 对比两种方式计算的xbow，应该是一样的
torch.allclose(xbow,xbow2)
Out[18]:
True
2.3 进一步我们利用 softmax 函数¶
softmax 可以把一堆实数映射到 [0,1] 区间，并保证映射后的和为1，非常适合用于计算概率
In [20]:
from torch.nn import functional as F
tril = torch.tril(torch.ones(T,T))
wei = torch.zeros((T,T))
# softmax 会把 -inf 映射为0，所以需要把希望为0的地方置为 -inf
wei = wei.masked_fill(tril == 0,float('-inf'))
# 对各个最后一维，即行，进行softmax
wei = F.softmax(wei,dim=-1)
wei
Out[20]:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],
[0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],
[0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],
[0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])
可见可以得到和前面用sum除一样的结果
In [22]:
# 同样方式计算xbow3
xbow3 = wei @ x
xbow3
Out[22]:
tensor([[[0.1720, 0.5959],
[0.4836, 0.6337],
[0.6523, 0.5880],
[0.5884, 0.5597],
[0.5420, 0.4713],
[0.5266, 0.5294],
[0.4623, 0.5671],
[0.5164, 0.4998]],
[[0.8279, 0.8667],
[0.8427, 0.5026],
[0.8541, 0.5418],
[0.6650, 0.4188],
[0.6219, 0.5140],
[0.6480, 0.4472],
[0.5668, 0.4282],
[0.5073, 0.4838]]])
In [23]:
# xbow3 也能和前两种方式结果一致
torch.allclose(xbow,xbow3)
Out[23]:
True
3. 用self-attention的方式¶
In [25]:
# 扩展一下输入数据的大小，重新设置输入数据
B, T, C = 4, 8, 32
x = torch.rand(B,T,C)
x.shape
Out[25]:
torch.Size([4, 8, 32])
3.1 实现不带v的单头注意力机制¶
In [27]:
import torch.nn as nn
from torch.nn import functional as F
# 注意头的大小设置
head_size = 16
# 每个头需要一个 key 矩阵，作用到x上以后提取x的特征信息
key = nn.Linear(C, head_size, bias = False)
# 每个头需要一个 query 矩阵，作用到x上以后提取想问的问题
query = nn.Linear(C, head_size, bias = False)
# 作用于x得到具体的key和query
k = key(x) # (B,T,16)
q = query(x) # (B,T,16)
In [28]:
k.shape
Out[28]:
torch.Size([4, 8, 16])
In [29]:
q.shape
Out[29]:
torch.Size([4, 8, 16])
In [30]:
# 让 k和q进行内积，获得key和query到底有多么匹配
# 为了求内积，需要转置一下后两维
# (B,T,16) @ (B,16,T) --> (B,T,T)
# 注：忽略B以后，k 和 q 都是由T个行向量组成的矩阵，内积有T*T组，对应T*T结果
# 这个 wei 类似之前的取平均的注意力矩阵类似，表征每个token应该有多么关心其他各个token
wei = q @ k.transpose(-2,-1)
wei.shape
Out[30]:
torch.Size([4, 8, 8])
In [31]:
# 查看第一个batch的wei的权重矩阵
wei[0]
Out[31]:
tensor([[ 0.6487, 0.7615, 0.1522, 0.4993, 0.2020, 0.5022, 0.2302, 0.4878],
[ 0.9861, 0.8410, 0.5973, 1.0108, 0.5287, 0.6795, 0.5102, 0.6981],
[ 0.7705, 0.7264, 0.2658, 0.6753, 0.0797, 0.3709, 0.1030, 0.3787],
[ 0.4662, 0.8602, 0.0690, 0.3791, -0.0987, 0.1380, 0.2878, 0.5163],
[ 0.2727, 0.4088, -0.1478, 0.2061, 0.0704, -0.0910, -0.2079, 0.0758],
[ 0.4846, 0.7829, 0.1513, 0.4147, -0.0211, 0.2098, 0.3993, 0.5004],
[ 0.3002, 0.3277, 0.0349, 0.2555, -0.2753, 0.0648, 0.1008, 0.2496],
[ 0.3715, 0.3061, 0.1405, 0.4753, -0.1130, 0.3157, 0.3514, 0.2744]],
grad_fn=<SelectBackward0>)
In [32]:
# 类似之前2.3的做法，我们让各个token仅关注其前面的token，并且用softmax归一化
tril = torch.tril(torch.ones(T,T))
wei = wei.masked_fill(tril == 0,float('-inf'))
wei = F.softmax(wei,dim=-1)
wei[0]
Out[32]:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.5362, 0.4638, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.3905, 0.3737, 0.2358, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2456, 0.3642, 0.1651, 0.2251, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2195, 0.2515, 0.1442, 0.2054, 0.1793, 0.0000, 0.0000, 0.0000],
[0.1866, 0.2514, 0.1337, 0.1740, 0.1125, 0.1417, 0.0000, 0.0000],
[0.1688, 0.1735, 0.1295, 0.1615, 0.0950, 0.1334, 0.1383, 0.0000],
[0.1372, 0.1285, 0.1089, 0.1522, 0.0845, 0.1297, 0.1345, 0.1245]],
grad_fn=<SelectBackward0>)
In [33]:
# 计算加权平均
out = wei @ x
out.shape
Out[33]:
torch.Size([4, 8, 32])
3.2 实现带v的单头注意力机制¶
刚才的做法下，仅基于key 和 query 的内积
self-attention 还进一步对 x 进行了一次变换，用另一个矩阵v
合在一起的思路是：
- query 是我想问我自己的一批问题
- key 是我基于我自己的信息的一批特征
- query @ key 得到我想关注我的哪些地方（权重表示）
- value 表示我各个位置想提供什么样的信息
整合输出的注意信息就是 out = (query(x) @ key(x).T)*value(x)
In [35]:
# 注意头的大小设置
head_size = 16
# 每个头需要一个 key 矩阵，作用到x上以后提取x的特征信息
key = nn.Linear(C, head_size, bias = False)
# 每个头需要一个 query 矩阵，作用到x上以后提取想问的问题
query = nn.Linear(C, head_size, bias = False)
# key,query作用于x得到具体的key和query
k = key(x) # (B,T,16)
q = query(x) # (B,T,16)
wei = q @ k.transpose(-2,-1)
tril = torch.tril(torch.ones(T,T))
wei = wei.masked_fill(tril == 0,float('-inf'))
wei = F.softmax(wei,dim=-1)
# 额外再加一个value矩阵直接对x进行作用
value = nn.Linear(C, head_size, bias = False)
# v 也直接作用于x
v = value(x) # (B,T,16)
out = wei @ v
In [36]:
out.shape
Out[36]:
torch.Size([4, 8, 16])
3.3 添加系数控制方差¶
在 attention is all your need 的论文中，计算out的时候还会除以 head_size^0.5 这么个系数
原因是为了控制矩阵的方差不因计算而膨胀
举例子：
In [38]:
k = torch.randn(B,T,head_size)
q = torch.randn(B,T,head_size)
wei = q @ k.transpose(-2,-1)
In [39]:
# torch.randn生成的都是均值0方差1的随机数，方差期望为1
k.var()
Out[39]:
tensor(0.9876)
In [40]:
# torch.randn生成的都是均值0方差1的随机数，方差期望为1
q.var()
Out[40]:
tensor(1.0327)
In [41]:
# 但计算后的wei的方差的期望却膨胀到了 head_size 倍
wei.var()
Out[41]:
tensor(15.7463)
方差变大的原因如下：
- 如果 X,Y随机变量都满足 均值0方差1，可计算 X*Y 也满足均值0方差1
- 由于计算 wei 即k和q的内积的时候相当于是计算了 head_size 组不同 X*Y 的和。
- 求和会导致方差为各随机变量的方差相加，于是方差的期望变为 head_size
方差变大会有什么后果：
- 而方差增加会导致每一层的权重在方差上不一致
- 特别当方差非常大时，对应到softmax函数（scale敏感），会产生尖锐集中的分布
- 这容易导致梯度消失，进而导致训练效果受限
- 同时我们也不希望每个 token 仅向少量几个其他 token 获取信息（wei的分布不应太尖锐和集中）
- 所以需要添加一个系数去控制方差
举个例子看softmax的尺度敏感性
In [43]:
# 生成一个一维张量
a = torch.randn(T)
a
Out[43]:
tensor([ 0.7613, 0.4432, 0.9386, -0.0056, -0.5113, -0.7695, 0.3200, -0.9199])
In [44]:
# 取softmax
b = torch.softmax(a,dim = -1)
b
Out[44]:
tensor([0.2122, 0.1544, 0.2534, 0.0986, 0.0594, 0.0459, 0.1365, 0.0395])
In [45]:
# 把所有元素扩大10倍以后取softmax
# 可以看到往a里面最大地那个值，占有了绝大部分b的结果，即集中度更高了
# 这是由于softmax是指数函数作用后的归一，而指数函数是 scale 敏感的
b = torch.softmax(a*10,dim = -1)
b
Out[45]:
tensor([1.4397e-01, 5.9852e-03, 8.4824e-01, 6.7290e-05, 4.2832e-07, 3.2393e-08,
1.7459e-03, 7.1994e-09])
由于计算已知方差扩大了 head_size 倍，所以仅需对结果除以 head_size^0.5 则能达到控制方差的目标。
所以修改后的 attention 代码如下
In [47]:
# 注意头的大小设置
head_size = 16
# 每个头需要一个 key 矩阵，作用到x上以后提取x的特征信息
key = nn.Linear(C, head_size, bias = False)
# 每个头需要一个 query 矩阵，作用到x上以后提取想问的问题
query = nn.Linear(C, head_size, bias = False)
# key,query作用于x得到具体的key和query
k = key(x) # (B,T,16)
q = query(x) # (B,T,16)
wei = q @ k.transpose(-2,-1)
tril = torch.tril(torch.ones(T,T))
wei = wei.masked_fill(tril == 0,float('-inf'))
wei = F.softmax(wei,dim=-1)
# 额外再加一个value矩阵直接对x进行作用
value = nn.Linear(C, head_size, bias = False)
# v 也直接作用于x
v = value(x) # (B,T,16)
# 计算注意后的结果，并且乘以系数稳定方差的期望
out = wei @ v * head_size ** -0.5