SlowGuess/NAACL_2022_OCR huggingface dataset:
Jiacheng Xu $^{\spadesuit}$  Siddhartha Reddy Jonnalagadda $^{\star}$  Greg Durrett $^{\spadesuit}$
$\clubsuit$ Department of Computer Science, The University of Texas at Austin
$^{\star}$ Alexa AI, Amazon, Seattle
{jcxu,gdurrett}@cs.utexas.edu sidjreddy@gmail.com
Conditional neural text generation models generate high-quality outputs, but often concentrate around a mode when what we really want is a diverse set of options. We present a search algorithm to construct lattices encoding a massive number of generation options. First, we restructure decoding as a best-first search, which explores the space differently than beam search and improves efficiency by avoiding pruning paths. Second, we revisit the idea of hypothesis recombination: we can identify pairs of similar generation candidates during search and merge them as an approximation. On both summarization and machine translation, we show that our algorithm encodes thousands of diverse options that remain grammatical and high-quality into one lattice. This algorithm provides a foundation for building downstream generation applications on top of massive-scale diverse outputs.<sup>1</sup>
Although pre-trained text generation models (Lewis et al., 2020; Raffel et al., 2020) have achieved impressive results across a range of tasks, these models do not always deliver what system developers want. Machine generated text may be non-factual (Kryscinski et al., 2020; Maynez et al., 2020; Goyal and Durrett, 2021) or toxic (Gehman et al., 2020). We might patch these problems by applying discriminators over the output (Holtzman et al., 2018; Yang and Klein, 2021) to enforce these properties post-hoc; we could, for instance, apply a secondary model as a reranker over a small collection of outputs. However, if the generator returns a homogeneous set of candidates, we may fail to find any usable generation output. What if generation models could return massive numbers of candidates rather than a few outputs with optimal score?
Figure 1: A lattice of outputs yielded by path recombination is a more efficient way to represent and explore related generation outputs compared to beam search.
With a large set of candidates, our secondary model could more easily find an acceptable one without having to take more extreme steps like re-training the initial generation model. Output diversity has separately been established as a useful goal for for applications such as dialogue and story generation (Li et al., 2016; Fan et al., 2019).
Standard approaches including beam search (BS) and sampling methods fall short of our goal. Beam search uses significant computational resources to explore similar hypotheses, and much of the computation in the search process is invested into paths that could be acceptable generation outputs, but are ultimately pruned. Sampling approaches like nucleus sampling (Holtzman et al., 2020), although achieving better diversity than beam search, often re-discover seen hypotheses and can be harder to control for quality. A central problem with both methods is that they do not handle very similar hypotheses efficiently.
In this paper, we present a decoding framework with two key components. First, we argue that a modified best-first search (BFS) is the right way to explore the search space. We augment standard best-first search with a depth-first path completion strategy: we eagerly expand each node until we reach an EOS token, thereby guaranteeing that each node is part of some completed path returned to
$^{1}$ Code, implementation guideline, and visualization are available at https://github.com/jiacheng-xu/lattice-generation.
Human Language Technologies, pages 4659 - 4676
the user. This generation strategy avoids exploring large numbers of states which end up being pruned. BFS is also more flexible than static beam search and can prioritize exploration in more uncertain parts of the generation.
Second, our algorithm returns a massive number of generation options encoded in a lattice, with different hypotheses recombined in an approximate fashion. Beam search preserves similar outputs such as "A Cardiff recycling company has gone into" and "A Cardiff waste management company has gone into" as different states. However, these prefixes actually have very similar distributions of following words under the model; if we identify states like this, we can recombine them (Figure 2) and treat them as the same from the perspective of future continuations. In Figure 1, we show an illustration of the lattice structure this recombination can form for document summarization. We broaden a recombination method used previously in beam search for machine translation (Och et al., 2001; Zhang et al., 2018), enabling us to compactly encode large number of generation candidates and achieve dense lattices.
We show results for both document summarization and machine translation in three language pairs. For each setting, we show that our lattice encodes a large number of high-quality candidates, including good matches with annotated reference generations. We further show that a variant of our method can still achieve strong results with a lower number of nodes expanded than the baselines, suggesting that this can be a path towards saving computational resources. We believe that computing thousands of high-quality generation candidates within a single compact data structure can provide a powerful starting point for various downstream purposes: diversity, factuality, customizability, and more.
2 Problem & Setup
We define our algorithm in the context of conditional text generation (Sutskever et al., 2014; Bahdanau et al., 2014). Conditional text generation is formulated as sequence transformation from a source input  $\mathbf{x}$  to target output  $\mathbf{y} = (y_{1},\dots ,y_{n})$  via a neural text generation model parameterized by  $\theta$ . Each  $y_{i}$  is a symbol in a vocabulary  $\nu$ . The probability of a decoded sequence is given by  $p(\mathbf{y}\mid \mathbf{x};\theta) = \prod_{t = 1}^{n}p(y_{t}\mid \mathbf{y}_{< t},\mathbf{x};\theta)$ . Decoding text from a model can be framed as a search problem, where the search objective is to find the
Figure 2: Our search algorithm. At each step, the algorithm pops a node from search frontier  $\mathcal{O}$ , checks for possible recombinations with existing nodes, and merges the nodes if a match is found. In this example, "has gone into" is the merge target to match. "waste company" and "recycling plant" are interchangeable paraphrases which do not affect the continuation from the model's perspective.
output sequence that maximizes the conditional probability under the model:  $\arg \max_{\hat{\mathbf{y}}}p(\hat{\mathbf{y}}\mid \mathbf{x};\theta)$ . Because  $p(\hat{y}_t\mid \hat{\mathbf{y}}_{< t},\mathbf{x};\theta)$  depends on the entire generated sequence so far, this decoding problem is intractable to solve exactly.
While typically the goal of decoding is to find the hypothesis with the highest possible model score, we instead focus on finding a large set of "good enough" hypotheses. That is, finding a set  $\mathcal{V}$ :
$$
\underset {\mathcal {Y}} {\arg \max } | \mathcal {Y} | \quad \text {s . t .} p (\mathbf {y} \mid \mathbf {x}; \theta) > \epsilon \text {f o r a l l} \mathbf {y} \in \mathcal {Y} \tag {1}
$$
for some threshold  $\epsilon$ .  $\epsilon$  emerges naturally by adjusting search hyperparameters to control the number of returned hypotheses. Our goal in this paper is to design an algorithm that can efficiently find  $\mathcal{V}$ .
Notation We encode predicted generation candidates  $\hat{\mathbf{y}}$  in a lattice. A lattice  $L = (N,E)$  is a directed graph where each node represents a token and paths defined by directed edges encode candidates. A path  $\pi$  in  $L$  from a unique start-of-sequence node  $n_{\mathrm{SOS}}$  to any node  $n$  represents a (partially) decoded string, consisting of the words in that path. All completed paths start with  $n_{\mathrm{SOS}}$  and end at (potentially different) end-of-sequence nodes  $n_{\mathrm{EOS}}$ . The search graph  $L$  is constructed iteratively through a search procedure. We maintain the closed graph  $\mathcal{C}$  with explored nodes and edges as well as a search frontier  $\mathcal{O}$ , a set consisting of successors to nodes currently in the graph. For each node, there are  $|\mathcal{V}|$  possible successors.
We define the search budget as the number of nodes expanded from the search frontier. Our ex
Figure 3: Correlation of ROUGE-2 and model score in beam search. For each example, we compare the hypothesis with the highest model score,  $h^*$ , with all other hypotheses.  $x$  and  $y$ -axis show the gaps of R2 and model score. The Pearson's  $\rho$  is 0.092 which suggests very low correlation between R2 and model score.
periments will seek to compare different methods using the same search budget. We will define this more precisely in Sec. 7.
3 Inadequacies of Beam Search
As we have alluded to, beam search is inadequate for our goal for several reasons. We show experiments on these aspects in this section and Appendix A.
Better Model Score  $\Rightarrow$  Better Hypothesis The most critical issue is that beam search is designed to efficiently approximate arg max  $\hat{\mathbf{y}} = p(\hat{\mathbf{y}}\mid \mathbf{x};\theta)$  but the optimal model score is neither our goal nor a guarantee of a good hypothesis. In Figure 3, we compare the correlation of model score and ROUGE under beam search for text summarization. The Pearson correlation between these two variables is very weak. Beyond ROUGE score, the example in Fig. 1 shows that the main differences between these summaries may be minor differences in surface realization that have little effect on our qualitative judgments of summary quality. The hypothesis with the best score under beam search is not substantially better quality than hypotheses with slightly lower scores.
Lack of Diversity in (Diverse) Beam Search Are the model outputs from BS and DBS diverse? We use Self-BLEU (sBL) (Zhu et al., 2018) to measure the BLEU score for randomly sampled pairs from each algorithm's output. The lower the self-BLEU, the more dissimilar the pairs are. On decoding summaries on XSum, the sBL for BS/DBS are 87/79 while a nucleus sampling method can achieve 57/50 depending on the configuration. Although DBS slightly improves the
Figure 4: Diversity vs. ROUGE of BS/DBS on XSum with varying beam size  $k$ , compared to a proposed model introduced later (blue star) with equivalent beam size  $k = 16$ . We consider sample ROUGE-2 lower than 13 as low relevance/quality generations. Diversity of BS does not scale well with  $k$  and DBS generations become low quality.
k
D	16
XSum	8
zh-en	fr-en	en-fr
BS	71.3%	63.3%	54.0%	59.2%
DBS	71.2%	56.1%	50.4%	55.7%
Table 1: Pruning ratio of BS and DBS on different tasks and datasets with beam size  $k$ . We report the average percentage of explored nodes getting pruned and not appearing in a finished hypothesis.
diversity compared to the original variant, the overlap of outputs from beam search based method is still very high, and the diversity remains a challenge.
Poor Efficiency from Pruning One final issue with beam search is that most of its computation is not even useful in producing finished hypotheses; that is, the set  $\mathcal{V}$  of answers produced does not contain most of the nodes expanded in the typical course of operation. We conduct an empirical pruning study on a summarization dataset and three translation datasets and show the results in Table 1. For all studied cases, beam search and diverse beam search prune over half of the expanded nodes. Many pruned hypotheses are not truly ungrammatical or low quality, but are merely slightly less likely than other nodes. How we can preserve more of the explored lattice and do so efficiently is addressed in next by our use of best-first search.
4 Modified Best-first Search
As established in the previous section, beam search prunes many paths that would potentially yield high-quality summaries and wastes computational resources expanding nodes that aren't included in a
final search graph. We tackle this issue by changing from beam search to best-first search (BFS) (Hart et al., 1968; Pearl, 1984). BFS prioritizes searching over nodes according to a scoring function, giving us more flexibility in how we explore the space. Our chief modification of the base algorithm is a heuristic we call depth-first completion.
Depth-first Path Completion Neural text generation is a search problem with large branching factor  $(\mathcal{V})$  and deep search depth (sequence length). As a result, applying BFS with the scoring function being the model score of a state often leads to a broad search that rarely returns a valid path. One solution to this problem is to incorporate a heuristic based on length. Model score is monotonically decreasing as a sequence grows in length, so prior work (Wu et al., 2016; Zhang et al., 2018; Meister et al., 2020b) has used a length reward term to alleviate this issue. We found that, even with a length heuristic, BFS will still have "dangling" nodes that are not part of any path to an EOS (goal) token, and it might return few or no valid hypotheses.
Recognizing our objective from Equation 1, we can take a simple step to ensure that every node ends up on some completed path: eagerly do a greedy search from each node until we reach  $n_{eOS}$
Figure 5: Depth-first completion. The red node is the current node being expanded. We depth-first expand a sequence of nodes (in blue) to get a completed path.
or exceed a maximum length. In Algorithm 1, we implement this by modifying the priority of the highest scored token with  $\infty$  (line 12), so it will be explored immediately after the current time step. In Figure 5, we illustrate depth-first completion.
Search Algorithm We describe BFS with depth-first completion in Algorithm 1. The algorithm is a modified best-first search algorithm applied to text generation.  $s(\cdot)$  is a function to evaluate the value of a path. Typically it is defined as  $s(\mathbf{y}) = \sum \log p(y_t \mid y_{<t})$ .  $b$  is the budget for total model calls to neural text generation model. Note that isRecomb and doRecomb do not invoke the neural generation model, so they do not count towards the computation budget we defined here. In practice, we only consider top 5 expansions rather than the whole vocabulary  $\mathcal{V}$  for line 10.
5 Path Recombination
Path recombination, also known as hypothesis recombination, was originally proposed and used in phrase-based machine translation (Och et al., 2001; Koehn et al., 2003; Zhang et al., 2018). The idea of path recombination is to combine similar paths if what the model predicts for them in the future is the same, reflecting a similar dynamic programming principle as the Viterbi algorithm. We focus on finding hypotheses which approximately exhibit this property, and show that merging them can yield high-quality outputs. Figure 2 shows an example of recombination. The two hypotheses being merged here roughly convey the same intent, and it turns out that the shared suffix "has gone into" is a strong indicator that the model will treat them similarly in the rest of the generation.
Prerequisites of Recombination Theoretically, two search states should only be recombined if they yield the exact same distribution over future generation decisions (see strong equivalence in Ap
2This can be considered a heuristic like in (weighted)  $\mathbf{A}^*$  search, but it is not necessarily admissible or consistent.
Figure 6: Illustration of two path recombination strategies, RCB and ZIP. Orange lines are the merging edges (MRG) built by recombination. Dotted lines and circles are removed after recombination. The key difference of RCB and ZIP is how much the recombination propagates, 1 step or  $n$  steps.
pendix B). However, this is intractable even to check approximately; we define a weaker criterion:
Definition 5.1 (Weak equivalence). Let  $\mathbf{a}$  and  $\mathbf{b}$  be two prefix strings starting with  $n_{SOS}$ .  $\mathbf{a}$  and  $\mathbf{b}$  are weakly equivalent if greedy completions of these two strings are the same:  $\arg \max_{\mathbf{y}} P(\mathbf{y} \mid \mathbf{a}) = \arg \max_{\mathbf{y}'} P(\mathbf{y}' \mid \mathbf{b})$ .
This criterion can be checked empirically, but it is not practical to do so during search.
To approximate equivalence, we define a similarity function isRecomb(h,  $\hat{\mathbf{h}}$  ) to determine if an expanded node  $\hat{\mathbf{h}}$  should be merged with an existing expanded node  $\mathbf{h}$  . A similar recombination idea was explored in Zhang et al. (2018). Following their work, we explore a family of rule-based heuristics for merging. There are two rules: (1) two strings share a common  $n$  -gram suffix, (2) the length difference of two strings is less than  $\alpha$  Assume that the canonical paths for  $\mathbf{h}$  and  $\hat{\mathbf{h}}$  are lengths  $l$  and  $\hat{l}$  , then isRecomb(h,h)  $= \mathbb{1}[\pi (\mathbf{h})_{l - n + 1,\dots,l} = \pi (\hat{\mathbf{h}})_{\hat{l} -n + 1,\dots,\hat{l}}\wedge |l - \hat{l}| <   \alpha ]$  where  $\alpha$  and  $n$  are hyper-parameters. For a large enough value of  $n$  note that the shared suffixes encourage hypotheses like this in Figure 6 that share large parts of the structure already.
Prior Work: BSZBEAM Zhang et al. (2018) use their merging criterion in the context of beam search for neural machine translation. If the merging criteria hold,  $\hat{\mathbf{h}}$  will be recombined with  $\mathbf{h}$ . However,  $\hat{\mathbf{h}}$  will not be considered as a future
merging candidate. We call this merging strategy ZBEAM. We implement this model together with its merging criteria and denote it as BSZBEAM. This strategy is tailored to beam search and explores a more limited set of merges than we do.
Canonical Paths After recombination, a single node may represent multiple different possible sentence prefixes. If an edge is created due to the extension of search graph via model's prediction, we call it a GEN edge. Otherwise, the edge is created due to path recombination, and we call it a MRG edge. We define the notion of a canonical path, which represents the single path used to score candidate expansions.
Definition 5.2 (Canonical Path). Let  $n$  be a node. The canonical path to  $n$  is defined as the unique path from  $n_{SOS}$  to  $n$  consisting only of GEN edges.
Theorem 5.1. For any node  $n$  in the graph except  $n_{SOS}$ , there exists exactly one canonical path.
We present the proof in Appendix. C. We define the path of a node  $n$ ,  $\pi(n)$ , as returning the sequence of words corresponding to the canonical path of that node. Expanding  $n$  computes  $P(y \mid \pi(n))$  under the neural model.
6 Recombination Mechanism
We illustrate the two major recombination techniques we introduce, RCB and ZIP, in Figure 6. These methods represent our two implementations of doRecomb in Algorithm 1.
RCB: Generalization of ZBEAM ZBEAM has a major limitation: a limited set of merging candidates. The potential merge candidates in ZBEAM are only nodes in the current beam hypotheses and their previous steps, so the method cannot merge with nodes from earlier timesteps. For example, "A waste plant has gone into" cannot be merged with the hypothesis with ending in node 4 shown in Figure 6. The proposed generalization, RCB, addresses this limitation. We index all of the nodes in the lattice across all timesteps by their  $n$ -grams using a hash table, making it  $O(1)$  time to look up an  $n$ -gram pattern and retrieve potential merge candidates if they exist.
ZIP: Recombining More If we take a closer look at RCB in Figure 6, we see that even in the merged structure, nodes 3 and 7 and nodes 2 and 6 are preserved as separate. They do not pass the recombination criterion themselves, but these nodes
3In Zhang et al. (2018), there is one extra constraint requiring  $P(\hat{\mathbf{h}} \mid \mathbf{x}) < P(\mathbf{h} \mid \mathbf{x})$ , which requires that the path getting recombined has lower model score than the existing path. However, we found that model score is not always a good indicator for merging, as suggested in Fig. 3, partially because it is challenging to calibrate scores across different sequence lengths, so we disregard this constraint.
are part of the suffix matched strings, still correspond to the same words, and have the same directly generated next word. There is reason to believe that these might be equivalent as well. Hence, we explore a variant called ZIP that propagates the merge backwards through the lattice. This change relaxes the merging criterion and up to  $n$  pairs of nodes are combined when a merge is identified, leading to a more compact lattice. We describe some of the details in Appendix D.
Our Methods In this work, we combine the two proposed techniques, the modified best-first search (BFS) and recombination methods, together. Hence, we name them as BFSRCB and BFSZIP.
To evaluate the proposed methods, we conduct experiments on abstractive text summarization and machine translation. Our evaluation focuses on two questions: (1) how large and diverse are our lattices? (2) are the candidates encoded in the lattices high quality and grammatical?
7.1 Datasets & Base Models
We obtain all the models and certain baseline decoding methods from the Transformers library (Wolf et al., 2020). Since our methods are inference techniques without learned components, we do not re-train any models. For summarization, we use XSum (Narayan et al., 2018), a popular English news summarization dataset. We sample 100 examples from the validation set. The base model we use is BART-large-XSum (Lewis et al., 2020). For machine translation, we study our models on the English-French (en-fr) pairs from WMT 2014 (Bojar et al., 2014) and Chinese-to-English (zhen) pair from WMT 2019 (Barrault et al., 2019). We use mBART (Liu et al., 2020), a state-of-the-art neural machine translation model. We set the max decoding length to be twice the input length, so it varies per example.
7.2 Search Budget
To fairly compare the resource usage of all methods, we define the search budget as the number of calls to the neural model, equivalent to the number of nodes expanded. With beam size  $k$  and maxi
mum length  $T$ , beam search methods are given a theoretical budget of  $kT$ . We could simply allow best-first search and sampling methods to expand this number of nodes. However, since hypotheses may terminate before they reach EOS, empirically there is a gap between effective length (the average generated hypothesis length) and max length for both beam search and sampling. To balance the computation used across the different methods, we apply a correction factor so that the different methods are expanding the same number of nodes in aggregate. We increase the beam size  $k$  by  $50\%$  for translation, from 8 to 12, and  $25\%$  for summarization, from 16 to 20, for our baseline methods:  $k$  to BS, DBS, NCLS, TEMP, and BSZBEAM. This correction was empirically determined to balance the number of nodes expanded between our method and the baselines. We emphasize that this correction improves the baseline performance relative to our methods.
7.3 Search Algorithms
We implemented GREEDY, BS, DBS, NCLS, and TEMP as baseline methods.  $\mathrm{NCLS}_{0.9}$  represents nucleus sampling method with  $p = 0.9$ . We refer to Appendix E for detailed descriptions. We also experiment with basic BFS without path recombination, but including our depth-first path completion technique to ensure that finished hypotheses are produced. BSZBEAM is our implementation of Zhang et al. (2018). We integrate RCB with nucleus sampling and best-first search as NCLSRCB and BFSRCB. We also test BFS with the ZIP strategy.  $\mathcal{B}\mathrm{FSZIP}$  is a resource-efficient version of BFSZIP where only  $25\%$  of the search budget is used, exploring what this method can achieve with a lower budget given its more aggressive merges.
7.4 Evaluation Metrics
We describe our metrics to evaluate both quality and diversity. Several of our methods build on ROUGE (including R1, R2, RL) (Lin, 2004) and BLEU (Papineni et al., 2002; Post, 2018) comparing the generated text to references.