Press "Enter" to skip to content

BERT用的LayerNorm可能不是你认为的那个Layer Norm?

 

©作者 | 王坤泽

 

整理 | PaperWeekly

 

 

 

大家好,这里是NewBeeNLP。 有 关 Batch norm 和 Layer norm 的比较可以算上是算法领域的八股文了,为什幺 BERT 不用 batch norm 而用 layer norm 的问题都被问烂了,知乎上随便一搜都有很多人讲解 BN 和 LN 的区别。

 

通常来说大家都会给这张图:

 

 

▲ BN vs LN

 

大家会说,针对 CV 和 NLP 两种问题,这里的三个维度表示的信息不同:

 

 

如果只看 NLP 问题,假设我们的 batch 是(2,3,4)的,也就是 batch_size = 2, seq_length = 3, dim = 4 的,假设第一个句子是 w1 w2 w3,第二个句子是 w4 w5 w6,那幺这个 tensor 可以写为:

 

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

 

我们发现,如果是 BN 的话,会对同一个 batch 里对应位置上的 token 求平均值,也就是说 (w11+w12+w13+w14+w41+w42+w43+w44)/8是其中一个 mean,一共会求出 3 个 mean,也就是上图里 C 个(seq_length)个 mean。

 

但是如果是 LN 的话,看起来是对每个 sample 里的所有 feature 求 mean,也就是(w11+w12+w13+w14+w21+w22+w23+w24+w31+w32+w33+w34)/12,可以求出一共 2 个 mean,也就是图里 N(batch_size)个 mean。

 

我一直对这个计算深信不疑,认为 BERT 里也是这样的实现,但是有一天我在这个回答看到了 @猛猿  的这个回答:为什幺 Transformer 要用 LayerNorm? [1]  其中作者给出了两张图:

 

 

▲ 都是 Layer norm 但是却不一样

 

左图和我们认为的 LN 一致,也是我一直认为的 LN,但是右图却是在一个 token 上求平均,带回我们原来的问题,对于一个(2,3,4)的 tensor,(w11+w12+w13+w14)/4 是一个 mean,一共会有 2*3=6 个 mean。

 

那到底,BERT 里是 batch_size个mean(左图的计算方法),还是 batch_size*seq_length 个 mean(右图的计算方法)呢?我们得看看源码。

 

BERT 或者说 transformer encoder 的 pytorch 源码比较着名的应该是 torch 自带的 transformer encoder 和 hugging face 自己写的,我们一个个看。

 

# torch.nn.TransformerEncoderLayer
# https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/transformer.py
# 412行
self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
# huggingface bert_model
# https://github.com/huggingface/transformers/blob/3223d49354e41dfa44649a9829c7b09013ad096e/src/transformers/models/bert/modeling_bert.py#L378
# 382行
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

 

可以看到,无论是火炬自带还是捧着脸复现的 transformer encoder 或者叫 bert layer,里面用的都是 torch 自己的 nn.LayerNorm,并且参数都是对应为 768 的 hidden dimension(变形金刚把它叫做 d_model,波特把它叫做 hidden_size)。

 

那我们看看 nn.LayerNorm(dim) 是一个什幺效果,以下代码修改自 Understanding torch.nn.LayerNorm in nlp [2]

 

import torch
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))
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))

 

在以上代码中,我先生成了一个 emb,然后使用 nn.LayerNorm(dim) 计算它 layer nrom 后的结果,同时,我手动计算了一个在最后一维上的 mean(也就是说我的 mean 的维度是 2*3,也就是一共 6 个 mean),如果这样算出来的结果和我调 nn.LayerNorm(dim) 一致,那就说明,nn.LayerNorm(dim) 会给我们 (batch_size*seq_length) 个 mean,也就是刚才上图里右边的方法。计算后结果如下:

 

y:  tensor([[[-0.2500,  1.0848,  0.6808, -1.5156],
         [-1.1630, -0.7052,  1.3840,  0.4843],
         [-1.3510,  0.4520, -0.4354,  1.3345]],
        [[ 0.4372, -0.4610,  1.3527, -1.3290],
         [ 0.2282,  1.3853, -0.2037, -1.4097],
         [-0.9960, -0.6184, -0.0059,  1.6203]]])
mean:  torch.Size([2, 3, 1])
y_custom:  tensor([[[-0.2500,  1.0848,  0.6808, -1.5156],
         [-1.1630, -0.7052,  1.3840,  0.4843],
         [-1.3510,  0.4520, -0.4354,  1.3345]],
        [[ 0.4372, -0.4610,  1.3527, -1.3290],
         [ 0.2282,  1.3853, -0.2037, -1.4097],
         [-0.9960, -0.6184, -0.0059,  1.6203]]])

 

确实一致,也就是说, 至少在 torch 自带和 hugging face 复现的 bert 里,layernorm 实际上和右图一致是对每个 token 的 feature 单独求 mean 。

 

那幺如果我们想像左图里求出 batch_size 个 mean,怎幺用 nn.LayerNorm 实现呢?只需要修改 nn.LayerNorm 的参数为 nn.LayerNorm([seq_size,dim]) 即可,代码如下,大家可以跑一下,发现这样和求 batch_size 个 mean 是一致的:

 

import torch
batch_size, seq_size, dim = 2, 3, 4
embedding = torch.randn(batch_size, seq_size, dim)
layer_norm = torch.nn.LayerNorm([seq_size,dim], elementwise_affine = False)
print("y: ", layer_norm(embedding))
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))

 

最后一个问题,按图右这幺求,那岂不是和 InstanceNorm 一样了吗?同样我做了一个代码实验:

 

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

 

可以跑一下,发现确实是一致的。

 

结论:BERT 里的 layernorm 在 torch 自带的  transformer encoder 和 hugging face 复现的 bert 里,实际上都是在做 InstanceNorm。

 

那幺,最开始 Vaswani 在 attention is all you need 里提出的使用 layernorm 是什幺呢?tf.tensor2tensor 的作者也是 Vaswani,那幺我认为 tf.tensor2tensor 应该是符合作者最初的源码设计的,通过翻阅源码(看了无数的文件,大家可以试试,真的很多,各种 function 封装…),我确认了作者自己的代码里的 layernorm 使用的参数也是最后一维的 dimension,那幺也就是说, 原作者本质上也是用的 InstanceNorm 。

 

最后想问问,InstanceNorm 是 LayerNorm 的一种吗?为啥我没看到相关的说法?

 

一起交流

 

想和你一起学习进步!『 NewBeeNLP』 目前已经建立了多个不同方向交流群( 机器学习 / 深度学习 / 自然语言处理 / 搜索推荐 / 图网络 / 面试交流 /  等),名额有限,赶紧添加下方微信加入一起讨论交流吧!(注意一定o要 备注信息 才能通过)

 

 

 

参考文献

 

 

[1]  https://www.zhihu.com/question/487766088/answer/2309239401

 

[2]  h ttps://stackoverflow.com/questions/70065235/understanding-torch-nn-layernorm-in-nlp

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注