BERT ese: Learning to Speak to BERT (2021)

Abstract: Large pre-trained language models have been shown to encode large amounts of world and commonsense knowledge in their parameters, leading to substantial interest in methods for extracting that knowledge. In past work, knowledge was extracted by taking manually-authored queries and gathering paraphrases for them using a separate pipeline. In this work, we propose a method for automatically rewriting queries into “BERTese”, a paraphrase query that is directly optimized towards better knowledge extraction. To encourage meaningful rewrites, we add auxiliary loss functions that encourage the query to correspond to actual language tokens. We empirically show our approach outperforms competing baselines, obviating the need for complex pipelines. Moreover, BERTese provides some insight into the type of language that helps language models perform knowledge extraction.

## introduction recent work has shown that large pre-trained language models (lm), trained with a masked language modeling (mlm) objective @xcite @xcite @xcite , encode substantial amounts of world knowledge in their parameters. this has led to ample research on developing methods for extracting that knowledge @xcite @xcite . the most straightforward approach is to present the model with a manually-crafted query such as "dante was born in @xcite " and check if the model predicts "florence" in the [mask] position. however, when this fails, it is difficult to determine if the knowledge is absent from the lm or if the model failed to understand the query itself. for example, the model might return the correct answer if the query is "dante was born in the city of [mask]". motivated by the above observation, we ask: can we automatically find the best way to "ask" an lm about its knowledge? we refer to this challenge as speaking "bertese". in particular, we ask how to rewrite a knowledge-seeking query into one that mlms understand better, where understanding is manifested by providing a correct answer to the query. prior work @xcite tackled this problem using a 2-step pipeline, where first a small list of paraphrase templates is collected using external resources, and then a model learns to extract knowledge by aggregating information from paraphrases of the input query. in this work, we propose a more general approach, where the model learns to rewrite queries, directly driven by the objective of knowledge-extraction. figure 1 provides an overview of our approach. our model contains a bert-based rewriter, which takes a query as input, and outputs for each input position a new token, which is its rewrite. this new query is fed into a different bert predictor from which the answer is extracted. importantly, the downstream predictor bert is a fixed pre-trained model, and thus the goal is to train the rewriter to produce queries for which the predictor outputs the correct answer. a technical challenge is that outputting discrete tokens leads to a non-differentiable model, which we tackle by adding a loss term that encourages the rewriter's output to be similar to bert token embeddings. moreover, we must guarantee that the bertese query contains the [mask] token from which the answer will be read. to achieve this, we first add an auxiliary loss term that encourages the model to output precisely one masked token in the query rewrite. we then add a layer that finds the token index that most closely resembles [mask], and this is where we expect the correct answer to be completed. training of this selection process is done using the straight-through estimator @xcite . we evaluate our approach on the lama dataset @xcite , and show that our model significantly improves the accuracy of knowledge extraction. furthermore, many of the rewrites correspond to consistent changes in query wording (e.g., changing tense), and thus provide information on the types of changes that are useful for extracting knowledge from bert. while we experiment on bert, our method is generic and can be applied to any mlm. taken together, our results demonstrate the potential of rewriting inputs to language models for both obtaining better predictions, and for potentially gaining insights into how knowledge is represented in these models. our code can be downloaded from https://github.com/ adihaviv/bertese . ## related work choosing the right language for extracting world knowledge from lms has attracted much interest recently. first, @xcite observed that mlms can complete simple queries with correct factual information. @xcite and heinzerling and inui (2020) then showed that in the zeroshot setting, small variations to such queries can lead to a drop in fact recall. orthogonally, another line of research focused on query reformulation for standard question answering (qa) tasks. @xcite demonstrated that even minor query modifications can lead to a significant decrease in performance for multiple qa models and tasks. @xcite showed that it is possible to train a neural network to reformulate a question using reinforcement learning (rl), optimizing the accuracy of a black-box qa system. similarly, @xcite used rl to create a query reformulation system that maximizes the recall of a black-box information retrieval engine. @xcite proposed an ensemble method for query reformulation from lms, that includes: (1) mining new queries, (2) using an off-the-shelf pre-trained translation model to collect additional paraphrased queries with back-translation, and (3) using a re-ranker to select one or more of the new queries. they then feed those queries to bert to get the masked token prediction. in this work, we take the idea of jiang et al. (2020) a step forward and train a model in an endto-end fashion to generate rephrased queries which are optimized to maximize knowledge extraction from the mlm. foot_0 ## the bertese model recall that our goal is to build a model that takes as input a query in natural language, and re-writes it into a query that will be fed as input to an existing bert model. we refer to the above re-writing model as the rewriter and the existing bert model as the predictor. we note that both input and output queries should include the token @xcite . for example the input could be "obama was born in [mask]" and the output "obama was born in the state of [mask]". we first describe the behaviour of our model at inference time (see figure 1 ). given a query, which is a sequence of tokens, @xmath0, . . . , s n ), we map s into a sequence of vectors q(s) ∈ r d×n using bert's embeddings of dimensionality d. this input is fed into a (bert-based) stack of transformer layers that outputs a new sequence of vectors q(s) ∈ r d×n . to obtain vectors that can be used as input to the predictor, we need to map the vectors in each position to their nearest neighbor in the set of bert embeddings. specifically, let b v be the set of bert embeddings, and let @xmath1. we map @xmath2 training this model involves two technical challenges. first, the nearest-neighbor operation is non-differentiable. second, to obtain the prediction of the [mask] token, we need to guarantee that the rewriter generates a [mask] token, and know its position (because this is where the ground-truth answer should be predicted). we overcome these by adding two auxiliary loss functions. the first encourages the model to output vectors that are similar to bert embeddings (thus reducing the loss in the nearest neighbor operation), and the second encourages the model to output one masked token. finally, we apply the straight-through estimator, which allows us to feed discrete word representations into the predictor and backpropagate the signal back to the rewriter. we next provide more details on the terms in our loss function used to train the rewriter. valid token loss: at training time we do not apply the non-differentiable nearest-neighbor operation. thus, we would like the vectors q(s) output by the rewriter to be as close as possible to valid bert embeddings. this loss is the average over tokens of the distance between a re-written query token and its nearest neighbor: single [mask] loss: the output of the rewriter must contain the [mask] token, so that the predictor can extract an answer from this token. to encourage the rewriter to output a [mask] we add a loss as follows. we define the following "softmin" distribution over @xmath3, . . . , | q(s)|}: prediction loss: the predictor should return the gold answer y when given q as input. without non-differentiability, we could find the index of the [mask] token in q, and use cross-entropy loss between the output distribution of the predictor in that index and the gold answer y. to remedy this, we use a differentiable formulation, combined with the straight-through estimator (ste) @xcite : let o i be the output distribution at the i th position of the predictor, and let ℓ(y, p) be the cross-entropy between the one-hot distribution corresponding to y and a distribution p. then, we use the loss: thus, if m is a one-hot on the index corresponding to @xcite , the loss will be the desired crossentropy between the gold answer and the predicted distribution. we optimize this objective using the ste. namely, in the forward pass, we convert m to a one-hot vector. our final training loss is the sum of the above three loss terms: the weights λ 1 , λ 2 are tuned using crossvalidation. to summarize, the main challenge is that the rewriter output needs to be optimized to predict the correct label for the [mask] token (eq. 4). however, the [mask] token needs to appear once in the rewriter output. in order to enforce the above, the "single [mask] loss" (eq. 3) is used. in addition, in order for the rewriter output to be a valid sentence, the "valid token loss" (eq. 1) is added. this encourages the model to output tokens that are close to bert input embeddings. this is done by minimizing the distance between each rewriter vector to some vector in the bert input embedding dictionary. rewriter pre-training we initialize the rewriter with a bert-based model, additionally fine-tuned to output the exact word embeddings it received as input (i.e., fine-tuned to the identity mapping). thus, when training for knowledge extraction, the rewriter is initialized to output exactly the query it received as input. for training our model, we use a separate training set, created by jiang et al. ( 2020 ), called t-rextrain. this dataset is constructed from wikidata and has no overlap with the original t-rex dataset. we evaluate our model on the complete t-rex dataset. implementation details both the rewriter and the predictor are based on bert base with the default settings from the huggingface @xcite platform. we optimize bertese using adamw with an initial learning rate of 1e-5. we train the model on a single 32gb nvidia v100 for 5 epochs with a batch size of 64. for the loss coefficients (see eq. ( 5 )) we set λ 1 = 0.3 and λ 2 = 0.5. baselines we compare our method to three baselines: (a) bert -a bert base model without any fine-tuning, as evaluated in @xcite . (b) lpaqa -the model proposed by @xcite , based on mining additional paraphrase queries. we report results on a single paraphrase. foot_1 (c) ft-bert: an end-to-end differentiable bert base model, explicitly fine-tuned on t-rex-train to output the correct answer. this model, like ours, is trained for knowledge extraction, but does this internally, without exposing an interpretable intermediate textual rewrite. ## conclusion we presented an approach for modifying the input to a bert model, such that factual information can be more accurately extracted. our approach uses a trained rewrite model that is optimized to maximize the accuracy of its rewrites, when used as input to bert. our rewriting scheme indeed turns out to produce more accurate results than baselines. interestingly, our rewrites are fairly small modifications, highlighting the fact that bert models are not invariant to these edits. our approach is not limited to knowledge extraction. it can, in principle, be applied to bert in general question answering datasets and even language modeling. in the former, we can change the predictor to a multiple-choice qa pretrained bert and exclude the single [mask] token loss. in the latter, we can for example envision a case where rewriting a sentence can make it easier to complete a masked word. our empirical setting focuses on the lama dataset, where a single mask token prediction is required. there are several possible extensions to multiple masks, and we leave these for future work. finally, it will be interesting to test the approach on other masked language models such as roberta @xcite and ernie @xcite , a mlm that is enhanced with external entity representations.