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: The authors proposed a variant of ensemble method in reinforcement learning for query reformulation. They train multiple specialized sub-agents on disjoint partitions of the training data, and use a meta-agent, which can see all the training data, to decide the final answer. This can speed up the training thanks to parallelization. They observed that this can improve the diversity of learnt reformulations and the overall performance in some cases.