Compositional Generalization via Neural-Symbolic Stack Machines

We design a stack machine that supports recursion, a key component to achieving compositional generalization. In particular, this machine supports general-purpose sequence-to-sequence tasks, where an input sequence in a source language is mapped into an output sequence in a target language. We present an overview of the machine operations in Table 1. Specifically, SHIFT is used for reading the input sequence, PUSH and POP are standard stack operations, REDUCE is used for output sequence generation, CONCAT_M and CONCAT_S concatenate the generated sequences to form a longer one, and the FINAL operation terminates the machine operation and produces the output. We provide an illustrative example in Figure 1, and defer more descriptions to the supplementary.

With the incorporation of a symbolic machine into our sequence-to-sequence generation framework, NeSS does not directly generate the entire output sequence. Instead, the neural network in NeSS acts as a controller to operate the machine. The machine runs the execution trace generated by the neural network to produce the output sequence. Meanwhile, the design of our machine allows the neural controller to make predictions based on the local context of the input, which is a key factor to achieving compositional generalization. We provide an overview of the neural controller in Figure 4, describe the key components below, and defer more details to the supplementary.

Machine status encoder. A key property of NeSS is that it enables the neural controller to focus on the local context that is relevant to the prediction, thanks to the recursion supported by the stack machine. Specifically, the input to the neural controller consists of three parts: (1) the next token in the input queue $\mathit{tok}$, e.g., the token “around” before executing step 4 in Figure 1; (2) the top 2 frames of the stack; and (3) the memory. Note that including the second stack frame from the top is necessary for determining the association rules of tokens, as discussed in [8]. Take arithmetic expression parsing as an example, when the top stack frame includes a variable “y” and the next input token is “*”, we continue processing this token when the previous stack frame includes a “+”, while we need a POP operation when the previous stack frame includes “*”. We use 4 bi-directional LSTMs as the machine status encoders to encode the input sequence, the top 2 stack frames and the memory respectively. Then we denote the 4 computed embedding vectors as $e_{\mathit{tok}}$, $e_{cur}$, $e_{pre}$ and $e_{M}$.

Operator predictor. The operator predictor is a multi-layer fully connected network with a $|Op|$-dimensional softmax output layer, where $|Op|=7$ is the number of operators supported by the machine, as listed in Table 1. Its input is the concatenation of $e_{\mathit{tok}}$, $e_{cur}$, $e_{pre}$ and $e_{M}$.

Argument predictors. We include three neural modules for argument prediction, which are used for REDUCE, CONCAT_M and CONCAT_S respectively. We design the REDUCE argument predictor as a standard LSTM-based sequence-to-sequence model with attention, which generates a token sequence in the target vocabulary as the arguments. For both CONCAT_M and CONCAT_S argument predictors, we utilize the pointer network architecture [46] to select indices from an element list as arguments, where the list contains the elements from the top stack frame and the memory.

Latent category predictors. We introduce two latent category predictors for source and target languages. The source category predictor, denoted as $p_{\text{sc}}(e_{\mathit{tok}})$, computes an embedding vector $ec_{\mathit{tok}}$ to indicate the categorical information given the input $e_{\mathit{tok}}$. Similarly, we denote the target category predictor as $p_{\text{tc}}(e_{s})$, where the input $e_{s}$ is the embedding of the token sequence $s$ in the target language. For the input to the operator predictor, we replace $e_{\mathit{tok}}$ with $ec_{\mathit{tok}}$ as the representation of the next input token $\mathit{tok}$, which encourages the neural controller to predict the same operator for tokens of the same category. Similarly, the categorical predictions for the target language are used for subsequent instruction predictions.

In the following, we discuss the expressiveness and generalization power of NeSS. In particular, NeSS preserves the same expressive power as sequence-to-sequence models, while our neural-symbolic design enhances its compositional generalization capability.

Expressive power. When we impose no constraints to regularize the machine, e.g., we do not restrict the argument length for REDUCE instruction, there is a degenerate execution trace that is valid for every input-output example. Specifically, this trace keeps running the SHIFT instruction until an $\verb|[EOS]|$ is observed as the next input token, then executes a REDUCE instruction with the entire output sequence as its argument. In this way, NeSS preserves the capacity of existing sequence-to-sequence models [3, 48, 19, 44, 13]. To leverage the recursion property of the machine, we could set a length limit for REDUCE arguments, so that the neural model mainly calls REDUCE instruction to generate phrases in the target language, and utilizes other instructions to combine the generated primitives to form a longer sequence. We call such compositional traces as non-degenerate traces hereafter.

Generalization. Some recent work proposes neural-symbolic architectures to achieve length generalization for program induction [6, 8, 49, 39]. The core idea is to incorporate a stack into the symbolic machine, so that the neural network model could restrict its attention to only part of the input important for the current decision. This recursion design principle is also crucial in achieving compositional generalization in our case. Meanwhile, capturing the operational equivalence enables NeSS to simultaneously obtain generalization capability and expressive power.

The SCAN benchmark has been widely used to evaluate the compositional generalization of neural networks, where the input sequence is a navigation command in English, and the output is the corresponding action sequence [26]. See Figure 1 for a sample usage of NeSS for the SCAN tasks.

Evaluation setup. Similar to prior work [27, 16, 38], we evaluate the following four settings. (1) Length generalization: the output sequences in the training set include at most 22 actions, while the output lengths in the test set are between 24 and 48. (2) Template generalization for “around right”: the phrase “around right” is held out from the training set; however, both “around” and “right” occurs in the training set separately. For example, the phrases “around left” and “opposite right” are included in the training set. (3) Primitive generalization for “jump”: all commands not including “jump” are included in the training set, but the primitive “jump” only appears as a single command in the training set. The test set includes commands combining “jump” with other primitives and templates, such as “jump twice” and “jump after walk”. (4) Simple split: randomly split samples into training and test sets. In this case, no compositional generalization is required.

Previous approaches. We compare NeSS with two classes of existing approaches on SCAN benchmark. The first class of approaches propose different neural network architectures, without additional data to provide training supervision. Specifically, sequence-to-sequence models (seq2seq) [26] and convolutional neural networks (CNN) [12] are standard neural network architectures, Stack LSTM learns an LSTM to operate a differentiable stack [18], while the equivariant sequence-to-sequence model [16] incorporates convolution operations into the recurrent neural networks, so as to achieve local equivariance discussed in Section 2.1.1. On the other hand, the syntactic attention model [42] and primitive substitution [30] learn two attention maps for primitives and templates separately. The second class of approaches design different schemes to generate auxiliary training data. Specifically, GECA [1] performs data augmentation by replacing fragments of training samples with different fragments from other similar samples, while the meta sequence-to-sequence model [27] and the rule synthesizer model (synth) [38] are trained with samples drawn from a meta-grammar with the format close to the SCAN grammar.

Results. Table 6 summarizes our results on SCAN tasks. In the top block, we compare with models trained without additional data. Among these approaches, NeSS is the only one achieving $100\%$ test accuracies on tasks that require compositional generalization, and the performance is consistent among 5 independent runs. In particular, the best generalization accuracy on the length split is only around $20\%$ for the baseline models. Note that the stack LSTM does not achieve better results, demonstrating that without a symbolic stack machine that supports recursion and sequence manipulation, augmenting neural networks with a stack alone is not sufficient. Meanwhile, without category predictors, NeSS still achieves $100\%$ test accuracy in 2 runs, but the accuracy drops to around $20\%$ for other 3 runs. A main reason is that existing models could not generalize to the input template “around left/right thrice”, when the training set only includes the template “around left/right twice”. Although NeSS correctly learns the parsing rules for different words, without category predictors, NeSS still may not learn that the parsing rule for “thrice” has the same priority as “twice”. For example, in the test set, there is a new pattern “jump around right thrice”. The correct translation is to parse “jump around right” first, then repeat the action sequence thrice, resulting in 24 actions. Without category prediction, NeSS could mistakenly parse “right thrice” first, concatenate the action sequences of “jump” and “right thrice”, then repeat it for four times, resulting in 16 actions. Such a model could still achieve $100\%$ training accuracy, because this pattern does not occur in the training set, but the test accuracy drops dramatically due to the wrong order for applying rules. Therefore, to achieve generalization, besides the parsing rules for each individual word, the model also needs to understand the order of applying different rules, which is not captured by the primitive equivalence [27, 30, 42] or local equivariance [16] studied in prior work. On the other hand, as shown in Table 2, the operational equivalence defined in Section 2.1.1 enables the model to learn the priorities of different parsing rules, e.g., categorizing “twice” and “thrice” together, which is crucial in achieving length generalization.

Next, we compare NeSS with models trained with additional data. In particular, the meta sequence-to-sequence model is trained with different permutations of primitive assignment, i.e., different one-to-one mapping of {run, look, jump, walk} to {RUN, LOOK, JUMP, WALK}, denoted as “(perm)”. We consider two evaluation modes of the rule synthesizer model (synth) [38], where the first variant performs greedy decoding, denoted as “(no search)”; the second one performs a search process, where the model samples candidate grammars, and returns the one that matches the training samples. We observe that even with additional training supervision, Synth (with search) is the only baseline approach that is able to achieve $100\%$ generalization on all these SCAN splits.

Although both Synth and NeSS achieve perfect generalization, there are key differences we would like to highlight. First, the meta-grammar designed in Synth restricts the search space to only include grammars with a similar format to the SCAN grammar [38]. For example, each grammar has between 4 and 9 primitive rules that map a single word to a single primitive (e.g., run $\rightarrow$ RUN), and 3 to 7 higher order rules that encode variable transformations given by a single word (e.g., x1 and x2 $\rightarrow$ [x1] [x2]). Therefore, Synth cannot be applied to other two benchmarks in our evaluation. Unlike Synth, NeSS does not make such assumptions about the number of rules nor their formats. Also, NeSS does not perform any search during the inference, while Synth requires a search procedure to ensure that the synthesized grammar satisfies the training samples.

Next, we evaluate on the few-shot learning benchmark proposed in [28], where the model learns to produce abstract outputs (i.e., colored circles) from pseudowords (e.g., “dax”). Compared to the SCAN benchmark, the grammar of this task is simpler, with 4 primitive rules and 3 compositional rules. However, while the SCAN training set includes over 10K examples, there are only 14 training samples in this benchmark, thus models need to learn the grammar from very few demonstrations. In [28], they demonstrate that humans are generally good at such few-shot learning tasks due to their inductive biases, while existing machine learning models struggle to obtain this capability.

Results. We present the results in Table 6, where we compare with the standard sequence-to-sequence model, the primitive substitution approach discussed in Section 4.1, and the human performance evaluated in [28]. We didn’t compare with the meta sequence-to-sequence model and the rule synthesizer model discussed in Section 4.1, because they require meta learning with additional training samples. Despite that the number of training samples is very small, NeSS achieves $100\%$ test accuracy in 5 independent runs, demonstrating the benefit of integrating the symbolic stack machine to capture the grammar rules.

Then we evaluate on the compositional machine translation benchmark proposed in [26]. Specifically, the training set includes 11,000 English-French sentence pairs, where the English sentences begin with phrases such as “I am”, “you are” and “he is”, and 1,000 of the samples are repetitions of “I am daxy” and its French translation, which is the only sentence that introduces the pseudoword “daxy” in the training set. The test set includes different combinations of the token “daxy” and other phrases, e.g., “you are daxy”, which do not appear in the training set. Compared to the SCAN benchmark, the translation rules in this task are more complex and ambiguous, which makes it challenging to be fully explained with a rigorous rule set.

Results. We present the results in Table 6, where we compare NeSS with the standard sequence-to-sequence model [26], and the primitive substitution approach discussed in Section 4.1. Note that instead of measuring the exact match accuracy, where the prediction is considered correct only when it is exactly the same as ground truth, we measure the semantic equivalence in Table 6. As discussed in [30], only one reference translation is provided for each sample in the test set, but there are 2 different French translations of “you are” that appear frequently in the training set, which are both valid translations. Therefore, if we measure the exact match accuracy, the accuracy of the Primitive Substitution approach is $62.5\%$, while NeSS achieves $100\%$ in 2 runs, and $62.5\%$ in 3 other runs. Although both NeSS and the Primitive Substitution approach achieves $100\%$ generalization, by preserving the sequence generation capability of sequence-to-sequence models with the REDUCE argument generator, NeSS is the only approach that simultaneously enables length generalization for rule learning tasks and achieves $100\%$ generalization on machine translation with more diverse rules, by learning the phrase alignment.

Finally we evaluate NeSS on the context-free grammar parsing tasks in [8]. Following [8], we mainly consider the curriculum learning setting, where we train the model with their designed curriculum, which includes 100 to 150 samples enumerating all constructors in the grammar. NeSS parses the inputs by generating the serialized parse trees, as illustrated in the supplementary material. The average input length of samples in the curriculum is around 10. This benchmark is mainly designed to evaluate length generalization, where the test samples are much longer than training samples.

Results. We present the main results in Table 6, where we compare NeSS with the sequence-to-sequence model [47], sequence-to-tree model [15], LSTM augmented with a differentiable stack structure [18], and the neural parser [8]. All these models are trained on the curriculum of the While language designed in [8], and we defer the full evaluation results of more setups and baselines to the supplementary material, where we draw similar conclusions. Again, NeSS achieves $100\%$ accuracy in 5 independent runs. We notice that none of the models without incorporating a symbolic machine generalizes to test inputs that are 500 $\times$ longer than training samples, suggesting the necessity of the neural-symbolic model design. Meanwhile, compared to the neural parser model, NeSS achieves the same capability of precisely learning the grammar production rules, while it supports more applications that are not supported by the neural parser model, as discussed in Section 2.3.