华为NEZHA预训练模型中的Relative Position与self_attention详解

---

前言

最近读了一下华为NAZHA预训练模型，该模型的主要创新之处在于在Transformer阶段加入了Relative Position Embedding(相对位置编码)，具体的是在Q*V计算得到的attention_scores后加入了Key的Relative Position，在计算得到的context 也加入了Value的Relative Position，详情请看代码。。小编已经对这部分代码进行了注释，主要注释了各个阶段的维度变化，这样对于小白来说代码读起来比较轻松。

---

直接上代码：

# coding=utf-8
# @Time:2021/4/2510:53
# @author: SinGaln
import math
import torch
import argparse
import numpy as np
import torch.nn as nn

def generation_relative_embedding(length, depth, max_relative_position):
    vocab_size = max_relative_position * 2 + 1  # 129
    range_vec = torch.arange(length)  #[1,129]
    range_mat = range_vec.repeat(length).view(length, length) # [512,512]
    distance_mat = range_mat - torch.t(range_mat) # [512,512]
    distance_mat_clipped = torch.clamp(distance_mat, -max_relative_position, max_relative_position) #[512,512]
    final_mat = distance_mat_clipped + max_relative_position  #[512,512]
    embeddings_table = np.zeros([vocab_size, depth]) # [129,12]
    for pos in range(vocab_size):
        for i in range(depth // 2):
            embeddings_table[pos, 2 * i] = np.sin(pos / np.power(10000, 2 * i / depth))
            embeddings_table[pos, 2 * i + 1] = np.cos(pos / np.power(10000, 2 * i / depth))

    embeddings_table_tensor = torch.tensor(embeddings_table).float()
    flat_relative_positions_matrix = final_mat.view(-1) #[262144]  512 * 512
    one_hot_relative_positions_matrix = torch.nn.functional.one_hot(flat_relative_positions_matrix,
                                                                    num_classes=vocab_size