Learning to Coordinate Multiple Reinforcement Learning Agents for Diverse Query Reformulation

1 INTRODUCTION . Reinforcement learning has proven effective in several language processing tasks , such as machine translation ( Wu et al. , 2016 ; Ranzato et al. , 2015 ; Bahdanau et al. , 2016 ) , question-answering ( Wang et al. , 2017a ; Hu et al. , 2017 ) , and text summarization ( Paulus et al. , 2017 ) . In reinforcement learning efficient exploration is key to achieve good performance . The ability to explore in parallel a diverse set of strategies often speeds up training and leads to a better policy ( Mnih et al. , 2016 ; Osband et al. , 2016 ) . In this work , we propose a simple method to achieve efficient parallelized exploration of diverse policies , inspired by hierarchical reinforcement learning ( Singh , 1992 ; Lin , 1993 ; Dietterich , 2000 ; Dayan & Hinton , 1993 ) . We structure the agent into multiple sub-agents , which are trained on disjoint subsets of the training data . Sub-agents are co-ordinated by a meta-agent , called aggregator , that groups and scores answers from the sub-agents for each given input . Unlike sub-agents , the aggregator is a generalist since it learns a policy for the entire training set . We argue that it is easier to train multiple sub-agents than a single generalist one since each sub-agent only needs to learn a policy that performs well for a subset of examples . Moreover , specializing agents on different partitions of the data encourages them to learn distinct policies , thus giving the aggregator the possibility to see answers from a population of diverse agents . Learning a single policy that results in an equally diverse strategy is more challenging . Since each sub-agent is trained on a fraction of the data , and there is no communication between them , training can be done faster than training a single agent on the full data . Additionally , it is easier to parallelize than applying existing distributed algorithms such as asynchronous SGD or A3C ( Mnih et al. , 2016 ) , as the sub-agents do not need to exchange weights or gradients . After training the sub-agents , only their actions need to be sent to the aggregator . We build upon the works of Nogueira & Cho ( 2017 ) and Buck et al . ( 2018b ) . Therefore , we evaluate the proposed method on the same tasks they used : query reformulation for document retrieval and question-answering . We show that it outperforms a strong baseline of an ensemble of agents trained on the full dataset . We also found that performance and reformulation diversity are correlated ( Sec . 5.5 ) . Our main contributions are the following : • A simple method to achieve more diverse strategies and better generalization performance than a model average ensemble . • Training can be easily parallelized in the proposed method . • An interesting finding that contradicts our , perhaps naive , intuition : specializing agents on semantically similar data does not work as well as random partitioning . An explanation is given in Appendix F . 2 RELATED WORK . The proposed approach is inspired by the mixture of experts , which was introduced more than two decades ago ( Jacobs et al. , 1991 ; Jordan & Jacobs , 1994 ) and has been a topic of intense study since then . The idea consists of training a set of agents , each specializing in some task or data . One or more gating mechanisms then select subsets of the agents that will handle a new input . Recently , Shazeer et al . ( 2017 ) revisited the idea and showed strong performances in the supervised learning tasks of language modeling and machine translation . Their method requires that output vectors of experts are exchanged between machines . Since these vectors can be large , the network bandwidth becomes a bottleneck . They used a variety of techniques to mitigate this problem . Anil et al . ( 2018 ) later proposed a method to further reduce communication overhead by only exchanging the probability distributions of the different agents . Our method , instead , requires only scalars ( rewards ) and short strings ( original query , reformulations , and answers ) to be exchanged . Therefore , the communication overhead is small . Previous works used specialized agents to improve exploration in RL ( Dayan & Hinton , 1993 ; Singh , 1992 ; Kaelbling et al. , 1996 ) . For instance , Stanton & Clune ( 2016 ) and Conti et al . ( 2017 ) use a population of agents to achieve a high diversity of strategies that leads to better generalization performance and faster convergence . Rusu et al . ( 2015 ) use experts to learn subtasks and later merge them into a single agent using distillation ( Hinton et al. , 2015 ) . The experiments are often carried out in simulated environments , such as robot control ( Brockman et al. , 2016 ) and video-games ( Bellemare et al. , 2013 ) . In these environments , rewards are frequently available , the states have low diversity ( e.g. , same image background ) , and responses usually are fast ( 60 frames per second ) . We , instead , evaluate our approach on tasks whose inputs ( queries ) and states ( documents and answers ) are diverse because they are in natural language , and the environment responses are slow ( 0.5-5 seconds per query ) . Somewhat similarly motivated is the work of Serban et al . ( 2017 ) . They train many heterogeneous response models and further train an RL agent to pick one response per utterance . 3 METHOD . 3.1 TASK . We describe the proposed method using a generic end-to-end search task . The problem consists in learning to reformulate a query so that the underlying retrieval system can return a better result . Following Nogueira & Cho ( 2017 ) and Buck et al . ( 2018b ) we frame the task as a reinforcement learning problem , in which the query reformulation system is an RL-agent that interacts with an environment that provides answers and rewards . The goal of the agent is to generate reformulations such that the expected returned reward ( i.e. , correct answers ) is maximized . The environment is treated as a black-box , i.e. , the agent does not have direct access to any of its internal mechanisms . Figure 1- ( b ) illustrates this framework . 3.2 SYSTEM . Figure 1- ( c ) illustrates the new agent . An input query q0 is given to the N sub-agents . A subagent is any system that accepts as input a query and returns a corresponding reformulation . Thus , sub-agents can be heterogeneous . Here we train each sub-agent on a partition of the training set . The i-th agent queries the underlying search system with the reformulation qi and receives a result ai . The set { ( qi , ai ) |0 ≤ i ≤ N } is given to the aggregator , which then decides which result will be final . 3.3 SUB-AGENTS . The first step for training the new agent is to partition the training set . We randomly split it into equal-sized subsets . For an analysis of how other partitioning methods affect performance , see Appendix F. In our implementation , a sub-agent is a sequence-to-sequence model ( Sutskever et al. , 2014 ; Cho et al. , 2014 ) trained on a partition of the dataset . It receives as an input the original query q0 and outputs a list of reformulated queries ( qi ) using beam search . Each reformulation qi is given to the same environment that returns a list of results ( a1i , .. , a K i ) and their respective rewards ( r1i , .. r K i ) . We then use REINFORCE ( Williams , 1992 ) to train the sub-agent . At training time , instead of using beam search , we sample reformulations . Note that we also add the identity agent ( i.e. , the reformulation is the original query ) to the pool of sub-agents . 3.4 META-AGENT : AGGREGATOR . The aggregator receives as inputs q0 and a list of candidate results ( a1i , .. a K i ) for each reformulation qi . We first compute the set of unique results aj and two different scores for each result : the accumulated rank score sAj and the relevance score s R j . The accumulated rank score is computed as sAj = ∑N i=1 1 ranki , j , where ranki , j is the rank of the j-th result when retrieved using qi . The relevance score sRj is the prediction that the result aj is relevant to query q0 . It is computed as : sRj = σ ( W2ReLU ( W1zj + b1 ) + b2 ) , ( 1 ) where zj = [ fCNN ( q0 ) ; fBOW ( aj ) ; fCNN ( q0 ) − fBOW ( aj ) ; fCNN ( q0 ) fBOW ( aj ) ] , ( 2 ) W1 ∈ R4D×D and W2 ∈ RD×1 are weight matrices , b1 ∈ RD and b2 ∈ R1 are biases . The brackets in [ x ; y ] represent the concatenation of vectors x and y . The symbol denotes the element-wise multiplication , σ is the sigmoid function , and ReLU is a Rectified Linear Unit function ( Nair & Hinton , 2010 ) . The function fCNN is implemented as a CNN encoder1 followed by average pooling over the sequence ( Kim , 2014 ) . The function fBOW is the average word embeddings of the result . At test time , the top-K answers with respect to sj = sAj s R j are returned . We train the aggregator with stochastic gradient descent ( SGD ) to minimize the cross-entropy loss : L = − ∑ j∈J∗ log ( sRj ) − ∑ j /∈J∗ log ( 1− sRj ) , ( 3 ) where J∗ is the set of indexes of the ground-truth results . The architecture details and hyperparameters can be found in Appendix B . 4 DOCUMENT RETRIEVAL . We now present experiments and results in a document retrieval task . In this task , the goal is to rewrite a query so that the number of relevant documents retrieved by a search engine increases . 4.1 ENVIRONMENT . The environment receives a query as an action , and it returns a list of documents as an observation/state and a reward computed using a list of ground truth documents . We use Lucene2 in its default configuration as our search engine , with BM25 as the ranking function . The input is a query and the output is a ranked list of documents . 4.2 DATASETS . To train and evaluate the models , we use three datasets : TREC-CAR : Introduced by Dietz et al . ( 2017 ) , in this dataset the input query is the concatenation of a Wikipedia article title with the title of one of its section . The ground-truth documents are the paragraphs within that section . The corpus consists of all of the English Wikipedia paragraphs , except the abstracts . The released dataset has five predefined folds , and we use the first four as a training set ( approx . 3M queries ) , and the remaining as a validation set ( approx . 700k queries ) . The test set is the same used evaluate the submissions to TREC-CAR 2017 ( approx . 1,800 queries ) . JEOPARDY : This dataset was introduced by Nogueira & Cho ( 2016 ) . The input is a Jeopardy ! question . The ground-truth document is a Wikipedia article whose title is the answer to the question . The corpus consists of all English Wikipedia articles . MSA : Introduced by Nogueira & Cho ( 2017 ) , this dataset consists of academic papers crawled from Microsoft Academic API.3 A query is the title of a paper and the ground-truth answer consists of the papers cited within . Each document in the corpus consists of its title and abstract .

Summary: In this paper, authors proposed an ensemble approach for query reformulation (QR).  The basic idea is that 1) train a bunch of models/sub-agents on subsets, e.g., randomly partitioned, of the training data; 2) and then train an additional meta model/meta-agent to aggregate the results from the step 1).  They conduct experiments on document retrieval and question answering tasks to show the effectiveness of the proposed model.