在读一篇关于模型蒸馏的知乎文章《文本 × 分类：简单网络可以替代BERT》的时候我有所启发，文中说蒸馏到了一个小模型中，小模型的开头是BERT的Embedding层直接拿过来的，并且删除了拖慢速度的 nn.LayerNorm 和 position_embedding 等。
所以到底有没有必要删除源代码中的这些部分呢？

原始BertEmbeddings

首先看看 🤗 的源码，其中 config 下的一些参数放在这里，当然也可以看源文档 BertConfig ：

  • vocab_size：21128，对于中文BERT来说是21128个token
  • hidden_size：768
  • pad_token_id：0
  • max_position_embeddings：512，也就是最大序列长度
  • type_vocab_size：2，也就是上下句的标记种类
  • layer_norm_eps：层归一化的 蒸馏时BERT的Embedding处理 - 图1 参数

  • hidden_dropout_prob：0.1 的 dropout 率

    1. class BertEmbeddings(nn.Module):
    2.   """Construct the embeddings from word, position and token_type embeddings."""
    4.   def __init__(self, config):
    5.       super().__init__()
    6.       self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
    7.       self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
    8.       self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
    10.       # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
    11.       # any TensorFlow checkpoint file
    12.       self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
    13.       self.dropout = nn.Dropout(config.hidden_dropout_prob)
    15.       # position_ids (1, len position emb) is contiguous in memory and exported when serialized
    16.       self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
    18.   def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
    19.       if input_ids is not None:
    20.           input_shape = input_ids.size()
    21.       else:
    22.           input_shape = inputs_embeds.size()[:-1]
    24.       seq_length = input_shape[1]
    26.       if position_ids is None:
    27.           position_ids = self.position_ids[:, :seq_length]
    29.       if token_type_ids is None:
    30.           token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
    32.       if inputs_embeds is None:
    33.           inputs_embeds = self.word_embeddings(input_ids)
    34.       position_embeddings = self.position_embeddings(position_ids)
    35.       token_type_embeddings = self.token_type_embeddings(token_type_ids)
    37.       embeddings = inputs_embeds + position_embeddings + token_type_embeddings
    38.       embeddings = self.LayerNorm(embeddings)
    39.       embeddings = self.dropout(embeddings)
    40.       return embeddings

    简化BertEmbeddings

    去掉position和token_type

    在通常任务中，我们是不需要 position_ids 和 token_type_ids 这两个输入参数的，因此可以直接干掉。
    经过化简以后的代码如下： ```python from torch import Tensor from torch import nn

class BertEmbeddings(nn.Module): “”” Construct the embeddings from word, position and token_type embeddings. 魔改一下，只保留word_embeddings和layer_norm “””

  1. def __init__(self):
  2.     super().__init__()
  3.     self.word_embeddings = nn.Embedding(21128, 768, padding_idx=0)
  5.     # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
  6.     # any TensorFlow checkpoint file
  7.     self.LayerNorm = nn.LayerNorm(768, eps=1e-12)
  8.     self.dropout = nn.Dropout(0.1)
  10. def forward(self, input_ids: Tensor = None) -> Tensor:
  11.     """
  12.     前向传播过程。参考输入如下：
  14.     >>> test_seq = [ 101,  704, 1290, 1957, 2094, 2110, 7368, 8038, 3315, 4906, 2231,
  15.                      3613,  788,  122,  683,  689, 2875, 4511, 4495,  102,    0,    0,
  16.                      0,    0,    0,    0,    0,    0,    0,    0,    0,    0]
  17.     >>> test_seq_inputs = torch.tensor(test_seq).clone().detach().unsqueeze(0)
  18.     >>> print(test_seq_inputs)
  19.     tensor([[ 101,  704, 1290, 1957, 2094, 2110, 7368, 8038, 3315, 4906, 2231, 3613,
  20.               788,  122,  683,  689, 2875, 4511, 4495,  102,    0,    0,    0,    0,
  21.               0,    0,    0,    0,    0,    0,    0,    0]])
  22.     >>> bert_embedding = BertEmbeddings()
  23.     >>> res = bert_embedding(input_ids=test_seq_inputs)
  24.     >>> print(res)
  25.     torch.Size([1, 32])
  26.     tensor([[[ 0.8166, -1.1873,  1.2912,  ...,  1.6106, -0.4729, -1.5941],
  27.                 [-0.2341, -0.1526,  0.3441,  ..., -1.6328,  1.7635,  0.4524],
  28.                 [ 1.2200,  1.1668, -0.1776,  ...,  0.5772, -0.5337, -1.8260],
  29.                 ...,
  30.                 [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
  31.                 [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
  32.                 [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],
  33.         grad_fn=<MulBackward0>)
  35.     >>> print(res.size())
  36.     torch.Size([1, 32, 768])
  37.     ```
  39.     Args:
  40.         input_ids (Tensor, optional): 句子的 token_ids. Defaults to None.
  42.     Raises:
  43.         ValueError: 如果输入不是ids则报错。注意这里修改为不接受 embedding 输入
  45.     Returns:
  46.         Tensor: 返回 embedding 句子结果
  47.     """
  48.     if input_ids is not None:
  49.         input_shape = input_ids.size()
  50.     else:
  51.         raise ValueError("检查input_ids是否正确，不接受传现成的embedding")
  52.     # print(input_shape)
  53.     inputs_embeds = self.word_embeddings(input_ids)
  55.     embeddings = inputs_embeds
  56.     embeddings = self.LayerNorm(embeddings)
  57.     embeddings = self.dropout(embeddings)
  58.     return embeddings
  1. <a name="aS5Ox"></a>
  2. ## 去掉LayerNorm
  3. ```python
  4. from torch import Tensor
  5. from torch import nn
  8. class BertEmbeddingsWithoutLayerNorm(nn.Module):
  9.     """
  10.     Construct the embeddings from word, position and token_type embeddings.
  11.     魔改一下，只保留word_embeddings，连LayerNorm也删掉
  12.     """
  14.     def __init__(self):
  15.         super().__init__()
  16.         self.word_embeddings = nn.Embedding(21128, 768, padding_idx=0)
  17.         self.dropout = nn.Dropout(0.1)
  19.     def forward(self, input_ids: Tensor = None) -> Tensor:
  20.         """
  21.         前向传播过程。参考输入如下：
  23.         >>> test_seq = [ 101,  704, 1290, 1957, 2094, 2110, 7368, 8038, 3315, 4906, 2231,
  24.                          3613,  788,  122,  683,  689, 2875, 4511, 4495,  102,    0,    0,
  25.                          0,    0,    0,    0,    0,    0,    0,    0,    0,    0]
  26.         >>> test_seq_inputs = torch.tensor(test_seq).clone().detach().unsqueeze(0)
  27.         >>> print(test_seq_inputs)
  28.         tensor([[ 101,  704, 1290, 1957, 2094, 2110, 7368, 8038, 3315, 4906, 2231, 3613,
  29.                   788,  122,  683,  689, 2875, 4511, 4495,  102,    0,    0,    0,    0,
  30.                   0,    0,    0,    0,    0,    0,    0,    0]])
  31.         >>> bert_embedding_no_layernorm = BertEmbeddingsWithoutLayerNorm()
  32.         >>> res2 = bert_embedding_no_layernorm(input_ids=test_seq_inputs)
  33.         >>> print(res2)
  34.             torch.Size([1, 32])
  35.             tensor([[[ 0.8470, -0.5573, -1.4137,  ...,  0.1017,  2.5658, -1.5135],
  36.                     [ 1.3631,  0.1439, -1.0117,  ...,  0.5226, -0.7416, -0.5690],
  37.                     [ 0.9139,  0.5850, -0.4989,  ...,  1.9164, -0.1690,  0.1694],
  38.                     ...,
  39.                     [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
  40.                     [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
  41.                     [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],
  42.                 grad_fn=<MulBackward0>)
  45.         >>> print(res2.size())
  46.         torch.Size([1, 32, 768])
  1.     Args:
  2.         input_ids (Tensor, optional): 句子的 token_ids. Defaults to None.
  4.     Raises:
  5.         ValueError: 如果输入不是ids则报错。注意这里修改为不接受 embedding 输入
  7.     Returns:
  8.         Tensor: 返回 embedding 句子结果
  9.     """
  10.     if input_ids is not None:
  11.         input_shape = input_ids.size()
  12.     else:
  13.         raise ValueError("检查input_ids是否正确，不接受传现成的embedding")
  14.     # print(input_shape)
  15.     inputs_embeds = self.word_embeddings(input_ids)
  17.     embeddings = inputs_embeds
  18.     embeddings = self.dropout(embeddings)
  19.     return embeddings

`` 在上述代码中，我进一步删掉了nn.LayerNorm` 。
注意，这里的输出维度是一样的。所以理论上要不要LayerNorm，至少程序上是不会报错的。

LayerNorm的作用

问题是，去掉了nn.LayerNorm 可以加速，但是效果会不会有所降低？
首先看看 nn.LayerNorm 的作用。在 NIPS19 的一篇论文《Understanding and Improving Layer Normalization》中有提到：
蒸馏时BERT的Embedding处理 - 图2

其中：
蒸馏时BERT的Embedding处理 - 图3 是 LayerNorm 的输出
蒸馏时BERT的Embedding处理 - 图4 是点乘操作
蒸馏时BERT的Embedding处理 - 图5 和 蒸馏时BERT的Embedding处理 - 图6 是输入的均值和标准差
蒸馏时BERT的Embedding处理 - 图7 偏差（bias）和 蒸馏时BERT的Embedding处理 - 图8 增益（gain）是与 蒸馏时BERT的Embedding处理 - 图9 维度一致的参数。

传统研究认为，正向归一化（forward normalization）是 LayerNorm 的决定性因素，而这篇论文的作者认为，正向归一化与有效性无关，而均值和方差的导数在LayerNorm中起重要作用。

Transformer结构中LayerNorm是一个重要的部分，但是在这里的Embedding层并不是不可替代的，所以我认为在这里采用更简单的、去掉了nn.LayerNorm 的 BertEmbeddingsWithoutLayerNorm 是完全可以的。

当然实际上还需要进一步实验检验。