对Transformer中Add&Norm层的理解

时间:2022-10-10 09:50:53

对Add&Norm层的理解

Add操作

首先我们还是先来回顾一下Transformer的结构:Transformer结构主要分为两大部分,一是Encoder层结构,另一个则是Decoder层结构,Encoder 的输入由 Input Embedding 和 Positional Embedding 求和输入Multi-Head-Attention,再通过Feed Forward进行输出。

由下图可以看出:在Encoder层和Decoder层中都用到了Add&Norm操作,即残差连接和层归一化操作。
对Transformer中Add&Norm层的理解
什么是残差连接呢?残差连接就是把网络的输入和输出相加,即网络的输出为F(x)+x,在网络结构比较深的时候,网络梯度反向传播更新参数时,容易造成梯度消失的问题,但是如果每层的输出都加上一个x的时候,就变成了F(x)+x,对x求导结果为1,所以就相当于每一层求导时都加上了一个常数项‘1’,有效解决了梯度消失问题。

Norm操作

首先要明白Norm做了一件什么事,从刚开始接触Transformer开始,我认为所谓的Norm就是BatchNorm,但是有一天我看到了这篇文章,才明白了Norm是什么。

假设我们输入的词向量的形状是(2,3,4),2为批次(batch),3为句子长度,4为词向量的维度,生成以下数据:

[[w11, w12, w13, w14], [w21, w22, w23, w24], [w31, w32, w33, w34]
[w41, w42, w43, w44], [w51, w52, w53, w54], [w61, w62, w63, w64]]

如果是在做BatchNorm(BN)的话,其计算过程如下:BN1=(w11+w12+w13+w14+w41+
w42+w43+w44)/8,同理会得到BN2和BN3,最终得到[BN1,BN2,BN3] 3个mean

如果是在做LayerNorm(LN)的话,则会进如下计算:LN1=(w11+w12+w13+w14+w21+
w22+w23+w24+w31+w32+w33+w34)/12,同理会得到LN2,最终得到[LN1,LN2]两个mean

如果是在做InstanceNorm(IN)的话,则会进如下计算:IN1=(w11+w12+w13+w14)/4,同理会得到IN2,IN3,IN4,IN5,IN6,六个mean,[[IN1,IN2,IN3],[IN4,IN5,IN6]]
下图完美的揭示了,这几种Norm
对Transformer中Add&Norm层的理解
接下来我们来看一下Transformer中的Norm:首先生成[2,3,4]形状的数据,使用原始的编码方式进行编码:

import torch
from torch.nn import InstanceNorm2d
random_seed = 123
torch.manual_seed(random_seed)

batch_size, seq_size, dim = 2, 3, 4
embedding = torch.randn(batch_size, seq_size, dim)

layer_norm = torch.nn.LayerNorm(dim, elementwise_affine = False)
print("y: ", layer_norm(embedding))

输出:

y:  tensor([[[ 1.5524,  0.0155, -0.3596, -1.2083],
         [ 0.5851,  1.3263, -0.7660, -1.1453],
         [ 0.2864,  0.0185,  1.2388, -1.5437]],
        [[ 1.1119, -0.3988,  0.7275, -1.4406],
         [-0.4144, -1.1914,  0.0548,  1.5510],
         [ 0.3914, -0.5591,  1.4105, -1.2428]]])

接下来手动去进行一下编码:

eps: float = 0.00001
mean = torch.mean(embedding[:, :, :], dim=(-1), keepdim=True)
var = torch.square(embedding[:, :, :] - mean).mean(dim=(-1), keepdim=True)

print("mean: ", mean.shape)
print("y_custom: ", (embedding[:, :, :] - mean) / torch.sqrt(var + eps))
mean:  torch.Size([2, 3, 1])
y_custom:  tensor([[[ 1.1505,  0.5212, -0.1262, -1.5455],
         [-0.6586, -0.2132, -0.8173,  1.6890],
         [ 0.6000,  1.2080, -0.3813, -1.4267]],
        [[-0.0861,  1.0145, -1.5895,  0.6610],
         [ 0.8724,  0.9047, -1.5371, -0.2400],
         [ 0.1507,  0.5268,  0.9785, -1.6560]]])

可以发现和LayerNorm的结果是一样的,也就是说明Norm是对d_model进行的Norm,会给我们[batch,sqe_length]形状的平均值。
加下来进行batch_norm,

layer_norm = torch.nn.LayerNorm([seq_size,dim], elementwise_affine = False)
eps: float = 0.00001
mean = torch.mean(embedding[:, :, :], dim=(-2,-1), keepdim=True)
var = torch.square(embedding[:, :, :] - mean).mean(dim=(-2,-1), keepdim=True)

print("mean: ", mean.shape)
print("y_custom: ", (embedding[:, :, :] - mean) / torch.sqrt(var + eps))

输出:

mean:  torch.Size([2, 1, 1])
y_custom:  tensor([[[ 1.1822,  0.4419, -0.3196, -1.9889],
         [-0.6677, -0.2537, -0.8151,  1.5143],
         [ 0.7174,  1.2147, -0.0852, -0.9403]],
        [[-0.0138,  1.5666, -2.1726,  1.0590],
         [ 0.6646,  0.6852, -0.8706, -0.0442],
         [-0.1163,  0.1389,  0.4454, -1.3423]]])

可以看到BN的计算的mean形状为[2, 1, 1],并且Norm结果也和上面的两个不一样,这就充分说明了Norm是在对最后一个维度求平均。
那么什么又是Instancenorm呢?接下来再来实现一下instancenorm

instance_norm = InstanceNorm2d(3, affine=False)
output = instance_norm(embedding.reshape(2,3,4,1)) #InstanceNorm2D需要(N,C,H,W)的shape作为输入
layer_norm = torch.nn.LayerNorm(4, elementwise_affine = False)
print(layer_norm(embedding))

输出:

tensor([[[ 1.1505,  0.5212, -0.1262, -1.5455],
         [-0.6586, -0.2132, -0.8173,  1.6890],
         [ 0.6000,  1.2080, -0.3813, -1.4267]],
        [[-0.0861,  1.0145, -1.5895,  0.6610],
         [ 0.8724,  0.9047, -1.5371, -0.2400],
         [ 0.1507,  0.5268,  0.9785, -1.6560]]])

可以看出无论是layernorm还是instancenorm,还是我们手动去求平均计算其Norm,结果都是一样的,由此我们可以得出一个结论:Layernorm实际上是在做Instancenorm!

如果喜欢文章请点个赞,笔者也是一个刚入门Transformer的小白,一起学习,共同努力。