Scalable Evaluation and Improvement of Document Set Expansion via Neural Positive-Unlabeled Learning (2021)

Abstract: We consider the situation in which a user has collected a small set of documents on a cohesive topic, and they want to retrieve additional documents on this topic from a large collection. Information Retrieval (IR) solutions treat the document set as a query, and look for similar documents in the collection. We propose to extend the IR approach by treating the problem as an instance of positive-unlabeled (PU) learning—i.e., learning binary classifiers from only positive (the query documents) and unlabeled (the results of the IR engine) data. Utilizing PU learning for text with big neural networks is a largely unexplored field. We discuss various challenges in applying PU learning to the setting, showing that the standard implementations of state-of-the-art PU solutions fail. We propose solutions for each of the challenges and empirically validate them with ablation tests. We demonstrate the effectiveness of the new method using a series of experiments of retrieving PubMed abstracts adhering to fine-grained topics, showing improvements over the common IR solution and other baselines.

## introduction we are interested in the task of focused document set expansion, in which a user has identified a set of documents on a focused and cohesive topic, and they wish to find more documents about the same topic in a large collection. this problem is also known as a "more like this" (mlt) query in web retrieval. a common way of modeling this problem is to consider the set of documents as a long query, with which information retrieval (ir) techniques can rank documents. ir literature on document similarity and ranking is vast (faloutsos and oard, 1995; mitra and chaudhuri, 2000, inter * most of the work done during internship at riken. alia)-beyond the scope of this work, and largely orthogonal to it, as will be explained later. current methods in document set expansion for very large collections are based on word-frequency or bag-of-words document similarity metrics such as term frequency-inverse document frequency (tf-idf) and okapi bm25 and its variants @xcite , considered strong due to their robustness to extreme class imbalance, corpus variance and variable length inputs, as well as their scalability and efficiency @xcite . however, the performance of such solutions is limited, as the models cannot capture local or global relationships between words. we examine methods to improve document set expansion by leveraging non-linear (neural) models under the setting of imbalanced binary text classification. to this end, we look to positive-unlabeled (pu) learning @xcite : a binary classification setting where a classifier is trained based on only positive and unlabeled data. in the standard document expansion setting, we indeed only possess positive (the document set) and unlabeled (the very large collection) data. pu learning has originally been employed for text classification by @xcite ; @xcite ; @xcite by using techniques such as em and svm. since then, the setting has been well studied theoretically @xcite @xcite , and recently objective functions have been developed to facilitate training of flexible neural networks from pu data @xcite . we discuss the pu learning setting in more detail in section 2, and relevant work on pu learning for text in section 8. we are, however, not interested in replacing traditional (term-frequency-based) ir solutions, but rather improve upon their results by further classifying the outputs of those models. there are two reasons for this approach: (1) traditional ir engines are based on word frequencies, and as a result, cannot capture features based on word order; (2) classification by the use of neural networks does not scale well to "extreme" imbalance 1 . following these observations, we see traditional ir engines and neural models as complementary to each other. our proposed solution is a two-step process: first, a bm25-based, mlt ir engine retrieves relevant candidates; then, a non-linear pu learning model is trained based on the subset of candidates. in this way, each method relieves the weakness of the other. as already discussed above, pu learning has recently become viable for deep neural network models. as a result, we are able to leverage it to train models that are able to capture higher order features between words. however, pu learning literature focused on theoretical analysis and experiments on small models and simple-notably, class-balanced-benchmarks such as mnist, ci-far10 and 20news @xcite @xcite . pu learning has not been extensively tested for imbalanced datasets. scaling pu solutions to high-dimensional, ambiguous and complex data is a significant challenge. one reason for this is that pu data is, by definition, difficult or sometimes impossible to be fully labeled for exhaustive, large-scale evaluation. for the purpose of document set expansion, and in particular for fine-grained topics, gathering fullylabeled data for an accurate benchmark is also a challenge. for this reason, we propose to simulate the scenario synthetically but realistically by using the pubmed collection of bio-medical academic papers. pubmed entries are manually assigned multiple terms from medical subject headings (mesh), a large ontology of medical terms and topics. we can treat a set of mesh terms as defining a finegrained topic, and use the mesh labels for deriving fully-labeled tasks (see examples of mesh topic conjunctions in the contributions of this work are thus: foot_2 ## the pubmed set expansion task in this section we discuss the method of generating an extensive benchmark for evaluating solutions of mlt document set expansion. we are inspired by the following scenario: a user has a set of documents which all pertain to a latent topic, and is interested in retrieving more documents about that topic from a large collection. while traditional term-frequency-based ir solutions scale well to extremely large collections of documents, they are imprecise, and contain a significant amount of noise. therefore, an additional step based on pu learning can be utilized to classify the output of the ir model, and improve the results. we are interested in gathering a task for evaluation of the second step. in other words, given an existing black-box ir solution, we would like to use it to produce a dataset for training and evaluation of models which should improve upon the black-box ir solution's performance. due to the varied nature of the setting, it is impractical to acquire full supervision for a large number of topics. therefore, we propose to generate synthetic tasks inspired by the real use-case application. ## experiment details the rest of this work will reference experiment results. unless otherwise noted, our base architecture is a single-layer cnn @xcite . the choice of cnn, over other recurrent-based or attentionbased models, is due to this architecture achieving the best performance in our experiments. test-set performance is reported as an average over multiple mesh topics (as many as our resources allowed). except for the experiments that use pretrained models, the inputs are tokenized by words, and word embeddings are randomly initialized and trained with the model. more details are available in appendix b. we stress that our intent in this work is not to report the very best scores possible, but rather to perform controlled experiments to test hypotheses. to this end, many orthogonally beneficial "tricks" in nlp literature were not utilized. additionally, nnpu-trained models generally required more diligent hyperparameter tuning due to an additional two hyperparameters. ## pu learning for document set expansion in pu classification literature, traditionally small (and in many cases, linear) models have been used on relatively simple tasks, such as cifar-10 and 20news. however, performance of existing methods does not scale well to very high-dimensional inputs and state-of-the-art neural models for text classification; applying the pu learning methods described in section 2 to a more practical setting results in several critical challenges that must be overcome-for example, pu learning methods often assume a known class prior, yet estimation of the class prior, particularly for text, is hard and inaccurate. in this section we discuss various challenges we have encountered in applying pu learning to the pubmed set expansion task, along with proposed, empirically validated solutions. ## effectiveness of pu learning in this section we evaluate the viability of our proposed solution. all experiments in this section use ber optimization and proportional batching (section 5), but no pre-trained embeddings. we refer to our proposed method as bm25+nnpu where the ir solution bm25 selects the unlabeled dataset for the pu solution, a cnn model with the nnpu loss. as an anchor for comparison, we use the following reference: upper-bound: an identical model, trained on the same training data with full supervision using the true labels. this reference can be regarded as the upper-bound performance in the ideal case. we directly compare against the current commonly-used and best-performing solution as ir (bm25) foot_6 : the top-k documents of the ir engine's output, for @xmath0, are selected as positive documents, while the rest are treated as negative. f1 mean and standard deviation are reported across k. this strong baseline serves as a reference to the state-of-the-art. we additionally compare against standard dse baselines all + (all positive): classifying all samples as the positive class; and naive: supervised learning between the labeled positive set (as p) and the unlabeled set (as n). finally, we compare against two additional baselines with the aim of validating the beneficial synergy between the ir step and the pu step. in the rand+pu baseline, we replace the ir step with a random selection of u data. in the bm25+copk baseline, we replace the pu step with a constrained k-means clustering @xcite solution, where we perform k-means clustering, @xmath1, under the constraint that all lp examples must be in the same cluster. to represent examples in embedding space for k-means, we encode the text with scibert. additional details of constrained clustering as a replacement to pu learning are discussed in appendix c. the ir baseline is the main alternative to our approach. the all-positive and naive baselines are very simplistic "lower-bound" models to be compared against, while the other two-step baselines evaluate the ir or pu steps separately, providing further justification to using the ir and pu solutions together. experiments in an interesting experiment in figure 1 shows the performance of the ir and pu models normalized by the performance of the upper-bound, as a function of the amount of labeled data. the reported values are the distance of f1 scores between each respective model with the upper-bound, normalized by the sum of scores. the figure shows that as more labeled data is added, the pu model (in truth ir+pu) increases in performance at a rate that is higher than the performance increase of the upper-bound. in comparison, the ir model improvement stays relatively constant beyond 300 labeled samples, while the upper-bound continues to increase, causing the disparity between them to grow. this experiment shows that the ir+pu system scales well with increase in lp data, increasing performance at a stronger pace than the fully-supervised reference, while the ir solution scales poorly. ## using negative data the document set expansion scenario may allow for cases where a limited amount of negative data can be collected. for example, the user may possess some number of relevant negative documents which were acquired alongside the positive documents, prior to training; alternatively, the user may label some documents from the model's output as they appear. therefore, it is of interest to augment the task with biased labeled negative data-i.e., negative documents which were not sampled from the true negative distribution, but were selected with some bias, such as their length, popularity (for example, the number of citations), or their placement within the ir engine's rankings. we consider a bias from document character length, randomly sampling abstracts that are below a certain amount of characters. alternative bias methods are dis ## conclusion we propose a two-stage solution to document set expansion-the task of retrieving documents from a large collection based on a small set of documents pertaining to a latent fine-grained topic-as a method of improving and expanding upon current ir solutions, by training a pu model on the output of a black-box ir engine. in order to accurately evaluate this method, we synthetically generated tasks by leveraging pubmed mesh term conjunctions to denote latent topics. finally, we discuss challenges in applying pu learning to this task, namely an unknown class prior, extremely imbalanced data and batch size restrictions, propose solutions (one of which-"proportional batching"applies in the general scope of pn imbalanced classification, as we empirically validate), and provide empirical evaluation against multiple baselines which showcase the effectiveness of the approach. future work. stronger class prior estimation, through additional task assumptions, may facilitate direct auc optimization. additionally, methods of increasing precision may be considered (such as data augmentation or adversarial training).