ERNIE程式碼解析

NLP論文解讀發表於2022-01-28

©原創作者 |瘋狂的Max

ERNIE程式碼解讀

考慮到ERNIE使用BRET作為基礎模型，為了讓沒有基礎的NLPer也能夠理解程式碼，筆者將先為大家簡略的解讀BERT模型的結構，完整程式碼可以參見[1]。

01 BERT的結構組成

BERT的程式碼最主要的是由分詞模組、訓練資料預處理、模型結構模組等幾部分組成。

1.1 分詞模組

模型在訓練之前，需要對輸入文字進行切分，並將切分的子詞轉換為對應的ID。這一功能主要由BertTokenizer來實現，主要在
/models/bert/tokenization_bert.py實現。

BertTokenizer 是基於BasicTokenizer和WordPieceTokenizer 的分詞器：

BasicTokenizer負責按標點、空格等分割句子，並處理是否統一小寫，以及清理非法字元。

WordPieceTokenizer在詞的基礎上，進一步將詞分解為子詞（subword）。

具有以下使用方法：

  • from_pretrained：從包含詞表檔案（vocab.txt）的目錄中初始化一個分詞器；
  • tokenize：將文字分解為子詞列表；
  • convert_tokens_to_ids：將子詞轉化為子詞對應的下標；
  • convert_ids_to_tokens ：將對應下標轉化為子詞；
  • encode：對於單個句子，分解詞並加入特殊詞形成“[CLS], x, [SEP]”的結構並轉換為詞表對應下標的列表；
  • decode：將encode的輸出轉換為句子。

1.2 訓練資料預處理

訓練資料的構建主要取決於預訓練的任務，由於BERT的預訓練任務包括預測上下句和掩碼詞預測是否為連續句，那麼其訓練資料就需要隨機替換連續的語句和其中的分詞，這部分功能由run_pretraining.py中的函式
create_instances_from_document實現。

該部分首先構建上下句，拼接[cls]和[sep]等特殊符號的id，構建長度為512的列表，然後根據論文中所使用的指定概率選擇要掩碼的子詞，這部分由函式
create_masked_lm_predictions實現。

1.3 模型結構

BERT模型主要由BertEmbeddings類、BertEncoder類組成，前者負責將子詞、位置和上下句標識（segment）投影成向量，後者實現文字的編碼。

編碼器BertEncoder又由12層相同的編碼塊BertLayer組成。每一層都由自注意力層BertSelfAttention和前饋神經網路層BertIntermediate以及輸出層BertOutput構成，在
/models/bert/modeling_bert.py中實現。

每一層編碼層的結構和功能如下：

  • BertSelfAttention：負責實現子詞之間的相互關注。注意，多頭自注意力機制的實現是通過將維度為hidden_size 的表示向量切分成n個維度為hidden_size / n的向量，再對切分的向量分別進行編碼，最後拼接編碼後的向量實現的；
  • BertIntermediate：將批次資料（三維張量）做矩陣相乘和非線性變化；
  • BertOutput ：實現歸一化和殘差連線；

工程小技巧: 如果模型在學習表示向量的過程中需要使用不同的編碼方式，以結合圖神經網路層和Transformer編碼層為例，筆者建議儘量使用相同的引數初始化方式，兩者都使用殘差連線，這能夠避免模型訓練時出現梯度爆炸的問題。

此外是否需要對注意力權重進行大小的變化，如Transformer會除以向量維度的開方，則取決於圖神經網路的層數，一般而言，僅使用兩層或以下的圖神經網路層，則無需對注意力權重做變化。

具體可以通過觀察圖神經網路層生成的表示向量的大小是否和Transformer編碼層生成的向量大小在同一個數量級來決定，如果在同一個數量級則無需改變注意力權重，如果出現梯度爆炸的現象，那麼則可以縮小注意力的權重。

02 從BERT到ERNIE

由於ERNIE是在BERT的基礎上進行改進，在資料層面需要構建與文字對應的實體序列，在預訓練層面加入了新的預訓練任務，那麼在程式碼上就對應著訓練資料預處理和模型結構這兩方面的改動。因此筆者也將重點針對這兩個方面進行講解，完整程式碼參見[2]。

其程式碼結構主要包含兩大模組，訓練資料預處理模組和模型構建模組。

2.1 訓練資料預處理模組

ERNIE模型的知識注入依賴於找到文字中存在的實體，這些實體是指具有意義的抽象或者具象的單個名詞或名詞短語，我們可以將其稱為文字指稱項（mention)。一個實體可以有多個別名，也就意味著一個實體可以對應著文字中的多個指稱項。

為了能夠找到文字語料中實體，作者使用維基百科作為ERNIE的訓練語料，將維基百科中具有超連結的名詞或者短語作為實體，利用這一現有資源能夠大大的簡化檢索實體的難度。

2.1.1 訓練資料構建

在利用現有抽取工具獲得語料和實體名檔案後，通過
pretrain_data/create_insts.py構建訓練資料。

我們知道在訓練之前，首先需要對語料進行分詞（tokenize），獲得子詞（tokens），然後根據詞典得到子詞的索引ID，模型在接收索引後將其投影成向量。從BERT的程式碼中我們可以知道，BERT首先構建用於下一句預測（next sentences prediction)所需要的上下句，並從中隨機選擇掩碼詞，生成用於自注意力階段的掩碼列表。

那麼為了能夠注入語句中對應的實體，ERNIE就需要在這一過程中建立和訓練語料等長的實體ID張量，以及對應的掩碼列表。

作者僅僅對文字指稱項第一個子詞所對應的位置標註實體ID，這也就意味模型僅使用第一個子詞向量預測實體。這種做法能夠直接複用BERT的程式碼，而無需單獨針對實體序列再構建訓練資料，減輕了工程實現的工作量。

for i, x in enumerate(vec):
    if x == "#UNK#":
        vec[i] = -1
    elif x[0] == "Q":
        if x in d:
           vec[i] = d[x]
           if i != 0 and vec[i] == vec[i-1]:
           # 以某個實體為例，Q123 Q123 Q123 -> d[Q123] -1 -1，僅在第一個子詞中記錄實體的ID，其他位置標誌為-1
               vec[i] = -1  
           else:
               vec[i] = -1
#函式 create_instances_from_document
    // 獲取句子a和b的實體和子詞   
    tokens = [101] + tokens_a + [102] + tokens_b + [102]
    entity = [-1] + entity_a + [-1] + entity_b + [-1]
    // 構造用於為資料構建索引的物件ds，並將對應的輸入語料id列表及掩碼列表，實體id列表和掩碼列表等訓練資料存入ds。
    ds.add_item(torch.IntTensor(input_ids+input_mask+segment_ids
            +masked_lm_labels+entity+entity_mask+[next_sentence_label]))

2.1.2 實體向量載入

BERT由於具有經過預訓練的向量表，子詞的ID值可以利用nn.embedding模組獲取投影向量。

那麼實體的向量是經過TransE表示學習獲得的，又應該如何讓模型獲取其投影向量呢？作者在code/iteration.py中自定義資料迭代器物件，該物件在返回資料時會呼叫
torch.utils.data.DataLoader，通過在該函式中傳入負責投影實體向量的函式collate_fn，能夠讓模型在載入資料時獲取實體的表示向量。

#類 EpochBatchIterator(object):
    return CountingIterator(torch.utils.data.DataLoader(
            self.dataset,
            # collate_fn是傳入實體向量的關鍵
            collate_fn=self.collate_fn,
            batch_sampler=batches,
        ))
#函式collate_fn:
def collate_fn(x):
    x = torch.LongTensor([xx for xx in x])
    entity_idx = x[:, 4*args.max_seq_length:5*args.max_seq_length]
    # embed = torch.nn.Embedding.from_pretrained(embed)
    # embed為載入了經過預訓練的二維實體張量
    uniq_idx = np.unique(entity_idx.numpy())
    ent_candidate = embed(torch.LongTensor(uniq_idx+1))

2.2 模型結構模組

在模型方面，作者依舊使用12層Transformer編碼層作為模型結構，與BERT所不同的是，在前6層沿用BERT的Transformer編碼層，但在第7層自定義知識融合層BertLayerMix，首次對經過對齊的實體向量和指稱項向量求和，並將其分別傳輸給知識編碼模組和文字編碼模組，在剩下5層自定義知識編碼層BertLayer，分別對經過融合了兩者資訊的實體序列和文字序列使用自注意力機制編碼。

模型的前5層就是論文所指的文字編碼器，後面的7層編碼層則構成了論文中的知識編碼器。

對於BERT的Transformer編碼層，由於第一部分已經介紹過，就不再贅述。下文主要針對作者自定義的編碼層做詳細解讀。

2.2.1 知識融合層BertLayerMix

具體來說，知識融合層BertLayerMix由自注意力層BertAttention_simple、融合層BertIntermediate以及輸出層BertOutput構成。

class BertLayerMix(nn.Module):
    def __init__(self, config):
        super(BertLayerMix, self).__init__()
        self.attention = BertAttention_simple(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)
     # 該編碼層僅針對文字進行自注意力操作、矩陣相乘和殘差連線
    def forward(self, hidden_states, attention_mask, hidden_states_ent, attention_mask_ent, ent_mask):
        attention_output = self.attention(hidden_states, attention_mask)
        attention_output_ent = hidden_states_ent * ent_mask
        # intermediate層負責實體和文字向量求和，並對求和向量非線性變化
        intermediate_output = self.intermediate(attention_output, attention_output_ent)
        # 然後通過輸出層output再次歸一化和殘差連線
        layer_output, layer_output_ent = self.output(intermediate_output, attention_output, attention_output_ent)
        return layer_output, layer_output_ent

自注意力層BertAttention_simple由BertSelfAttention和BertSelfOutput構成，前者負責對文字進行自注意力操作，實現上與BERT的自注意力操作相同，就不再展示程式碼。後者則用於對向量進行矩陣變化和殘差連線，生成attention_output

class BertAttention_simple(nn.Module):
    def __init__(self, config):
        super(BertAttention_simple, self).__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)


    def forward(self, input_tensor, attention_mask):
        self_output = self.self(input_tensor, attention_mask)
        attention_output = self.output(self_output, input_tensor)
        return attention_output
class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)


    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states

前饋神經網路層BertIntermediate負責將兩者進行線性變化轉換為同一維度，求和並做非線性變化。

class BertIntermediate(nn.Module):
    def __init__(self, config):
        super(BertIntermediate, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        self.dense_ent = nn.Linear(100, config.intermediate_size)
        self.intermediate_act_fn = ACT2FN[config.hidden_act] \
            if isinstance(config.hidden_act, str) else config.hidden_act
    def forward(self, hidden_states, hidden_states_ent):
        # 線性變化轉換為同一維度
        hidden_states_ = self.dense(hidden_states)
        hidden_states_ent_ = self.dense_ent(hidden_states_ent)
        # 求和並使用intermediate_act_fn做非線性變化
        hidden_states = self.intermediate_act_fn(hidden_states_+hidden_states_ent_)
        return hidden_states

最終使用BertOutput分別對文字向量和實體向量做矩陣相乘，將經過融合的向量和兩者殘差連線，並做歸一化操作。

class BertOutput(nn.Module):
    def __init__(self, config):
        super(BertOutput, self).__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dense_ent = nn.Linear(config.intermediate_size, 100)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.LayerNorm_ent = BertLayerNorm(100, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
    def forward(self, hidden_states_, input_tensor, input_tensor_ent):
        # 針對文字向量矩陣相乘
        hidden_states = self.dense(hidden_states_)
        hidden_states = self.dropout(hidden_states)
        # 針對文字向量殘差連線和歸一化
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        # 針對實體向量的矩陣相乘、殘差連線和歸一化
        hidden_states_ent = self.dense_ent(hidden_states_)
        hidden_states_ent = self.dropout(hidden_states_ent)
        hidden_states_ent = self.LayerNorm_ent(hidden_states_ent + input_tensor_ent)
        return hidden_states, hidden_states_ent

2.2.2 知識編碼層BertLayer

該編碼層針對融合後的實體向量和文字向量分別進行自注意力編碼，從而使實體序列中的所有實體也能夠實現相互關注。

再次基礎上實體向量將和對應位置的文字向量求和，將實體資訊傳遞給文字向量，從而使整個文字序列在下一個編碼層中實現對實體序列的關注。

class BertLayer(nn.Module):
    def __init__(self, config):
        super(BertLayer, self).__init__()
        self.attention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)


    def forward(self, hidden_states, attention_mask, hidden_states_ent, attention_mask_ent, ent_mask):
        attention_output, attention_output_ent = self.attention(hidden_states, attention_mask, hidden_states_ent, attention_mask_ent)
        attention_output_ent = attention_output_ent * ent_mask
        intermediate_output = self.intermediate(attention_output, attention_output_ent)
        layer_output, layer_output_ent = self.output(intermediate_output, attention_output, attention_output_ent)
        # layer_output_ent = layer_output_ent * ent_mask
        return layer_output, layer_output_ent

這一編碼層自定義了自注意力層，其中針對實體的自注意力層僅使用4個注意力頭。

class BertAttention(nn.Module):
    def __init__(self, config):
        super(BertAttention, self).__init__()
        self.self = BertSelfAttention(config)
        self.output = BertSelfOutput(config)
        config_ent = copy.deepcopy(config)
        config_ent.hidden_size = 100
        config_ent.num_attention_heads = 4
        self.self_ent = BertSelfAttention(config_ent)
        self.output_ent = BertSelfOutput(config_ent)
    def forward(self, input_tensor, attention_mask, input_tensor_ent, attention_mask_ent):
        # BertSelfAttention對文字向量進行自注意力操作
        self_output = self.self(input_tensor, attention_mask)
        self_output_ent = self.self_ent(input_tensor_ent, attention_mask_ent)
        # BertSelfAttention對實體向量進行自注意力操作
        attention_output = self.output(self_output, input_tensor)
        attention_output_ent = self.output_ent(self_output_ent, input_tensor_ent)
        return attention_output, attention_output_ent

輸出層同知識融合層一樣，都是使用BERToutput實現歸一化和殘差連線。

03 原始碼參考

[1] https://github.com/google-research/bert

[2] https://github.com/thunlp/ERNIE

 

私信我領取目標檢測與R-CNN/資料分析的應用/電商資料分析/資料分析在醫療領域的應用/NLP學員專案展示/中文NLP的介紹與實際應用/NLP系列直播課/NLP前沿模型訓練營等乾貨學習資源。

相關文章