Title: SUPERBLOOM: BLOOM FILTER MEETS TRANSFORMER

Abstract: We extend the idea of word pieces in natural language models to machine learning tasks on opaque ids. This is achieved by applying hash functions to map each id to multiple hash tokens in a much smaller space, similarly to a Bloom filter. We show that by applying a multi-layer Transformer to these Bloom filter digests, we are able to obtain models with high accuracy. They outperform models of a similar size without hashing and, to a large degree, models of a much larger size trained using sampled softmax with the same computational budget. Our key observation is that it is important to use a multi-layer Transformer for Bloom filter digests to remove ambiguity in the hashed input. We believe this provides an alternative method to solving problems with large vocabulary size.

1. INTRODUCTION:

In natural language processing, one recent development, made popular by Wu et al. (2016) is to use a smaller sub-word vocabulary (Sennrich et al., 2016) , or so called word piece model. In such a model, only frequent words and word pieces are kept in the vocabulary. Each word is then segmented as a sequence of word pieces. Both the input and the prediction are then represented in the smaller word piece space.

The word piece model has multiple benefits. Besides its generalizability and compact size, one crucial benefit is that we can afford to compute the full softmax loss on its much smaller vocabulary. This leads to more precise predictions, (measured e.g. using recall at k for small values of k), compared to alternative approaches such as the sampled softmax method (Bengio & Sénécal, 2003; 2008) or the hierarchical softmax (Morin & Bengio, 2005) . Word pieces have been shown to work well for natural language understanding (NLU) tasks. For example, the recent break-through of BERT (Devlin et al., 2018) uses a vocabulary of about 30K word pieces. The goal of this paper is to extend this idea to machine learning tasks where we have to model a large number of categorical values, which are represented by opaque ids (e.g. product ids, video ids) or named entities (e.g. Wikipedia or Knowledge Graph entities).

While word pieces are a natural way for breaking up words, it is unclear how this could be done for a set of arbitrary categorical values (referred to as vocabulary throughout the paper). We adopt a technique proposed by Serrà & Karatzoglou (2017) of using multiple hashing to reduce the vocabulary size while still keeping each entity identifiable. The idea is to map each id to multiple hash tokens in a smaller space, similarly to a Bloom filter (Bloom, 1970) , and then embed and predict at the token, instead of the id, level. This way, the vocabulary is reduced to a much smaller size, similar to the effect of word pieces. However, hashing introduces much ambiguity due to random collisions. To solve this problem, we propose to use hashing in combination with a multi-layer Transformer (Vaswani et al., 2017) , based on observations that the Transformer can disambiguate word meanings well using context. A hashed token can be viewed as a word piece with many different meanings. We hope that a Transformer model is also able to remove the ambiguity of hash tokens using the context, i.e. the set of other input tokens.

With these motivations, we build Superbloom in which we apply a Transformer model to the Bloom filter digest of the input. On the output side, the predictions are on the hashed tokens, similar to (Serrà & Karatzoglou, 2017) . We demonstrate, through experiments, that Superbloom works well for tasks with a large vocabulary size -it can be efficiently trained and outperforms nonhashed models of a similar size, and larger models trained with sampled softmax with the same computational budget.

The key insight from our experiments is that the multiple-layer Transformer is effective for resolving the ambiguity in the hashed input, hence works particularly well for Bloom filter digests. For instance, we find that the model quality gap between a one layer and a twelve layer Transformer model is significantly larger when using Bloom filter digests, compared to that when the vocabulary is not hashed. This capability of the Transformer to "unhash" the Bloom digest is a key difference to earlier work on feature hashing (Weinberger et al., 2009) and multiple hashing (Serrà & Karatzoglou, 2017; Svenstrup et al., 2017; Daniely et al., 2017) . In addition, we propose an efficient approximate inference algorithm, by applying "beam search" in the much smaller token space. Such fast inference is another useful property of Superbloom.

The Superbloom model is trained using BERT's masked language model task, i.e. on predicting hash tokens masked out from the input. Since the hashes have a much smaller vocabulary, this aspect is closely related to the error-correcting output codes (ECOC) (Dietterich & Bakiri, 1994; Berger, 1999) 1 . In the ECOC model, a multi-class classification is converted into many binary classification problems where redundancy is added to improve the robustness of the model. The prediction task of Superbloom can be viewed as reducing a large multi-class classification problem to a few smaller multi-class classification problems. However, unlike the ECOC models, we do not reduce the problem all way down to binary predictions. This might be the additional reason that Superbloom is able to achieve high accuracy.

1.1. RELATED WORK:

Learning with a large vocabulary is a well-studied but still open research problem. Weinberger et al. (2009) proposed feature hashing which uses random hashing to reduce the input vocabulary size, and then learns embeddings for hashed ids in the smaller vocabulary. Several follow-up works propose to better resolve collisions by using multiple hashes: Svenstrup et al. (2017) proposed to learn a weighted sum of hashed embeddings; Shu & Nakayama (2018) used an unweighted sum, but proposed instead to learn the hash function itself; and Chen et al. (2018) proposed to learn both the hash function and the combiner, for which they use either a linear function or an LSTM. A key difference with the aforementioned work is that we do not resolve the hashing early at the input of the model; instead, we feed all hashed embeddings to the Transformer and let it learn to resolve the hashing collisions using the context. Our experiments show that multi-layer Transformer models indeed have the capacity to resolve hashing collisions while learning a high quality model.

Besides reducing the input space and memory usage, another set of related work focuses on dealing with large output vocabularies and improving training efficiency. A commonly used method is sampled softmax (Bengio & Sénécal, 2003; 2008) where for each gradient update, only a subset of the output vocabulary is considered. Another line of work is hierarchical softmax where classes are organized in clusters (Goodman, 2001) or in a tree structure (Morin & Bengio, 2005) to allow for efficient pruning of the output vocabulary. Through our experiments, we show that Superbloom, which allows us to train a full softmax on the hashed vocabularies, can lead to more accurate results than using sampled softmax on the larger output vocabulary. Serrà & Karatzoglou (2017) proposed to use Bloom filters as a general tool in deep models, for both the input and output. Our work demonstrates the efficiency of a multi-layer Transformer-like architecture to use contextual information to resolve hash ambiguity. Indeed, we show that shallow models, even with attention, fail.

2. SUPERBLOOM MODEL ARCHITECTURE:

Given discrete sets S I , S O , representing respectively the input and output spaces (e.g. word tokens or entities), the goal is to model a function that maps a sequence of n elements2 in S I , to a sequence of probability distributions over S O . The space of probability distributions over a set S will be denoted by

∆(S) = {p ∈ R |S| + : s∈S p s = 1}.

The input and output entities are typically represented using embedding matrices E I ∈ R |S I |×d and E O ∈ R |S O |×d , which map each entity to an embedding vector of dimension d. This makes training and inference expensive if the number of entities is very large. In order to reduce the model size and We decompose the Superbloom model architecture into ) and inference (Sec. 2.5) need to be adapted accordingly.

M = O • (T L • • • • • T 1 ) • I,

2.1. INPUT LAYER:

I : (S I ) n → R mn×d

The input layer consists of m hash functions3 h j : S I → H I , j ∈ {1, . . . , m} and an embedding matrix E I ∈ R |H I |×d . The input sequence (s 1 , . . . , s n ) is mapped to the sequence of embeddings (E h1(s1) , . . . , E hm(s1) , E h1(sn) , . . . , E hm(sn) ). For ease of notation, we will write x i,j = E hj (si) ∈ R d , and denote the sequence by (x i,j ) i=1,...,n j=1,...,m . Throughout, we use subscripts i, j to denote the j-th hash of the i-th element.

2.2. TRANSFORMER LAYERS:

T : R mn×d → R mn×d

The Transformer is an attention-based model that was initially proposed for sequence transduction tasks, and that has been used in various other settings such as BERT. For the intermediate layers of Superbloom, we use the same architecture as the original transformer model (Vaswani et al., 2017) , which we briefly summarize in Appendix A. Each transformer layer is a function T : R mn×d → R mn×d which maps a sequence of mn embeddings in R d to another sequence in the same space. We will denote by (y i,j ) i=1,...,n j=1,...,m the output sequence of the last transformer layer, where each y i,j ∈ R d .

2.3. OUTPUT LAYER:

: O : R mn×d → ∆(H O ) mn

Similarly to the input layer, we have m hash functions η j : S O → H O , j ∈ {1, . . . , m} for the output space. We modify the original goal of predicting distribution over S O to predicting distributions over H O , as follows. If (y i,j ) i=1,...,n j=1,...,m is the output of the last transformer layer, then the output layer O maps each y i,j to

p i,j := σ(y i,j (E O ) ) ∈ ∆(H O ), where E O ∈ R |H O |×d

is an output embedding matrix, and σ is the softmax function. Note that in some problems, the input and output spaces coincide, so it can be advantageous to use identical input and output hash functions, h j = η j , and the same embedding matrices (p i,j , η j (t i )).

E I = E O . 2.4 TRAINING If the target sequence in S O is (t 1 , . . . , t n ), then the corresponding target sequence in H O is (η j (t i )) i=1,...,

where p i,j ∈ ∆(H O ) are the model's output distributions, and

: ∆(H O ) × H O → R is a loss

function, e.g. the cross-entropy loss. Note that we can pre-process the training data to map the elements in the original spaces (S I ) n , (S O ) n to the hash spaces (H I ) mn , (H O ) mn , and training proceeds entirely in the hash spaces.

Model size and efficiency Compared to a model trained on the original space, the main advantage of Superbloom is a reduction in the size of the embedding matrices E I , E O . For instance, if a α-toone hashing is used (i.e., each hash bucket contains α elements), then |H| = |S|/α and the size of the input matrices is reduced by a factor α. This not only reduces the memory cost, but may also improve the efficiency of gradient updates during training. Consider a cross-entropy loss, for each training example, all elements in the output space have a non-zero gradient due to the partition function in softmax, and thus the full matrix E O needs to be updated at each step, unless approximate methods such as sampled softmax (Bengio & Sénécal, 2003) For each j, compute S b j = {s ∈ S O : p j (η j (s)) ≥ p b j }.

6: Score all candidates in S b = S b 1 ∪ • • • ∪ S b m . Let s = arg max s∈S b γ(s). 7: if γ(s ) ≥ j log p b j then break.

This guarantees γ(s) ≤ γ(s ) for all s.

8: return s .

2.5. INFERENCE:

For each position i in the sequence, the model outputs m distributions (p i,j ) j=1,...m ∈ ∆(H O ) m , and our goal is to use these distributions to rank the elements of S O . To simplify notation, we assume in this section that i is fixed and will omit it by writing p j instead of p i,j .

One simple way to rank items, as proposed by Serrà & Karatzoglou (2017) , is to compute, for each s, γ(s) := m j=1 log p j (η j (s)). When S O is very large, this can be expensive, so instead of exhaustively scoring all items s, we propose an iterative beam-search, given in Algorithm 1 and illustrated in Figure 2 , that can be used to compute the top-k elements5 , either exactly, or approximately by fixing a maximum iteration number.

Let us fix a beam width b, and let p b j be the b-th largest value of p j , and let S b j = {s ∈ S O : p j (η j (s)) ≥ p b j }. In words, S b j are elements whose hash is in the top b values according to p j . S b j is obtained by pre-computing and storing inverse look-up tables6 η -1 j for each hash function, and observing that S b j = ∪ h:pj (h)≥p b j η -1 j (h). This defines a set of candidates to score

S b := S b 1 ∪ • • • ∪ S b m ,

and guarantees the following upper-bound: for all s / ∈ S b , γ(s) ≤ j log p b j . If the best scored element s := arg max s∈S b γ(s) satisfies γ(s ) ≥ j log p b j , then we have a certificate that s is the best element over the entire set and the algorithm terminates. Otherwise, we increase the beam width and score a new set of candidates.

An example is illustrated in Figure 2 for m = 2, a beam width B = 2, and hash functions with α = 4 (four elements share the same hash value along each dimension). The subset of candidates to score during the first iteration is highlighted in blue. The top element s is circled, and the solid line shows its γ level set. In the left figure, the level set does not intersect the shaded area (unscored elements), thus we have a certificate that s is the exact maximizer. In the right figure, the level set does intersect the shaded area, so to find the exact maximizer, a second iteration is performed where the search region is extended (highlighted in red). For example, with the parameter setting described in Section 3.3, approximate inference is 10 times faster than exhaustive scoring. In Appendix B, we study the effect of the beam width on the quality of the model.

Computational complexity:

Remark 1. While we describe a particular choice of ranking function γ, it is possible to generalize the algorithm to other ranking functions that are increasing, in a sense described in Appendix C.

3. WIKIPEDIA ENTITY PREDICTION:

We apply Superbloom to the Wikipedia entity prediction task, in which we use surrounding links on a Wikipedia page to predict a held-out link. This task is derived from the same data set as many NLU tasks, but uses entities instead of natural language. We believe this study is complementary to previous NLU models trained on Wikipedia, that focus on modeling language. Indeed, we show through examples that the model can learn entity relations well and demonstrates a strong use of contextual information.

The task needs to model about 5.3 million entity pages on Wikipedia. This vocabulary size is two orders of magnitude larger than in previous work that applies a Transformer model with full softmax loss (Devlin et al., 2018; Zhang et al., 2018; Sun et al., 2019) . Other works, such as Zhang et al. (2019) and Soares et al. (2019) , train a Transformer model with a large number of entities using sampled softmax, with either in-batch or in-example negative sampling. But as we shall show, sampled softmax, even with a large number of 128K negative samples, results in much worse quality.

3.1. TASK:

We take all the entity pages on the website en.wikipedia.org. For each page, we obtain the URL links to other Wikipedia entity pages. We only use "raw" links, i.e. links that explicitly appear on the page. We obtain 5,281,889 pages and 462,588,415 links. Since the Wikipedia site usually removes duplicates of links on each page, the distribution of pages is rather long tail. For example, the top 100 most frequent pages represent only 3.8% of the total links, and the top 10% most frequent pages represent about 60% of the total links. We hold out 10% random entity pages for testing. For the training data, we apply a masking similar to BERT -from each page, we take a random contiguous segment of entities, of length up to n = 32, and mask 15% of the segment. The task is then to predict the masked entities. We also apply the same input perturbation, where for the input, each masked out link is either replaced with a special [MASK] entity (with 80% probabilty), replaced with a random entity (with 10% probability), or left unchanged (with 10% probability). For evaluation, we hold out (i.e. replace with the [MASK] token) one random entity from a random segment on a test page. For quality evaluation, we use recall at k metric (abbreviated as rec@k below), which represents the chance the held out entity is in one of the top k predictions.

3.2. MODEL:

To apply Superbloom, we first create m hash maps from entities to hash tokens with a given hash density α. Each hash map is obtained by applying a random permutation to the vocabulary and map every consecutive α entities to the same token. This way we guarantee each hash token to have the same number of collisions α.7 Special tokens [CLS], [MASK], [SEP] , are each mapped to m tokens with no collisions. For example we create

[MASK 1 ], .., [MASK m ] tokens corresponding to [MASK].

We apply the hashing to the input and target, to map each entity to m tokens as described in Section 2. We then apply the Transformer model to the input to predict the masked tokens. Unlike in BERT, we do not use position embeddings, in other words, we treat the input as a set instead of a sequence. Since the input and output spaces coincide, we use the same hash functions and the same embedding matrices in the input and output layer.

We carry out experiments on both the full vocabulary as well as a smaller subset consisting of the top 500K entity pages. On the smaller vocabulary, we are able to train a baseline model with large capacity, with no hashing and no sampling, which is useful for understanding the best achievable model quality.

We train all of our models on 16 Cloud TPUs. We use a batch size of 1024 for experiments with full vocabulary and 4096 for experiments with 500K vocabulary. All the experiments use the Adam optimizer (Kingma & Ba, 2014), and use a decreasing learning rate sequence with inverse square root decay, and initial learning rate 1e-4 for the full vocabulary and 2e-4 for the 500K vocabulary. All the experiments have been run for more than 1 million steps to reach near convergence.

3.3. SUPERBLOOM IS MORE ACCURATE:

We experiment with two models of similar size: one is a baseline model (baseline) with full vocabulary of size N equal to the number of entities; the other is a Superbloom model (superbloom) with a heavy 50 to 1 hashing. We set other hyper-parameters (such as the embedding dimension) so both models have a similar size. We also compare to a large model (sampled-softmax) trained using sampled softmax. The Superbloom model clearly outperforms, to a large extent, both the baseline and the sampledsoftmax model. We note that the sampled-softmax model has much worse rec@k than the other two models, and this gap is larger for smaller k. This is not surprising given the relatively small percentage (2.5%) of negative examples we can afford to sample.

While the Superbloom model performs well overall, there is a possibility that it devotes most of the embedding capacity to the top entities, so it loses accuracy on the less frequent entities. To test this, we plot the rec@1 value as a function of label frequency. In Figure 3 , we show the mean rec@1 for every 10 percentile bucket in terms of the label frequency. We can observe that Superbloom is more accurate than the baseline in all the buckets. Another interesting phenomenon is that the most challenging labels are those in the 20 and 30 percentile. One possible reason is that they lack the higher predictability of the most frequent labels, and also the strong regularity of less frequent labels. Figure 3 : Rec@1 with respect to label frequency, starting from the most frequent labels.

Besides the high predictive accuracy, the prediction from the model shows strong semantic ability and context dependency. We show some examples of predictions in Figure 4 in Appendix E. In one set of examples, we pair "Copenhagen" with different entities, and observe that the predictions change accordingly, depending on the context. Another observation is that despite the heavy hashing, there are almost no unrelated entities in the top 10 predictions. The model even exhibits an ability to perform certain analogy tasks (without being trained on such tasks) -for example, given "Tunisia Tunis Thailand", it predicts "Bangkok" as the top result.

3.4. MULIT-LAYER TRANSFORMER IS IMPORTANT FOR SUPERBLOOM:

Intuitively, given the large noise introduced by hashing, it is more important for Superbloom to use multiple attention layers in Transformer to "remove" the noise. To test this intuition, we run experiments with a smaller vocabulary size of the top 500K entity pages (about 60% of the links). On this smaller vocabulary size, we can afford to run a full softmax model with a larger embedding dimension. We consider different embedding dimensions and model complexity. Table 3 lists the model parameters well as the recall metrics for each model. We observe that for the baseline models, the quality difference is small between models of different complexity. For example, rec@1 of baseline-l12 (55.0%) is about 8% better than baseline-l1 (51.0%). Since a one layer Transformer is close to a bag-of-words (BOW) model, one may argue that it may be unnecessary to use a Transformer in this case -instead one can use a larger dimension BOW model to achieve a similar accuracy.

EQUATION

However, for Superbloom models, the quality improves significantly with more layers. When increasing the number of layers from 1 (superbloom-d256l1) to 12 (superbloom-d256l12), rec@1 increases from 17.8% to 43.4%. The multi-layer model also performs much better than the single layer model with the same size (superbloom-d384l1). Note that previous work on hashed vocabularies relies on BOW models, which are less expressive than even a single-layer transformer. This highlights one of our key observations that multi-layer Transformer models are more effective for working with hashed vocabularies.

4. EXPERIMENTS ON NATURAL LANGUAGE DATA:

In this section, we apply Superbloom to natural language data. We consider a large vocabulary that contains frequent unigrams and bigrams and use it to tokenize the text, then apply a Bloom filter to reduce the vocabulary size. We show that despite high hash collisions, the model can achieve high accuracy on natural language data. Since many named entities appear in the large vocabulary, we observe that the model seems to make better predictions of named entities than the BERT model.

While each hash id can be regarded as a word piece in an NLU model, there are important differences between hash ids and word pieces. First, hashing causes random collisions, while wordpiece tokenization can be viewed as a special hashing scheme based on the spelling -there is often coherence between words that share a word piece. As suggested by the experiments in Appendix D, random hashing with Superbloom digests may outperform coherent hashing. In addition, as every token in the large vocabulary is hashed, we do not have unambiguous anchors (such as the exact word pieces) to help bootstrap the disambiguation process. Despite these differences, our experiments suggest that even with high hashing collision α = 40, the Transformer is capable of resolving, or unhashing, the Bloom filter digest effectively and produces highly accurate predictions and meaningful embeddings.

We construct a vocabulary of size 1M by taking the union of standard BERT word piece vocabulary (∼ 30K) with the most frequent unigrams and bigrams, and follow the same procedure in BERT to create training examples. For Superbloom, we apply random hash maps to the 1M vocabulary similar to the approach described in Section 3.2 to ensure an even number of collisions. The Superbloom architecture is chosen to have a comparable model size to the baseline BERT model.

We compare four models: For the non-hashed baselines, we have a large model with embedding dimension d = 256, and a small model with d = 64. And we have two Superbloom models with similar model sizes. We list the parameters in Table 4 . In models. We observe that with comparable model size, Superbloom outperforms the baseline model in all the recall metrics, and the improvement is more significant for smaller model size.

model name rec@1 rec@10 rec@20 model name rec@1 rec@10 rec@20 baseline-h64 

5. CONCLUSION:

Our experiments show that the multi-layer Transformer is effective for achieving high accuracy on hashed inputs, represented using Bloom filter digests. Besides applying it to tasks with large vocabularies, it also points to a few interesting future research directions.

The Transformer model has been mostly studied in natural language settings and for sequence data.

In our setup, we show that it can work effectively with sets of hashed entities. We hope that by investigating this simpler setup, it can help us better understand the properties of the Transformer. For example, due to hashing, each token is similar to words with multiple meanings, so its embedding can be viewed as a combination, possibly linear (Arora et al., 2018) , of the embeddings of multiple entities. A multi-layer Transformer model may provide a mechanism for iteratively filtering such noisy representations, using the context. It would be interesting to further study this mechanism.

While hashing adds noise to the learned representations, it can also increase the flexibility of these representations -when we hash multiple entities to the same token, the model is free to allocate the corresponding embedding unevenly among entities, which results in a different effective embedding dimension for each entity. Such learned capacity allocation might be more efficient than using a fixed embedding size or frequency-based allocation. Of course, an effective "denoising" model is a pre-requisite for such an approach to work. Perhaps Superbloom, with its strong denoising ability, can help further realize the potential of embedding models on hashed vocabularies. c(ρ) = j log ρ j (used in Section 2.5), but also includes a much larger class, e.g. c(ρ) = min j ρ j or c(ρ) = max j ρ j . The beam search Algorithm 1 can be immediately generalized to any such function. The only required change is to replace the condition on line 7 (γ(s ) ≥ j log(p b j )) with γ(s ) ≥ c(p b ). To prove correctness, one needs to verify that the optimality certificate holds in this general case, i.e., Lemma 1. Let γ(s) = c((p j (η j (s))) j ), where c is increasing in the sense defined above, and let p b , S b j be as defined in Algorithm 1. Then,

∀s / ∈ S b , γ(s) ≤ c(p b ).

It follows from the lemma that the algorithm can terminate whenever γ(s ) ≥ c(p b ), since γ(s ) ≥ γ(s) for all s ∈ S b (by definition of s ) and for all s / ∈ S b (by the lemma).

Proof. Since S b = ∪ j {s : p j (η j (s)) ≥ p b }, then the complement of S b is the set {s : p j (η j (s)) < p b j ∀j}. Thus, since c is increasing, it follows that for all s / ∈ S b , c((p j (η j (s))) j ) ≤ c(p b ), as claimed.

D COMPARISON OF DIFFERENT HASHING SCHEMES:

We have used random hashing functions in Superbloom. One natural alternative is "coherent" hashing, in which we map similar entities to the same hash bucket. A potential benefit of coherent hashing is that it may use embedding capacity more effectively by sharing it among similar entities. However, the downside is that it becomes difficult to distinguish those similar entities.

To create a coherent hashing function, we first run a co-occurrence factorization algorithm and then group similar entities together using the following procedure, designed to guarantee equal-sized hash buckets. For each entity, in decreasing frequency order, we compute the nearest neighbors (scored using cosine similarity), then create a hash bucket that includes the elements and its α -1 nearest neighbors which have not been already assigned a bucket. When creating a second coherent hash function, we add the constraint that any pair of elements that share a bucket for the first hash function cannot be assigned to the same bucket in the second hash. This ensures that no two elements have the same collision in both hash functions.

We carry out the experiments on the data set with smaller vocabulary (500K). We train different models that all use two hash functions, with the following configurations: both random, one random and one coherent; and both coherent. We also use different hashing densities α = 10 and α = 20. All the models have the same hyper-parameters as the superbloom-l12 model in Section 3.4. The results are given in the following We observe that with coherent hashing, we get higher accuracy for predicting hash tokens but lower accuracy for predicting entities. And the entity recall@1 is significantly lower when both hash functions are coherent. This indicates that with higher coherence, it becomes increasingly difficult for the model to make finer distinctions between similar items.

F EXAMPLES OF NATURAL LANGUAGE ENTITY PREDICTIONS:

Miles Davis is a Jazz musician, he is similar to [MASK] .

BERT: jazz himself beethoven him davis chopin bowie williams jones baseline-h256: miles davis john coltrane bill evans charlie parker louis armstrong sonny rollins keith jarrett thelonious monk jazz duke ellington hash20-h1024: miles davis john coltrane charlie parker thelonious monk dizzy gillespie bill evans billie holiday duke ellington humans is louis armstrong Empire state building is an iconic site of [MASK1] , it is close to [MASK2] .

[MASK1] BERT: architecture chicago manhattan downtown pittsburgh art philadelphia history washington america baseline-h256: architecture modern art contemporary art modern architecture national significance new york art its day historical significance the city hash20-h1024: the city new york lower manhattan manhattan the neighborhood downtown wall street the area harlem architecture

[MASK2] BERT: downtown it chicago philadelphia rome london broadway manhattan chinatown campus baseline-h256: downtown downtown pittsburgh city hall new york the city times square columbia university san francisco philadelphia the pentagon hash20-h1024: central park city hall times square wall street union station broadway lower manhattan the pentagon fifth avenue carnegie hall 4 .

:

This connection is suggested by an anonymous reviewer.

We assume a fixed sequence length for simplicity. This is also a useful assumption for practical implementation on a TPU, which requires fixed input dimensions.

The hash functions and their inverse mappings are randomly generated and stored as look-up tables. When generating the hash functions, we make sure that each hash bucket is evenly sized, and that there are no complete collisions.

Note that unlike ECOC(Dietterich & Bakiri, 1994), the task is not to predict the individual bits in the output Bloom digest, but rather to predict (a probability distribution over) the index of the m non-zero bits.

For simplicity, we describe the algorithm for k = 1, but we apply it for larger k in our experiments.

The cost of storing inverse look-up tables is dominated by that of storing embedding tables as a long as mα < d for an α-to-one hashing scheme, since the inverse lookups have total size O(mα|H O |), while the embedding tables have size O(d|H O |). This is always the case in our experiments.

The procedure described here is for simplicity. If we are concerned with space, we may use some space efficient methods, for example a perfect hash function(Fredman et al., 1984).

A minor difference with the original Transformer model is that we operate on R mn×d instead of R n×d , since we have m embeddings for each element in the sequence.

Figures and Tables:

are used. Our experiments (see Section 3.3) show that the cost of updating E O dominates that of training, and a reduction in vocabulary size allows us to significantly reduce training time without resorting to negative sampling. Algorithm 1 Approximate and exact inference in Superbloom 1: Input: Beam width B, a maximum iteration number N , model outputs p j ∈ ∆(H O ) and hash inverse look-up tables η -1 j , for j = 1, . . . , m. 2: For each j, sort p j 3: for b = B, . . . , N B do 4:Let p b j be the b-th largest value in p j .5:

Figure 2: Illustration of approximate and exact inference, with a number of hashes m = 2, a fourto-one hashing scheme, and a beam width B = 2.

Consider an α-to-one hashing scheme. Sorting the vectors p j (line 2) costs O(m|H O | log |H O |). Each iteration consists of computing S b j (line 5) then scoring candidates in S b (line 6) which costs O(m 2 Bα). The total cost for N iterations is O(m|H O | log |H O | + m 2 N Bα) which can be significantly cheaper than scoring all candidates O(|S O |).

Figure 5: Natural language fill-in-the-blank examples. BERT is the base BERT model in Devlin et al. (2018); baseline-h256 and hash20-h1024 are the Superbloom models with 1M vocabulary, with model parameters listed in Table4.

Given m hash functions h j , j ∈ {1, . . . , m}, each element s is represented by the hashes (h 1 (s), . . . , h m (s)), which we refer to as a Bloom digest, given its similarity to the Bloom filter data structure. The set of values the hashes can take is much smaller than the original spaces S I , S O , which allows us to reduce the vocabulary size and thus the size of embedding matrices.

Table content: <table><tr><td/><td/><td colspan="2">Input layer</td><td/><td colspan="3">Transformer layers</td><td/><td colspan="2">Output layer</td><td/><td colspan="3">Output hashing</td></tr><tr><td colspan="2">Input hashing</td><td/><td colspan="2">Input embedding</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>Loss</td><td/><td/><td/><td/></tr><tr><td>s 1 . . . s n</td><td>h</td><td>h1(s1) . . . hm(s1) . . . h1(sn) . . .</td><td>E I</td><td>x 11 . . . x 1k . . . x n1 . . .</td><td>L layers Attention layer norm</td><td>Feed Forward</td><td>layer norm</td><td>y 11 . . . y 1k . . . y n1 . . .</td><td>E O</td><td>softmax</td><td>p 11 . . . p 1k . . . p n1 . . .</td><td>η1(t1) . . . ηm(t1) . . . η1(tn) . . .</td><td>η</td><td>t 1 . . . t n</td></tr><tr><td/><td/><td>hm(sn)</td><td/><td>x nk</td><td/><td/><td/><td>y nk</td><td/><td/><td>p nk</td><td>ηm(tn)</td><td/><td/></tr><tr><td>(S I ) n</td><td colspan="2">(H I ) mn</td><td/><td>R mn×d</td><td/><td/><td colspan="2">R mn×d</td><td/><td colspan="2">∆(H O ) mn</td><td>(H O ) mn</td><td/><td>(S O ) n</td></tr><tr><td/><td/><td/><td/><td colspan="7">Figure 1: Superbloom model architecture</td><td/><td/><td/><td/></tr><tr><td colspan="8">improve efficiency, we hash each element as follows.</td><td/><td/><td/><td/><td/><td/><td/></tr></table>

as illustrated in Figure1: an input layer (Sec. 2.1) I : (S I ) n → R mn×d which maps each item in the input sequence to m embeddings of dimension d; L transformer layers (Sec. 2.2) T i : R mn×d → R mn×d which apply transformations in the embedding space; and an output layer (Sec. 2.3) O : R mn×d → ∆(H O ) mn mapping each embedding to a probability distribution. Since the model predicts distributions over H

Table content: <table/>

n j=1,...,m . We define the training objective 4 as

Table content: <table><tr><td>n</td><td>m</td></tr><tr><td>i=1</td><td>j=1</td></tr></table>

Table1lists the hyper-parameters of each model. Recall that α denotes the number of collisions (1 if there is no hashing), d the embedding dimension, n A the number of attention heads, d F the dimension of intermediate hidden layers, and L the number of transformer layers. In all of our experiments, we use two hash functions for Superbloom models. Hence their vocabulary size is 2N/α. Model parameters. "#samples" lists the number of samples in the softmax loss computation. For baseline and superbloom, since there is no sampling, this number corresponds to the full vocabulary, 5.3M and 200K, respectively. For sampled-softmax, we use 128K samples.Table2shows the recall metrics of the models. For the Superbloom model, we set the beam width to B = 20 (our experiments suggest that it is sufficient to set B = k in order to achieve the best rec@k metric, see Appendix B for details).

Table content: <table><tr><td>model</td><td>α</td><td colspan="2">d n A</td><td>d F</td><td colspan="2">L #parameters #samples</td></tr><tr><td>baseline</td><td>1</td><td>48</td><td colspan="3">4 1024 12</td><td>248M</td><td>5.3M</td></tr><tr><td>sampled-softmax</td><td colspan="2">1 512</td><td colspan="3">8 2048 12</td><td>2.6G</td><td>128K</td></tr><tr><td>superbloom</td><td colspan="5">50 768 12 3072 12</td><td>229M</td><td>200K</td></tr><tr><td colspan="2">model</td><td/><td colspan="4">rec@1 rec@10 rec@20</td></tr><tr><td colspan="2">baseline</td><td/><td colspan="2">36.2%</td><td>63.1%</td><td>68.2%</td></tr><tr><td colspan="3">sampled-softmax</td><td/><td>3.1%</td><td>36.2%</td><td>55.1%</td></tr><tr><td colspan="2">superbloom</td><td/><td colspan="2">51.1%</td><td>72.3%</td><td>76.5%</td></tr></table>

Recall metrics for different models.

Table content: <table/>

Model parameters and recall metrics.

Table content: <table><tr><td/><td>256</td><td>1 1024</td><td>1</td><td>123M 51.0%</td><td>70.4%</td><td>75.5%</td></tr><tr><td>baseline-l12</td><td>1 256</td><td colspan="2">8 1024 12</td><td>132M 55.0%</td><td>73.7%</td><td>77.3%</td></tr><tr><td>superbloom-d256l1</td><td>20 256</td><td>1 1024</td><td>1</td><td>13M 17.8%</td><td>35.8%</td><td>42.6%</td></tr><tr><td>superbloom-d384l1</td><td>20 384</td><td>1 1536</td><td>1</td><td>21M 30.6%</td><td>52.9%</td><td>58.7%</td></tr><tr><td colspan="2">superbloom-d256l12 20 256</td><td colspan="2">8 1024 12</td><td>21M 43.4%</td><td>60.1%</td><td>64.0%</td></tr></table>

Table 5 we list the recall metrics for the The model parameters.

Table content: <table><tr><td>model</td><td>α</td><td colspan="2">d n A</td><td>d F</td><td colspan="2">L #parameters</td></tr><tr><td>baseline-h64</td><td>1</td><td>64</td><td>4</td><td colspan="2">256 12</td><td>62.6M</td></tr><tr><td>baseline-h256</td><td>1</td><td>256</td><td colspan="3">8 1024 12</td><td>254.4M</td></tr><tr><td>hash40-h512</td><td>40</td><td>512</td><td colspan="3">8 2048 12</td><td>62.3M</td></tr><tr><td colspan="6">hash20-h1024 20 1024 16 4096 12</td><td>246.3M</td></tr></table>

Recall metrics.Since many named entities are included in the larger vocabulary, the Superbloom model shows that it may have better "understanding" or representation of those entities. We show some anecdotal evidence in Appendix E by comparing predictions of pretrained BERT and Superbloom model on some fill-in-the-blanks examples. The BERT model often predicts generic words, seemingly ignoring other named entities in the sentence. The Superbloom model, on the other hand, can often fill in the blank with related entities.

Table content: <table><tr><td>28.4%</td><td>44.9%</td><td>48.6% baseline-h256 37.2%</td><td>57.4%</td><td>63.3%</td></tr><tr><td>hash40-h512 31.7%</td><td>48.3%</td><td>52.9% hash20-h1024 39.2%</td><td>58.5%</td><td>64.5%</td></tr></table>

table. Random hashing versus coherent hashing.

Table content: <table><tr><td>model</td><td colspan="4">α #coherent hashing token rec@1 entity rec@1</td></tr><tr><td colspan="2">hash10-00 10</td><td>0</td><td>36.32%</td><td>52.50%</td></tr><tr><td colspan="2">hash10-01 10</td><td>1</td><td>38.19%</td><td>50.20%</td></tr><tr><td colspan="2">hash10-11 10</td><td>2</td><td>38.55%</td><td>34.70%</td></tr><tr><td colspan="2">hash20-00 20</td><td>0</td><td>33.39%</td><td>43.70%</td></tr><tr><td colspan="2">hash20-01 20</td><td>1</td><td>36.98%</td><td>41.10%</td></tr><tr><td colspan="2">hash20-11 20</td><td>2</td><td>37.65%</td><td>30.20%</td></tr></table>

Peer review (iclr2020):
The authors propose to learn a Transformer in embedded spaces defined via hash functions, with application to (fixed length) sequence-to-sequence classification for NLP. The method differentiates itself, in large part, by formulating model outputs in terms of distinct softmax layers, each corresponding to a hash of output space . The relative low dimensionality of each hashed output space helps to avoid computational bottlenecks associated with wide softmax layers. For a given labelled pair , the model outputs a matrix , whose -th row is trained using the -th hash . At test-time, predictions are made by finding tokens whose hashes best *align* with model outputs . By way of example, target alignment might be defined in terms of the aggregated log likelihood , where can be understood as the element of corresponding to the -th hash of token . At this time, the work is significantly hindered by a lack of clarity. This issue begins with the chosen notation, which routinely obscures otherwise simple points. Similarly, the section on test-time predictions (Sect 2.5) is needlessly hard to follow. Given its pivotal role, this section feels strangely rushed. Some unanswered questions I had include: a) What happens if two tokens* hashes collide all times? b) How were hash functions inverted and what was the cost of doing so? c) How does the test-time throughput of the proposed compare with that of alternatives? Questions: - Did you compare against different approaches to sparse softmax, such as LSH-based methods [1]? - What was the impact of approximate vs. exact inference and which was used during experiments? - How important were embedding matrices ? What happens if you directly feed hashes into the Transformer or use random projections for ?. [1] *Deep Networks With Large Output Spaces*, Vijayanarasimhan et al, 2015