Inducing Systematicity in Transformers by Attending to Structurally Quantized Embeddings (2024)

Abstract: Transformers generalize to novel compositions of structures and entities after being trained on a complex dataset, but easily overfit on datasets of insufficient complexity. We observe that when the training set is sufficiently complex, the model encodes structurally equivalent sentences using a systematic attention pattern. Inspired by this observation, we propose SQ-Transformer (Structurally Quantized) that explicitly encourages systematicity in the embeddings and attention layers even with low-complexity data. At the embedding level, we introduce Structure-oriented Vector Quantization (SoVQ) to cluster word embeddings into several classes of structurally equivalent entities. At the attention level, we devise the Systematic Attention Layer (SAL) and an alternative, Systematically Regularized Layer (SRL) that operate on the quantized word embeddings so that sentences of the same structure are encoded with invariant or similar attention patterns. Empirically, we show SQ-Transformer achieves stronger compositional generalization than the vanilla Transformer on multiple low-complexity semantic parsing and machine translation datasets. In our analysis, we show SoVQ indeed learns a syntactically clustered embedding space, and SAL/SRL induces generalizable attention patterns, altogether leading to improved systematicity.

## introduction natural languages demonstrate compositionality, which states that the meaning of a complex expression is determined by its syntactic structure and the meanings of its lexical constituents @xcite . it leads to humans' algebraic capacity to systematically understand a potentially infinite number of novel combinations of known structures and entities. for example, someone who understands "the cat is asleep" and "the dog is awake" must simultaneously understand "the dog is asleep" and "the cat is awake". early works argued that neural networks are associative devices that cannot capture compositionality @xcite and are supported by the empirical results that a transformer @xcite trained to parse "walk twice", "walk around left", and 'jump' fails to parse "jump twice" and "jump around left" in scan addjump @xcite . later works presented a more promising picture: for example, @xcite found that transformers trained on an augmented, high-complexity dataset with more examples and diverse entities/structures can systematically generalize to novel compositions in scan addjump. studies on large pretrained models @xcite also reveal their ability to systematically generalize. however, data augmentation requires domain-specific knowledge and pretraining on large datasets is also prohibitively expensive. therefore, how to induce systematicity with lowcomplexity data has a significant value for improving the model's data efficiency, and remains an open and important research question. to understand the emergence of systematicity in transformers, we start by analyzing the attention maps from models trained on scan addjump data of different complexities. first, we demonstrate that a transformer trained on the original, low-complexity training set (with only 4 primitives) uses different attention weights to encode in-distribution training sentences like "walk around left" (fig. 1a ) and an unobserved sentence "jump around left" (fig. 1b ). it only achieves 3.7% test accuracy in unobserved sentences. in contrast, the same model trained on large augmented data (with 84 distinct primitives like 'walk' and 'jump') uses highly similar attention patterns to encode these two sentences (see fig. 1c and fig. 1d ). this model atterns for the same structure different attention patterns for the same structure same attention patterns for the same structure different attention patterns for t figure 1 : attention maps encoding a training example "walk around left" and a test example "jump around left". the maps in (a) and (b) are from the transformers trained on the original scan addjump training set; the maps in (c) and (d) are from the model trained on 20x augmented training set @xcite with 20 times more primitives like 'walk1' (mutated from 'walk') and more examples like "walk1 around left". in all attention maps, brighter yellow squares correspond to an attention value close to 1 and darker purple squares correspond to an attention value close to 0. we highlight the attention maps in (b) that differ from (a) in red boxes. when trained on 20x augmented training set, the model encodes the two examples with highly similar attention maps across all layers and heads (c and d). we show the attention maps on other examples of the structure "@xmath0x around left", sal is hard-invariant for any $x in a structural equivalence class c established by sovq. it thus enables the transformer to systematically represent common structural relations among those quantized classes of equivalently cognizable entities. to retain the attention's ability to represent nonstructural relations that commonly exist in natural languages, we also introduce an alternative to sal: systematically regularized layer (srl). it inherits the architecture of a regular attention layer, but additionally minimizes the l2 distance between the layers' outputs computed from word embeddings and the layers' outputs computed from quantized word embeddings (fig. 2b ). therefore, unlike sal, srl encourages attention's soft invariance to structurally equivalent entities: sentences with common structures are processed with similar but not necessarily the same attention pattern. overall, we name this model with sovq and sal/srl as (structurally quantized) sq-transformer. to demonstrate that predictive clustering in embeddings and invariance in attention can lead to systematicity in the model's predictions, we train and evaluate sq-transformer from scratch on multiple low-complexity semantic parsing and machine translation datasets requiring compositional generalization. in semantic parsing, sq-transformer improves upon transformer on scan addjump x2 @xcite (40%→99.4%), aroundright @xcite (69.5%→99.6%), and cogs (kim and linzen, 2020) (82.6%→83.4%). in machine translation, sq-transformer achieves higher bleu scores (60.5→62.8) and lower novel compound translation error (29.6%→18.1%) on cognition @xcite . importantly, it also shows generalizability to higher-complexity, natural datasets that do not have a significant distribution shift between training and test sets: in wmt en↔de and en↔fr, sq-transformer with srl obtains significantly higher bleu scores. we further analyze sq-transformer and present two findings: (1) sovq can more effectively cluster word embeddings based on their syntactic functions compared to vq; (2) sal and srl learn attention patterns that can systematically encode unseen compositions of structure and entities. these analyses explain the working mechanism of sq-transformer and verify our insights in designing these modules. in summary, sq-transformer quantizes word embeddings based on their syntactic functions and learns generalizable attention for sentences of the same structure. as a result, it can correctly parse and translate more sentences with unseen compositions of syntactic structures and lexical constituents. we hope this work sheds light on the inner mechanism of transformers' generalization and inspire future work in architecture design. ## background vector quantization (vq) @xcite is a compression technique that represents a set of representations e x of the variable x by a small, fixed number of code embeddings z. the code is inferred with the nearest neighbor look-up on a codebook @xmath1: previous works @xcite @xcite have shown that vq-vae can generate high-fidelity, continuous signals like images and speech. the exploration of vq on languages (lingle, 2023) remains limited. in this work, we use vq to cluster words based on their syntactic function (sec. 3.1). brown clustering @xcite is a word clustering algorithm that divides a vocabulary v into m mutually exclusive classes by maximizing the mutual information i(z 1 , z 2 ) between the classes of a random bigram (x 1 , x 2 ) in a sentence: ## sq-transformer in this section, we introduce the components of sq-transformer and discuss their technical details. 4 we use the smallest index when there is a tie in argmin. notations. we denote the source and target sequences as [x i ] and [y j ]. the seq2seq framework consists of an encoder with word embeddings e x and a decoder with word embeddings e y . for quantizing e x and e y , we define two codebooks z x and z y with k x and k y code embeddings respectively. ## analysis in this section, we conduct an ablation study (sec. 5.1) and analyze the embeddings (sec. 5.2) as well as the attention patterns (sec. 5.3). ## discussion in this part, we (1) discuss why the systematic attention layer (sal) can only work on small, synthetic datasets like scan; and (2) further motivate the systematically regularized layer (srl) that achieves a balance between compositionally encoding structures and maintaining flexibility in incorporating non-compositional relationships that commonly exist in natural languages. recall sal represents the "common structural relations" by computing attention weights in the quantized, syntactic space. this ensures that sentences that have (1) the same structure and (2) equivalently cognizable entities at all positions foot_9 are processed with the same attention weights across all heads and layers. however, natural languages are only approximately compositional, indicating that sometimes the meaning of a piece of text does not only depend on its structure and its lexical constituents. we discuss two situations where sal is overly strict and thus prevents the model from encoding non-structural linguistic features. situation 1: semantics. we use the classic winograd challenge @xcite to demonstrate how lexical semantics, or more specifically, commonsense knowledge affects the comprehension of a sentence. here, the model sees two sentences that differ only in one or two words, which are often of the same syntactic role but contain a referential ambiguity resolved in opposite directions. for example: the trophy. here, it is necessary to use commonsense semantics to resolve the ambiguous anaphora. however, the attention maps learned by sal for these two sentences would be the same, given that "big" and "small" are most likely quantized into the same class. therefore, a model with sal cannot encode this commonsense knowledge into the attention. situation 2: pragmatics. other than the inability to incorporate commonsense to resolve the coreference within a sentence, sal with "hard invariance" over structural equivalence classes is only designed for capturing local dependency within a sentence. it cannot scale to represent the complex relationships between words across sentences. previously, @xcite bias the attention with the sentence parse tree, and reported deteriorated performance on document-level language modeling. similarly, we observe that using quantized embeddings to compute cross-attention weights would result in degenerated performance on scan tasks, and thus opt for the original word embeddings as the queries and keys (appendix a.5). in both situations, it is necessary to compute the attention using the word embeddings rather than their quantized code embedding, so that the attention can capture other non-structural relations (e.g., commonsense) as well. empirically, we also show that using the sal results in degenerated performance on cogs and cognition datasets ( previous works have also proposed many novel methods to improve the compositional generalization of neural models. such methods include novel architectures @xcite @xcite @xcite , grammarbased approaches @xcite , task decomposition @xcite , data augmentation @xcite @xcite , careful architecture selection @xcite , and novel learning methods @xcite @xcite . structures captured by transformer. researchers have long been studying how the at-tention maps of transformers encode the syntactic structure (e.g., dependency parse) of a sentence @xcite @xcite @xcite . recently, jian and reddy (2023) substituted words in a sentence with words from the same syntactic category and then averaged the attention maps of a bert @xcite model across these "syntactically invariant sentences". our analysis in fig. 1 also makes use of such syntactically invariant sentences. we further reveal the correlation between the emergence of systematic attention patterns and the model's generalization ability in a transformer trained from scratch, which leads to our effort to inject linguistic structure into the model, as discussed below. ## conclusion in this work, we propose sq-transformer with structure-oriented vector quantization and two types of attention layers that use the quantized embeddings as the keys and values. our experiments show that sq-transformer can better generalize to unseen expressions in multiple semantic parsing and machine translation tasks. we conduct multiple analyses and show that sovq can cluster word embeddings based on their syntactic roles and the model learns systematic attention patterns in processing sentences of the same syntactic structure.