MG - BERT : Multi-Graph Augmented BERT for Masked Language Modeling (2021)

Abstract: Pre-trained models like Bidirectional Encoder Representations from Transformers (BERT), have recently made a big leap forward in Natural Language Processing (NLP) tasks. However, there are still some shortcomings in the Masked Language Modeling (MLM) task performed by these models. In this paper, we first introduce a multi-graph including different types of relations between words. Then, we propose Multi-Graph augmented BERT (MG-BERT) model that is based on BERT. MG-BERT embeds tokens while taking advantage of a static multi-graph containing global word co-occurrences in the text corpus beside global real-world facts about words in knowledge graphs. The proposed model also employs a dynamic sentence graph to capture local context effectively. Experimental results demonstrate that our model can considerably enhance the performance in the MLM task.

## introduction in recent years, pre-trained models have led to promising results in various natural language processing (nlp) tasks. recently, bidirectional encoder representations from transformers (bert) @xcite has received much attention as a pre-trained model that can be easily fine-tuned for a wide range of nlp tasks. bert is pre-trained using two unsupervised tasks, masked language modeling (mlm) and next sentence prediction (nsp). in @xcite , some psycholinguistic diagnostics are introduced for assessing the linguistic capacities of pre-trained language models. these diagnostic tests consist of commonsense and pragmatic inferences, role-based event prediction, and negation. @xcite observes some shortcomings in bert's results and demonstrates that although bert sometimes predicts the first candidate for the masked token almost correctly, some of its top candidates contradict each other. besides, in the tests targeting commonsense and pragmatic inference, it is illustrated that bert can not precisely fill the gaps based on just the input context @xcite . in this paper, we incorporate co-occurrences and global information about words through graphs describing relations of words along with local contexts considered by bert. the intention is to find more reliable and meaningful embeddings that result in better performance in mlm task. utilizing external information about the corpus and the world in the form of graphs helps the model fill the gaps in the mlm task more easily and with more certainty. we take advantage of the rich information source accessible in knowledge graphs and also condensed information of words co-occurrence in graphs using relational graph convolutional network (r-gcn) to enrich the embedding of tokens. we also utilize the words in the current context as a dynamic complete graph using an attention mechanism. these graphs can considerably influence the performance of bert in the mlm task as shown in the experiments. ## related work knowledge graphs (kgs) are valuable sources of facts about real-world entities. many studies have been recently introduced to utilize knowledge graphs for various purposes, such as recommender systems @xcite or link prediction @xcite @xcite . recently, using bert along with knowledge graphs has also been attended for knowledge graph completion and analysis. @xcite employ kg-bert in triple classification, link prediction, and relation prediction tasks. furthermore, knowledge graphs are used in nlp tasks such as text classification @xcite @xcite , named entity recognition @xcite , and language modeling @xcite . ernie @xcite ) is an enhanced language representation model incorporating knowledge graphs. in addition to bert's pre-training objectives, it uses an additional objective that intends to select appropriate entities from the knowledge graph to complete randomly masked entity alignments. moreover, named entity mentions in the text are recognized and aligned to their corresponding entities in kgs. other types of graphs have also been utilized in nlp tasks in some studies. for instance, text gcn @xcite applies graph convolutional network (gcn) to the task of text classification. this paper's employed graph is a text graph created based on token co-occurrences and document-token relations in a corpus. moreover, vgcn-bert @xcite enriches the word embeddings of an input sentence using the text graph inspired by text gcn @xcite and examines the obtained model in fire hate language detection tasks @xcite . in this paper, we aim to improve bert's performance (in the mlm task) by incorporating a static multi-graph that includes both the knowledge graph and global co-occurrence graphs derived from the corpus as well as a dynamic graph including input sentence tokens. static text graphs have been recently employed in vgcn-bert @xcite via a modified version of gcn that extends the input by a fixed number of embeddings. however, the modification of embeddings in this work is only based on input tokens. neither other vocabularies in the static text graphs nor real-world facts (available in kgs) affect the final embeddings of tokens. on the other hand, while enrie @xcite and kepler @xcite utilize kgs to reach an improved model, they do not employ other graphs derived from the corpus. also, ernie does not learn graph-based embedding during representation learning and only adopts embeddings trained by transe @xcite . however, in our model, since we incorporate a multi-graph by extending bert architecture and providing a graph layer of an r-gcn module and attention mechanism, a multi-graph augmented representation learning model is obtained. ## preliminaries gcn @xcite is one of the most popular models for graph node embedding. r-gcn @xcite extends gcn to provide node embedding of multi-relational graphs: ## methodology this section presents the overall architecture of our model, called multi-graph augmented bert (mg-bert). mg-bert takes advantage of bert's power in capturing context of an input text as well as a graph module including an r-gcn layer over a static multi-graph and a graph attention layer over a dynamic sentence graph. this static multi-graph includes global information about words available as facts in kgs in addition to dependencies between tokens of the input text and other words in the vocabulary which are discovered by computing co-occurrence statistics in the corpus. two graphs are used to condense co-occurrences of words in the corpus inspired by text gcn @xcite that are also employed by vgcn-bert @xcite . one of these graphs includes local co-occurrences of terms that is computed based on point-wise mutual information (pmi) of terms i and j which is calculated by: in the above equations, #w (i) and #w (i, j) denote the number of fixed size windows containing term i and both of the terms i and j, respectively. #w is the whole number of windows in the corpus. the other graph includes the document level co-occurrence of tokens in the corpus computed based on term frequency-inverse document frequency (tf-idf). the knowledge graph is also incorporated in this multi-graph. formally, the weighted edges between token i and token j for three types of relations @xmath0, pmi, tf-idf} in the multi-graph are: furthermore, we use a graph attention mechanism to capture local information via a dynamic and complete graph in which nodes represent all tokens of the input sentence. the complete dynamic graph is used in order to obtain context-dependent new embeddings while the r-gcn layer itself provides the same new embeddings for a specific token even if the token appears in different contexts. this happens because the single r-gcn layer always performs on the same static multi-graph. as shown in fig. 1 , the whole graph module is placed immediately after the bert token embeddings layer since the hidden states of the whole vocabulary are available in this layer. we pass the entire multi-graph to the r-gcn module so that the global dependencies would affect embeddings of tokens properly using eq. 3. we also use an attention mechanism as in eq. 4 to consider the local context. the new embedding of token i in sentence s is computed as: next, we aggregate the obtained tokens' embeddings by the graph module with position embeddings and segment embeddings (similar to bert). afterward, we feed these representations to bert encoders to find final embeddings. the proposed model architecture is shown in figure 1 . in the training phase, a token from each sentence is randomly masked, and the model is trained to predict the masked token based on both the context and the incorporated static multi-graph. ## experiments in this section, we explain the details of training mg-bert and conduct experiments to evaluate and compare our model with the related methods recently proposed. datasets. during training, we use the wn18 knowledge graph, derived from wordnet, as an unlabeled graph @xcite . we also experiment mg-bert and other recent models on cola, sst-2, and brown datasets @xcite @xcite . the detailed description of these datatsets is given in appendix a. parameter setting. in order to capture word co-occurence statistics of the corpus, we use the bert's tokenizer on sentences and set the sliding window size to 20 when calculating the pmi value. the whole bert module in mg-bert is first initialized with the pre-trained bert-base-uncase version of bert in pytorch and the model is trained on the mlm task with cross entropy loss @xcite . regarding eq. 2, different hyper-parameter settings have been used for each dataset. λ k , λ p , and λ t are set to 0.01, 0.001, and 0.001, respectively in both cola and brown datasets and 0.001, 1.0, and 0.001 in sst-2 dataset. the hyperparameter λ dyn is also set to 0.8. the graph attention mechanism is performed with 12 heads. the r-gcn and graph attention layers' output dimension are also set to 768 that equals to the dimension of the bert token embeddings layer to substitute easily bert's token embeddings with the embeddings derived from the graph module. we also employ the normalization trick introduced in gcn @xcite to normalize each adjacency matrix in the multi-graph. compared methods. to assess our model, we compare it with bert as the baseline. moreover, ernie and vgcn-bert are being compared as the recent methods utilizing knowledge graph and text graph, respectively @xcite . we also compare mg-bert with mg-bert(base) which doesn't use the dynamic graph incorporating the context according to eq. 4. all these models are fine-tuned on the text datasets for a fair evaluation. results. we evaluate our model using hits@1 and hits@5 metrics. hits@k shows the proportion of correct tokens appearing in the top k results for each sample. in the reason to our superiority over vgcn-bert @xcite is that it doesn't take advantage of real-world facts (available in kgs). moreover, as opposed to mg-bert, it modifies initial embeddings of tokens only based on input tokens of each sentence and other vocabularies in the text graphs don't influence the final embeddings of tokens. on the other hand, ernie @xcite doesn't take full advantage of graphs since it doesn't use graphs derived from the corpus. be-sides, it does not learn graph-based embeddings during representation learning. it is worth mentioning that the entity embedding model used in ernie has been trained on a huge subset of wikidata foot_0 , which is almost 120 times bigger than wn18 knowledge graph employed in our method. the superiority of mg-bert over mg-bert(base) demonstrates the importance of the dynamic sentence graph and the results of mg-bert(base) itself shows that utilizing the static multi-graph has been useful. ## conclusion in this paper, we proposed a language representation learning model that enhances bert by augmenting it with a graph module (i.e. an r-gcn layer over a static multi-graph, including global dependencies between words, and a graph attention layer over a dynamic sentence graph). the static multi-graph utilized in this work consists of a knowledge graph as a source of information about real-world facts and two other graphs built based on word co-occurrences in local windows and documents in the corpus. therefore, the proposed model utilizes the local context, the corpuslevel co-occurence statistics, and the global word dependencies (through incorporating a knowledge graph) to find the input tokens' embeddings. the results generally show the superiority of the proposed model in the masked language modeling task compared to both the bert model and the recent models employing knowledge or text graphs.