Exploring Compositional Generalization (in COGS/ReCOGS_pos) by Transformers using Restricted Access Sequence Processing (RASP)

TL;DR

The paper investigates whether Transformers can achieve true compositional generalization on COGS and ReCOGS tasks by using Restricted Access Sequence Processing (RASP), a flat, non-hierarchical, Transformer-compatible language. By compiling RASP programs into Transformer weights, the authors demonstrate near-perfect performance on structural generalization—including prepositional phrase recursion and sentential complements—without requiring tree-structured representations. They show that a carefully designed, word-level POS/verb-type embedding plus 19 pattern-matching rules, coupled with decoder loop unrolling, can yield exact semantic and string matches across both COGS and ReCOGS generalization sets. The work also analyzes errors in a strong baseline Transformer (Wu2023), identifying attraction errors and illustrating how masking in a flat pattern framework mitigates them, thereby offering new insights into Transformer inductive biases and the role of data/inductive design in grammatical generalization. Overall, the study argues that compositional generalization in these tasks can be achieved with flat, non-recursive strategies and highlights RASP as a valuable tool for diagnosing Transformer capabilities and guiding future research.

Abstract

Humans understand new combinations of words encountered if they are combinations of words recognized from different contexts, an ability called Compositional Generalization. The COGS benchmark (Kim and Linzen, 2020) arXiv:2010.05465 reports 0% accuracy for Transformer models on some structural generalizations. We use (Weiss et al., 2021) arXiv:2106.06981's Restricted Access Sequence Processing (RASP), a Transformer-equivalent programming language, to demonstrate that a Transformer Encoder-Decoder can perform COGS and the semantically equivalent ReCOGS_pos (Wu et al., 2024) arXiv:2303.13716 systematically and compositionally: Our RASP models attain near perfect scores on structural generalization splits on COGS (exact match) and ReCOGS_pos (semantic exact match). Our RASP models show the (Re)COGS tasks do not require a hierarchical or tree-structured solution (contrary to (Kim and Linzen, 2020) arXiv:2010.05465, (Yao and Koller, 2022) arXiv:2210.13050, (Murty et al., 2022) arXiv:2211.01288, (Liu et al., 2021) arXiv:2107.06516): we use word-level tokens with an "embedding" layer that tags with possible part of speech, applying just once per encoder pass 19 attention-head compatible flat pattern-matching rules (easily identified with specific training examples), shown using grammar coverage (Zeller et al., 2023) to cover the non-recursive aspects of the input grammar, plus masking out prepositional phrases ("pp noun") and/or sentential complements (cp) when recognizing grammar patterns and extracting nouns related to the main verb in the sentence, and output the next logical form (LF) token (repeating until the LF is complete). The models do not apply recursive, tree-structured rules like "np_det pp np -> np_pp -> np", but score near perfect semantic and string exact match on both COGS and ReCOGS pp recursion, cp recursion using the decoder loop.

Figures (13)

• Figure 1: Introducing parse trees, logical form, and semantic graphs. Two semantically identical but syntactically distinct (i) sentences (a) "a boy painted the girl" and (b) "the girl was painted by a boy" are shown with (ii) their distinct parse tree (parsed into the shared (Re)COGS input grammar), (iii) the string form of their semantics (ReCOGS logical form; differs in indices and ordering), and (c) the graph representation of their logical form (semantic graph, not different at all between the two examples). Note the (iii) logical forms (LFs) differ by String Exact Match but not Wu2023's Semantic Exact Match (order and indices do not match but nouns, normalized verbs, and relationships between nouns and verbs are same). Note the "agent", "theme" order in the logical form string is not required to match for Semantic Exact Match. COGS and ReCOGS tasks require extracting the semantics (c) encoded in LFs (iii) of sentences (i). RASP-for-(Re)COGS shows decoding the LF (Figures \ref{['rasp-for-recogs-decoder-loop-supplementary-figure-incl-encoder-and-decoder-and-grammar-vertical']} (ReCOGS), \ref{['rasp-for-cogs-encoder-decoder-with-grammar-patterns']} (COGS)) does NOT require representing the parse tree.
• Figure 2: RASP-for-(Re)COGS Encoder determines an example contains "np v_dat_p2 np np" via flat pattern matching (maps to agent, recipient, theme noun-verb relation ordering (\ref{['RASP-model-flat-patterns-after-masking-to-nv-relationships-table']})). See Figure \ref{['rasp-model-flat-pattern-match-example-v_dat_p2_match_with_middle_pp']} for middle-noun pp modification, Figure \ref{['rasp-model-flat-pattern-match-example-v_dat_p2_non_matching']} for the non-matching case, and Figure \ref{['rasp-for-recogs-decoder-loop-supplementary-figure-incl-encoder-and-decoder-and-grammar-vertical']} (ReCOGS), \ref{['rasp-for-cogs-encoder-decoder-with-grammar-patterns']} (COGS) for the full Encoder-Decoder models translating input to logical form. 1d embed. equiv. to higher dim orthogonal vector idx (A\ref{['rasp-word-level-model-design']}).
• Figure 3: Wu2023 Encoder-Decoder Transformer trained from scratch generalizing to new v_dat_p2 pp moved to recipient (from theme) split is as hard as the previously reported hardest generalization split consistent with the flat/non-recursive/non-tree representation hypothesis ((d) rejects $H_0$). Figure \ref{['rasp-model-flat-pattern-match-example-v_dat_p2_match_with_middle_pp']} shows how a flat RASP model can recognize (b).
• Figure 4: Example of the RASP-for-(Re)COGS model flat grammar pattern matching, for np v_dat_p2 np np, for a matching sentence.
• Figure 5: Example of the RASP-for-(Re)COGS model flat grammar pattern matching, for np v_dat_p2 np np, for a matching sentence, despite pp modification of middle recipient noun. This is in the encoder (shared for COGS and ReCOGS). See also Figures \ref{['attraction_errors_figure']}, \ref{['rasp-for-recogs-decoder-loop-supplementary-figure_pp_depth_2']}, \ref{['rasp-for-cogs-decoder-loop-figure_pp_depth_2']} for how the RASP model avoids attraction errors in assigning agents, recipients, themes due to prepositional phrase modification in the decoder.