\section{Introduction}
Using neural models for representing and processing textual data has become a de-facto standard in research and practice. Recent approaches to information retrieval (IR) and natural language processing (NLP) tasks adopt contextual language models (e.g., BERT \cite{devlin2019bert}) that gives promising results with richer representations of texts, ameliorating not only synonymy but also polysemy problems in IR \cite{zhan2020repbert,khattab2020colbert,ding2020rocketqa}. In these approaches, queries and documents are first encoded into contextual embeddings, either independently \cite{zhan2020repbert} or with interactions \cite{ding2020rocketqa}, which are usually low-dimensional dense representations. Documents are retrieved based on a similarity metric, such as cosine similarity, defined for two vectors.

Despite the positive results of the dense embedding approaches, every dimension is too polysemous to be used as a discriminating feature because every dimension is heavily re-used because only a few dimensions are available (low-dimensional), and the dimensions are always filled with some information that requires inspection, incurring inefficiency (dense) \cite{conneau-etal-2018-cram}. For this reason, it's usually not possible to build an efficient first-stage neural ranker without relying on external modules such as the Approximate Nearest Neighbors (ANN) engines \cite{8733051}. Most notable drawbacks of the dense embedding approaches to IR include: 1) The process of matching query and document embeddings for a large number of documents is inherently inefficient, making it hardly useful for full ranking without requiring extra efficiency-enhancing resources; 2) it is not straightforward to adopt the term-based symbolic IR techniques, such as pseudo-relevance feedback (PRF) and inverted index, that have been proven to be effective and/or useful for industrial applications, and 3) The retrieval results are hardly interpretable like other neural network approaches, making it difficult to improve the design through failure analyses or implement conditional/selective processing \cite{belinkov-glass-2019-analysis}.

Another line of recent research on neural IR attempts to improve term-based symbolic IR by taking advantage of contextualized representations of words for document expansion \cite{nogueira2019document,Nogueira2019FromDT}, term weighting \cite{10.1145/3397271.3401204}, and learn new term-based document representations \cite{bai2020sparterm}. The approaches along with this line address the three drawbacks mentioned above, lending itself to first-stage ranking with improved results over traditional term-based approaches like BM25. However, these sparse approaches are symbolic in nature and thus their capabilities are confined to the representational power (or their lack) of the original vocabulary space, still suffering relying on lexical matching.

Our focus in this paper is to propose a full-blown neural IR model that is amenable to an inverted index while alleviating the aforementioned shortcomings of dense neural IR and yet achieving compatible effectiveness. The main thrust of our approach is to utilize ultra-high dimensional (UHD) embeddings\footnote{The dimension size is close to half a million.} with high sparsity (i.e. a small fraction of non-zero dimensions), which possesses superior expressive and discriminative power required for effective and efficient IR. We construct UHD embeddings with BERT so that they are not only disentangled for fine-grained semantics, compared to dense embeddings, but also context-sensitive, addressing the word polysemy problem.

In order to obtain sparse representations (i.e. sparsification), we adopt the $Winner\text{-}Take\text{-}All$ (WTA) method for high dimensional sparse representations \cite{Ahmad2015PropertiesOS} with a new learning objective for IR. It employs a linear layer that only preserves top-$k$ activation and set the others to zero. WTA is chosen partly because we can precisely control the outputs' sparsity by explicitly setting the parameter $k$ whereas the method relying on regularization for sparsity (e.g., minimizing L1-norm) \cite{faruqui2015sparse, sun2016sparse, DBLP:journals/corr/abs-1903-11257} is neither straightforward nor precise in controlling the degree of sparsity. The method prevents us failing to model the samples that comprise only relatively weak information. We also observe that it is possible to increase the sparsity at inference time without retraining the model, which is useful for flexibly adjusting the trade-off between effectiveness and efficiency.

Besides the practical goals in terms of efficiency and effectiveness, our approach aims at a versatile neural IR model that can represent different aspects of linguistic properties of documents and queries. Instead of the common approach of using the final layer of BERT for document and query representations, we attempt to make use of multiple layers, each of which emits token representations for different linguistic properties \cite{jo-myaeng-2020-roles}. We devise a bucketing mechanism to select and merge different parts of a document/query representation, which participate in the final matching process. Buckets can be constructed for a single layer horizontally or arranged vertically for different layers. Finally, our proposed approach has the potential to make individual dimensions of the document/query embeddings interpretable \cite{faruqui2015sparse, sun2016sparse} and to enable mutually non-interfering aggregation \cite{Ahmad2015PropertiesOS}.

We evaluate our fully-blown neural IR model as the first-stage ranker on the MSMARCO passage retrieval task by comparing it against both the sparse approaches and dense approaches. The results show that our model outperforms the sparse models by a large margin, and is competitive to dense models especially in terms of recall.

In summary, our key contributions are:

\begin{itemize}
    \item We propose UHD-BERT, which generates an ultra-high dimensional (UHD) sparse representation designed to support efficient full ranking by adding the benefits of symbolic IR to neural representations. A novel bucketed learning method is employed for effective training of the model\footnote{The code will be publicly available upon acceptance.}.

    \item We show that our choice of the WTA model for sparsification enables us to reliably generate sparse representations, at the same time giving us the flexibility to control the sparsity (i.e., efficiency) at inference time depending on the need of the application or task.

    \item We prove that our model outperforms previous sparse models and is competitive to dense models especially in terms of recall, which is advantageous for the first-stage ranker.
\end{itemize}

\begin{figure*}[ht!]
  \centering
  \includegraphics[width=0.6\linewidth]{overall_architecture}
  \caption{Overall architecture with the training scheme for building UHD sparse representations.}
  \label{fig:overall_architecture}
\end{figure*}

\section{Related Work}
\label{sec::2_related_work}
With the rapid development of neural models, dense neural retrieval models have become dominant for IR tasks. However, many state-of-the-art neural retrieval models, such as RocketQA \cite{ding2020rocketqa}, utilize a cross-encoder that shows high effectiveness while known to be impractical for large-scale retrieval \cite{luan2020sparse}. This is due to the inherent problem of the cross-encoder where queries and documents must be encoded together, which limits many neural models to handle only re-ranking tasks.

To overcome this problem, some dense neural models have been proposed with first-stage retrieval capabilities in mind. RepBERT \cite{zhan2020repbert} encodes queries and documents separately to avoid heavy computation required for a cross-encoder and uses inner products for ranking. It achieves an efficiency of fast enough retrieval by requiring a GPU for a query for using Faiss, a billion-scale similarity search method \cite{8733051}, which would be unrealistic in a usual environment. ColBERT \cite{khattab2020colbert} proposes the concept of late-interaction, which is to calculate relevance at the token-level after encoding a query and a document. Although it shows excellent retrieval effectiveness without using the cross-encoder, it is questionable whether the first-stage retrieval is computationally feasible at a large scale \cite{bai2020sparterm}.

As dense neural retrieval models are usually not feasible for the first-stage retrieval, an alternative is to use a sparse vector representation, like the bag of word (BOW) approaches, where only part of the dimensions is activated for the IR task. Sparse models are amenable to inverted index and hence efficient retrieval. Two possibilities have been explored: One is to extend the same vocabulary space (as in symbolic BOW approaches) with neural models to enhance the symbolic representation; The other is to build sparse embeddings for documents/queries in a new space.

For the former category, DeepCT \cite{10.1145/3397271.3401204} and its extended work, HDCT \cite{dai2020context} were proposed for context-aware term-weighting using pre-trained language models such as BERT \cite{devlin2019bert}. Doc2query \cite{nogueira2019document} and DocTTTTTquery \cite{Nogueira2019FromDT} predict a query given a document, and expand the document by using the predicted query. Also included in the category is SparTerm \cite{bai2020sparterm} that generates term-weighted and term-expanded sparse vectors that belong to the whole vocabulary space. These approaches are successful as an efficient first-stage retriever with enhanced effectiveness. However, they inherit the same shortcomings of symbolic IR because they are limited to the representation power of the vocabulary space, without being able to retrieve documents based on non-symbolic latent semantics that can be captured by neural models.

SNRM \cite{zamani2018neural} creates a latent sparse representation using a fully connected layer of an hourglass structure similar to an autoencoder. While it is a pioneering work for its sparse semantic representation designed for IR, its simple word embedding method and lower dimensionality (20K) exhibit the lack of versatility in capturing a variety of linguistic properties obtainable from language models, suffering from the polysemy problem.
Besides, it's not clear how the regularization method using L1 norm can be adjusted and controlled in a transparent way.
SOLAR \cite{Medini2020SOLARSO} proposes to use randomly initialized high-dimensional and ultra-sparse embeddings for large scale book classification. Since it ignores the content of documents (books), it is not suitable for the IR task at hand.

\section{Sparsified and Bucketized Embeddings}

\subsection{Design Objectives}
Our goal is to build an efficient first-stage neural ranker without relying on an additional efficiency-enhancing device like ANN \cite{8733051}, which is external to the retrieval model. To achieve this, we posit that it is crucial to represent the distinct textual signals in the embeddings for discriminative processing for IR. By doing go, we hope to achieve the effect of adding symbol-like characteristics to the neural approaches.

As such, our design for embeddings is guided by the lessons learned from the benefits and problems of the low-dimensional dense representations (e.g., BERT embedding \cite{devlin2019bert}). While welcoming the contextualized language modeling capabilities of recent developments in transformer architectures, we note the shortcomings of the embedding representations for IR: 1) All the dimensions must always be accessed during document-query matching; 2) Dimensions are so highly overloaded and polysemous that they hardly serve as useful discriminators (e.g., \cite{arora-etal-2018-linear}); 3) As a corollary, various linguistic properties or semantics are entangled. In order to circumvent the problems, we attempt to enforce the following properties for desirable representations:

\begin{itemize}
    \item The embeddings need to be sufficiently high dimensional and sparse with most of the dimensions being zero, so that they can be processed efficiently for matching while enjoying the high expressive power. Ahmed et al. \cite{DBLP:journals/corr/abs-1903-11257} authors theoretically prove that the high-dimensional vectors with sufficient non-zero dimensions can be extremely robust to random matches (i.e., false positives).

    \item Each dimension represents more specific concepts, making it suitable for precise semantic computation. It is reported that the dimensions of high-dimensional sparse embeddings can be interpreted as the clusters of the word \textit{senses} \cite{arora-etal-2018-linear,murphy-etal-2012-learning}, implying that the semantics of the dimensions is highly concentrated.

    \item Different aspects of a text, i.e., queries and documents, should be captured for versatile representations as IR queries are notoriously ambiguous with diverse intents \cite{azad2019query}. Recent studies show that different layers of the transformer-based architectures tend to deal with distinct linguistic properties and that they can contribute to improving NLP tasks \cite{hendrycks2016bridging,conneau-etal-2018-cram}.
\end{itemize}

Our model shares a similar goal with Zamani et al. \cite{zamani2018neural}, in which the authors suggest a number of requirements that a responsive first-stage neural ranker should satisfy. Most importantly, the ranker must be able to encode a query and document independently to support offline encoding of documents. Also, the outputs should be high-dimensional sparse representations to support inverted index. It also stipulates that the matching function should return zero if none of the non-zero dimensions of a query and document are matched. This is an essential property to avoid unnecessary computation. Our model largely follows this with an added requirement that the representation should be ultra-high dimensional (UHD) to achieve effectiveness comparable to dense models \cite{khattab2020colbert,zhan2020repbert}, not losing the efficiency coming from sparsity.

\subsection{Overall Architecture}
Figure \ref{fig:overall_architecture} depicts the overall architecture of our proposed approach. Like most sparse models (e.g., \cite{zamani2018neural,bai2020sparterm}), it comprises a query encoder, a document encoder, and a scoring function. While the query and document encoders are run separately, unlike interaction models (e.g. \cite{ding2020rocketqa}), the weights are shared for the query and document encoders. In other words, a single encoder is used to produce both query and document embeddings. When the final query and document representations are formed with sparsification and bucketization to be explained below, they are matched with a dot product as the scoring function.

The encoder is composed of three modules: 1) the BERT module to convert text into dense token embeddings, 2) the WTA module that sparsifies the BERT module's outputs, converting a dense embedding to sparse one, and 3) the Max-pool module that performs \textit{non-interfering aggregation} of sparse token embeddings. We call the final output of from the WTA module a \textit{bucket} (Figure \ref{fig:encoder}), implying that an ultra-high sparse representation to be used for matching contains multiple buckets. The BERT and encoder are trained jointly with an objective to maximize the similarity between a query bucket and individual relevant document buckets. An in-depth explanation of the encoder is in the next section.

Conceptually, our UHD representations are bucketed or segmented into multiple parts so that they represent different aspects of a query or document. From the procedural point of view, multiple buckets are produced for a query or document by different versions of WTA's and concatenated to build the final UHD representation. Buckets can be generated either \textit {horizontally} from a particular layer of BERT or \textit {vertically} from multiple layers as in Figure \ref{fig:multi_buckets}.

It's been shown that exploiting token-level information is a good strategy for accurate retrieval \cite{khattab2020colbert}. However, storing every token embedding for indexing is impractical because of the length of an ordinary document and the size of each embedding. Our high-dimensional representations for tokens can help achieve the effect of storing individual token embeddings because they are spacious and sparse enough to encode distinct signals and max-pooled in our architecture. In other words, a high-dimensional bucket for a query or a document is formed by max-pooling the token-level embeddings of the dimension size obtained from BERT and WTA (see Figure \ref{fig:encoder}). An in-depth explanation is in Section \ref{ssec:encoder}.

\subsection{Query \& Document Encoder}
\label{ssec:encoder}

\begin{figure}[ht]
  \centering
  \includegraphics[width=0.8\linewidth]{encoder}
  \caption{The Encoder structure.}
  \label{fig:encoder}
\end{figure}

Our approach leverages pre-trained BERT to construct contextualized dense embeddings. Let $q = \{q_1, q_2, ..., q_{|q|}\}$ be a query and $d = \{d_1, d_2, ..., d_{|d|}\}$ be a document. Then, we obtain BERT dense representations for the tokens in $q$ and  $d$. Although we propose our method of using bucketed representations as in Section \ref{ssec:multi-buckets}, we here explain a common option of using only the last layer for a query or document representation. An embedding for a query or a document is obtained as follows:

\begin{equation}
    \begin{aligned}
        E_q = BERT(q) \in \mathbb{R}^{|q| \times h}
        \\
        E_d = BERT(d) \in \mathbb{R}^{|d| \times h}
    \end{aligned}
    \label{eq:bert}
\end{equation}
where $h$ is the hidden size of BERT. Queries and documents are encoded separately while the same encoder is shared.

To sparsify each dense token embedding, we use a  WTA layer \cite{makhzani2015winner}, which is an $n$-dimensional linear layer in which only top-$k$ activations are preserved while the others are set to zero. We usually set $k \ll n$ to generate a high-dimensional sparse representation $S_{t_i}$.
More formally, let $t_i$ be the $i^{th}$ token of a query or a document and $E_{t_i}$ be its dense embedding. Then $S_{t_i}$ is build as follows:
\begin{equation}
    z_{t_i} = E_{t_i} \cdot W + b,\ W \in \mathbb{R}^{h \times n},\ b \in \mathbb{R}^{n}
\end{equation}
\begin{equation}
    S_{t_i}[dim] =
    \begin{cases}
      z_{t_i}[dim], & \text{if}\ z_{t_i}[dim] \in \text{top-}k(z_{t_i}) \\
      0, & \text{otherwise}
    \end{cases}
\end{equation}
where $dim$ $\in [1, n]$ is a dimension of the vectors $z_{t_i}$ and $S_{t_i}$.

In order to concentrate on the learning signals, the WTA module considers only the gradients flowing back through the \textit{winning} dimensions. In addition, we impose $weight sparsity$ constraint proposed in Ahmad et al. \cite{DBLP:journals/corr/abs-1903-11257}, which is like applying dropout \cite{JMLR:v15:srivastava14a} to the weights $W$ between the input and the output. As a result, each dimension of WTA can see only sub-samples of the input dimensions. In addition to the usual effect of avoiding overfitting, we expect this constraint to promote the dimensions of WTA to extract diverse aspects of the input, which collectively can store a complete view of the input.

Some benefits from adopting WTA are: 1) We can control sparsity of a resulting embedding precisely and conveniently by adjusting $k$, an ability considered important for generating sparse representations reliably without outliers (e.g., a sample with only zeros); 2) We can modify $k$ even at inference time so that we can control the output's sparsity for an application's need without re-training the model (although only increasing the sparsity is recommended).

Sparse token embeddings generated for a given query or document by the encoder is aggregated with token-wise max-pool followed by L2-normalization to produce a single sparse embedding $B_t$, a \textit{bucket}.
\begin{equation}
    B_t = {max\text{-}pool}({S_{t_1}, S_{t_2}, ..., S_{t_{|t|}}}) \in \mathbb{R}^{n}
    \label{eq:bucket}
\end{equation}
Note that dense representations are destined to lose much information when they are aggregated, because most, it not all, the dimensions of one token is likely to be \textit{interfered} by those of the other tokens being aggregated. As a result, the representation for a query/document becomes ambiguous and overly generalized. However, our sparse embeddings can be merged with minimal information interference because non-zero dimensions are not likely to overlap among the token embeddings, resulting in \textit{non-interfering} max-pooling. This method is inspired by Ahmad et al. \cite{Ahmad2015PropertiesOS}, in which the authors prove that high-dimensional \textit{binary} vectors could be merged into a single vector using OR operation without increasing the false positive ratio much.

\subsection{Representations with Multiple Buckets}
\label{ssec:multi-buckets}
\begin{figure}[h]
  \centering
  \includegraphics[width=0.9\linewidth]{multi_buckets}
  \caption{Illustration of multiple buckets. Multiple buckets can be grounded either horizontally on a single layer or vertically on multiple layers.}
  \label{fig:multi_buckets}
\end{figure}

Multi-buckets are formed to generate different variants of $B_t$ (Eq. \ref{eq:bucket}) in two different ways: \textit{vertically} and \textit{horizontally} as shown in Figure \ref{fig:multi_buckets}.
The \textit{vertical} buckets encode information with different levels of abstraction coming from different layers of BERT, while the \textit{horizontal} buckets are expected to model multi-aspect information from a single layer.
As in Jo et al. \cite{jo-myaeng-2020-roles}, \textit{vertical} buckets or the outputs from different layers are deemed to contain different linguistic properties of the input text. Therefore a UHD representation with \textit{vertical} buckets contains richer information about the text than a particular bucket from a layer (e.g., the last layer) often used for a downstream task. We first extract $V$ representations corresponding to the numbers of the BERT layers for a sequence of tokens $t$.
\begin{equation}
    E^j_t = BERT^j(t) \in \mathbb{R}^{|t| \times h}
\end{equation}
where $j$ is the index of BERT layer $\in [1,\ V]$.
WTA layers are then independently applied to all BERT layers as in Section \ref{ssec:encoder} to obtain $V$ buckets.
\begin{equation}
    B^j_t = WTA^j(E^j_t) \in \mathbb{R}^{n}
\end{equation}

\textit{Horizontal} buckets are produced by projecting $E^j_t$ on to space for the purpose of obtaining a particular aspect of the $E^j_t$. By applying multiple randomly initialized $H$ WTA layers, we attempt to produce multiple versions of $B^j_t$ representing different latent aspects.
\begin{equation}
    B^{jm}_t = WTA^{jm}(E^j_t) \in \mathbb{R}^{n}
\end{equation}
where $m \in [1,\ H]$. This is under the assumption that a UHD sparse representation of a query/document captures lexical, syntactic, and semantic aspects and encode them.

After the bucketing mechanism is equally applied to a query and a document to produce $B^{jm}_q$ and $B^{jm}_d$, a relevance score for a pair of a query $q$ and a document $d$ is computed through a bucket-wise dot product as follows.
\begin{equation}
    \begin{aligned}
    Rel(q, d) = \sum_{j,\ m}B^{jm}_q \cdot B^{jm}_d
    \end{aligned}
    \label{eq:relevance_score}
\end{equation}

Note that placing multiple buckets with the same dimensionality $n$ makes each bucket smaller and more entangled, possibly leading to performance degradation. For example, if 10 buckets are placed with $n$ set to 81,920, then each bucket's dimensionality becomes $n$ = 8,192, making each dimension overloaded and less discriminative than those in a larger single bucket. In Section \ref{sec::4.4_multiBucket}, we describe various bucketing strategies and compare them.

\subsection{Training}
\label{ssec:training}

Figure \ref{fig:overall_architecture} shows how we jointly train the underlying BERT and WTA layers altogether. The entire model is trained to make a query similar to the relevant documents and dissimilar to the irrelevant ones. Given a query $q$, a set of relevant (positive) documents $D^{p}$, and a set of irrelevant (negative) documents $D^{n}$, we calculate the loss as shown in Equation \ref{eq:loss}. For efficient negative sampling, for each query in a mini-batch, we regard the positives of other queries in the same mini-batch as the negative samples (i.e., in-batch negatives), following Zhan et al. \cite{zhan2020repbert}. We use the following loss function:
{\small
\begin{equation}
    \mathcal{L}(q, D^{p}, D^{n}) = \sum_{d^p \in D^{p},\ d^n \in D^{n}} {\rm max}(0, 1-Rel(q, d^p)+Rel(q, d^n))
    \label{eq:loss}
\end{equation}
}

\subsection{Ranking}

\begin{figure}[h]
  \centering
  \includegraphics[width=\linewidth]{ranking}
  \caption{Ranking scheme using a bucketed inverted index.}
  \label{fig:ranking}
\end{figure}

Our model allows for exploiting an inverted index.
We regard each dimension in our bucketed sparse representations, which is indexable, as a signal or a latent term. For instance, if $n$ (dimensionality) is 81,920, a document is represented with a combination of a few latent terms out of 81,920. Only the non-zero dimensions of a document enter the inverted index with their weights. The level of efficiency in symbolic IR can be achieved since only a small fraction of the dimension in our UHD representation contains a non-zero value.

As shown in Figure \ref{fig:ranking}, we first encode all documents in a collection using the trained encoder to construct an inverted index. Each bucket conceptually has its own independent inverted index. Note that this process is needed only once offline. At inference (retrieval) time, we encode a given query to make $|B|$ representations so that we calculate bucket-wise dot products and sum up the scores for the final relevance score as in Eq. \ref{eq:relevance_score}. We note that this bucket-wise index construction and scoring can be easily distributed in practice for added efficiency.

\section{Experiments}
We conducted a series of experiments to show the validity of the proposed retrieval model as well as the associated methods, comparing it primarily against recently proposed sparse methods. We also juxtapose our results with those of the most recent dense models although they are not geared toward full ranking with sufficient efficiency, without resorting to additional computational resources. Since it has been primarily used by the previous work, we opted for the MSMARCO Passage Retrieval dataset \cite{bajaj2018ms} (henceforth, MSMARCO), consisting of 8.84M passages collected from the Web and 1.01M queries generated from real-world users of Bing. Given that our model is based on BERT, which is limited in terms of the number of tokens per document, passage retrieval is most appropriate for our experiments, too.

\subsection{Settings}
\label{sec::4.1_setting}

\paragraph{Training} We use 25\% of MSMARCO train triples, (\textit{query}, \textit{positive passage}, \textit{negative passage}), containing a total of 39,782,779 triples, in all our experiments unless otherwise specified. Note that we use only queries and positive passages with in-batch negatives for training (see Section 3.5), ignoring the negative passages. We sample 7,000 queries that are not in the above 25\%  to construct a re-ranking train subset with the corresponding top 1000 passages, provided by MSMARCO. This train subset is used for the bucket weight tuning in Section \ref{sec::4.4_multiBucket}. Training a single-bucket model from the last layer of BERT with $n$ = 81,920 (dimensionality) and $k$ = 400 (sparsity) takes about 40 hours on a GeForce RTX 3090.

\paragraph{Evaluation} MSMARCO provides two settings: re-ranking and full ranking. In the re-ranking setting, given a query and 1000 passages (top-1000) retrieved by BM25, the model should re-rank the passages. In full ranking, on the other hand, the model should retrieve relevant passages from the entire collection for a given query. Our main goal is to validate our model in the full ranking setting, we utilize the re-ranking setting when a large number of model variants are compared, as a way to save computational cost. In both settings we use the MSMARCO dev set, which contains 6,980 queries. MRR@10 (mean reciprocal rank) is used as the evaluation metric as in the official evaluation. We use recall at various levels as additional metric.

\paragraph{Model implementation} The language model (LM) used in our proposed architecture is pre-trained BERT (\textit{bert-base-uncased}) \cite{devlin2019bert} implemented by Hugging Face \cite{wolf-etal-2020-transformers}. For the buckets on top of the LM, we modified Winner\text{-}Take\text{-}All (WTA) layer \cite{makhzani2015winner} by adding a weight sparsity proposed by Ahmad et al. \cite{DBLP:journals/corr/abs-1903-11257}. For the multi-bucket setting, it is possible to jointly train all WTA layers at once, but it is also possible to separately train the model for each WTA layer and combine the result later. We discuss the difference between those two in Section \ref{sec::4.4_multiBucket}.

\paragraph{Hyperparameters}
We use the Adam optimizer with the learning rate of 5e-6, $\beta_1$ = 0.9, $\beta_2$ = 0.999 and epsilon =  1e-8, and L2 weight decay of 0.01. The learning rate is linearly warmed up over the first 2,000 steps, and then linearly decayed. The batch size is 32, and the epoch is 0.25 (25\% of MSMARCO Train Triples Small set). For the BERT LM, we use GELU \cite{hendrycks2016bridging} as an activation and set the maximum length to 32 for queries and 180 for passages (truncated if exceeded). For the WTA layer, we select dimensionality $n$ from \{8,192, 24,576, 49,152, 81,920\} and the number of non-zero dimensions $k$ from \{4, 8, 16, 20, 40, 80, 200, 400\}. We set weight sparsity to 0.3, meaning that 30\% of the weights connecting each node to input is set to zero.

\paragraph{Baselines} We compare our proposed approach with both dense and sparse methods. Among the six below, the first two are based on dense embeddings, while the rest are based on sparse representations.

\begin{itemize}
  \item \textbf{RepBERT} \cite{zhan2020repbert} is a dense IR model that encodes queries and documents separately, with mean-pooling on BERT token embeddings.
  \item \textbf{ColBERT} \cite{khattab2020colbert} proposes token-level late-interactions between queries and passages, which are encoded separately.
  \item \textbf{BM25} \cite{robertson1994some} is a BoW retrieval method with the term weighting based on the query/document frequency signals.
  \item \textbf{Doc2query} \cite{nogueira2019document} \& \textbf{DocTTTTTquery} \cite{Nogueira2019FromDT} expands documents in the vocabulary space by predicting queries using BERT (Doc2query) or T5 \cite{raffel2020exploring} (DocTTTTTquery).
  \item \textbf{DeepCT} \cite{10.1145/3397271.3401204} is a deep contextualized term weighting framework, which maps BERT representations to context-aware term weights.
  \item \textbf{SparTerm} \cite{bai2020sparterm} is a sparse representation-based IR model that generates term-weighted \& term-expanded sparse vectors belonging to the vocabulary space.
\end{itemize}

\begin{table*}[ht!]
\caption{Comparisons on the MSMARCO Passage Retrieval dev set for full ranking. The sparse representation approaches are amenable to an inverted index and other BoW-based symbolic IR techniques whereas the embedding approaches are not.}
\begin{center}
\begin{tabular}{lcccccc}
\toprule
\textbf{Model} & \textbf{MRR@10} & \textbf{Recall@50} & \textbf{Recall@100} & \textbf{Recall@200} & \textbf{Recall@500} & \textbf{Recall@1000} \\ \midrule
\multicolumn{7}{c}{{\textbf{Dense Embedding Approaches}}} \\
\hline
RepBERT & 30.4 & - & - & - & - & 94 \\
ColBERT & 36.0 & 82.9 & - & 92.3 & - & 96.8 \\
\hline
\multicolumn{7}{c}{{\textbf{Sparse Representation Approaches}}} \\
\hline
BM25 & 18.6 & 60 & 69 & 75 & 82 & 85.71 \\
Doc2query & 21.5 & - & - & - & - & 89.1 \\
DeepCT & 24.3 & 69 & 76 & 82 & 86 & 91 \\
DocTTTTTquery & 27.7 & 75.6 & 81.89 & 86.88 & 91.64 & 94.7 \\
SparTerm & 27.94 & 72.48 & 78.95 & 84.05 & 89.5 & 92.45 \\
\hline
\textbf{UHD-BERT} & 29.56 & 77.7 & 84.1 & 89.1 & 94.1 & 96.1 \\
\bottomrule
\end{tabular}
\end{center}
\label{table:overall_performance}
\end{table*}

In the following experiments, we use the performances of the dense embedding approaches reported in the corresponding papers. For the sparse representation approaches, the baseline performances are provided by SparTerm \cite{bai2020sparterm}. While SNRM \cite{zamani2018neural} would be a suitable baseline for its sparse model, it was not possible to reproduce the model since its accurate hyperparameters were not provided. Besides, there is no experimental result for MSMARCO.

\subsection{Overall Performance}

Table \ref{table:overall_performance} shows the full ranking performance of the proposed model (UHD-BERT) against the baselines on MSMARCO Passage Retrieval dataset. In this experiment, we train a UHD-BERT model with vertical buckets on the layer \{$2,4,6,8,10,12$\}, where dimensionality $n$ and the number of non-zero dimensions $k$ are 81,920 and 400, respectively. The reason we chose to use only vertical buckets without horizontal buckets is described in Section \ref{sec::4.4_multiBucket}. Bucket weight tuning was also performed in the same manner as in Section \ref{sec::4.4_multiBucket}, and the retrieved layer weights were ($0.3:0.0:0.2:0.6:0.7:0.9$). Compared to sparse approaches, UHD-BERT outperforms all the baseline models, from traditional term-based methods to recent neural approaches. The result is attributed to the capabilities of our model that can represent a query and a document with a sufficient amount of information. Unlike the previous sparse approaches that mainly focus on efficiency at the expense of losing information, our approach achieves high-quality, rich representations using bucketed UHD sparse representations while maintaining the necessary sparsity for efficiency.

As can be seen in Table \ref{table:overall_performance}, the two dense approaches are more effective, especially the highest performance of ColBERT \cite{khattab2020colbert}. It is evident that this is due to its complex matching method compared to all the other baselines that use practical similarity-based matching. While the proposed late interaction outside of BERT is promising compared to the full interaction-based approaches, it still requires relatively heavy computation, which might be impractical for industrial applications. We notice that UHD-BERT is just slightly worse than RepBERT \cite{zhan2020repbert} in all but Recall@1000, for which UHD-BERT is superior. Note that dense models need to employ external resources (i.e. approximate nearest neighbor search or GPUs) for comparable efficiency. This result highlights the advantages of our model that is effective and efficient. Furthermore, our work is amenable to other symbolic IR techniques such as PRF that have been proven for additional improvements.

\subsection{Analysis on dimensionality and sparsity}

\begin{figure}[h]
  \centering
  \includegraphics[width=\linewidth]{NK.pdf}
  \caption{Impact of $n$ (dimensionality) and $k$ (non-zero dimensions).}
  \label{fig:NK}
\end{figure}

To understand the roles of $n$ (dimensionality) and $k$ (non-zero dimensions), we tested our model with two different settings. In figure \ref{fig:NK}, we show the performance trends:  the impact of varying $k$ = (4, 8, 16, 20, 40, 80, 200, 400)\footnote{We trained with the k of 400 and changed it to above settings at test time. There was no significant difference between the setting where k matches and the setting that does not match at training and test times.} on the left with the fixed $n=81,920$ and the impact of varying $n$ = (8,192, 24,576, 49,152, 81,920) on the right with the fixed $k= 400$.

First, given the fixed $n$, it is evident that the higher $k$, the better the score. This is because the absolute amount of information for representing a query and a document increases. Given the time constraint, however, it is not desirable to set an exorbitant $k$ (e.g., 8,192) as it would yield high computational costs to get a desirable relevance score. We see the score is reasonable even with $k$ being 16 (4\%). While it keeps increasing as $k$ becomes higher, the degree of improvement rapidly reduces after a certain point. This suggests there is an optimal $k$ which achieves both effectiveness and efficiency.

Next, we observe that the score improves as $n$ increases, even with $k$ (the absolute amount of information) being fixed. This is basically due to the increased number of latent terms to be expressed specifically. Moreover, even though the same $k$ is used to represent each query and document, the discrimination power can become increased for matching since the number of cases to get activated increases exponentially as $n$ becomes larger. Note that a large $n$ also means sparsity increases. This result supports our motivation of UHD sparse representations that endow the discrimination power (for effectiveness) while increasing the sparsity (for efficiency).

\begin{figure}
  \centering
  \includegraphics[width=0.9\linewidth]{images/Figure_sparsity_vertical.pdf}
  \caption{Density of the generated sparse representation of queries and passages. Queries and passages are sorted according to their token lengths (x-axis).}
  \label{fig:sparsity}
\end{figure}

\subsection{Analysis on Multiple Buckets}
\label{sec::4.4_multiBucket}

\begin{table}
\caption{Comparison of various multi-bucket strategies on the MSMARCO Passage Retrieval dev set re-ranking task. The total dimensionality ($n$) and non-zero dimension size ($k$) are the same across the strategies. $b$ stands for the number of buckets.}
\begin{center}
\begin{tabular}{lc}
\toprule
\textbf{Strategy} & \textbf{MRR@10}   \\ \midrule
Single-Bucket ($b=1$, $n=49,152$, $k=480$) & 26.18 \\
\hline
Multi-Bucket ($b=6$, $n=8,192$, $k=80$) &  \\
\hspace{0.15cm}  Horizontal bucket & 23.19 \\
\hspace{0.15cm}  Vertical bucket (Joint) & 25.58 \\
\hspace{0.15cm}  Vertical bucket (Separate)& 27.08 \\
\hspace{0.15cm}  + bucket weight tuning & \textbf{27.41} \\
\hline
\hspace{0.15cm}  Ideal query-layer predictor & \textbf{39.62} \\
\bottomrule
\end{tabular}
\end{center}
\label{table:multilayer}
\end{table}

As mentioned in Section \ref{ssec:multi-buckets}, we adopt the multi-bucket strategies to capture and represent multiple aspects of information. In order to verify the effectiveness of these multi-bucket strategies, we compare the performance of various strategies on the MSMARCO dev re-ranking set in Table \ref{table:multilayer}.

For fair comparisons among them, we set the total dimensionality ($n$) and the number of non-zero dimensions ($k$) to be the same for all the strategies. In this experiment, we set $n$ to 49,152, $k$ to 480, and the number of buckets ($b$) for the multi-bucket strategies to 6. As a result, $n$ and $k$ are set to 8,192 and 80, respectively. For the horizontal bucket setting, we create each bucket on the 2nd, 4th, 6th, 8th, 10th, and 12th layers of BERT, whereas we create buckets on the 12th layer of BERT for the single-bucket and vertical bucket settings. For training the models, we only use 3\% of the MSMARCO Train Triples Small set. For the bucket weight tuning, we conducted a grid search on the top-1000 re-ranking dataset we built from the MSMARCO dataset (See Section \ref{sec::4.1_setting}). About the vertical bucket setting, we evaluate two types of implementation (Joint/Separate) introduced in \textit{Model Implementation}, Section \ref{sec::4.1_setting}.

In Table \ref{table:multilayer}, we can observe that various strategies yield very different results. Comparing the single-bucket and multi-bucket strategies, the vertical bucket (Separate) setting of the multi-bucket outperforms the single-bucket setting, meaning that the multi-bucket is successful in capturing more useful information. However, the joint setting shows lower performance than the single-bucket. We believe that this is because the WTA layers in the joint setting interfere with each other during the training process. A special mechanism that can minimize the interference between WTA layers will solve the problem of the joint model, but we leave this as future work.

Also, the horizontal setting of the multi-bucket shows very poor performance, which implies: 1) Given the same information source, it is advantageous to have one large bucket (e.g. $n=49,152$) than several smaller buckets (e.g. $n=8,192 * 6$), due to the disadvantages of the multiple buckets mentioned in Section \ref{ssec:multi-buckets}. Inherently, this eliminates the need of combining vertical and horizontal buckets. 2) In order to maximize the utility of the multi-bucket strategy, it is necessary to carefully consider what information each bucket extracts. By extracting information from different BERT layers, the vertical bucket setting outperforms the single bucket setting, despite the disadvantages of the multi-bucket strategy (small $n$ for each bucket). This is also evident in the result of the bucket weight tuning. By controlling the weight of each bucket, we can further increase the performance of the vertical bucket setting. The bucket weight tuning yields the searched weights ($0.66:0.0:0.33:1.00:0.33:1.00$).

The weight of each layer (bucket) according to the above bucket weight tuning gives us another intuition about the selection of BERT layers suitable for the IR task. We believe it is worthwhile to analyze why certain layers are more useful than others. As many previous studies including Jo et al. \cite{jo-myaeng-2020-roles} have discovered, each BERT layer models information of a different granularity; a lower layer focuses more on syntactic while the upper layer focuses more on semantic information. By incorporating this, we can interpret the above weight as following: crucial to IR tasks are not only the last layer that contains well-refined semantics (which makes a good use of the powerful transformer model) but also the low- and mid-level lexical and syntactic information.

Going further from this analysis, one can imagine that each query may have a more suitable set of layer weights. For example, certain queries may require syntactic contents while others may require topical (semantic) contents. From this perspective, we measured the performance when a layer with the highest MRR was ideally selected for each query (Ideal query-layer predictor in Table\ref{table:multilayer}). The result clearly suggests that selecting different layers or bucket weight distributions for each query has the potential to improve performance. This implies the necessity of a bucket weight predictor which takes a query as an input, but we leave it as future work.

\subsection{Analysis on UHD Embeddings}
\label{sec::4.5_analysis_embedding}

To better understand the characteristics of our proposed model, we conducted an in-depth analysis of the generated embeddings. In this experiment, we use a model with a single bucket on the 12th layer of BERT LM, where dimensionality $n$ and the number of non-zero dimensions $k$ are 81,920 and 40, respectively.

In order to check whether the generated embedding properly contains the information of queries and passages, in Figure \ref{fig:sparsity}, we measured the density (ratio of non-zero dimension) of the queries (left-hand side) and passages (right-hand side) with respect to their length. Intuitively, we can assume that the more words (or tokens) the text contains, the more information it contains. Because each dimension of sparse embedding is activated only when the corresponding feature exists, the density of the generated embeddings must correlate with the amount of information (the number of tokens) in the query/passage, as shown in Figure \ref{fig:sparsity}. However, interestingly, the density continues to increase according to the number of tokens in the query, while it converges in the passage. Our interpretation is that since the query reflects the intention of the user, it is highly likely that each query term represents a unique feature that narrows down the search space, which does not overlap with other query terms. For the example \textit{'how much do psychologists earn uk'}, a user has put \textit{'uk'} to reduce the search space. As such, increasing the length of the query has a high probability of adding new features, which increases the density. Unlike queries, passages usually deal with limited topics, which makes density converges even if the length increases.

\begin{figure}[h]
  \centering
  \includegraphics[width=0.9\linewidth]{images/Figure_dimension_frequency.pdf}
  \caption{Active (non-zero) dimension frequency in queries and passages.}
  \label{fig:frequency}
\end{figure}

Besides the above query/passage-level embedding analysis, dimension-level analysis is also crucial to understand the characteristics of our proposed model. In Figure \ref{fig:frequency}, we measured how many times each dimension was activated (non-zero) in the query and passage, similar to traditional word frequency analysis. The analysis suggests that the embeddings generated by our model work in a similar fashion to words in natural languages (Zipf's law). It is also worth noting that similar to stop words in natural language words, there are dimensions that appear in many passages (8.84M in total) and queries (7K in total). Dealing with these "stop-dimensions" can also be expected to have a big impact on performance, but we leave it as future work.

\subsection{Utilization of Sparse Characteristics}

\begin{table}[ht!]
\footnotesize
\caption{Interpretation of the each dimension in the UHD embedding by analyzing terms of queries that activates a specific dimension. The query terms that support the interpretation are indicated in \textit{italic}.}
\begin{center}
\begin{tabular}{c|l|c}
\toprule
\multirow{2}{*}{\textbf{Dim}} & \multirow{2}{*}{\textbf{Appeared Query Terms}} & \textbf{Inter-} \\
 & & \textbf{pretation} \\
\midrule
68,604 & \textit{costs}, \textit{paid}, \textit{fee}, \textit{prices}, \textit{fees} & Price \\
29,938 & \textit{tall}, paid, \textit{square}, \textit{feet}, \textit{miles}, \textit{words} & Unit \\
8,797 & \textit{soon}, \textit{longest}, heal, \textit{wait}, \textit{been}, \textit{length}, \textit{weeks}, & Period \\
66,589 & \textit{growth}, \textit{movement}, \textit{development}, degree, theory & Advancement \\
3,740 & \textit{treatment}, \textit{heal}, \textit{surgery}, \textit{doctors}, \textit{doctor} & Treatment \\
24,673 & \textit{medication}, normal, \textit{constipation}, \textit{doctor}, \textit{medicine} & Medical \\
37,089 & \textit{problems}, \textit{diseases}, \textit{signs}, \textit{syndrome}, \textit{fever} & Disease \\
476 & \textit{die}, \textit{killed} & Death \\
47,873 & \textit{technology}, \textit{science}, \textit{chemical}, \textit{scientific}, union & Science \\
374 & science, technology, fastest, \textit{java}, \textit{language}, \textit{software} & Programming \\
68,111 & lower, low, clear, \textit{deductible}, \textit{irs}, depression & Tax \\
21,822 & \textit{party}, beach, \textit{amendment}, tissue, \textit{constitution} & Politic \& Law \\
22,627 & \textit{founded}, \textit{invented}, \textit{begin}, \textit{established}, \textit{started}, released & Founding \\
3,266 & \textit{numbers}, \textit{customer}, \textit{telephone}, \textit{zip}, \textit{contact}, union, npi & E-commerce \\
46,843 & \textit{washington}, \textit{il}, \textit{ma}, \textit{mn}, \textit{indiana}, \textit{south} & Location \\
61,183 & \textit{child}, \textit{kids}, \textit{children}, \textit{disney}, paint, \textit{parents} & Child \\
48,698 & \textit{education}, \textit{student}, training, \textit{degree}, workout & Education \\
26,433 & \textit{synonym}, \textit{description}, \textit{synonyms}, \textit{defined}, english & Definition \\
1,377 & \textit{sang}, \textit{played}, \textit{sings}, requirements, \textit{cats}, \textit{plays} & Play (theatre)\\
\bottomrule
\end{tabular}
\end{center}
\label{table::interpretability}
\end{table}

One of the advantages of the sparse representation is its interpretability, which is possible in our case due to the disentangled feature of the UHD embedding. To prove this, we conduct an analysis about the meaning of each dimension of the generated UHD embedding. In this analysis, we use the same UHD-BERT model with Section \ref{sec::4.5_analysis_embedding} to generate UHD embeddings.

Table \ref{table::interpretability} shows interpretation results on a few sampled dimensions of the UHD embedding. We use a simple method that validates the interpretation, which is to inspect terms of queries that have non-zero values (activated) on a specific dimension. For example of the dimension index "61,183" in Table \ref{table::interpretability}, the terms that activated queries frequently have were \textit{child}, \textit{kids} and \textit{disney}, implying that the dimension is specialized for the "child" information. For clarity, query terms that have appeared at least five times are listed, along with dimensions that are relatively easy to interpret. We observe that there are dimensions representing abstract latent terms such as "Unit" or "Period". To summarize the analysis, our UHD embedding is not only interpretable (advantage of sparse representations), but also has the power of the latent representation (advantage of dense representations).

\section{Conclusion}
In this work, we present UHD-BERT, a novel sparse retrieval method empowered by extremely high dimensionality and controllable sparsity. To fully utilize the abundant capacity of our representation scheme, we propose bucketed learning, which is a method of learning from the same or different layers of BERT for different aspects or abstraction levels of information. We also introduce and justify our choice of WTA model as our sparsification model by proving the benefits of reliable representation generation and the ability to control the sparsity at inference time. We show that our model outperforms previous sparse models significantly and delivers competitive performance compared to dense models especially for recall, which is an important measure for a first-stage ranker. We expect greater performance gain can be achieved for industrial systems by smartly incorporating symbolic techniques, which will be our future work.

\bibliographystyle{ACM-Reference-Format}