MICO : Selective Search with Mutual Information Co-training (2022)

Abstract: In contrast to traditional exhaustive search, selective search first clusters documents into several groups before all the documents are searched exhaustively by a query, to limit the search executed within one group or only a few groups. Selective search is designed to reduce the latency and computation in modern large-scale search systems. In this study, we propose MICO, a M utual I nformation CO -training framework for selective search with minimal supervision using the search logs. After training, MICO does not only cluster the documents, but also routes unseen queries to the relevant clusters for efficient retrieval. In our empirical experiments, MICO significantly improves the performance on multiple metrics of selective search and outperforms a number of existing competitive baselines.

## introduction in information retrieval (ir), searching over all the documents is quite costly at a large scale @xcite . selective search @xcite divides documents into a number of shards (clusters) to allow an incoming query to search over only a small number of these shards that are most relevant to this query. it tries to preserve the quality of search results similar to that of searching over the whole corpus, with less computation and shorter retrieval latency, achieved by limiting an incoming query to search within only a few chosen document shards. sharding is the process of breaking a massive collection of documents into smaller chunks called shards. randomly assigning documents to shards, a popular approach, cannot guarantee documents relevant to a specific query restrained within one shard or only a few shards. with the observation that semantically similar documents tend to appear together in the search results of the same query @xcite , researchers in recent years have devised multiple machine learning-based algorithms for selective search. these algorithms are tailored to benefit the document retrieval process by dividing a large volume of the corpus in a certain way so that semantically related documents are allocated in the same cluster and then routing new queries to the most relevant clusters subsequently. in this way, the retrieval system can return the search results with both efficiency and accuracy, as shown in figure 1 . these approaches significantly differ from the traditional sharding methods that either randomly split the collection into several groups @xcite or allocate them based on simple handcrafted rules @xcite , and have achieved comparable performance to exhaustive search over the entire document corpus. nevertheless, having led giant strides in selective search, these approaches fail to address a few critical practical issues when scalable ir systems for a large number of documents in the industry are needed. a major concern is that these approaches heavily rely on a multi-stage process: firstly, complex clustering procedures are applied to allocate documents into different shards, while later additional elaborate algorithms are used to choose a small number of shards for incoming queries. the multi-stage process needs not only additional infrastructure support to bridge the output of the previous step as the input of the next step but also is prone to error propagation. a natural attempt to mitigate the aforementioned difficulties is to construct a model for both document allocation and subsequent query routing, with an objective function to be optimized in an end-toend paradigm. this approach ensures applying gradient-based learning to modules or components of the system coherently. to this end, we introduce mutual information co-training (mico), a novel end-to-end approach for selective search. we argue that by maximizing the mutual information of the query routing variable and document assignment variable of a relevant query-document pair, we can leverage co-training in an end-to-end training fashion to overcome the aforementioned difficulties. this approach first creates a bipartite graph by exploiting the query-document relations and then treats query and document as two different views of the same example while maximizing the mutual information between the cluster variable of these two views. we show an overview of the mico framework in figure 2 and summarize our contributions: ## mico: mutual information co-training our assumption is that for a relevant querydocument pair in the search, the cluster index assignments z of the query q and z of the document d are two views of that pair (q, d), thus z is equal to z . mico, consisting of a document allocation module and a query routing module, is trained based on this assumption. after training, mico assigns the document to the cluster based on the prediction made by the document allocation module and routes new incoming queries to the cluster based on the prediction by the query routing module. our goal is to learn a model from the doubleview examples in the search logs. the search logs are historical data that contains the relevance information between the query and its relevant documents, but does not carry explicit information of how documents shall be allocated into different shards, neither the query routing information of which shard the query shall be dispatched to. thus, we have minimal supervision signal when training the model to learn to cluster documents and queries. the desired model maximizes the document coverage (recall) and minimizes the search cost simultaneously. we force the cluster selection of the query and the relevant document to be coherent during training, formally as p(z|q) ≈ p(z |d). here we force the two distributions to be consistent if the query and the document are relevant to reflect the coherent cluster assignment. meanwhile, to avoid certain resulting clusters being oversized, causing the search in these shards to be inefficient, we consider balancing the document shard sizes. this concept is illustrated in figure 3 . this balance is achieved by assigning the documents to clusters as uniform as possible (formally p(z ) ≈ uniform where p(z ) is the distribution of cluster indices z ). mico achieves cluster selection coherency between queries and documents by minimizing the mutual information between p(z|q) and p(z |d), and balances shard size by adding an entropy regularization term p(z ). ## conclusion in this study, we present mico, a mutual information co-training framework for document sharding and query routing in selective search with minimal supervision. our contributions are: first, during training, mico maximizes the mutual information between the cluster assignment variables of the query and the document, forcing the prediction of two views of the query-document pair coherently. second, this design enables it to be trained end-toend for both document sharding (clustering) and subsequent query routing, featuring its practical use when vast volumes of documents are required to be searched simultaneously. third, we show improved performance of mico empirically on multiple important metrics in selective search. our future research direction includes further investigation of the regularization for more balanced sizes among all the shards, the utilization of complex neural language models to enrich the document and the query representation for higher query coverage via better clustering results, the reduction of potential noise between a query and its associated document in the data set, and the detection policy to decide when the model needs to be retrained if constant update of the model is needed when deployed. we are also interested in potential extension of this approach to multi-modal data, e.g., image search, music search and software search where the non-traditional documents are in the non-textual formats.