SlowGuess/NAACL_2022_OCR huggingface dataset:
Linlu Qiu $^{1*†}$  Peter Shaw $^{1*}$  Panupong Pasupat $^{1}$  Paweł Krzysztof Nowak $^{1}$  Tal Linzen $^{1,2}$  Fei Sha $^{1}$  Kristina Toutanova $^{1}$
<sup>1</sup> Google Research
2 New York Univeristy
{linluqiu,petershaw,ppasupat,pawelnow,linzen,fsha,kristout}@google.com
Generic unstructured neural networks have been shown to struggle on out-of-distribution compositional generalization. Compositional data augmentation via example recombination has transferred some prior knowledge about compositionality to such black-box neural models for several semantic parsing tasks, but this often required task-specific engineering or provided limited gains.
We present a more powerful data recombination method using a model called Compositional Structure Learner (CSL). CSL is a generative model with a quasi-synchronous context-free grammar backbone, which we induce from the training data. We sample recombined examples from CSL and add them to the fine-tuning data of a pre-trained sequence-to-sequence model (T5). This procedure effectively transfers most of CSL's compositional bias to T5 for diagnostic tasks, and results in a model even stronger than a T5-CSL ensemble on two real world compositional generalization tasks. This results in new state-of-the-art performance for these challenging semantic parsing tasks requiring generalization to both natural language variation and novel compositions of elements.
Compositional generalization refers to the ability to generalize to novel combinations of previously observed atoms. For example, we may ask a model to interpret the instruction "jump twice", when the atoms "jump" and "twice" were each observed separately during training but never in combination with each other (Lake and Baroni, 2018).
Improving compositional generalization is seen as important for approaching human-like language understanding (Lake et al., 2017; Battaglia et al.,
Figure 1: An overview of our method for compositional data augmentation with CSL, a generative model with a QCFG backbone, which is automatically induced from the training data. We show a notional set of original and synthetic examples mapping utterances to programs.
2018) and is practically significant for real world applications, where models deployed in the wild often need to interpret new combinations of elements not well-covered by expensive and potentially skewed annotated training data (Herzig and Berant, 2019; Yin et al., 2021).
Generic neural sequence-to-sequence models have improved substantially and reached high levels of performance, particularly when combined with large-scale unsupervised pretraining and sizable in-distribution labeled data. However, these models often perform poorly on out-of-distribution compositional generalization tasks (Lake and Baroni, 2018; Furrer et al., 2020; Shaw et al., 2021).
$\dagger$  Work done as part of the Google AI Residency program.
Also commonly referred to as elements or concepts.
Human Language Technologies, pages 4341 - 4362
In contrast, specialized architectures with discrete latent structure (Chen et al., 2020; Liu et al., 2020; Nye et al., 2020; Herzig and Berant, 2021; Shaw et al., 2021) have made strides in compositional generalization, but without task-specific engineering or ensembling, the gains have been limited to synthetic semantic parsing tasks. Although following SCAN (Lake and Baroni, 2018) some increasingly realistic synthetic tasks such as CFQ (Keysers et al., 2020) and COGS (Kim and Linzen, 2020) have been created, and several approaches achieve good performance on these tasks, the out-of-distribution generalization ability of state-of-the-art models on real-world, nonsynthetic tasks is still far from sufficient (Shaw et al., 2021; Yin et al., 2021).
Given their different strengths and weaknesses, it is compelling to combine the compositional bias of such specialized models with the greater flexibility and ability to handle natural language variation that characterizes generic pre-trained neural sequence-to-sequence models. One method for this is data augmentation. For example, Jia and Liang (2016) generate new training examples using example recombination via induced high-precision synchronous grammars, resulting in improvements on in-distribution and compositional splits of semantic parsing tasks. Another example is GECA (Andreas, 2020), a more general data augmentation approach that does not require task-specific assumptions. GECA achieved further gains on a larger variety of tasks, but provided limited improvements on some compositional generalization challenges.
We present a compositional data augmentation approach that generalizes these earlier methods. Training examples are recombined using the Com-positional Structure Learner (CSL) model, a generative model with a (quasi-)synchronous context-free grammar (QCFG) backbone, automatically induced from the training data. As illustrated in Figure 1, CSL is used to sample synthetic training examples, and the union of original and synthesized examples is used to fine-tune the T5 sequence-to-sequence model (Raffel et al., 2020). CSL is more generally applicable than the method of Jia and Liang (2016), employing a generic grammar search algorithm to explore a larger, higher-coverage space of possible grammars. Unlike GECA, CSL can re-combine examples recursively and also defines a probabilistic sampling distribution over input-output pairs.
CSL builds on the NQG model of Shaw et al. (2021), a discriminative parsing model over an induced QCFG backbone, which Shaw et al. (2021) proposed to ensemble with T5. Like NQG, CSL can, on its own, address a variety of compositional generalization diagnostic tasks on synthetic datasets and achieves high precision (but limited recall) on non-synthetic compositional generalization tasks, leading to overall gains when ensembled with T5. However, CSL offers several significant improvements over NQG, allowing it to efficiently address a wider range of datasets (see §3.1). Additionally, unlike NQG which is a discriminative model assigning probabilities to outputs  $y$  given inputs  $x$ , CSL is a generative model which admits sampling from a joint probability distribution  $p(x,y)$ . This enables the creation of new input-output training examples.
Empirically, augmenting the training data for T5 with samples from CSL transfers most of CSL's compositional bias to T5 for diagnostic tasks (SCAN and COGS), and outperforms a  $\mathrm{T}5+\mathrm{CSL}$  ensemble on non-synthetic compositional generalization tasks defined by compositional splits of GeoQuery (Zelle and Mooney, 1996) and SM-CalFlow (Andreas et al., 2020; Yin et al., 2021), resulting in new state-of-the-art performance on these splits.
In this section, we discuss the problem setting common to the compositional generalization evaluations we study, and propose some general assumptions that motivate our proposed method.
Problem Setting Consider a training dataset  $\mathcal{D}$  consisting of input-output pairs  $\langle x,y\rangle \in \mathcal{X}\times \mathcal{Y}$ , where  $\mathcal{X}$  is the set of valid inputs and  $\mathcal{Y}$  is the set of valid outputs. We assume that  $\langle x,y\rangle \in \mathcal{D}$  are sampled from a source distribution  $p_{s}(x,y)$ . Our model will be evaluated on inputs from a target distribution  $p_t(x,y)$ . We make an assumption that the conditional distribution of  $y$  given  $x$  is unchanged between source and target distributions; i.e.,  $p_s(y|x) = p_t(y|x)$ , which is also a standard assumption for domain adaptation evaluations under covariate shift. Any or all of the following may be true:  $p_{s}(x,y)\neq p_{t}(x,y)$ ,  $p_{s}(x)\neq p_{t}(x)$ ,  $p_{s}(y)\neq p_{t}(y)$ , and  $p_{s}(x|y)\neq p_{t}(x|y)$ .
2Our code is available at https://github.com/google-research/language/tree/master/language/compgen/csl.
Figure 2: An example derivation that derives the string pair  $\langle \text{jump twice and walk}, \text{JUMP JUMP WALK} \rangle$ .
What differentiates our setting from other forms of distribution shift is the added assumption that the source and target distributions share common "atoms" (see §1). In order to translate this intuitive notion of atom sharing into formal conditions, we define a general class of models termed derivational generative models, based on representing atoms as functions which can be recombined via function application. As a modeling hypothesis, we will assume that the training and evaluation distributions can be modeled by derivational generative models that share a common set of functions, but may vary in how they assign probability to derivations formed by recombining these functions.
Derivational Generative Models A derivational generative model defines a distribution  $p(x,y)$  over input-output pairs. The model contains a set of functions,  $\mathcal{G}$ , and a distribution over derivations. A derivation  $z$  can be viewed as a tree of functions from  $\mathcal{G}$  which derives some element  $[[z]] = \langle x,y\rangle \in \mathcal{X}\times \mathcal{Y}$  determined by recursively applying the functions in  $z$ . An example derivation is shown in Figure 2.
Given  $\mathcal{X}, \mathcal{Y}$ , and  $\mathcal{G}$ , we can generate a set  $\mathcal{Z}^{\mathcal{G}}$  of possible derivations. We define some shorthands for important subsets of  $\mathcal{Z}^{\mathcal{G}}$  for given  $x$  and  $y$ :
$$
\begin{array}{l} \mathcal {Z} _ {\langle x, y \rangle} ^ {\mathcal {G}} = \left\{z \in Z ^ {\mathcal {G}} \mid [   [ z ]   ] = \langle x, y \rangle \right\} \\ \mathcal {Z} _ {\langle x, \ast \rangle} ^ {\mathcal {G}} = \{z \in Z ^ {\mathcal {G}} \mid \exists y ^ {\prime} \in \mathcal {Y}, [   [ z ]   ] = \langle x, y ^ {\prime} \rangle \} \\ \end{array}
$$
A derivational generative model also consists of some probability distribution  $p_{\theta}(z)$  over the set of derivations  $\mathcal{Z}^{\mathcal{G}}$ , which we assume to be parameterized by  $\theta$ . We define  $p_{\mathcal{G},\theta}(x,y)$  in terms of  $p_{\theta}(z)$  as:
$$
p _ {\mathcal {G}, \theta} (x, y) = \sum_ {z \in \mathcal {Z} _ {\langle x, y \rangle} ^ {\mathcal {G}}} p _ {\theta} (z), \tag {1}
$$
$$
p _ {\mathcal {G}, \theta} (y | x) = \frac {\sum_ {z \in \mathcal {Z} _ {\langle x , y \rangle} ^ {\mathcal {G}}} p _ {\theta} (z)}{\sum_ {z \in \mathcal {Z} _ {\langle x , * \rangle} ^ {\mathcal {G}}} p _ {\theta} (z)}, \tag {2}
$$
for  $p_{\mathcal{G},\theta}(x) > 0$
Discussion In general, we are interested in a set of functions that captures some knowledge of how the parts of inputs correspond to parts of outputs. If we can recover some approximation of the underlying set of functions,  $\mathcal{G}$ , given  $\mathcal{D}$ , then we could sample derivations consisting of new combinations of functions that are not observed in  $\mathcal{D}$ . This could potentially help us improve performance on the target distribution, since we assume that the set of functions is unchanged between the source and target distributions, and that what is varying is the distribution over derivations.
However, even assuming  $\mathcal{G}$  can be exactly recovered given  $\mathcal{D}$  is not sufficient to ensure that we can correctly predict the most likely  $y$  given  $x$  according to the true  $p(y|x)$  (shared between source and target distributions) for  $x \sim p_t(x)$ .<sup>4</sup> We must also assume that there exists a parameterization of  $p_{\theta}(z)$  such that when we estimate  $\hat{\theta}$  given  $\mathcal{D}$ ,  $p_{\mathcal{G},\hat{\theta}}(y|x)$  sufficiently approximates the true  $p(y|x)$  for  $x \sim p_t(x)$ . We hypothesize that conditional independence assumptions with respect to how  $p_{\theta}(z)$  decomposes across the function applications in  $z$  can be helpful for this purpose. In particular, such assumptions can enable "reusing" conditional probability factors across the exponential space of derivations, potentially improving transfer to the target distribution.
With this intuition in mind, in §3 we propose a specific class of functions for  $\mathcal{G}$  based on (quasi-)synchronous context-free grammars, as well as a factorization of  $p_{\theta}(z)$  with strong conditional independence assumptions.
As shown in Figure 1, our method consists of two stages. First, we induce our generative model, CSL, from training data (§3.1). Second, we sample synthetic examples from the generative model and use them to augment the training data for a sequence-to-sequence model (§3.2).
4 One special case is where  $|\mathcal{Z}_{\langle x,* \rangle}^{\mathcal{G}}| = 1$  for all  $x$ . In this case, every  $x$  has exactly one unique derivation and  $p_{\mathcal{G}}(y|x)$  is deterministic given  $\mathcal{G}$  and does not depend on  $\theta$ , and therefore recovering  $\mathcal{G}$  is sufficient.
Formally, let  $\mathcal{G}$  be a set of partial functions over some set of elements. If  $f\in \mathcal{G}$  is a constant function, then  $f$  is a ground term. If  $f\in \mathcal{G}$  has arity  $k$ , and  $a_1,\dots ,a_k$  are ground terms, then  $f(a_{1},\dots ,a_{k})$  is a ground term. We define a derivation as a ground term that generates an element  $\in \mathcal{X}\times \mathcal{Y}$ .
Figure 3: The example derivation of Figure 2 using QCFG notation.
3.1 Compositional Structure Learner (CSL)
CSL can be viewed as a derivational generative model, as defined in §2, where the set  $\mathcal{G}$  of recursive functions is defined by a (quasi-)synchronous context free grammar (QCFG).<sup>5</sup>
We first describe the grammar formalism and the parameterization of our probabilistic model. Then we describe our two-stage learning procedure for inducing a grammar and learning the model parameters. CSL builds on the NQG model of Shaw et al. (2021), with several key differences discussed in the following sections:
- Unlike NQG, which is discriminative, CSL is a generative model that admits efficient sampling from the joint distribution  $p(x,y)$ .
- CSL enables a more expressive set of grammar rules than NQG.
- CSL offers a more computationally efficient and parallelizable grammar induction algorithm and a more efficient parametric model, allowing the method to scale up to larger datasets such as SMCalFlow-CS.
See section 4.4 for experiments and analysis comparing the components of CSL and NQG.
3.1.1 Grammar Formalism
An example QCFG derivation is shown in Figure 3, and the notation for QCFGs is reviewed in Appendix B.1. Notably, the correspondence between rules over input and output strings in QCFGs is akin to a homomorphism between syntactic and semantic structures, commonly posited by formal theories of compositional semantics (Montague, 1970; Janssen and Partee, 1997). We restrict our grammars to have only a single unique nonterminal symbol,  $NT$ . In constraint to standard synchronous context-free grammars (SCFGs), our grammars can
be quasi-synchronous (Smith and Eisner, 2006) because we allow a one-to-many alignment between non-terminals. Unlike the formalism of Shaw et al. (2021), which limited rules to contain  $\leq 2$  nonterminals, in the current work the maximal number of non-terminals is a configurable parameter, which enables inducing grammars with higher coverage for certain datasets.
3.1.2 Probabilistic Model
We factorize the probability of a derivation in terms of conditional probabilities of sequentially expanding a rule from its parent. Formally, let  $r$  denote a rule expanded from its parent rule  $r_p$ 's  $NT_{[i]}$  nonterminal (or a special symbol at the root of the derivation tree). We assume conditional independence and factorize the probability of  $z$  as:
$$
p _ {\theta} (z) = \prod_ {r, r _ {p}, i \in z} p _ {\theta} (r | r _ {p}, i) \tag {3}
$$
This non-terminal annotation with context from the tree is akin to parent annotation or other structure conditioning for probabilistic context-free grammars (Johnson, 1998; Klein and Manning, 2003).
Using independent parameters for each combination of a rule, its parent, and non-terminal index may lead to overfitting to the training set, limiting our ability to generalize to new combinations of rule applications that are needed for compositional generalization. We therefore factor this distribution using a soft clustering into a set of latent states  $S$  representing parent rule application contexts:
$$
p _ {\theta} (r | r _ {p}, i) = \sum_ {s \in \mathcal {S}} p _ {\theta} (r | s) p _ {\theta} (s | r _ {p}, i) \tag {4}
$$
where  $p_{\theta}(s|r_p,i)\propto e^{\theta_{r_p,i,s}}$  and  $p_{\theta}(r|s)\propto e^{\theta_{s,r}}$  and the  $\theta$ s are scalar parameters. The number of context states  $|\mathcal{S}|$  is a hyperparameter. While these conditional independence assumptions may still be too strong for some cases (see Appendix C.2), we find them to be a useful approximation in practice.
${}^{5}$  QCFG rules can be interpreted as functions which are limited to string concatenation. For notational convenience, we will therefore treat  $\mathcal{G}$  as a set of QCFG rules in §3.
6Concretely, for a rule  $NT\to \langle \alpha ,\beta \rangle$  , a non-terminal in  $\alpha$  can share an index with more than one non-terminal in  $\beta$  . This is important for datasets such as SCAN, as it allows rules such as NT  $\rightarrow \langle \mathrm{NT}_{[1]}$  twice,  $\mathrm{NT}_{[1]}\mathrm{NT}_{[1]}\rangle$  which enable repetition.
<sup>7</sup>We find that 4 is a computationally tractable choice for the datasets we study.
Using the example derivation from Figure 3, for the rule application  $r = \mathrm{NT}\rightarrow \langle \mathrm{walk},\mathrm{WALK}\rangle$  we have  $r_p = \mathrm{NT}\rightarrow$ $\langle \mathrm{NT}_{[1]}$  and  $\mathrm{NT}_{[2]}$ $\mathrm{NT}_{[1]}\mathrm{NT}_{[2]}\rangle$  and the expansion probability for that rule application is  $p(r|r_p,2)$
The full definition of these terms is in Appendix B.2.
We also optionally consider a task-specific output CFG, which defines valid output constructions. $^{10}$
3.1.3 Learning Procedure
A principled method to estimate  $\mathcal{G}$  and  $\theta$  given  $\mathcal{D}$  would be to find the MAP estimate based on some prior,  $p(\mathcal{G},\theta)$ , that encourages compositionality:
$$
\underset {\mathcal {G}, \theta} {\arg \max } p (\mathcal {G}, \theta) \times \prod_ {\langle x, y \rangle \in \mathcal {D}} p _ {\mathcal {G}, \theta} (x, y) \tag {5}
$$
However, since optimizing  $\mathcal{G}$  and  $\theta$  jointly is computationally challenging, we adopt a two-stage process similar to that of Shaw et al. (2021). First, we learn an unweighted grammar using a surrogate objective for the likelihood of the data and a compression-based compositional prior that encourages smaller grammars that reuse rules in multiple contexts, inspired by the Minimum Description Length (MDL) principle (Rissanen, 1978; Grunwald, 2004). We describe the induction objective and algorithm in §3.1.4. Second, given an unweighted grammar  $\mathcal{G}$ , we optimize the parameters  $\theta$  by maximizing the log-likelihood of  $p_{\mathcal{G},\theta}(x,y)$ , as defined by Eq. 1, using the Adam optimizer (Kingma and Ba, 2015).
3.1.4 Grammar Induction Algorithm
Our method for inducing a QCFG is based on that of Shaw et al. (2021), but with several modifications, which improve the computational scalability of the algorithm as well as the precision and coverage of the induced grammar. We analyze the relative performance of the two algorithms in §4.4.
Objective The main idea of the grammar induction objective,  $L_{\mathcal{D}}(\mathcal{G})$  , is to balance the size of the grammar with its ability to fit the training data:
$$
L _ {\mathcal {D}} (\mathcal {G}) = \sum_ {N T \rightarrow \langle \alpha , \beta \rangle \in \mathcal {G}} | \alpha | + | \beta | - c _ {\mathcal {D}} (\alpha , \beta), \tag {6}
$$
where  $|\cdot|$  is a weighted count of terminal and nonterminal tokens (the relative cost of a terminal vs. nonterminal token is a hyperparameter) and
$$
c _ {\mathcal {D}} (\alpha , \beta) = k _ {\alpha} \ln \hat {p} _ {\mathcal {D}} (\alpha | \beta) + k _ {\beta} \ln \hat {p} _ {\mathcal {D}} (\beta | \alpha) (7)
$$
where  $k_{\alpha}$  and  $k_{\beta}$  are hyperparameters and  $\hat{p}_{\mathcal{D}}(\alpha |\beta)$  is equal to the fraction of examples  $\langle x,y\rangle \in \mathcal{D}_{train}$  where  $\alpha$  "occurs in"  $x$  out of the examples where  $\beta$  "occurs in"  $y$ , and vice versa for  $\hat{p}_{\mathcal{D}}(\beta |\alpha)$ . The correlation between  $\alpha$  and  $\beta$  as measured by the  $\hat{p}$  terms provides a measure related to how well the rule fits the training data. We use sampling to optimize the computation of  $\hat{p}_{\mathcal{D}}$  for larger datasets. While conceptually similar to the objective used by NQG, we found that CSL's objective is more efficient to compute and can be more effective at penalizing rules that lead to lower precision.
Initialization To initialize  $\mathcal{G}$  we add a rule  $NT\to \langle x,y\rangle$  for every  $\langle x,y\rangle \in \mathcal{D}$ . We also optionally add a set of seed rules, such as  $NT\rightarrow \langle x,x\rangle$  where a terminal or span  $x$  is shared between the input and output vocabularies. For details on seed rules used for each dataset, see Appendix A.
Greedy Algorithm Following the initialization of the set of rules  $\mathcal{G}$ , we use an approximate parallel greedy search algorithm to optimize  $L_{\mathcal{D}}(\mathcal{G})$ , while maintaining the invariant that all examples in  $\mathcal{D}$  can be derived by  $\mathcal{G}$ .
At each iteration, the algorithm considers each rule  $r$  in the current grammar in parallel. The algorithm determines a (potentially empty) set of candidate actions for each  $r$ . Each candidate action consists of adding a new rule to the grammar that can be combined with an existing rule to derive  $r$ , enabling  $r$  to be removed. Certain candidate actions may enable removing other rules, too. The algorithm then selects the candidate action that leads to the greatest improvement in the induction objective, if any action exists that leads to an improvement. The selected actions are then aggregated and executed, resulting in a new set of rules. The algorithm continues until no further actions are selected, or a maximum number of steps is reached. The detailed implementation of the greedy algorithm is detailed in Appendix B.
3.1.5 Inference Procedure
While the primary goal of CSL is to be used to sample new examples for data augmentation (discussed
10 The outputs for several of the tasks we study consist of executable programs or logical terms, for which we can assume the availability of a CFG for parsing. Details are in Appendix A.
To optimize  $\theta$  efficiently, we use a variant of the CKY algorithm (Cocke, 1969; Kasami, 1965; Younger, 1967) to determine the set of derivations, represented as a parse forest, and use dynamic programming to efficiently sum over this set.
<sup>12</sup>By  $\alpha$  "occurs in"  $x$ , we mean that there exists some substitution for any non-terminals in  $\alpha$  such that it is a substring or equal to  $x$ .
<sup>13</sup>For example, CSL's objective enables our algorithm to induce a "clean" 20 rule grammar for SCAN, while using our algorithm with the objective of NQG leads to grammars with additional spurious rules for SCAN.
next), we can also use CSL as a discriminative parsing model, by using a variant of the CKY algorithm to find the highest scoring derivation  $z$  that maximizes Eq. 3 for a given input  $x$ . We then output the corresponding  $y$  if it can be derived by the given output CFG, or if no output CFG is provided.
3.2 Data Augmentation
We synthesize a configurable number of examples by sampling from the learned generative model, CSL. To generate a synthetic example  $(x,y)$ , we use forward sampling: we start from the single  $NT$  symbol and sample recursively to expand each non-terminal symbol with a rule, based on  $p_{\theta}(r|r_p,i)$  defined by Eq. 4.
Given that generic sequence-to-sequence models perform poorly on length extrapolation (Newman et al., 2020), we optionally bias our sampling to favor deeper derivations. We achieve this by adding a bias  $\delta > 0$  to  $\theta_{t,r}$  for any rule  $r$  that contains more nonterminals than our configurable threshold.