文本的操作在github上都有Shirley-Xie/pytorch_exercise · GitHub,且有运行结果。
本文代码从下往上写,从小往大的写。
首先我们看到整个模型由不同的组件组成,其中最重要的部分是多头注意力部分,我们首先从这部分下手。
1. Multi-Head Attention
将多头拆解为自注意力和多头两部分。在代码方面不仅会对文本序列进行讲解还会带上图像的讲解,本质是一样的,只是处理细节会有不一样。
1.1 缩放点积注意力
注意力函数可以描述为将查询(query)和一组键值(key-value)对映射到输出,其中查询、键、值和输出都是向量。输出被计算为值的加权和,其中分配给每个值的权重由查询与相应键的兼容性函数计算。将这个特别的注意力称为“Scaled Dot-Product Attention”。输入是query和key的d_k,和vlue的维度d_v。我们用所有键计算查询的点积,每个除以d_k,并应用softmax函数来获得这些值的权重。
首先我们确认输入的三个变量,query,key, value,每个张量的维度都是[batch_size, seq_length, features],那么公式的右侧d_k也就是query和key的features维度。代码如下:
d_k = query.size(-1)
# 右边括号里的部分,转置最后两维
scores = query@key.tanspose(-2,-1)/d_k**0.5
#另一种写法 torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
# softmax的部分
p_attn = scores.softmax(dim=-1)
# 最后乘以value
p_attn@value
需要的注意的是,Transformer的多头注意力在Decoder那里还有Masked Mulit-head Attention,所以要考虑进去。将为mask==0的位置将scores替换为-1000000000。所以完整版为:
def ScaledDotProductAttention(query, key, value, mask=None, dropout=None):
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1)
scores = query@key.transpose(-2, -1)) / d_k**0.5
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = scores.softmax(dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
1. 2 多头注意力
计算完缩放点积注意力,看看多头注意力如何处理。
,Q,K,V维度是bsd_model,其中
,
。这部分使用nn.Linear进行操作。
,
1. 在缩放注意力之前需要批量进行所有线性投影,并将输入数据进行切分为多头。
# 创建线性层,因为输入三个线性层,输出一个,总共四个
module = nn.Linear(d_model,d_model)
N = 4
d_k = d_model//h
linears = nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
# 将三个输入向量分别输入到线性层,一一对应
nbatches = query.size(0)
query, key, value = [l(x).view(nbatches,-1,h,d_k).transpose(1,2) for l, x in zip(linears, (query, key, value))]
输入的维度的变化为:bsd_model -> bshd_k -> bhsd_k,其中b:nbatches, s:seq_length。
从维度的变化可以看出,将输入的维度d_model切分为h个头d_k大小,然后将头和句子长度变换一下位置,数量维度(batch,head)优先的原则。
2. 将线性投影后的张量输入到缩放注意力中。
SDPAttn = ScaledDotProductAttention()
x, sttn = SDPAttn(query, key, value)
3. 连接之后进行线性处理。
进行注意力计算之后将维度变换回来,bhsd_k -> bshd_k ->bsd_model。就是注意力之前的样子,一切的变化只为了注意力的计算。变换后进行线性运算结束。
x = x.transpose(1, 2).contiguous().view(nbatches, -1, h * d_k)
linears[-1](x)
将代码进行汇总:
class MultiHeadAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
"Take in model size and number of heads."
super(MultiHeadAttention, self).__init__()
assert d_model % h == 0
# We assume d_v always equals d_k
self.d_k = d_model // h
self.h = h
self.linears = nn.ModuleList([deepcopy(nn.Linear(d_model, d_model)) for _ in range(4)])
self.attn = None
self.dropout = nn.Dropout(p=dropout)
self.attention = ScaledDotProductAttention()
def forward(self, query, key, value, mask=None):
"Implements Figure 2"
if mask is not None:
# Same mask applied to all h heads.
mask = mask.unsqueeze(1)
nbatches = query.size(0)
# 1) Do all the linear projections in batch from d_model => h x d_k
query, key, value = \
[l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))]
# 2) Apply attention on all the projected vectors in batch.
x, self.attn = self.attention(query, key, value, mask=mask,
dropout=self.dropout)
# 3) "Concat" using a view and apply a final linear.
x = x.transpose(1, 2).contiguous() \
.view(nbatches, -1, self.h * self.d_k)
return self.linears[-1](x)
图块序列(VIT)的代码如下:
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
# NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.attn_gradients = None
self.attention_map = None
def save_attn_gradients(self, attn_gradients):
self.attn_gradients = attn_gradients
def save_attention_map(self, attention_map):
self.attention_map = attention_map
def forward(self, x, register_hook=False):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
# BN(3*C)->BN3HC'->3BHNC'
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
#BHNC'
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
if register_hook:
self.save_attention_map(attn)
attn.register_hook(self.save_attn_gradients)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
可以看出线性处理方面不一样,他将三个张量qkv和在一起进行运算。这个没考虑mask ,最后进行一个dropout运算,其他的地方一样。维度变换是,Attention之前是BNC,qkv合起来之后线性变换为BN(3*C)->BN3HC'->3BHNC,最后为q维度为BHNC',也是数据维度(3BH)优先排列。输入到Attention之后原路返回,BHNC'-> BNHC' -> BNC。
补充:
这里对attn使用了register_hook,这里的钩子是为了可以反向传播的时候可以访问这层的梯度,甚至是修改,进而影响到浅层的梯度传播。
官方示例,将梯度放大了两倍:
import torch
v = torch.tensor([0., 0., 0.], requires_grad=True)
h = v.register_hook(lambda grad: grad * 2) # double the gradient
v.backward(torch.tensor([1., 2., 3.]))
print(v.grad)
h.remove()
# tensor([2., 4., 6.])
咱们可以在cache中获取到梯度了。因为即使是hook住的张量梯度,也会在运行结束后立刻释放掉,所以想要保存梯度,一定要在func函数里进行。比如
cache = []
def func(grad):
cache.append(grad)
return grad
1.3 Attention问答
1.3.1 Self-Attention表达式是?一定要这样表达吗?
。
不一定,只要可以建模相关性就好了。最好能够高速计算(矩阵乘法),并且表达能力强(query可以主动去关注到其他的key并在value上进行强化,并且忽略不相关的其他部分)。模型容量够(引入了project_q/k/v,att_out,多头)。
1.3.2 为什么要对QK进行scaling?为什么要除以
?有其他方法不用除以
吗?
为了避免
变得很大时softmax函数的梯度趋于0。scaling可以视为将QK进行归一化,使得输入的数据的分布变得更好。从而进行softmax的时候不至于QK极大或极小的时候落在softmax的极端区域而导致梯度趋于0,scaling可以防止梯度消失,让模型能够更容易训练。 有不用除以
,只要能缓解梯度消失的问题就可以。详情可以了解Google T5的Xavier初始化。 (具体而言,假设Q和K中的取出的两个向量q和k的每个元素值都是正态随机分布,那么两个向量做点积,会得到d_k个正态随机变量的和,其方差是原来的d_k倍,标准差是原来的
倍。如果不做scale, 当d_k很大时,求得的
元素的绝对值容易很大,导致落在softmax的极端区域(趋于0或者1),极端区域softmax函数的梯度值趋于0,不利于模型学习。除以
,恰好做了归一,不受d_k变化影响。)
1.3. 3 MultiHeadAttention的参数数量和head数量有何关系?
MultiHeadAttention的参数数量和head数量无关。多头注意力的参数来自对QKV的三个变换矩阵以及多头结果concat后的输出变换矩阵。假设嵌入向量的长度是d_model, 一共有h个head. 对每个head,
这三个变换矩阵的尺寸都是 d_model×(d_model/h),所以h个head总的参数数量就是3×d_model×(d_model/h)×h = 3×d_model×d_model。它们的输出向量长度都变成 d_model/h,经过attention作用后向量长度保持,h个head的输出拼接到一起后向量长度还是d_model,所以最后输出变换矩阵的尺寸是d_model×d_model。因此,MultiHeadAttention的参数数量为 4×d_model×d_model,和head数量无关。
1.3.3 transformer为什么要用三个不一样的QKV?
是为了增强网络的容量和表达能力。更极端点,如果完全不要project_q/k/v,就是输入x本身来做,当然可以,但是表征能力太弱了(x的参数更新得至少会很拧巴)
1.3.4 为什么要多头?举例说明多头相比单头注意力的优势。
进一步增强网络的容量和表达能力。你可以类比CV中的不同的channel(不同卷积核)会关注不同的信息,事实上不同的头也会关注不同的信息,不同头关注的内容是很抽象的。
你当然可以就用一个头同时做这个事,但是还是这个道理,我们的目的就是通过增加参数量来增强网络的容量从而提升网络表达能力。
经过多头之后,我们还需要att_out线性层来做线性变换,以自动决定(通过训练)对每个头的输出赋予多大的权重,从而在最终的输出中强调一些头的信息,而忽视其他头的信息。这是一种自适应的、数据驱动的方式来组合不同头的信息。
2. Feed Forward
也就是全连接前向网络,对每个position的向量分别进行相同的操作,包括两个线性变换和一个ReLU激活输出。
class PositionwiseFeedForward(nn.Module):
"Implements FFN equation."
def __init__(self, d_model, d_ff, dropout=0.1):
super(PositionwiseFeedForward, self).__init__()
self.w_1 = nn.Linear(d_model, d_ff)
self.w_2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.w_2(self.dropout(self.w_1(x).relu()))
图块序列,ViT版本,只是最后多加了一层dropout。
class Mlp(nn.Module):
""" MLP as used in Vision Transformer, MLP-Mixer and related networks
"""
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
就此核心的两个部分注意力和前馈全连接层已经讲完,接下来可以用Add & Norm 进行拼接,便可形成EncoderLayer和DecoderLayer。
3. Add & Norm
3.1 LayerNorm
代码如下:
class LayerNorm(nn.Module):
"Construct a layernorm module (See citation for details)."
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
self.eps = eps
def forward(self, x):
mean = x.mean(-1, keepdim=True) # batchNorm的维度是0
std = x.std(-1, keepdim=True)
return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
3.2 SublayerConnection
从图中可以看出所有的结构都是子层加上Add &Norm。还有一种是掩码注意力,可以包含在注意力中。
因后来证明Norm放在前面效果更好,所以放在子层的前面。代码如下:
class SublayerConnection(nn.Module):
"""
A residual connection followed by a layer norm.
Note for code simplicity the norm is first as opposed to last.
"""
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, sublayer):
"Apply residual connection to any sublayer with the same size."
return x + self.dropout(sublayer(self.norm(x)))
3.3 Layer Norm问答
3.3.1 为什么transformer用Layer Norm?有什么用?
任何norm的意义都是为了让网络输入数据分布变的更好,也就是转换为标准正态分布。数值进入敏感度区间,以减缓梯度消失,从而更容易训练。当然,这也意味着舍弃了除此维度之外其他维度的其他信息。为什么能舍弃呢?请看下一题。
3.3.2 为什么不用BN?
如果在同一维度内进行normalization,那么在这个维度内,相对大小有意义的,是可以比较的;但是在normalization后的不同的维度之间,相对大小这是没有意义的。
BN(batch normalization)广泛应用于CV,针对同一特征,以跨样本的方式开展归一化,也就是对不同样本的同一通道间的所有像素值进行归一化,因此不会破坏不同样本同一特征之间的关系,毕竟“减均值,除标准差”只是一个平移加缩放的线性操作。这一性质进而决定了经过归一化操作后,样本之间仍然具有可比较性。但是,特征与特征之间的不再具有可比较性,也就是上一个问题中我所说的“舍弃了除此维度之外其他维度的其他信息”。
NLP中不用BN,而用LN呢?
- 对不同样本同一特征的信息进行归一化没有意义:
- 三个样本(为中华之崛起而读书;我爱中国;母爱最伟大)中,“为”、“我”、“母”归一到同一分布没有意义。
- 舍弃不了BN中舍弃的其他维度的信息,也就是同一个样本的不同维度的信息:
- “为”、“我”、“母”归一到同一分布后,第一句话中的“为”和“中”就没有可比性了,何谈同一句子之间的注意力机制?
CV中:
- 对不同样本同一特征(channel)的信息进行归一化有意义:
- 因为同一个channel下的所有信息都是遵循统一规则下的大小比较的,比如黑白图中越白越靠近255,反之越黑越靠近0
- 可以舍弃其他维度的信息,也就是同一个样本的不同维度间(channel)的信息:
- 举例来说,RGB三个通道之间互相比较意义不大
4. 编码器
4.1 EncoderLayer
现在个层已经就位,开始拼接。这个EncoderLayer可以分为上下两个部分,两个子层的结构是一直的,只是中间核心层不同。
# 定义一个clones函数,来更方便的将某个结构复制若干份
def clones(module, N):
"Produce N identical layers."
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
class EncoderLayer(nn.Module):
"Encoder is made up of self-attn and feed forward (defined below)"
def __init__(self, size, self_attn, feed_forward, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 2)
self.size = size
def forward(self, x, mask):
x = self.sublayer[0](x, lambda x:self.self_attn(x,x,x,mask))
return self.sublayer[1](x,self.feed_forward)
4.2 Encoder
class Encoder(nn.Module):
"""
Encoder
The encoder is composed of a stack of N=6 identical layers.
"""
def __init__(self, layer, N):
super(Encoder, self).__init__()
# 调用时会将编码器层传进来,我们简单克隆N分,叠加在一起,组成完整的Encoder
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def forward(self, x, mask):
"Pass the input (and mask) through each layer in turn."
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
5.解码器
5.1 DecoderLayer
class DecoderLayer(nn.Module):
"Decoder is made of self-attn, src-attn, and feed forward (defined below)"
def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = self_attn
self.src_sttn = src_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3)
self.size = size
def forward(self, x, mermory, src_mask, tgt_mask):
#forward函数中的参数有4个,分别是来自上一层的输入x,来自编码器层的语义存储变量memory,以及源数据掩码张量和目标数据掩码张量,将memory表示成m之后方便使用。
m = memory
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
#接着进入第二个子层,这个子层中常规的注意力机制,q是输入x;k,v是编码层输出memory,同样也传入source_mask,但是进行源数据遮掩的原因并非是抑制信息泄露,而是遮蔽掉对结果没有意义的padding。
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
#最后一个子层就是前馈全连接子层,经过它的处理后就可以返回结果,这就是我们的解码器结构
return self.sublayer[2](x, self.feed_forward)
5.2 Encoder
class Decoder(nn.Module):
"Generic N layer decoder with masking."
def __init__(self, layer, N):
super(Decoder, self).__init__()
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def forward(self, x, memory, src_mask, tgt_mask):
#memory是编码器层的输出,source_mask,target_mask代表源数据和目标数据的掩码张量,
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)
6. 输入和输出
输入部分包含两个模块,Embedding
和 Positional Encoding
。
6.1 Embeddings层
Embedding层的作用是将某种格式的输入数据,例如文本,转变为模型可以处理的向量表示,来描述原始数据所包含的信息。
Embedding层输出的可以理解为当前时间步的特征,如果是文本任务,这里就可以是Word Embedding,如果是其他任务,就可以是任何合理方法所提取的特征。
class Embeddings(nn.Module):
def __init__(self, d_model, vocab):
super(Embeddings, self).__init__()
self.lut = nn.Embedding(vocab, d_model)
self.d_model = d_model
def forward(self, x):
return self.lut(x) * math.sqrt(self.d_model)
6.2 Positional Encodding
Positional Encodding位置编码的作用是为模型提供当前时间步的前后出现顺序的信息。因为Transformer不像RNN那样的循环结构有前后不同时间步输入间天然的先后顺序,所有的时间步是同时输入,并行推理的,因此在时间步的特征中融合进位置编码的信息是合理的。
位置编码可以有很多选择,可以是固定的,也可以设置成可学习的参数。
这里,我们使用固定的位置编码。具体地,使用不同频率的sin和cos函数来进行位置编码,如下所示:
其中pos代表时间步的下标索引,向量 也就是第pos个时间步的位置编码,编码长度同Embedding层,这里我们设置的是512。上面有两个公式,代表着位置编码向量中的元素,奇数位置和偶数位置使用两个不同的公式。
class PositionalEncoding(nn.Module):
"Implement the PE function."
def __init__(self, d_model, dropout, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# Compute the positional encodings once in log space.
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
def forward(self, x):
x = x + self.pe[:, : x.size(1)].requires_grad_(False)
return self.dropout(x)
6.3 Softmax
#将线性层和softmax计算层一起实现,因为二者的共同目标是生成最后的结构
#因此把类的名字叫做Generator,生成器类
class Generator(nn.Module):
"Define standard linear + softmax generation step."
def __init__(self, d_model, vocab):
#初始化函数的输入参数有两个,d_model代表词嵌入维度,vocab.size代表词表大小
super(Generator, self).__init__()
#首先就是使用nn中的预定义线性层进行实例化,得到一个对象self.proj等待使用
#这个线性层的参数有两个,就是初始化函数传进来的两个参数:d_model,vocab_size
self.proj = nn.Linear(d_model, vocab)
def forward(self, x):
#前向逻辑函数中输入是上一层的输出张量x,在函数中,首先使用上一步得到的self.proj对x进行线性变化,然后使用F中已经实现的log_softmax进行softmax处理。
return F.log_softmax(self.proj(x), dim=-1)
最后整个模型代码:
def make_model(
src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
"Helper: Construct a model from hyperparameters."
c = copy.deepcopy
attn = MultiHeadedAttention(h, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
model = EncoderDecoder(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
Generator(d_model, tgt_vocab),
)
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
return model
7. Transformer面试八股
来到这里,你认为你懂了transformer了吗?让我来个灵魂考问。
7.1 Transformer是如何解决长距离依赖的问题的?
Transformer是通过引入Scale-Dot-Product注意力机制来融合序列上不同位置的信息,从而解决长距离依赖问题。以文本数据为例,在循环神经网络LSTM结构中,输入序列上相距很远的两个单词无法直接发生交互,只能通过隐藏层输出或者细胞状态按照时间步骤一个一个向后进行传递。对于两个在序列上相距非常远的单词,中间经过的其它单词让隐藏层输出和细胞状态混入了太多的信息,很难有效地捕捉这种长距离依赖特征。但是在Scale-Dot-Product注意力机制中,序列上的每个单词都会和其它所有单词做一次点积计算注意力得分,这种注意力机制中单词之间的交互是强制的不受距离影响的,所以可以解决长距离依赖问题。
7.2 Self-Attention相关问题
7.2.1 Self-Attention表达式是?一定要这样表达吗?
。
不一定,只要可以建模相关性就好了。最好能够高速计算(矩阵乘法),并且表达能力强(query可以主动去关注到其他的key并在value上进行强化,并且忽略不相关的其他部分)。模型容量够(引入了project_q/k/v,att_out,多头)。
7.2.2 为什么要对QK进行scaling?为什么要除以
?有其他方法不用除以
吗?
为了避免
变得很大时softmax函数的梯度趋于0。scaling可以视为将QK进行归一化,使得输入的数据的分布变得更好。从而进行softmax的时候不至于QK极大或极小的时候落在softmax的极端区域而导致梯度趋于0,scaling可以防止梯度消失,让模型能够更容易训练。 有不用除以
,只要能缓解梯度消失的问题就可以。详情可以了解Google T5的Xavier初始化。 (具体而言,假设Q和K中的取出的两个向量q和k的每个元素值都是正态随机分布,那么两个向量做点积,会得到d_k个正态随机变量的和,其方差是原来的d_k倍,标准差是原来的
倍。如果不做scale, 当d_k很大时,求得的
元素的绝对值容易很大,导致落在softmax的极端区域(趋于0或者1),极端区域softmax函数的梯度值趋于0,不利于模型学习。除以
,恰好做了归一,不受d_k变化影响。)
7.3