Title: CONTEXTUALIZED SPARSE REPRESENTATION WITH RECTIFIED N-GRAM ATTENTION FOR OPEN-DOMAIN QUESTION ANSWERING

Abstract: A sparse representation is known to be an effective means to encode precise lexical cues in information retrieval tasks by associating each dimension with a unique ngram-based feature. However, it has often relied on term frequency (such as tf-idf and BM25) or hand-engineered features that are coarse-grained (document-level) and often task-specific, hence not easily generalizable and not appropriate for finegrained (word or phrase-level) retrieval. In this work, we propose an effective method for learning a highly contextualized, word-level sparse representation by utilizing rectified self-attention weights on the neighboring n-grams. We kernelize the inner product space during training for memory efficiency without the explicit mapping of the large sparse vectors. We particularly focus on the application of our model to phrase retrieval problem, which has recently shown to be a promising direction for open-domain question answering (QA) and requires lexically sensitive phrase encoding. We demonstrate the effectiveness of the learned sparse representations by not only drastically improving the phrase retrieval accuracy (by more than 4%), but also outperforming all other (pipeline-based) open-domain QA methods with up to x97 faster inference in SQUAD OPEN and CURATEDTREC. 1

1. INTRODUCTION:

Retrieving text from a large collection of documents is an important problem in several natural language tasks such as question answering (QA) and information retrieval. In the literature, sparse representations have been successfully used in encoding text at sentence or document level, capturing precise lexical information that can be sparsely activated by n-gram based features. For instance, frequency-based sparse representations such as tf-idf map each text segment to a vocabulary space where the weight of each dimension is determined by the associated word's term and inverse document frequency. However, these sparse representations are mainly coarse-grained, task-specific, and not suitable for word-level representations since they statically assign the identical weight to each n-gram and do not change dynamically depending on the context.

In this paper, we introduce an effective method for learning a word-level sparse representation that encodes precise lexical information. Our model, COSPR, learns a contextualized sparse phrase representation leveraging rectified self-attention weights on the neighboring n-grams and dynamically encoding important lexical information of each phrase (such as named entities) given its context. Consequently, in contrast to previous sparse regularization techniques on dense embedding of at most few thousand dimensions (Faruqui et al., 2015; Subramanian et al., 2018) , our method is able to produce more interpretable representations with billion-scale cardinality. Such large sparse vector space is prohibitive for explicit mapping; we leverage the fact that our sparse representations only interact through inner product and kernelize the inner-product space for memory-efficient training, inspired by the kernel method in SVMs (Cortes & Vapnik, 1995) . This allows us to handle extremely large sparse vectors without worrying about computational bottlenecks. The overview of our model is illustrated in Figure 1 .

We demonstrate the effectiveness of our model in open-domain question answering (QA), the task of retrieving answer phrases given a web-scale collection of documents. Following Seo et al. (2019) , We propose to learn contextualized sparse phrase representations which are also very interpretable.

Overview Figure:

we concatenate both sparse and dense vectors to encode every phrase in Wikipedia and use maximum similarity search to find the closest candidate phrase to answer each question. We only substitute (or augment) the baseline sparse encoding which is entirely based on frequency-based embedding (tf-idf) with our contextualized sparse representation (COSPR). Our empirical results demonstrate its state-of-the-art performance in open-domain QA datasets, SQUAD OPEN and CURATEDTREC. Notably, our method significantly outperforms DENSPI (Seo et al., 2019) , the previous end-toend QA model, by more than 4% with negligible drop in inference speed. Moreover, our method achieves up to 2% better accuracy and x97 speedup in inference compared to pipeline (retrievalbased) approaches. Our analysis particularly shows that fine-grained sparse representation is crucial for doing well in phrase retrieval task. In summary, the contributions of our paper are:

1. we show that learning sparse representations for embedding lexically important context words of a phrase can be achieved by contextualized sparse representations, 2. we introduce an efficient training strategy that leverages the kernelization of the sparse inner-product space, and 3. we achieve the state-of-the-art performance in two open-domain QA datasets with at least x97 faster inference time.

2. RELATED WORK:

Open-domain question answering Most open-domain QA models for unstructured texts use a retriever to find documents to read, and then apply a reading comprehension (RC) model to find answers (Chen et al., 2017; Wang et al., 2018a; Lin et al., 2018; Das et al., 2019; Lee et al., 2019; Yang et al., 2019; Wang et al., 2019) . To improve the performance of the open-domain question answering, various modifications have been studied to the pipelined models which include improving the retriever-reader interaction (Wang et al., 2018a; Das et al., 2019) , re-ranking paragraphs and/or answers (Wang et al., 2018b; Lin et al., 2018; Lee et al., 2018; Kratzwald et al., 2019) , learning end-to-end models with weak supervision (Lee et al., 2019) , or simply making a better retriever and a reader model (Yang et al., 2019; Wang et al., 2019) . Due to the pipeline nature, however, these models inevitably suffer error propagation from the retrievers.

To mitigate this problem, Seo et al. (2019) propose to learn query-agnostic representations of phrases in Wikipedia and retrieve phrases that best answers a question. While Seo et al. (2019) have shown that encoding both dense and sparse representations for each phrase could keep the lexically important words of a phrase to some extent, their sparse representations are based on static tf-idf vectors which have globally the same weight for each n-gram.

Phrase representations:

In NLP, phrase representations can be either obtained in a similar manner as word representations (Mikolov et al., 2013) , or by learning a parametric function of word representations (Cho et al., 2014) . In extractive question answering, phrases are often referred to as spans, but most models do not consider explicitly learning phrase representations as these answer spans can be obtained by predicting only start and end positions in a paragraph (Wang & Jiang, 2017; Seo et al., 2017) . Nevertheless, few studies have focused on directly learning and classifying phrase representations (Lee et al., 2017) which achieve strong performance when combined with attention mechanism. In this work, we are interested in learning query-agnostic sparse phrase representations which enables the precomputation of re-usable phrase representations (Seo et al., 2018) .

Sparse representations One of the most basic sparse representations is the bag-of-words modeling, and recent works often emphasize the use of bag-of-words models as strong baselines for sentence classification and question answering (Joulin et al., 2017; Weissenborn et al., 2017) . tf-idf is another good example of sparse representations that is used for document retrieval, and is still widely adopted both in IR and QA community. Our work could be seen as an attempt to build a trainable tf-idf model for phrase (n-gram) representations, which should be more fine-grained than paragraphs or documents. There have been some attempts in IR that learn sparse representations of documents for duplicate detection (Hajishirzi et al., 2010) , inverted indexing (Zamani et al., 2018) , and more. Unlike previous works, however, our method does not require hand-engineered features for sparse n-grams while keeping the original vocabulary space. In NLP, there are some works on training sparse representations specifically designed for improving interpretability of word representations (Faruqui et al., 2015; Subramanian et al., 2018) , but they lose an important role of sparse represenations, which is keeping the exact lexical information, as they are trained by sparsifying dense representations of a higher (but much lower than V ) dimension.

3. BACKGROUND: OPEN-DOMAIN QA THROUGH PHRASE RETRIEVAL:

Open-Domain QA We primarily focus on open-domain QA on unstructured data where the answer is a text span in the corpus. Formally, given a set of K documents x 1 , . . . , x K and a question q, the task is to design a model that obtains the answer â by â = arg max

x k i:j F (x k i:j , q),

where F is the score model to learn and x k i:j is a phrase consisting of words from the i-th to the j-th word in the k-th document. Typically, the number of documents (K) is in the order of millions for open-domain QA (e.g. 5 million for English Wikipedia), which makes the task computationally challenging. Pipeline-based methods typically leverage a document retriever to reduce the number of documents to read, but they suffer from error propagation when wrong documents are retrieved and can be slow if the reader model is computationally cumbersome.

Open-domain QA with Phrase Encoding and Retrieval As an alternative, phrase-retrieval approaches (Seo et al., 2018; 2019) mitigate this problem by directly accessing all phrases in K documents by decomposing F into two functions, â = arg max

x k i:j F (x k i:j , q) = arg max x k i:j H x (x k i:j ) • H q (q)

where • denotes inner product operation. Unlike running a complex reading comprehension model in pipeline-based approaches, H x query-agnostically encode (all possible phrases of) each document just once, so that we just need to compute H q (which is very fast) and perform similarity search on the phrase encoding (which is also fast). Seo et al. (2019) have shown that encoding each phrase with a concatenation of dense and sparse representations is effective, where the dense part is computed from BERT (Devlin et al., 2019) and the sparse part is obtained from the tf-idf vector of the document and the paragraph of the phrase. We briefly describe how the dense part is obtained below.

Dense Representation Assuming that the document x has N words as x 1 , . . . , x N , Seo et al.

(2019) use BERT (Devlin et al., 2019) to compute contextualized representation of each word as h 1 , . . . , h N = BERT(x 1 , . . . , x N ). Based on the contextualized embeddings, we obtain dense phrase representations as follows: We split each

h i ∈ R d into four parts [h 1 i , h 2 i , h 3 i , h 4 i ] where h 1 i , h 2 i ∈ R dse and h 3 i , h 4 i ∈ R dc (hence d = 2d se + 2d c ).

Then each phrase x i:j is densely represented as follows:

h i:j = [h 1 i , h 2 j , h 3 i • h 4 j ] ∈ R 2dse+1 (1)

where • denotes inner product operation. h 1 i and h 2 j are start/end representations of a phrase, and the inner product of h 3 i and h 4 j is used for computing coherency of the phrase. Refer to Seo et al. (2019) for details; we mostly reuse its architecture.

4. SPARSE ENCODING OF PHRASES:

Sparse representations are often suitable for keeping the precise lexical information present in the text, complementing dense representations that are often good for encoding semantic and syntactic information. While the sparsification of dense word embeddings has been explored previously (Faruqui et al., 2015) , its main limitations are that (1) it starts from the dense embedding which might have already lost rich lexical information, and (2) its cardinality is in the order of thousands at max, which we hypothesize is too small to encode sufficient information. Our work hence focuses on creating the sparse representation of each phrase which is not bottlenecked by dense embedding and is capable of increasing the cardinality to billion-scale without much computational cost.

4.1. WHY DO WE NEED TO LEARN SPARSE REPRESENTATIONS?:

Suppose we are given a question "How many square kilometres of the Amazon forest was lost by 1991?" and the target answer is in the following sentence, Between 1991 and 2000, the total area of forest lost in the Amazon rose from 415,000 to 587,000 square kilometres (160,000 to 227,000 sq mi).

To answer the question, the model should know that the target answer (415,000) corresponds to the year 1991 while the (confusing) phrase 587,000 corresponds to the year 2000, which requires syntactic understanding of English parallelism. The dense phrase encoding is likely to have a difficulty in precisely differentiating between 1991 and 2000 since it needs to also encode several different kinds of information. Window-based tf-idf would not help either because the year 2000 is closer (in word distance) to 415,000. This example illustrates the strong need to create an n-gram-based sparse encoding that is highly syntax-and context-aware.

4.2. CONTEXTUALIZED SPARSE REPRESENTATIONS:

Our sparse model, unlike pre-computed sparse embeddings such as tf-idf, dynamically computes the weight of each n-gram that depends on the context. Intuitively, for each phrase, we want to compute a positive weight for each n-gram near the phrase depending on how important the n-gram is to the phrase.

Sparse Representation Sparse representation of each phrase is also obtained as the concatenation of its start word's and end word's sparse embedding, i.e. s i:j = [s start i , s end j ]. This way, similarly to how the dense phrase embedding is obtained, we can efficiently compute them without explicitly enumerating all possible phrases. We obtain each (start and end) sparse embedding in the same way (with unshared parameters), so we just describe how we obtain the start sparse embedding here and omit the superscript 'start'. Given the contextualized encoding of each document H = [h 1 , . . . , h N ] ∈ R N ×d , we obtain its (start or end) sparse encoding S = [s 1 , . . . ,

s N ] ∈ R N ×F by S = ReLU QK √ d F ∈ R N ×F (2)

where Q, K ∈ R N ×d are query, key matrices obtained by applying a (different) linear transformation on H (i.e., using W Q , W K : R N ×d → R N ×d ), and F ∈ R N ×F is an one-hot n-gram feature representation of the input document x. That is, for instance, if we want to encode unigram (1-gram) features, F i will be simply a one-hot representation of the word x i , and F will be equivalent to the vocabulary size. Note that F will be very large, so it should always exists as an efficient sparse matrix format (e.g. csc) and one should not explicitly create its dense form. We also discuss how we can leverage it during training in an efficient manner in Section 4.3.

Note that Equation 2 is similar to the scaled dot-product self-attention (Vaswani et al., 2017) , with two key differences that (1) ReLU instead of softmax is used, and (2) a sparse matrix F instead of (dense) value matrix is used. In fact, our sparse embedding is related to attention mechanism in that we want to compute how important each n-gram is for each phrase's start and end word. Intuitively, s i contains a weighted bag-of-ngram representation where each n-gram is weighted by its relative importance on each start or end word of a phrase. Unlike attention mechanism whose role is to summarize the target vectors via weighted summation (thus softmax is used, which sums up to 1), we do not perform the summation and directly output the unnormalized attention weights to a large sparse embedding space. In fact, we experimentally observe that ReLU is more effective than softmax for this objective.

Since we want to handle several different sizes of n-grams, we create the sparse encoding S for each n-gram and concatenate the resulting sparse encodings. That is, let S n be the sparse encoding for n-gram, then the final sparse encoding is the concatenation of different n-grams we consider.

In practice, we experimentally find that unigram and bigram are sufficient for most use cases; in this case, the sparse vector for the (start) word x i will be

s i = [s 1 i , s 2 i ].

We do not share linear transformation parameters for across different n-grams. Note that this is also analogous to multiple attention heads in Vaswani et al. (2017) .

Question Encoding For a question q = [CLS], . . . , q M where [CLS] denotes a special token for BERT inputs, contextualized question representations are computed in a similar way (h [CLS] , . . . h M = BERT([CLS], . . . , q M )). We share the same BERT used for the phrase encoding. We compute sparse encodings on the question side (s ∈ R F ) in a similar way to the document side, with the only difference that we use the [CLS] token instead of start and end words to represent the entire question. Linear transformation weights are not shared.

4.3. TRAINING:



Kernel function:

As training phrase encoders on the whole Wikipedia is computationally prohibitive, we use training examples from an extractive question answering dataset (SQuAD) to train our encoders. Given a pair of question q and a golden document x (a paragraph in the case of SQuAD), we first compute the dense logit of each phrase x i:j by l i,j = h i:j • h . Unlike Seo et al. (2019) , each phrase's sparse embedding is also trained, so it needs to be considered in the loss function. We define the sparse logit for phrase x i:j as l sparse i,j

= s i:j • s [CLS] = s start i •s start [CLS] +s end j •s end [CLS]

. For brevity, we describe how we compute the first term

s start i •s start [CLS]

corresponding to the start word (and dropping the superscript 'start'); the second term can be computed in the same way.

s start i • s start [CLS] = ReLU QK √ d i F ReLU Q K √ d [CLS] F where Q , K ∈ R M ×d , F ∈ R M ×F

denote the question side query, key, and n-gram feature matrices, respectively. We can efficiently compute it if we precompute FF ∈ R N ×M . Note that FF can be considered as applying a kernel function, i.e. K(F, F ) = FF where its (i, j)-th entry is 1 if and only if n-gram at i-th position of the context is equivalent to j-th n-gram of the question, which can be efficiently computed as well. One can also think of this as kernel trick (in the literature of SVM (Cortes & Vapnik, 1995) ) that allows us to compute the loss function without explicit mapping.

The loss to minimize is computed from the negative log likelihood over the sum of the dense and sparse logits:

EQUATION

where i * , j * denote the true start and end positions of the answer phrase. While the loss above is an unbiased estimator, in practice, we adopt early loss summation as suggested by Seo et al. (2019) for larger gradient signals in early training. Additionally, we also add dense-only loss that omits the sparse logits (i.e. original loss in Seo et al. (2019) ) to the final loss, in which case we find that we obtain higher-quality dense phrase representations.

Negative sampling We train our model on SQuAD v1.1 which always has a positive paragraph that contains the answer. To learn robust phrase representations, we concatenate negative paragraphs to the original SQuAD paragraphs. To each paragraph x, we concatenate the paragraph x neg which was paired with the question whose dense representation h neg is most similar to the original dense question representation h , following Seo et al. (2019) . Note the difference, however, that we concatenate the negative example instead of considering it as an independent example with noanswer option Levy et al. (2017) . During training, we find that adding tf-idf matching scores on the word-level logits of the negative paragraphs further improves the quality of sparse representations as our sparse models have to give stronger attentions to positively related words in this biased setting.

5. EXPERIMENTS:



5.1. DATASETS:

We use two popular open-domain QA datasets to evaluate our model: SQUAD OPEN and CURAT-EDTREC. SQUAD OPEN is the open-domain version of SQUAD (Rajpurkar et al., 2016) . We use 87,599 examples with the golden evidence paragraph to train our encoders, and use 10,570 QA pairs from development set to test our model, as suggested by Chen et al. (2017) . CURATEDTREC consists of question-answer pairs from TREC QA (Voorhees et al., 1999) constructed based on real user queries. We use 694 test set QA pairs for testing our model. Note that we only train on SQuAD and test on both SQuAD and CuratedTREC, relying on the generalization ability of our model for zero-shot inference on CuratedTREC. This is in a clear contrast to previous work that utilize distant supervision (Chen et al., 2017) or weak supervision (Lee et al., 2018; Min et al., 2019) on CuratedTREC.

5.2. IMPLEMENTATION DETAILS:

We use and finetune BERT LARGE for our encoders. We use BERT vocabulary which has 30522 unique tokens based on byte pair encodings. As a result, we have F = 30522 when using unigram feature for F, and F ≈ 1B when using both uni/bigram features. We do not finetune the word embedding during training. We pre-compute and store all encoded phrase representations of all documents in Wikipedia (more than 5 million documents). It takes 600 GPU hours to index all phrases in Wikipedia. Each phrase representation has 2d se + 2F + 1 dimensions. We use the same storage reduction and search techniques by Seo et al. (2019) . For storage, the total size of the index is 1.3 TB including unigram and bigram sparse representations. For search, we either perform dense search first and then rerank with sparse scores (DFS) or perform sparse search first and rerank with dense scores (SFS), and also consider a combination of both (Hybrid).

Constraining sparse encoding We find that injecting some prior domain knowledge to avoid spurious matchings allows us to learn better sparse representations. For example, in machine reading comprehension, answer phrases do not occur as an exact phrase in questions (e.g., the phrase "August 4, 1961" would not appear in the question "When was Barack Obama born?"). Therefore, we can zero the elements in the diagonal axis of QK ∈ R N ×N in Equation 2 which correspond to the attention values of target phrase itself. Additionally, we mask special tokens in BERT such as [SEP] or [CLS] to have zero weights as matching of these tokens means nothing.

5.3. RESULTS:

We evaluate the effectiveness of COSPR by augmenting DENSPI (Seo et al., 2019) with contextualized sparse representations (DENSPI+COSPR). We extensively compare the model with the original DENSPI and previous pipeline-based QA models.

Model SQUADOPEN CURATEDTREC EM F1 Exact Match s/Q DrQA (Chen et al., 2017) 29.8 ** -25.4 * 35 R 3 (Wang et al., 2018a) 29.1 37.5 28.4 * -Paragraph Ranker (Lee et al., 2018) 30.2 -35.4 * -Multi-Step-Reasoner (Das et al., 2019) 31.9 39.2 --BERTserini (Yang et al., 2019) 38.6 46.1 -115 ORQA (Lee et al., 2019) 20.2 -30.1 -Multi-passage BERT † † (Wang et al., 2019) 53.0 60.9 --DENSPI (Seo et al., 2019) 36 (Yang et al., 2019) while being almost two orders of magnitude faster. We expect much bigger speed gaps between ours and other pipeline methods as most of them put additional complex components to the original pipelined methods.

On CURATEDTREC, which is constructed from real user queries, our model also achieves the stateof-the-art performance. Even though our model is only trained on SQUAD (i.e. zero-shot), it outperforms all other models which are either distant-or semi-supervised with at least 29x faster inference. Note that our method is orthogonal to end-to-end training (Lee et al., 2019) or weak supervision (Min et al., 2019) methods and future work can potentially benefit from these.

Ablation study Table 2 (left) shows the effect of contextualized sparse representations by comparing different variants of our method on SQUAD OPEN . For these evaluations, we use a subset of Wikipedia dump (1/100, 1/10 scale, approximately 50K, 500K documents, respectively). We evaluate the effect of removing (1) COSPR, (2) tf-idf sparse representations, (3) tf-idf bias on negative sample training (section 4.3), and (4) and adding trigram features to our sparse representation (i.e., s i = [s 1 i , s 2 i , s 3 i ]). One interesting observation is that adding trigram features in our sparse representations is worse than using uni-/bigram representations only. We suspect that the model becomes too dependent on trigram features, which means we might need a stronger regularization for high-order n-gram features.

In Table 2 (right), we show how we consistently improve over DENSPI when COSPR is added in different search strategies. Note that on SQUAD OPEN , SFS is better than DFS as questions in SQUAD are created by turkers given particular paragraphs, whereas on CURATEDTREC where the questions more resemble real user queries, DFS outperforms SFS showing the effectiveness of dense search when not knowing which documents to read. Similar observation was pointed by Lee et al. (2019) who showed different performance tendencies between two datasets, but as we are using

:

Our code will be publicly available upon acceptance.

Figures and Tables:

Figure1: Overview of our model. Using the billions of precomputed phrase representations, we perform a maximum inner product search between the phrase vectors and an input question vector. We propose to learn contextualized sparse phrase representations which are also very interpretable.

Results on open-domain question answering datasets.

Table content: <table><tr><td/><td/><td colspan="2">.2 44.4</td><td colspan="2">31.6  †</td><td>0.81</td></tr><tr><td colspan="2">DENSPI+COSPR(Ours)</td><td colspan="2">40.7 49.0</td><td colspan="2">35.7  †</td><td>1.19</td></tr><tr><td colspan="3">* Trained on distantly supervised training data.</td><td/><td/></tr><tr><td colspan="2">** Trained on multiple datasets</td><td/><td/><td/></tr><tr><td colspan="3">† No supervision using target training data.</td><td/><td/></tr><tr><td colspan="2">† † Concurrent work.</td><td/><td/><td/></tr><tr><td>Model Ours (DFS) -COSPR -doc./para. tf-idf -tf-idf bias training + trigram sparse</td><td>SQUAD 1/100 SQUAD 1/10 60.0 51.6 55.9 (-4.1) 48.4 (-3.2) 58.1 (-1.9) 50.9 (-0.7) 58.1 (-1.9) 49.1 (-2.5) 58.0 (-2.0) 49.8 (-1.8)</td><td>Model SFS DFS HYBRID</td><td colspan="2">SQUAD OPEN DENSPI + COSPR 33.3 36.9 (+3.6) 28.5 34.4 (+5.9) 36.2 40.7 (+4.5)</td><td>CURATEDTREC DENSPI + COSPR 28.8 30.0 (+1.2) 29.5 34.3 (+4.8) 31.6 35.7 (+4.1)</td></tr></table>

Ablations of our model. On the left, we show how each sparse representation contribute to the performance, and on the right, we show how much gain we obtain from COSPR in different search strategies. Exact match scores are reported.Open-domain QA experiments Table1shows experimental results on two open-domain question answering datasets. For DENSPI and ours, we use HYBRID search strategy. On both datasets, our model with contextualized sparse representations (DENSPI + COSPR) significantly improves the performance of the phrase-indexing baseline model (DENSPI). Moreover, our model outperforms BERT-based pipeline approaches such as BERTserini

Table content: <table/>

Peer review (iclr2020):
The authors propose a method for learning contextualized, word-level sparse representations for the phrase retrieval problem. Their model dynamically computes the weight of each n-gram depending on the context. The authors show that their approach outperforms most competing methods on SQUAD (open) and CURATEDTREC. The biggest problem I see is that a BERT-based approach (Wang et al., 2019) outperforms the proposed approach by a large margin and doesn*t seem to be discussed. I understand that it*s concurrent work, but have the authors investigated if there might be any setting where their method would be preferable? This does seem important. The authors further perform an ablation study of different components of the model. The subsequent analysis is short, but I like it, and I think it is helpful in the context of the paper.