Year: 2021

Abstract:
This paper presents first successful steps in designing agents that learn metastrategies for iterative query refinement. Our approach uses machine reading to guide the selection of refinement terms from aggregated search results. Agents are then empowered with simple but effective search operators to exert fine-grained and transparent control over queries and search results. We develop a novel way of generating synthetic search sessions, which leverages the power of transformerbased language models through (self-)supervised learning. We also present a reinforcement learning agent with dynamically constrained actions that learns interactive search strategies from scratch. We obtain retrieval and answer quality performance comparable to recent neural methods using a traditional term-based BM25 ranking function. We provide an in-depth analysis of the search policies.

heading: 1 INTRODUCTION
text: Can machines learn to use a search engine as an interactive tool for finding information? Web search is the portal to a vast ecosystem of general and specialized knowledge, designed to support humans in their effort to seek relevant information and make well-informed decisions. Utilizing search as a tool is intuitive, and most users quickly learn interactive search strategies characterized by sequential reasoning, exploration, and synthesis (Hearst, 2009; Rutter et al., 2015; Russell, 2019). The success of web search relies on machines learning human notions of relevance, but also on the users’ ability to (re-)formulate appropriate queries, grounded in a tacit understanding of strengths and limitations of search engines. Given recent breakthroughs in language models (LM) (Vaswani et al., 2017; Devlin et al., 2019; Brown et al., 2020) as well as in reinforcement learning (RL) (Mnih et al., 2013; Silver et al., 2016; Berner et al., 2019), it seems timely to ask whether, and how, agents can be trained to interactively use search engines. However, the lack of expert search sessions puts supervised learning out of reach, and RL is often ineffective in complex natural language understanding (NLU) tasks. The feasibility of autonomous search agents hence remains an open question, which inspires our research.
We pursue a design philosophy in which search agents operate in structured action spaces defined as generative grammars, resulting in compositional, productive, and semantically transparent policies. Further domain knowledge is included through the use of well-known models and algorithms from NLU and information retrieval (IR). Most notably, we develop a self-supervised learning scheme for generating high-quality search session data, by exploiting insights from relevance feedback (Rocchio, 1971), used to train a supervised LM search agent based on T5 (Raffel et al., 2020). We also build an RL search agent based on MuZero (Schrittwieser et al., 2020) and BERT (Devlin et al., 2019), which performs planning via rule-constrained Monte Carlo tree search and a learned dynamics model.
We run experiments on an open-domain question answering task, OpenQA (Lee et al., 2019). Search agents learn diverse policies leading to deep, effective explorations of the search results. The MuZero agent outperforms a BM25 (Robertson & Zaragoza, 2009) search function running over a Wikipedia index, on both retrieval and answer quality metrics. Thus, providing novel evidence for the potential of knowledge-infused RL in hard NLU tasks. The T5 agent can more easily leverage large pre-trained encoder-decoders and proves superior to MuZero. Furthermore, a straightforward ensemble of agents is comparable in performance to a state of the art neural retrieval system, DPR (Karpukhin et al., 2020), while relying solely on interpretable, symbolic retrieval operations. This suggests new challenges for future work; e.g., involving hybrid architectures and policy synthesis.1
1We open-source the code and trained checkpoints for both agents: anonymized during review.

heading: 2 SELF-SUPERVISED LEARNING FOR INTERACTIVE SEARCH
text: It has been a powerful vision for more than 20 years to design search engines that are intuitive and simple to use. Despite their remarkable success, search engines are not perfect and may not yield the most relevant result(s) in one shot. This is particularly true for rare and intrinsically difficult queries, which may require interactive exploration by the user to be answered correctly and exhaustively. Contextual query refinement is a common technique (Jansen et al., 2009), even among children (Rutter et al., 2015), used to improve search by combining evidence from previous results and background knowledge (Huang & Efthimiadis, 2009). Such refinements often rely on inspecting result snippets and titles or on skimming the content of top-ranked documents. This process is iterative and may be repeated to produce a sequence of queries q0, q1, . . . , qT until (optimistically) a satisfactory answer is found. It seems natural to mimic this interactive process by a search agent, which learns the basic step of generating a follow-up query from previous queries and their search results.
Furthermore, it is noteworthy how power users apply dedicated search operators and sophisticated investigative strategies to solve deep search puzzles (Russell, 2019). In particular, unary operators offer a great deal of fine-grained control and transparency and as such are highly effective in expert hands. We concentrate on three operators: ‘+’, which limits results to documents that contain a specific term, ‘-’ which excludes results that contain the term, and ‘∧i’ which boosts a term weight in the BM25 score computation by a factor i ∈ R. For instance – see also Figure 1 – the query ’who won the us open’, may lead to both tennis and golf-related results. An expert searcher could zero-in on the tennis intent by excluding golf-related terms and boosting tennis-related ones. As we show in this paper, these operators are also pivotal in designing interactive search agents.

heading: 2.1 RESULT AGGREGATION FOR CONTEXTUAL REFINEMENT
text: Web searchers expect the best answer to be among the top two hits on the first results page (Hearst, 2009, §5) and pay marginal attention to the bottom half of the 10 blue links (Granka et al., 2004; Joachims et al., 2005; Nielsen & Pernice, 2009; Strzelecki, 2020). Likewise, a search agent considers only the top k documents returned by the search engine at every step, where k = 5. During a search session the agent maintains a list of the top-k documents overall, which is returned at the end.
To aggregate results we use a machine reader (MR, cf. (Rajpurkar et al., 2016)). Specifically, we use a DPR-like reader/passage scorer (Karpukhin et al., 2020), which builds upon a pre-trained BERT model. Beyond identifying the most promising answer span within each result document d, the system also estimates the probability of d containing the (unspecified) answer P(d 3 ANSWER | q) ∈ [0; 1]. This probability can be viewed as a Passage Score (PS) that induces a calibrated ranking across all result documents within a session.
An observation representing the session at any given step is built by extracting a fixed-length token window centered at the answer span predicted by the reader for each document. In addition, we include the document titles. Finally, the query tokens and refinements describing qt are also included. This leads to a segmented observation token sequence ot which is truncated to length ≤ 512 , a common input length for pre-trained transformer-based LMs (cf. Appendix B for details, examples). We then use BERT or T5 to produce an embedding st from which the search agent will generate the next query. If we denote the result set for qt by Dt, then we get diagrammatically q0, . . . , qtsearch ↓ engine
D0, . . . ,Dt  MR/PS7−→ ot︸ ︷︷ ︸ observation LM7−→ st︸ ︷︷ ︸ encoding agent7−→ qt+1︸ ︷︷ ︸ generation
(1)
We focus on the case where qt+1 is obtained from qt through augmentation. This may add a keyword w ∈ Σidx, where Σidx is the search index vocabulary, with the usual disjunctive search engine semantics or a structured search term formed by the use of unary operators (‘+’,‘-’,‘∧’) and fields.

heading: 2.2 ROCCHIO QUERY EXPANSIONS
text: In the absence of training sessions from human expert users, we propose to generate synthetic search sessions in a self-supervised manner, making use of a set of question-answer pairs (q, a). We initialize q0=q and aim to find a sequence of refinements that make progress towards identifying documents containing the answer a, based on a reward function qt 7→ Dt 7→ rt ∈ [0; 1] (cf. §4). A query is not further refined, if either t=20 (maximal length) or if no score increasing refinement can be found.
To create candidate refinements, we make use of the idea of relevance feedback as suggested in Rocchio (1971). An elementary refinement – called a Rocchio expansion – then takes the form
qt+1 := qt ∆qt,∆qt := [+| − | ∧i TITLE | CONTENT ] wt, wt ∈ Σt := Σqt ∪ Στt ∪ Σαt ∪ Σβt (2) where i is the boosting coefficient and Σt refers to a set of terms accessible to the agent. By that we mean terms that occur in the top PS-ranked session documents. We use superscripts to refer to the vocabulary of the question (q), titles (τ ), answers (α) or bodies (β) of documents in ot. Note that adding terms 6∈ Σt would make refinements difficult to reproduce for an agent and thus would provide supervision of low utility.
Another aspect of creating sessions as described above has to do with the search complexity of finding optimal sequences of Rocchio expansions. We consider q∗ = q + a as the “ideal” query, whose results define the vocabulary Σ∗. For efficiency reasons, we further constrain the terms to be added via exact matches, term boosting or term exclusions by defining respective constrained dictionaries
Σ↑t = Σt ∩ Σ∗, Σ↓t = Σt − Σ∗ . (3) This means it is possible to upgrade accessible terms, wt, to exact matches, or weight boosting, if they also occur in the ideal result set (wt ∈ Σ↑t ); and to exclude accessible terms if they are not present in the ideal results (wt ∈ Σ↓t ). We have found experimentally that this leads to a good trade-off between the quality of Rocchio expansions and the search effort to find them. The search for sequences of Rocchio expansions is done heuristically. More details, pseudo-code illustrating the procedure and examples can be found in §5, Appendix A Appendix G.

heading: 2.3 SELF-SUPERVISED T5 AGENT
text: We suggest to train a generative search agent in a supervised manner by making use of synthetic search sessions generated by Rocchio expansions. We use T5, a pretrained transformer encoder-decoder model which achieves state-of-the-art results on multiple NLU tasks. As a search agent, T5 predicts a single new search expansion from an observed state. In the spirit of everything-is-string-prediction, both state and expansions are represented as plain strings. See Appendix B for a full example.
Our T5 agent is trained via Behavioral Cloning (BC) (Michie, 1990). We treat each step in a Rocchio session as a single training example. As is common in sequence prediction tasks, we use the crossentropy loss for optimization. BC is perhaps the simplest form of Imitation Learning (IL), and has been proven effective in a variety of application domains (Sharma et al., 2018; Rodrı́guez-Hernandez
et al., 2019). In our query refinement task, it allows to inherit the expressive power of the Rocchio query expansions and, differently from other IL approaches (Ross et al., 2011; Ho & Ermon, 2016; Ding, 2020), requires only offline interactions with the search engine. Crucially, this enables scaling to the large action spaces and model sizes typical of recent LMs. Our T5 agent can also be described as a Decision Transformer with fixed max return (Chen et al., 2021).
At test time, we start with the initial query and incrementally add new expansions, querying the trained T5 model in every step. We then use the refined query to retrieve new documents and continue until either the set of new documents is empty or we reach the maximum number of steps. Throughout the session, we maintain the top-5 documents among all those retrieved.

heading: 3 REINFORCEMENT LEARNING: MUZERO AGENT
text: MuZero (Schrittwieser et al., 2020) is a state-of-the-art agent characterized by a learnable model of the environment dynamics. This allows the use of Monte Carlo tree search (MCTS) to predict the next action, in the absence of an explicit simulator. In our use case, MuZero aims to anticipate the latent state implied by each action with regard to the results obtained by the search engine. For instance, in the example of Figure 1, it may learn to predict the effect of boosting the term ’tennis’. This approach to planning is intuitive for search, as searchers learn to anticipate the effect of query refinements while not being able to predict specific results. Furthermore, this offers a performance advantage of many orders of magnitude against executing queries with the real search engine.

heading: 3.1 GRAMMAR-GUIDED SEARCH
text: To map observations to states, the MuZero agent employs a custom BERT with additional embedding layers to represent the different parts (cf. Appendix B for details). Compared to T5, MuZero has a more challenging starting point: its BERT-based representation function is pre-trained on less data, it has fewer parameters (110M vs. billions) and no cross-attention: predictions are conditioned on a single vector, [CLS]. Moreover, it cannot as easily exploit supervised signals. However, it can more openly explore the space of policies, e.g. independent of the Rocchio expansions. Through many design iterations, we have identified it to be crucial to structure the action space of the MuZero agent and constrain admissible actions and refinement terms dynamically based on context. This provides a domain-informed inductive bias that increases the statistical efficiency of learning a policy via RL.
We take inspiration from generative, specifically context-free, grammars (CFGs) (Chomsky, 1956) and encode the structured action space as a set of production rules, which will be selected in (fixed) top-down, left-to-right order. A query refinement is generated as follows
Q⇒ U Q |W Q, U ⇒ Op Field W, Op⇒ + | − |∧i, Field⇒ TITLE |CONTENT, (4)
which allows for adding plain or structured keywords using unary operators. The selection of each refinement term W proceeds in three steps, the first two can be described by the rules
W ⇒W qt |W τt |W βt |Wαt |W idx, W xt ⇒ w ∈ Σxτ , x ∈ {q, τ, β, α} , W idx ⇒ w ∈ Σidx (5)
which means that the agent first decides on the origin of the refinement term, i.e., the query or the different parts of the top-scored result documents, and afterwards selects the term from the corresponding vocabulary. As the term origin correlates strongly with its usefulness as a refinement term, this allows to narrow down the action space effectively. The agent is forced to pick a term from the larger vocabulary (1.6M terms) of the search index Σidx during MCTS, as there is no observable context to constrain the vocabulary.
The third level in the action hierarchy concerns the selection of the terms. We have found it advantageous to make use of subword units; specifically, BERT’s 30k lexical rules involving word pieces, to generate terms sequentially, starting from a term prefix and adding one or more suffixes. Note that this part of the generation is context-sensitive, as we restrict the generation to words present in the vocabulary. We make use of tries to efficiently represent each Σxτ and amortize computation. The grammar-guided MCTS is explained in detail in Appendix F.

heading: 4 THE OPENQA ENVIRONMENT
text: We evaluate search agents in the context of open-domain question answering (Open-QA) (Voorhees, 2000; Chen et al., 2017). Given a question q, we seek documents D that contain the answer a using a search engine, the environment. Following common practice, we use Lucene-BM25 with default settings on the English Wikipedia. BM25 has provided the reference probabilistic IR benchmark for decades (Robertson & Zaragoza, 2009), only recently outperformed by neural models (Lee et al., 2019). The Lucene system provides search operators comparable to commercial search engines.
Exploration-based learning is vulnerable to discovering adversarial behaviors. As a safeguard we design a composite reward. The score of a results set D, given q, interpolates three components. The first is the Normalized Discounted Cumulative Gain (NDCG) at k. See Eq. 6a, where wi = log2(i+1) −1/ ∑k j=1 log2(j + 1) −1 are normalizing weights, and rel(d|q) = 1, if a ∈ d, 0 otherwise:
a) NDCGk(D|q) = k∑ i=1 wi rel(di|q), b) NDCEMk(D|q) = k∑ i=1 wi em(di|q). (6)
NDCG is a popular metric in IR as it accounts for rank position, it is comparable across queries, and it is effective at discriminating ranking functions (Wang et al., 2013). NDCG alone can have drawbacks: on “easy” questions a score of 1 can be achieved in short meritless episodes, while on “hard” ones it may be impossible to find a first valid step, since Eq. 6a takes discrete values. Hence, we introduce a second component, NDCEMk (Eq. 6b) where em(d|q) = 1 if the answer extracted from d by the reader exactly matches a, 0 otherwise. NDCEMk helps validate results by promoting high-ranking passages yielding correct answer spans. Finally, to favour high-confidence result sets we add the normalized Passage Score of the top k results, leading to the following scoring function
Sk(D|q) := (1−λ1−λ2)·NDCGk(D|q)+λ2·NDCEMk(D|q)+λ1· 1
k k∑ i=1 PS(di|q) ∈ [0, 1] (7)
Based on (7), we define the search step reward
rt = S5(Dt|q0)− S5(Dt−1|q0). (8) We train the MuZero agent directly on the reward. The reward is sparse, as none is issued in between search steps. The T5 agent is trained indirectly on it via the induction of Rocchio sessions (cf. §2.2).

heading: 5 EXPERIMENTS
text: We use the OpenQA-NQ data (Lee et al., 2019), derived from Natural Questions (Kwiatkowski et al., 2019), consisting of Google queries paired with answers extracted from Wikipedia. The data includes 79,168 train questions, 8,757 dev questions and 3,610 for test. We use the provided partitions and Wikipedia dump. Following Lee et al. (2019) we pre-process Wikipedia into blocks of 288 tokens, for a total of 13M passages. We evaluate each system on the top-5 288-token passages returned. Model selection and data analysis are performed on NQ Dev, using the reward (Eq. 8) as the objective.
We generate synthetic search sessions using Rocchio expansions for 5 grammars: G0 (only simple terms), G1 (only term boosting), G2 (‘+’ and ‘-’), G3 (G0+G2) and G4 (G0+G1+G2). We use the reward above with λ1=0.2, λ2=0.6, after a search against the quality metrics (cf. Appendix C). We select 5 possible values, i ∈ {0.1, 2, 4, 6, 8}, for term boosting weights. At each step, we attempt at most M BM25 searches, based on the current observation ot, to find an expansion that improves the reward. To further improve efficiency, we sort the terms in ot by Lucene’s IDF score and keep the top N . We set N=M=100. For NQ Train, and G4, we find 298,654 Rocchio expansions from 77,492 questions. Similarly, we generate sessions for the NQ Dev and Test to estimate headroom.

heading: 5.1 AGENTS TRAINING AND INFERENCE
text: The DPR reader/selector and MuZero’s hθ function use 12-layer BERT systems.2 To train the former, we generate for each query in NQ Train 200 candidate passages from our BM25 system, picking
2BERT-base, initialized from https://tfhub.dev/google/bert_uncased_L-12_H-768_ A-12/1.
one positive and 23 negative passages for each query at random whenever the query is encountered during training. The reader/scorer is not trained further.
The MuZero implementation is scaled and distributed via an agent-learner setup (Espeholt et al., 2018) in the SEED RL (Espeholt et al., 2020) framework allowing for centralized batching of inference for effective use of accelerators. MuZero is trained on NQ Train for a total of 1.6 million steps (≈10 days) using 500 CPU-based actors and 4 Cloud TPU v2 for inference and training on the learner.3 For each step, 100 simulations are performed. During training, we limit sessions to a maximum of 20 steps. The agent also can decide to stop early by selecting a dedicated stop action. Training of MuZero can be improved by providing advice to the actors. An actor may receive information about which terms wt should be promoted, wt ∈ Σ↑t , or demoted, wt ∈ Σ↓t . The probability of an episode receiving advice starts at 0.5 and decays linearly to 0 in one million steps.
For the T5 agent we start from the pretrained T5-11B (11 billion parameters) public checkpoint and continue training on the NQ Train expansions. Training took about 5 days using 16 Cloud TPU v3. At inference time, we found that fixing the sessions to 20 steps worked best for both T5 and MuZero. We report detailed training configurations and ablations in Appendix D.

heading: 5.2 RESULTS
text: Table 1 summarizes the results on NQ Test. We evaluate passage retrieval quality by means of ranking (NDCG@5) and precision (Top-1, Top-5) metrics. We also report Exact Match (EM) to evaluate answer quality. Our baseline is Lucene’s BM25 one-shot search. Reranking the same BM25 documents by the PS score (PS(D0)) is easy and improves performance on all metrics, particularly noticeable on Top-1 and EM,4 providing a more informative comparison. The last column (RQE-G4) reports the metrics of the episodes generated by the Rocchio Query Expansion process (§2.2) using the grammar with all operators (G4). RQEs use the gold answer and can be seen as a, possibly conservative, estimate of the performance upper bound. As the external benchmark we use DPR (Karpukhin et al., 2020), a state-of-the-art OpenQA neural retriever based on dual encoders, the dominant architecture for deep learning-based ad hoc retrieval (Craswell et al., 2020).
The MuZero agent (MZ) outperforms both BM25 and PS(D0). While this result may seem trivial, it marked a milestone that required many iterations to achieve. The challenges for RL in IR, and NLU, are extreme in terms of state and action space dimensionality, data sparsity etc. (Zhang et al., 2021). We propose ideas for tackling some of these key challenges by fitting out agents with domain knowledge in principled ways, with the grammar-guided MCTS as the centerpiece, The best MuZero converges to a policy which uses only term boost actions with a weight of 2 (Figure 2a). This agent is not able to find better-performing, diverse policies. This is an extreme case of a more general pattern. Different sub-grammars represent different tactics; e.g., ‘+’ and ‘-’ affect the accessible documents in irreversible ways, while boosting only affects ranking. It is challenging for all agents, and particularly MuZero, to modulate effectively multiple sub-policies.
Training T5 is less involved than MuZero. This allows us to evaluate all 5 grammars. The best one, ‘T5-G1’ in Table 1, is limited to term boosting (G1), but it learns to use all available weight values (Figure 2a). In terms of Top-1 this agent outperforms the published and the most recently posted DPR results5 but has worse Top-5 than both. Results for all five T5 agents are found in Table A.5.
3For details, see https://cloud.google.com/tpu. 4Top-5 is identical to BM25 since the documents are the same. 5https://github.com/facebookresearch/DPR.
In the last experiment we combine all trained agents, the five T5 agents and the MuZero one, in one ensemble. We simply rank the union of all the documents returned by the ensemble, by means of the PS score on each document. This ensemble (‘MZ+T5s’ in Table 1) has slightly better precision than the recent DPR in top position, and slightly worse for the Top-5. This results indicates that the ability to orchestrate diverse sub-policies may indeed be the key to future progress for search agents. The current SOTA for Top-5 is 74.0 (Qu et al., 2021).
We conclude by discussing answer quality. Agents routinely produce answer spans, as predicted by the reader/scorer, to build observations. The MR/PS component is trained once, before the agents, on the output of BM25. However, agents deeply change results. As Figure 2b shows, they dig deep in the original BM25 ranking. This is positive, as behavior discovery is one of the main motivations for researching exploratory methods like RL. As a consequence, though, the MR/PS component effectively operates out of distribution and the EM numbers of the internal reader are not competitive with recent methods, Table A.6 reports all the numbers including on NQ Dev.
Ideally, one would co-train the observation builder with the search agent. However, combining the two would introduce significant engineering complexity in the current architecture. For instance, one could interleave training the two as in DQNs (Mnih et al., 2013). A simpler alternative is to add the answer prediction task to the T5 agent. Retrieval-augmented answer generation is known to produce strong results (Izacard & Grave, 2021). Multitasking would simplify the design of the generative agents and possibly produce better models. We make a first step in this direction by training one single, dedicated T5 agent. The system uses as training input the top-5 documents of the RQE-G4 episodes, but its task is to generate the gold answer, instead of the query expansion. By using the output of the ‘T5-G1’ and ‘MZ+T5s’ agents, the EM performance of the answer generation T5 is comparable to methods that build on DPR, such as RAG (Lewis et al., 2020b) (44.5 EM). Although not as good as FID (Izacard & Grave, 2021) that condition on many more (100) documents. The performance of the MuZero agent is lower here, we believe this is due its more idiosyncratic results.

heading: 5.3 DISCUSSION AND FUTURE WORK
text: Table 2 illustrates an example where the T5-G4 agent switches policy mid-session. This agent is trained on the full grammar (G4) and can use all search operators. The question is about basketball records and BM25 does not find good results. In the first three steps the agent focuses on re-ranking by boosting terms like ‘per’ (from the phrase ‘per game’ in the results for q0) and ‘scorer’. This produces a good hit and answer span (‘Pete Maravich’) at position 1 of step 3. The agent then switches to filtering mode, to zero in on documents containing the predicted answer term. This is a clear instance of successful policy synthesis. However it is a gamble, as predicted answers can be incorrect and filtering actions are not reversible. This suggests that the action space may benefit by including more control actions, e.g. to undo or go back to a specific state, to better support safe exploration and the emergence of meta policies. We plan to investigate this in future work.
The previous point extends to the agents’ architecture. It is reasonable to hypothesise that the superior performance of T5 is due to two main factors. T5s are bigger models, trained on more data, and have a more expressive prediction process based on encoder-decoders. In addition, they are finetuned on a self-supervised tasks which provides significant headroom. While T5-like LMs seem the obvious choice forward there are open questions concerning exploration. It is not clear how much the model can generalize being trained offline and never being exposed to its own prediction. This moves the learning problem back towards RL. We plan to investigate approaches like Decision Transformers (Chen et al., 2021) next, as a natural framework for expressing the search task within LMs, while at the same time bringing back in key RL concepts such as expected returns and learning also from negative experiences, possibly produced by different sources; e.g., MuZero.
We would like to note that pre-trained language models of the kind used here have been shown to capture societal biases (Tan & Celis, 2019; Webster et al., 2020), which motivates a broad discussion about potential harms and mitigations (Blodgett et al., 2020; Bender et al., 2021). We have no reason to believe our architectures would exacerbate biases, but the overall problems may persist. We also hope that end-to-end optimization methods based on composite rewards, as in this proposal, can contribute to solve some of these challenges; e.g., by providing means of adversarial testing, and by including relevant metrics directly in the objective design.
While our agents yield performance comparable to neural retrievers they rely solely on interpretable, symbolic retrieval operations. Exploration and transparency are core objectives of our framework.

heading: 6 RELATED WORK
text: Query optimization is an established topic in IR. Methods range from hand-crafted rules (Lawrence & Giles, 1998) to data-driven transformation patterns (Agichtein et al., 2001). Narasimhan et al. (2016) use RL to query the web for information extraction. Nogueira & Cho (2017) and Buck et al. (2018) use RL-trained agents to seek good answers by reformulating questions with seq2seq models. These methods are limited to one-step episodes and queries to plain natural language. This type of modeling is closely related to the use of RL for neural machine translation, whose robustness is currently debated (Choshen et al., 2020; Kiegeland & Kreutzer, 2021). Montazeralghaem et al. (2020) propose a feature-based network to score potential relevance feedback terms to expand a query. Das et al. (2019) propose to perform query reformulation in embedding (continuous) space and find that it can outperform the sequence-based approach. Xiong et al. (2021) successfully use relevance feedback by jointly encoding the question and the text of its retrieved results for multi-hop QA. Other work at the intersection of IR and RL concerns bandit methods for news recommendation (Li et al., 2010) and learning to rank (Yue & Joachims, 2009). Recently, interest in Deep RL for IR has grown (Zhang et al., 2021). There, the search engine is the agent, and the user the environment. In contrast, we view the search problem from the user perspective and thus consider the search engine as the environment.
The literature on searchers’ behavior is vast, see e.g. Strzelecki (2020) for an overview of eye-tracking studies. While behavior evolves with interfaces, users keep parsing results fast and frugally, attending to just a few items. From a similar angle, Yuan et al. (2020) offer promising findings on training QA agents with RL for template-based information gathering and answering actions. Most of the work in language-related RL is otherwise centered on synthetic navigation/arcade environments (Hu et al., 2019). This line of research shows that RL for text reading can help transfer (Narasimhan et al., 2018) and generalization (Zhong et al., 2020) in synthetic tasks but skirts the challenges of more realistic language-based problems. On the topic of grammars, Neu & Szepesvári (2009) show that Inverse RL can learn parsing algorithms in combination with PCFGs (Salomaa, 1969).
Current work in OpenQA focuses on the search engine side of the task, typically using dense neural passage retrievers based on a dual encoder framework instead of BM25 (Lee et al., 2019; Karpukhin et al., 2020). Leveraging large pre-trained language models to encode the query and the paragraphs separately led to a performance boost across multiple datasets, not just in the retrieval metrics but also in exact-match score. While Karpukhin et al. (2020) use an extractive reader on the top-k returned paragraphs, Lewis et al. (2020b) further improves using a generative reader (BART (Lewis et al., 2020a)). This design combines the strengths of a parametric memory – the pre-trained LM – with a non-parametric memory – the retrieved Wikipedia passages supplied into the reader’s context. This idea of combining a dense retriever with a generative reader is further refined in Izacard & Grave (2021), which fuses multiple documents in the decoding step. A recent line of work is concerned with constraining the model in terms of the number of parameters or retrieval corpus size while remaining close to state-of-the-art performance (Min et al., 2021). This effort led to a synthetic dataset of 65 million probably asked questions (Lewis et al., 2021) used to do a nearest neighbor search on the question – no learned parameters needed – or train a closed-book generative model.

heading: 7 CONCLUSION
text: Learning to search sets an aspiring goal for AI, touching on key challenges in NLU and ML, with far reaching consequences for making the world’s knowledge more accessible. Our paper provides the following contributions. First, we open up the area of search session research to supervised language modeling. Second, we provide evidence for the ability of RL to discover successful search policies in a task characterized by multi-step episodes, sparse rewards and a high-dimensional, compositional action space. Lastly, we show how the search process can be modeled via transparent, interpretable machine actions that build on principled and well-established results in IR and NLU.
Our findings seem to agree with a long-standing tradition in psychology that argues against radical behaviorism – i.e., pure reinforcement-driven learning, from tabula rasa – for language (Chomsky, 1959). RL agents require a remarkable share of hard-wired domain knowledge. LM-based agents are easier to put to use, but they need abundant task-specific data for fine tuning. Supplied with the right inductive bias, LM and RL search agents prove surprisingly effective. Different architectures learn different policies, suggesting broad possibilities in the design space for future work.
Reproducibility statement The code for this paper has already been open sourced (reference omitted for anonymity) and included with the supplementary material. We will make sure results are fully reproducible by updating the open source code based on the final version of the paper after completion of the review process.

heading: A ROCCHIO QUERY EXPANSIONS
text: Algorithm 1 provides a schematic summary of the procedure for generating Rocchio sessions. We omit the terms source for simplicity and readability, but it should be straightforward to reconstruct. Table A.7 shows an example of a session produced by Rocchio expansions using the full grammar.
Algorithm 1: Rocchio Query Expansions input :A question-answer pair (q, a), k = 5, num steps = 20, N = 100, M = 100 output :A set of observation-query expansion pairs for training a T5 agent RQE = {(ot,∆qt)} RQE← ∅ qt ← q; Dt ← ∅ ; // Unique documents found in the session q∗ ← q +(contents:”a”) ; // The ideal query D∗ ← LuceneBM25(q∗) ; // Use search to get the top k documents // Use the agent PS and MR components to rerank the documents, extract
answer spans to compute the snippets from the top k results, and compile the observation (cf. also Appendix B)
o∗ ← ComputeObservation(q, q∗,D∗, k); Σ∗ ← ComputeDictionary(o∗) ; // Collect good search terms for t← 1 to num steps do Dt ← Dt ∪ LuceneBM25(qt); ot ← ComputeObservation(q, qt,Dt, k); Σt ← ComputeDictionary(ot); Σ↑ ← Σ∗ ∩ Σt, Σ↓ ← Σt − Σ∗; st ← ComputeScore(q,Dt, k); // Compute the score using Eq.(7) max score← st; best action← ∅, num tries← 0; // Evaluate all available operators/fields, and also without for op, field ∈ {{+,−,∧0.1,∧2,∧4 ∧ 6 ∧ 8} × {contents:, title:}} ∪ {(′′,′′ )} do
for w ∈ TopNTermsByLuceneIDF(Σt,N) do if (op == ′−′ ∧ w ∈ Σ↓) ∨ (op 6= ′−′ ∧ w ∈ Σ↑) ∧ (num tries < M) then
∆qt ← op(field:′w′); q′ ← qt + ∆qt; D′ ← Dt ∪ LuceneBM25(q′); s′ ← ComputeScore(q,D′, k); num tries← num tries + 1; if s′ > max score then
max score← s′; best action← ∆qt;
end else
continue; end
end end if max score > st then
// If the best action improves the score, add this step to the data, and continue the session qt ← qt + best action; RQE← RQE ∪ (ot, best action);
else return RQE;
end end return RQE

heading: B OBSERVATION BUILDING DETAILS
text: This section provides more details and examples about the encoding of observations for both the MuZero and the T5 agent. As described in Section 2.1, the main part of the observation consists of the top-5 documents from all results retrieved so far, ∪ti=0Di. The documents are sorted according to the PS score and reduced in size by extracting fixed-length snippets around the DPR reader’s predicted answer. Moreover, the corresponding Wikipedia article title is appended to each document snippet. In addition to the top documents, the observation includes the original question and information about any previous refinements. While the main part of the observation is shared between the MuZero and the T5 agent, there are differences in the exact representation. The following two paragraphs give a detailed explanation and example for both agents.
MuZero Agent’s State (cf. §2.1) The MuZero agent uses a custom BERT (initialized from BERTbase) with additional embedding layers to represent the different parts of the observation. It consists of four individual embedding layers as depicted in Figure A.1. At first, the standard layer for the tokens of the query, the current tree, and the current top documents D5. The second layer assigns a type ID to each of the tokens representing if a token is part of the query, the tree, the predicted answer, the context, or the title of a document. The last two layers add scoring information about the tokens as float values. We encode both the inverse document frequency (IDF) of a word and the documents’ passage selection (PS) score. Figure A.1 shows a concrete example of a state used by the MuZero agent.
T5 Agent’s State (cf. §2.3) T5 represents the state as a flat string. The input is a concatenation of the original query, zero or more expansions, and five results. For each result, we include the answer given by the reader, the document’s title, and a span centered around the answer. The prediction target is simply the next expansion. See Table A.2 for a full example.

heading: C REWARD DETAILS
text: We investigate the effects of the three score components in the definition of the composite scoring function of Eq. 7. As mentioned in Section 4, in our experiments we have observed that using only the NDCGk score as reward signal (i.e., setting λ1 = λ2 = 0 in Eq. 7) has several limitations. This motivated us to introduce the NDCEMk and PS components with the intent of: 1) providing further guidance to the agent (whenever NDCGk cannot be increased, the agent can further refine the query by increasing NDCEMk or PS), and 2) regularizing the search episodes by making the score more robust with respect to exploratory behaviors that could lead to drift.
We run a grid search over the reward coefficients λ1, λ2 and, for each of their values, we evaluate the performance of the Rocchio query expansions on NQ Dev (for a high throughput, we select grammar G3 and set N = M = 20). Figure A.2 shows the respective end-to-end performance in terms of our three main quality metrics: NDCG@5, EM, and Top-5.
The results in Figure A.2 support our intents: by introducing NDCEMk and PS scores in the reward (i.e., setting λ1, λ2 > 0), the Rocchio expansions can achieve significantly higher performance, in all the three metrics, with respect to using only an NDCGk score (λ1 = λ2 = 0) (notably, it improves also the NDCG@5 itself, meaning that the agent is not trading-off performance metrics but it is indeed producing higher quality sessions). It is also worth pointing out the role of the NDCEMk score component, weighted by coefficient λ2. Notice that good NDCG@5 and Top-5 performance could be achieved also setting λ2 = 0 (see, e.g., the bottom-right corner λ1 = 1, λ2 = 0). However, this leads to definitely worse EM results compared to when λ2 > 0. Intuitively, a NDCEMk component λ2 > 0 ensures that the returned documents, in addition to containing the gold answer (thus having high NDCG@5 and Top-5), are also relevant for the query (thus reaching a high EM). Hence, it is crucial to prevent semantic drifts. Based on these results we set λ1 = 0.2, λ2 = 0.6, which is a sweet spot in Figure A.2.

heading: D MODEL AND TRAINING CONFIGURATION
text: MuZero The MuZero agent learner, which performs both inference and training, runs on a Cloud TPU v2 with 8 cores which is roughly equivalent to 10 Nvidia P100 GPUs in terms of TFLOPS.6 One core is allocated for training and 7 cores are allocated for inference. We use 500 CPU based actors along with 80 actors dedicated to evaluation. Each agent is trained for 1.6 million steps, with 100 simulations per step, at an approximate speed of 10,000 steps per hour. In total, training takes about 10 days. Hyperparameters are listed in Table A.3.
T5 The T5 agent is trained on 16 Cloud TPU v3, starting from the pre-trained T5-11B checkpoint. We mostly follow the suggested hyperparameter choices given by Raffel et al. (2020), see Table A.4. We train for about 2 days at roughly 3 steps per second for a total of 50k steps and select a checkpoint based on dev performance.
6The Cloud TPU v2 has a peak performance of 180 TFLOPS (https://cloud.google.com/tpu), whereas the Nvidia P100 GPU goes up to 18.7 TFLOPS depending on precision (https://cloud.google. com/compute/docs/gpus).

heading: E RESULTS

heading: F DETAILS AND EXAMPLES FOR THE GRAMMAR-GUIDED MCTS
text: Figure A.3a lists the detailed rules schemata for the query grammar used by the MuZero agent – explained in Section 3.1. An optional STOP rule allows the agent to terminate an episode and return the results collected up to that point. Using the BERT sub-word tokens as vocabulary allows us to generate a large number of words with a total vocabulary size of ∼30k tokens.
Our implementation of MuZero modifies the MCTS to use the query grammar for efficient exploration. Figure A.3b shows the different network components used during MCTS. Each node expansion is associated with a grammar rule. When the simulation phase is complete, the visit counts collected at the children of the MCTS root node provide the policy π from which the next action at+1 is sampled.
Each simulation corresponds to one or more hypothetical follow-up queries (or fragments) resulting from the execution of grammar rules. The MCTS procedure executes Depth-First node expansions, guided by the grammar, to generate a query top-down, left-to-right, in a forward pass. To control the process, we add two data structures to MCTS nodes: a stack γ, and an output buffer ω: γ contains a list of unfinished non-terminals, ω stores the new expansion. The stack is initialized with the start symbol γ = [Q]. The output buffer is reset, ω = [], after each document search. When expanding a node, the non-terminal symbol on the top of γ is popped, providing the left-hand side of the rule associated with the new edge. Then, symbols on the right-hand side of the rule are pushed right-to-left onto γ. When a terminal rule is applied, the terminal symbol is added to ω. The next time γ contains only Q, ω holds the new query expansion term ∆qt to be appended to the previous query qt for search.
Figure A.4 illustrates the process. Nodes represent the stack γ and output buffer ω. Edges are annotated with the rule used to expand the node. We illustrate the left-branching expansion. Starting from the top, the symbol ”Q” is popped from the stack, and a compatible rule, ”Q ⇒ W Q”, is sampled. The symbols ”W” and ”Q” are added to the stack for later processing. The agent expands the next node choosing to use the document content vocabulary (W ⇒ Wβ), then it selects a vocabulary prefix (‘dial’), adding it to the output buffer ω, followed by a vocabulary suffix (‘ects’). At that point, the stack contains only Q, and the content of ω contains a new expansion, the term ‘dialects’. A latent search step is simulated through MuZero’s gθ sub-network. Then the output buffer ω is reset.
After the latent search step, the simulation is forced to use the full trie (W ⇒Widx), which includes all terms in the Lucene index. This is necessary since there is no observable context that can be used to restrict the vocabulary. Instead of simply adding an OR term (Q⇒W Q), the right branch of the example selects an expansion with unary operator and field information (Q⇒ U Q).

heading: G SEARCH SESSION EXAMPLES
text: Table A.7 shows an example of a session produced by Rocchio expansions using the full grammar. Table A.8 shows a session generated by the MuZero agent.

Review summaries:
This work designs an RL (customized BERT MuZero) agent for query reformulation, where the underlying observation is defined as an aggregation of windows of answer spans found with a machine reading module, document titles, query tokens, and refinements. The action space involves adding keywords to the query or form search terms with three operators that constrain the search space. Due to the lack of training data, synthetic search sessions are generated with a pretrained transformer model and imitation learning. The reward is a linear combination of NDCG (well-established IR ranking for non-binary relevance), a revised NDCG-like metric for exact matches (NDCEM) and a top-k Passage Score. Results on an OpenQA dataset are *close* to neural passage retrieval methods.   

This paper proposes an approach based on MuZero to learn strategies to enhance the search query. It borrows the idea of pseudo relevance feedback of Rocchio, selects terms from the top-5 retrieved results, and incorporate a term using a predefined grammar. The expanded query is used by BM25 to find documents. The approach is tested on NQ dataset. It is shown to achieve a performance comparable to that of DPR. The key contribution of the exploration of MuZero for query enhancement, combined with the idea of Rochio. This exploration is new.

This paper investigates a vital direction: How can RL agents learn to use search engines to find information. The authors propose a method that leverages pre-trained models conditioned on machine reading results to guide the selection of query refinement from aggregated search documents. The method performs comparably to neural retrievers on OpenQA-NQ but operates in a more interpretable way.