BERT:输入嵌入的权重作为掩蔽语言模型的一部分

BERT: Weights of input embeddings as part of the Masked Language Model

我查看了 BERT 掩码语言模型的不同实现。 预训练有两个个常见版本:

  1. 解码器将简单地采用 [MASK]ed 令牌的最终嵌入并将其传递通过线性层(无需任何修改):
    class LMPrediction(nn.Module):
        def __init__(self, hidden_size, vocab_size):
            super().__init__()
            self.decoder = nn.Linear(hidden_size, vocab_size, bias = False)
            self.bias = nn.Parameter(torch.zeros(vocab_size))
            self.decoder.bias = self.bias
        def forward(self, x):
             return self.decoder(x)
  1. 一些实现会使用输入嵌入的权重作为解码器线性层的权重:
    class LMPrediction(nn.Module):
        def __init__(self, hidden_size, vocab_size, embeddings):
            super().__init__()
            self.decoder = nn.Linear(hidden_size, vocab_size, bias = False)
            self.bias = nn.Parameter(torch.zeros(vocab_size))
            self.decoder.weight = embeddings.weight ## <- THIS LINE
            self.decoder.bias = self.bias
        def forward(self, x):
             return self.decoder(x)

哪个是正确的?大多数情况下,我看到的是第一个实现。然而,第二个也有道理 - 但我找不到任何论文中提到它(我想看看第二个版本是否优于第一个)

对于那些感兴趣的人,它被称为权重绑定联合输入输出嵌入。有两篇论文争论这种方法的好处: