SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dhawal Gupta*
University of Massachusetts dgupta@cs.umass.edu
Yinlam Chow
Google Research  
yinlamchow@google.com
Aza Tulepbergenov
Google Research atuplep@google.com
Mohammad Ghavamzadeh†
ghavamza@amazon.com
Craig Boutilier
cboutilier@google.com
Reinforcement learning (RL) has shown great promise for developing agents for dialogue management (DM) that are non-myopic, conduct rich conversations, and maximize overall user satisfaction. Despite the advancements in RL and language models (LMs), employing RL to drive conversational chatbots still poses significant challenges. A primary issue stems from RL's dependency on online exploration for effective learning, a process that can be costly. Moreover, engaging in online interactions with humans during the training phase can raise safety concerns, as the LM can potentially generate unwanted outputs. This issue is exacerbated by the combinatorial action spaces facing these algorithms, as most LM agents generate responses at the word level. We develop various RL algorithms, specialized in dialogue planning, that leverage recent Mixture-of-Expert Language Models (MoE-LMs)—models that capture diverse semantics, generate utterances reflecting different intents, and are amenable for multi-turn DM. By exploiting the MoE-LM structure, our methods significantly reduce the size of the action space and improve the efficacy of RL-based DM. We evaluate our methods in open-domain dialogue to demonstrate their effectiveness with respect to the diversity of intent in generated utterances and overall DM performance.
Natural language processing (NLP) has made significant strides in recent years, notably in the field of language generation. Thanks to advances in language modeling, particularly with the use of the transformer architecture (Vaswani et al., 2017), NLP models can now generate text that is often difficult to distinguish from that written by a human. However, despite these advancements, these models still fall short when it comes to having rich conversations. Current NLP models lack effective dialogue management: they are good at generating individual sentences, but struggle with maintaining coherent and engaging conversations. Whereas most compelling conversations generally span numerous topics, are rather open-ended, and often have an underlying goal (e.g., customer success, task completion, recommendation). This requires dialogue agents to understand the context of the conversation and respond appropriately while maintaining the ability to achieve goals.
Reinforcement learning (RL) is a natural approach to learning a policy for a dialogue management (DM) agent. Earlier work on RL-based dialogue systems relies on specific, hand-crafted semantic states (Levin and Pieraccini, 1997; Singh et al., 2002; Walker, 2000) or partially observable belief
*The work was done as a student researcher at Google Research (dhawgupta@google.com).
†The work was done prior to joining Amazon, while the author was at Google Research.
states (Williams and Young, 2007; Young et al., 2010), in which the agent encodes conversations and chooses the best-structured dialogue action at each turn. Applications include relational reasoning (Shah et al., 2018), task completion (Shi and Yu, 2018), and query fulfillment (Serban et al., 2017), whose action spaces are structured enough to be represented by hand-crafted features. To handle more complex dialogues, recent approaches use language models (LMs) to extract semantic representations from conversation histories, treat them as dialogue states, and apply RL to learn a word-level generative DM agent (Jaques et al., 2019; Li et al., 2016, 2017; Shin et al., 2020).
However, unlike supervised learning approaches, where one can train imitation agents with offline conversation data, RL-based DM algorithms require online exploration. Unfortunately, constant interactions with real users are often expensive and time-consuming. While one can potentially address the DM problem using offline RL, issues such as model exploitation leading to distribution shift on the state and action spaces, when training on static datasets are of paramount concern (Levine et al., 2020). Moreover, the myriad variations of language make incorporating all possible conversation histories and bot utterances into the state and action spaces of an RL formulation of the DM problem impractical due to the combinatorics at play. As a result, naive application of RL to DM may result in poorly-performing agents that generate incomprehensible utterances (Zhao et al., 2019).
We tackle the above issues related to the use of offline RL in DM systems by leveraging recent advances in Mixture-of-Expert Language Models (MoE-LMs) (Chow et al., 2022). Specifically, we develop a suite of offline RL algorithms specialized in dialogue planning that exploit the structure of MoE-LMs. Our methods consist of three main components: 1) a primitive LM, which uses a probabilistic encoder and decoder and is capable of generating diverse semantic intents; 2) a number of specialized expert LMs, each of which generates utterances corresponding to a specific intent; and 3) a compositional DM that, at each turn, given the encoded conversation history, selects an utterance from a set of candidate utterances suggested by the experts and pass it to the DM agent to execute.
Our contributions to offline RL adapted for MoE-based DM agents are four-fold. First, we exploit the hierarchical structure of MoE-LMs, allowing our offline RL methods to work with a significantly smaller, finite action space, making the RL problem more tractable. Second, by leveraging pre-trained MoE-LMs—which generate coherent utterances—and regularization techniques from offline RL that align the DM's behavior with that of the primitive LM—the proposed RL algorithms can focus on higher-level dialogue planning. The proposed combination results in higher data efficiency than standard RL methods by delegating the responsibility of language fluency to be handled by the MoE-LMs. Third, by using the diverse semantic representations of MoE-LMs, our methods operate at the sentence embedding space and have much simpler critic and actor updates. This circumvents the word-level credit-assignment issue, particularly challenging in long conversations (Saleh et al., 2020). Fourth, in contrast to the findings of Verma et al. (2022), where offline RL agents tend to lack utterance diversity (due to potential reward hacking and optimization of a single objective), our MoE-based DM agents by design are adept at generating utterances that reflect different intents.
We begin with a brief introduction of LMs, the MoE-LM architecture, and the use of Markov decision processes (MDPs) in DM in Section 2. We then describe the pre-training procedure for MoE-LMs—which encode diverse semantics and generate fluent utterances capturing specific intents—in Section 3. We derive state-of-the-art (SOTA) offline RL algorithms for training MoE-LMs in Section 4, and three MoE-LM specialized offline RL algorithms in Section 5. Finally, in Section 6, we evaluate our algorithms in open-domain dialogues against their ability to generate utterances with diverse intents and their overall DM performance.
Language Models (LMs) In this work, we employ seq2seq LMs (Sutskever et al., 2014) to generate the next utterances in a dialogue. We assume access to a dataset of the form  $\mathcal{D} = \{(\mathbf{X}^{(k)},Y^{(k)})\}_{k = 1}^{|D|}$ , where each  $\mathbf{X}$  is an  $L$ -turn conversation history  $\mathbf{X} = \{X_{l}\}_{l = 0}^{L - 1}$ , wherein  $X_{l}$  is the utterance in a conversation at turn  $l$  and  $Y$  is the next utterance. We define  $N_{\mathbf{X}}$  to be an upper bound on the length (number of tokens) of each utterance  $X_{l}$  in  $\mathbf{X}$ . The role of a LM is to predict the probability of the next utterance  $Y$ , consisting of  $N$  tokens, conditioned on the conversation history  $\mathbf{X}$ , i.e.,  $\operatorname*{Pr}\left(Y = \{y_n\}_{n = 1}^N\mid \mathbf{X}\right)$ . In the transformer architecture (Wolf et al., 2019), a LM first encodes the conversation history  $\mathbf{X}$  using an encoder  $\Phi$  to a  $(L\times N_{\mathbf{X}})$ -length sequence of embeddings
<sup>3</sup>If  $X_{l}$  has fewer tokens than  $N_{\mathbf{X}}$ , the remaining spaces will be padded by a specific token and masked.
$\{(z_{l,0},\dots ,z_{l,N_{\mathbf{X}} - 1})\}_{l = 0}^{L - 1}$ , where each  $z_{l,n}$  is a vector in the latent space induced by the encoder  $\Phi$ . For notational convenience, we concatenate these embeddings into a single embedding  $z\in \mathcal{Z}\subseteq \mathbb{R}^d$  where  $d$  is the overall dimension of the latent space. The next utterance  $\widehat{Y} = \{\widehat{y}_n\}_{n = 1}^N$  is then sampled, token-by-token, from a decoder  $\Psi$ , i.e.,  $\widehat{Y}\sim \Psi (\cdot |z)\coloneqq \prod_{n = 1}^{N}\Psi (\widehat{y}_n|\widehat{y}_0,\ldots ,\widehat{y}_{n - 1};z)$ , where  $\widehat{y}_0$  is a fixed initial (start-of-sentence) token (Chien and Kuo, 2019) and the latent state is denoted as  $z = \Phi (\mathbf{X})$ .
Markov Decision Processes (MDPs) have been used to model dialogue management (DM) problems in a variety of settings (Li et al., 2016; Asadi and Williams, 2016; Jaques et al., 2019). In such MDPs, denoted by  $\mathcal{M} = (\mathcal{S},\mathcal{A},P,r,s_0,\gamma)$ , the state space  $\mathcal{S}$  represents the tokenized conversation history and the initial state  $s_0\in \mathcal{S}$  is the initial user's query. The action space  $\mathcal{A}$  is the tokenized language space with each action  $a\in \mathcal{A}$  represents one possible next utterance of the agent. The transition kernel  $P$  models the distribution over the user's response to the action taken by the agent (bot) and current conversational context. Finally, the reward function  $r$  measures the user's satisfaction as a function of the conversation until the most recent step. In these MDPs, we can think of the LM as a policy that maps conversation histories to the next utterances. The goal is to find a policy  $\pi^{*}$  with maximum expected discounted return, i.e.,  $\pi^{*}\in \arg \max_{\pi}J_{\pi}$ , where  $J_{\pi}\coloneqq \mathbb{E}[\sum_{k = 0}^{\infty}\gamma^{t}r_{t}|P,s_{0},\pi ]$ . Note that the size of the tokenized state and action spaces grows exponentially with the vocabulary size. This makes it intractable to solve MDPs of this type, even for a medium-sized vocabulary.
Mixture-of-expert Language Models (MoE-LMs). Chow et al. (2022) recently demonstrated promising results using MoE-LMs to enrich a bot's utterances and improve DM (see Figure 1 for a sketch of the MoE-LM architecture). These results were achieved mainly due to (i) learning a language representation that captures different semantics (primitive discovery), (ii) a machinery, called expert construction, that embeds different intents into sub-models of this LM, so that they can behave appropriately when prompted, and (iii) a compositional dialogue manager module that comprehends the conversation and determines which response deems most appropriate.
For primitive discovery, one first learns a language model  $\mathbf{L}\mathbf{M}_0\coloneqq (\Phi ,\mathcal{G}_0,\Psi)$  that consists of a stochastic encoder  $\mathcal{G}_0\circ \Phi$  , where  $\Phi$  is an encoder mapping tokenized conversation histories  $\mathbf{X}$  to a latent space  $\mathcal{Z}\subseteq \mathbb{R}^d$  and  $\mathcal{G}_0(z'|z)\coloneqq \mathcal{N}\bigl (\mu_0(z),\sigma_0^2 (z)\mathbf{I}_{d\times d}\bigr)$  is a Gaussian distribution, and a decoder  $\Psi$  .  $\mathrm{LM_0}$  predicts the next utterance  $\widehat{Y}_0$  (token-by-token) conditioned on  $z^{\prime}$  sampled from the latent distribution  $\Psi (\widehat{Y}_0|z')$  , i.e.,  $z^{\prime}\sim \mathcal{G}_{0}(\cdot |z)$  . We overload our notation to denote the primitive by  $\mathrm{LM}_0(Y|\mathbf{X})\coloneqq \mathbb{E}_{z'\sim \mathcal{G}_0(\cdot |z),z = \Phi (\mathbf{X})}[\Psi (Y|z')]$  , which predicts the next utterance accurately and also has strong generalization in  $\mathcal{Z}$  over a diverse set of possible utterances.
Figure 1: (Left) The MoE-LM architecture (Chow et al., 2022). Step 1:  $\Phi$  encodes conversation history. Step 2:  $\Psi \circ \mathcal{G}_i$ ,  $\forall i$ , generate candidate bot utterances. Step 3: The compositional dialogue manager  $\mu$  selects the bot response by  $Q$ -score ranking and post-processing. (Right) Sample utterance workflow generated by a MoE-LM trained with Reddit data.
Given a primitive  $\mathbf{L}\mathbf{M}_0$ , the algorithm learns  $m$  expert distributions  $\{\mathcal{G}_i\}_{i=1}^m$ ,  $\mathcal{G}_i(z'|z) = \mathcal{N}(\mu_i(z), \sigma_i^2(z)\mathbf{I}_{d\times d})$ , each corresponds to a particular personality/intent and generates samples in specific parts of the latent space  $\mathcal{Z}$ . This results in  $m$  LMs,  $\{\mathbf{L}\mathbf{M}_i\}_{i=1}^m$ ,  $\mathbf{L}\mathbf{M}_i := (\Phi, \mathcal{G}_i, \Psi)$ , each serving as an expert that generates one or more candidate next utterances  $\widehat{Y}_i$  that are relevant to the
conversation  $\mathbf{X}$ , and also compatible with its respective personality/intent. The compositional DM  $\mu$  takes the encoded conversation history  $z = \Phi(\mathbf{X})$  and candidate action utterances generated by the experts  $\{\widehat{Y}_i\}_{i=0}^m$  as input, and selects one of them to execute, i.e.,  $Y \sim \mu(\cdot \mid z, \{\widehat{Y}_i\}_{i=0}^m)$ . Given the state  $s = \mathbf{X}$  and action  $a = Y$ , the MoE-LM policy that optimizes the DM MDP can be expressed as
$$
\pi_ {\mathrm {M o E}} (a | s) = \int_ {\left\{\hat {a} _ {i}, z _ {i} ^ {\prime} \right\} _ {i = 0} ^ {m}} \mu \left(a | \Phi (s), \left\{\hat {a} _ {i} \right\} _ {i = 0} ^ {m}\right) \prod_ {i = 0} ^ {m} d \Psi \left(\hat {a} _ {i} \mid z _ {i} ^ {\prime}\right) d \mathcal {G} _ {i} \left(z _ {i} ^ {\prime} \mid \Phi (x)\right). \tag {1}
$$
3 Warmstarting the Mixture-of-Expert LM
The MoE-LM approach effectively reformulates the reinforcement learning (RL) dialogue management problem, resulting in considerably smaller action spaces. This is achieved as the DM agent is now re now required to select the most appropriate utterance for presentation to the user from a finite, predefined set of candidate utterances, rather than generating responses word by word, hence allowing the agent to focus on optimizing the specific goal of the conversation task (as candidate utterances are separately optimized to follow particular bot-based characteristics/intents). Recall that the DM is a policy conditioned on both the latent state and the actions suggested by the experts. Before introducing the different RL methods for DM (Sections 4 and 5), in the following, we outline (i) the learning of diverse semantics (primitive LM) for conversation histories, which allows the agent to generate a wide variety of utterances, and (ii) the construction of specialized LMs (experts), which generate utterances with different intents.
Following the primitive discovery procedure in Chow et al. (2022), we learn the primitive LM,  $\mathrm{LM}_0$ , by solving the following KL-constrained optimization problem that aims at capturing diverse semantics:
$$
\min  _ {(\Phi , \mathcal {G} _ {0}, \Psi), \rho} \widehat {\mathbb {E}} _ {z ^ {\prime} \sim \rho (\cdot | z, Y), z = \Phi (\mathbf {X})} [ - \log \Psi (Y | z ^ {\prime}) ] \quad \text {s . t .} \widehat {\mathbb {E}} _ {z = \Phi (\mathbf {X})} [ \mathrm {K L} (\rho (z ^ {\prime} | z, Y) | | \mathcal {G} _ {0} (z ^ {\prime} | z)) ] \leq \epsilon_ {\mathrm {K L}}. \tag {2}
$$
In (2),  $\widehat{\mathbb{E}}$  is the empirical expectation over  $(\mathbf{X}, Y)$  in the dataset  $\mathcal{D}$ ,  $\rho$  is a distribution over the latent space conditioned on the encoded conversation history  $z$  and the target utterance  $Y$ , and  $\epsilon_{\mathrm{KL}}$  is a positive real-valued threshold. The distribution  $\rho(\cdot | z, Y)$  is a Gaussian  $\mathcal{N}(\mu_{\rho}(z, \Phi_{\rho}(Y)), \sigma_{\rho}^{2}(z, \Phi_{\rho}(Y)) \mathbf{I}_{d \times d})$  in which  $\Phi_{\rho}$  is a pre-trained encoder for the target utterance  $Y$ , and the mean  $\mu_{\rho}(\cdot, \cdot)$  and variance  $\sigma_{\rho}^{2}(\cdot, \cdot)$  are trainable models. Note that by solving (2), we maximize the log-likelihood of sentence  $Y$  for a context and latent generation, while enforcing consistency between the latent variable  $z'$  predicted by  $\mathcal{G}_0(\cdot | z)$  and  $\rho(\cdot | z, Y)$  via the KL constraint. In practice, we implement the KL constraint in (2) as a penalty weighted by a chosen coefficient.
To complete the MoE framework, one needs to train a set of experts  $\{\mathrm{LM}_i\}_{i=1}^m$ , each generating candidate utterances of different intents. By viewing each expert as a distribution of particular behaviors in conversation data  $\mathcal{D}$ , we leverage the results in Chow et al. (2022) and adopt a universal encoder-decoder  $(\Phi, \Psi)$  among all the experts. In this view, each expert  $i$  is a distribution  $\mathcal{G}_i(\cdot|z)$  over certain regions of the latent space  $\mathcal{Z}$ . We train each  $\mathcal{G}_i(\cdot|z)$  by solving
$$
\min  _ {\mathcal {G} _ {i}} \widehat {\mathbb {E}} _ {z ^ {\prime} \sim \mathcal {G} _ {i} (\cdot | z), z = \Phi (\mathbf {X}), Y \sim \Psi (\cdot | z ^ {\prime})} \left[ - \ell_ {i} (\mathbf {X}, Y) \right], \tag {3}
$$
where  $\ell_i(\mathbf{X}, Y) \in \mathbb{R}$  is a real-valued label that characterizes the intent of expert  $i$ . We can think of  $\ell_i(\mathbf{X}, Y)$  as a score assigned to  $Y$  resembling how strongly  $Y$  exhibits the trait expert  $i$  is meant to represent. Each expert is learned via reward-maximization, where  $\ell_i$  is treated as an intent-aligned reward signal for expert  $i$ . Note that there is a connection between the above approach and contextual bandits (Chu et al., 2011), where both the context and action spaces are the latent space  $\mathcal{Z}$ , and the bandit policy is the latent distribution  $\mathcal{G}_i$ . The choice of greedy reward maximization is to encourage a particular behavior in the expert's immediate utterance rather than controlling its future utterances.
Long-term dialogue planning is handled by the compositional dialogue manager. For example, with Gaussian experts  $\{\mathcal{G}_i\}_{i=1}^m$ , we can use the standard REINFORCE algorithm (Sutton et al., 1999a), where the model parameters  $(\mu_i, \sigma_i)$  are updated in the direction  $\alpha \cdot \mathbb{E}_{z' \sim \mathcal{G}_i(\cdot|z)}, Y \sim \Psi(\cdot|z') [\ell_i(\mathbf{X}, Y) \cdot \nabla_{\{\mu_i, \sigma_i\}} \log \mathbb{P}_{\mathcal{G}_i}(z'|z)]$  with learning rate  $\alpha > 0$ . To reduce the variance of these estimates, we can also adopt the baseline variance reduction technique (Greensmith et al., 2004).
4 RL for Mixture-of-Expert DM
In offline RL, the policy is learned from the collected conversations  $\mathcal{D}$ , without further online interactions. This potentially allows RL DM methods to leverage the abundance of offline conversational data for policy learning. Let  $(\mathbf{X}, Y, X_{+}) \sim \mathcal{D}$  be a tuple sampled from the offline conversation data  $\mathcal{D}$ , with  $X_{+}$  being the follow-up user response. We can formulate this tuple as a MDP data by defining  $s := \mathbf{X}$ ,  $a := Y$ ,  $r(X_{+})$ , and  $s_{+} := (\mathbf{X}, Y, X_{+})$  as the state, action, reward (w.r.t. the follow-up user response), and next state. A standard offline RL algorithm is  $Q$ -learning (Watkins and Dayan, 1992) that solves:  $\min_{Q} \mathbb{E}_{(s,a,r,s_{+}) \sim \mathcal{D}}[(r + \gamma \max_{a_{+}} Q(s_{+}, a_{+}) - Q(s, a))^{2}]$ .
However, with a large action space, the inner maximization (also termed as greedification) in  $Q$ -learning is generally computationally intractable. Furthermore, since one cannot ensure that the greedy  $a_{+}$  is sampled from the same action distribution as in the offline RL dataset (an issue worsened by the large action set), such a covariate shift in the sampling distribution can cause an overestimation bias for the  $Q$  estimate. To alleviate these issues, we propose to leverage the warm-started MoE-LM (Section 3), where the diverse semantic representation and the expert LMs are learned separately. This is crucial to make our offline RL DM problem tractable as the language fluency is captured by the MoE-LM, while our RL-based DM focuses on higher-level planning strategies. In the following, we describe how this can be achieved via different offline RL algorithms.
Offline RL Methods for MoE-LMs: One approach to address the aforementioned offline RL issues is entropy regularization (Haarnoja et al., 2018; Carta et al., 2021), which regularizes the greedification step to ensure the learned policy is either diverse enough or close to the behavior (data-generation) policy (e.g., through Shannon entropy or KL divergence between these policies). Recall that the primitive LM,  $\mathsf{LM}_0$ , models the utterance distribution in  $\mathcal{D}$  and the state-action-reward next-state tuple  $(s,a,r,s_{+})$  of the DM MDP. With the following latent states generated by the primitive LM:  $z = \Phi(s)$ ,  $z_a = \Phi((s,a))$ , and  $z_{+} = \Phi(s_{+})$ , we define the latent conversation data  $\Phi(\mathcal{D})$  as a collection of  $(z,z_a,r,z_{+})$  tuples. With Shannon-entropy regularization, we can utilize the soft actor critic framework (Haarnoja et al., 2018) to develop RL updates for the value function  $V(z)$ , state-action value function  $Q(z_a)$ , and latent generator  $\mathcal{G}(z'|z)$ , which is initialized with the primitive latent expert  $\mathcal{G}_0$ . This framework minimizes the following losses:
$$
L _ {Q} = \mathbb {E} _ {\left(z, z _ {a}, r, z _ {+}\right) \sim \Phi (\mathcal {D})} \left[ \left(r + \gamma V _ {\operatorname {t a r}} \left(z _ {+}\right) - Q \left(z _ {a}\right)\right) ^ {2} \right] \tag {4}
$$
$$
L _ {V} = \mathbb {E} _ {z \sim \Phi (D), (\hat {a}, z ^ {\prime}) \sim \Psi \circ \mathcal {G} (\cdot | z)} \left[ Q _ {\operatorname {t a r}} \left(z _ {\hat {a}}\right) - \alpha \log \mathcal {G} \left(z ^ {\prime} \mid z\right) - V (z) ^ {2} \right] \tag {5}
$$
$$
L _ {\mathcal {G}} = \mathbb {E} _ {z \sim \Phi (\mathcal {D}), (\hat {a}, z ^ {\prime}) \sim \Psi \circ \mathcal {G} (., | z)} [ Q (z _ {\hat {a}}) - \alpha \log \mathcal {G} (z ^ {\prime} | z) ], \tag {6}
$$
where the critic,  $(V,Q)$ , takes any encoded conversation history as input and predicts the corresponding cumulative return,  $\alpha >0$  is the entropy temperature,  $(V_{\mathrm{tar}},Q_{\mathrm{tar}})$  are the target value networks,  $z^{\prime}\sim \mathcal{G}(.|z)$  is the latent sample generated by  $\mathcal{G}$ ,  $\hat{a}\sim \Psi (z^{\prime})$  is the utterance sampled from  $\Psi \circ \mathcal{G}$ , and finally  $z_{\hat{a}} = \Phi ((\mathbf{X},\hat{a}))$  is the corresponding latent state.
From a hierarchical RL perspective (Sutton et al., 1999b; Saleh et al., 2020), the latent generator behaves like a high-level policy, whose latent sample  $z'$  is used to generate a bot utterance via  $\Psi$ -decoding (with the primitive decoder  $\Psi$  acting as the low-level policy). Extending the above RL updates to the case of relative-entropy (KL) regularization can be straightforwardly done by replacing the term  $\log \mathcal{G}(z'|z)$  with  $\log (\mathcal{G}(z'|z) / \mathcal{G}_0(z'|z))$ , since the primitive LM approximates the behavior policy and the encoder-decoder pair  $(\Phi, \Psi)$  is shared among the experts.