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 lowcomplexity 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 SQTransformer 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.1

heading: 1 Introduction
text: 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 (Chomsky, 1957; Montague, 1970). It leads to humans’ algebraic capacity to systematically understand a potentially infinite number of novel combinations of known structures and entities. For example,
1Our code is publicly available at https://github. com/jiangycTarheel/SQ-Transformer.
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 (Fodor and Pylyshyn, 1988; Marcus, 1998) and are supported by the empirical results that a Transformer (Vaswani et al., 2017) trained to parse “walk twice”, “walk around left”, and ‘jump’ fails to parse “jump twice” and “jump around left” in SCAN ADDJUMP (Lake and Baroni, 2018). Later works presented a more promising picture: for example, Zhou et al. (2023) 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 (Furrer et al., 2020; Drozdov et al., 2023) 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
8360
Same attention patterns for the same structure Different attention patterns for the same structure Same attention patterns for the same structure Different attention patterns for the same structure
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(a) Attention maps encoding “walk around left”, trained on original ADDJUMP. (b) Attention maps encoding “jump around left”, trained on original ADDJUMP.
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(c) Attention maps encoding “walk around left”, trained on 20x augmented ADDJUMP.
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(d) Attention maps encoding “jump around left”, trained on 20x augmented ADDJUMP.
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 (Zhou et al., 2023) 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 “$x around left” in Fig. 5.
achieves 100% test accuracy and encodes the syntactic structure “$x around left” with a unified attention pattern invariant to the choice of $x, as long as $x has the same syntactic function (e.g., being a verb). Over the entire test set, we observe many such reused attention.2 Therefore, we test the hypothesis that this ability to systematically reuse learned attention patterns on novel sentences is critical for a Transformer to systematically generalize.
To this end, we propose SQ-Transformer with two improvements to the embeddings and attention layers respectively, so as to induce the same systematicity seen above, even with low-complexity training data. The first improvement brings linguistic categorization to the word embeddings: as Johnson (2004) argues “the claim that natural languages are systematic presupposes a natural nonoverlapping linguistic categorization of all the expressions.” For example, for the model to generalize to the unseen “jump twice”, it first has to learn that ‘jump’ belongs to the same category as other primitives like ‘walk’ and ‘run’. Motivated by this theory and the finding, we propose Structureoriented Vector Quantization (SoVQ) to actively cluster all word embeddings into a fixed number of structural equivalence classes, and quantize each word into a code embedding shared among an entire class. We introduce a variational and generalized Brown Clustering (Brown et al., 1992) objective. This unsupervised objective encourages a “predictive clustering” such that the class of a token can be predicted by the classes of its context
2We discuss a quantified analysis in Sec. 5.3.
tokens, and hence ensures that words of the same syntactic function are in the same class. After being trained on examples like “walk”, “walk around left”, and “jump”, SoVQ can quantize ‘jump’ and ‘walk’ into the same class with a code embedding encoding their shared role in a sentence structure.
Our second improvement encourages a unified attention pattern for encoding sentences of a common syntactic structure. The general belief in cognitive science states that systematicity involves a capacity to represent common structural relations among the equivalently cognizable entities (Phillips and Wilson, 2016). That is, a systematic mind can always represent a structure even if one or more of its entities is substituted with any equivalently cognizable entity.3 Since SoVQ has quantized each class of equivalently cognizable entities into a code embedding, we then propose the Systematic Attention Layer (SAL) that uses these code embeddings as the queries and keys, and the word embeddings as the values (Fig. 2a). When encoding sentences with a common syntactic structure like “$x 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
3The notion of “equivalently cognizable entities” generally refers to entities within an equivalence class with respect to certain structural equivalence (e.g., all proper nouns).
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 (Jiang et al., 2022) (40%→99.4%), AROUNDRIGHT (Loula et al., 2018) (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 (Li et al., 2021). 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.

heading: 2 Background
text: Vector Quantization (VQ) (Agustsson et al., 2017; Van Den Oord et al., 2017) is a compression technique that represents a set of representations ex 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 Z ∈ RK×D made up of K embeddings of the dimension D:
q(zk|x) = { 1 if k = argminj f(ex, zj) 0 otherwise
VQ(x) = zk where q(zk|x) = 1 (1)
where f is a distance function (e.g., negative cosine similarity).4 The discrete code embeddings are updated using exponential moving averages of ex. Previous works (Van Den Oord et al., 2017; Razavi et al., 2019; Ramesh et al., 2021) 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 (Brown et al., 1992) is a word clustering algorithm that divides a vocabulary V into m mutually exclusive classes by maximizing the mutual information I(Z1, Z2) between the classes of a random bigram (X1, X2) in a sentence:
max Z:V→[m]
= ∑
z1,z2
#(z1, z2)
N log(
#(z1, z2)N #(z1)#(z2) ) (2)
where #(z, z′) denotes the number of occurrences of the cluster pair (z, z′) for any bigram in [x1...xN ].5 This algorithm can cluster a vocabulary based on the syntactic functions of words by promoting “predictive clustering”: the class of a token must be predictable from the class of its context token. However, it requires nontrivial combinatorial optimization and is difficult to scale and generalize for modern neural networks. In Sec. 3.1.2, we propose a variational Brown Clustering objective that can be optimized with gradient descent.

heading: 3 SQ-Transformer
text: In this section, we introduce the components of SQ-Transformer and discuss their technical details.
4We use the smallest index when there is a tie in argmin. 5By assuming a uniform distribution over consecutive word pairs (xi−1, xi), Brown et al. (1992) approximate p(z1, z2) and p(z) using #(z1,z2)N and #(z) N to derive Eqn. 2.
Notations. We denote the source and target sequences as [xi] and [yj ]. The seq2seq framework consists of an encoder with word embeddings Ex and a decoder with word embeddings Ey. For quantizing Ex and Ey, we define two codebooks Zx and Zy with Kx and Ky code embeddings respectively.

heading: 3.1 Structure-oriented Vector Quantization
text: Same as the original VQ, Structure-oriented Vector Quantization (SoVQ) clusters the (sub)word embeddings into several classes and quantizes each class into a shared code embedding z (Eqn. 1). We discuss a previous MMI objective and then propose variational Brown Clustering that better cluster words based on their syntactic functions.

heading: 3.1.1 Variational MMI objective
text: Stratos (2019) proposed an unsupervised part-ofspeech tagging method by maximizing the mutual information (MMI) between the inferred class Z of a token X and its surrounding context X̂ . It defines q(z|x) that directly infers the class of x (posterior) and p(z|x̂) that predicts the cluster of x based on its context x̂ (prior). It maximizes the variational lower bound of the mutual information I(X̂, Z):
I(X̂, Z) = H(Z)−H(Z|X̂) ≥ H(Z)−H(q, p)
(3)
where H(q, p) is the cross entropy over samples. We show more details about the derivation of this lower bound in Appendix A.1. As we can see in Eqn. 3, maximizing this ELBo is equivalent to (1) minimizing cross-entropy between the cluster inference posterior q(z|x) and cluster prediction prior p(z|x̂), so that words appearing in the same context are assigned to the same class (we introduce Theorem. 1 to demonstrate this); and (2) maximizing the entropy H(Z) of the cluster distribution to avoid assigning all words to the same cluster.

heading: 3.1.2 Variational Brown Clustering
text: In this work, we propose another MMI objective that marries the original Brown Clustering objective I(Z1, Z2) and the variational MMI (Eqn. 3). First, we redefine the cluster prediction distribution as p(z|ẑ), where ẑ are the quantized codes of all context tokens x̂ inferred from q(z|x̂) (Eqn. 1). This differs from the p(z|x̂) that predicts the cluster of x directly from its context x̂. Then, instead of maximizing the ELBO of I(X̂, Z), we maximize
the ELBO of I(Ẑ, Z):
I(Ẑ, Z) = H(Z)−H(Z|Ẑ) ≥ H(Z)−H(q(z|x), p(z|ẑ)) (4)
This inequality is still valid6 even though we replaced X̂ and p(z|x̂) in Eqn. 3 with Ẑ and p(z|ẑ). The objective I(Ẑ, Z) becomes the exact Brown Clustering objective if we set x̂ as a random context token rather than all of them. We show the implementation of this variational loss in Appendix A.4.
We argue that this variational Brown Clustering objective can better cluster words based on their syntactic functions than the lower bound of I(X̂, Z) (Eqn. 3). This is because, according to Theorem. 1, maximizing I(X̂, Z) can only cluster words that appear in similar contexts into the same class. However, some words having the same syntactic function might rarely occur in the same context due to semantics. For example, ‘police’ and ‘professor’ usually appear in very different contexts: “The police arrested a thief.” and “The professor appraised a student.” Therefore, maximizing I(X̂, Z) might not push the model to assign them to the same cluster. In comparison, maximizing I(Ẑ, Z) (Brown Clustering objective) can encourage clustering structurally equivalent words that appear in various contexts together: even though ‘police’ and ‘professor’ have different X̂ , they share the same Ẑ given a structure-oriented word cluster.7 We support this claim with ablations in Sec. 4.

heading: 3.2 Systematic Attention Layer
text: Now that we have quantized each class of words into a code z encoding structural information, we then use z as the queries and keys in computing the attention weights. Here we show the encoder’s self-attention module (visualized in Fig. 2a):
MHAttn(Q,K, V ) = softmax(QKT )V
zl+1 = MHAttn(q = zl, k = zl, v = zl)
xl+1 = MHAttn(q = zl, k = zl, v = xl)
where x0 and z0 are the non-contextualized word embeddings and their quantized code embeddings respectively. The two MHAttn modules share all parameters. We call it the Systematic Attention Layer (SAL) because this modified attention module promotes the systematic reusing of attention
6We show the derivation in Appendix A.3. 7The necessary clustering scheme that can achieve the
purpose is [{arrested,appraised}, {thief,student}].
2
2
patterns: as words of the same syntactic function (e.g., ‘cat’ and ‘dog’, ‘asleep’ and ‘awake’) are in the same cluster, the transformer would process two sentences of the same syntactic structure (“The cat is asleep” and “The dog is awake”) using the same attention pattern. As a result, a model that understands one sentence is more likely to generalize to the other one, which is the key ability stemming from a systematic language understanding. Similarly, we also use SAL with regular cross attention in the decoder (see Appendix A.5 for details). In summary, SAL enforces hard attention invariance among sentences of the same syntactic structure, but at the cost of the flexibility of encoding non-structural relations that commonly exist in natural languages (e.g., idioms, commonsense, etc). We discuss these cases in Sec. 6.

heading: 3.3 Systematically Regularized Layer
text: To encourage systematicity in attention while keeping its ability to encode non-structural relations, we instead use the attention outputs zl computed from the quantized embeddings to regularize the attention outputs xl computed from word embeddings, by minimizing the squared L2 distances (MSE loss) between zl and xl for all layers l:
zl+1 = MHAttn(q = zl, k = zl, v = zl)
xl+1 = MHAttn(q = xl, k = xl, v = xl)
where x0 are the word embeddings. We name it Systematically Regularized Layer (SRL) and visualize it in Fig. 2b. Unlike SAL, SRL demonstrates “soft invariance” so that sentences of a common structure are processed with similar (because of the L2 regularization loss) but not necessarily the same
attention pattern. During inference, we do not need to perform SoVQ and compute zl since it is only used to regularize xl in training. Therefore, SRL does not incur any computation or memory overhead than a vanilla Transformer layer in inference.

heading: 4 Experiments

heading: 4.1 Datasets
text: We use three semantic parsing tasks including SCAN ADDJUMP (Lake and Baroni, 2018), AROUNDRIGHT (Loula et al., 2018), COGS (Kim and Linzen, 2020), and the CoGnition (Li et al., 2021) En→Zh translation task, all of which require OOD compositional generalization to test examples. For example, SCAN ADDJUMP tests the models’ ability to parse syntactic structures (e.g., “$x twice”) combined with a novel entity ($x = ‘jump’), which is never associated with other structures during training. We also use WMT17 English↔German (Bojar et al., 2017) and WMT14 English↔French (Bojar et al., 2014) to test models for generalizability. We introduce each dataset in detail in Appendix B.1.

heading: 4.2 Results
text: Semantic Parsing results. We show the experimental setup in Appendix B.2. We report the results on the two SCAN tasks and COGS in Table 1. Specifically, SQ-Transformer with SAL achieves significant8 improvements over the baseline on SCAN ADDJUMP9 and AROUNDRIGHT. With SRL, SQ-Transformer manages to outperform the
8Bootstrapped test with α < 0.01. 9Both models are trained with 2x augmented SCAN AD-
DJUMP from Jiang et al. (2022).

heading: PREVIOUS MODELS

heading: OUR MODELS
text: baseline on the larger, more natural COGS dataset. This shows the effectiveness of SQ-Transformer in generalizing to unseen combinations of syntactic structure and lexical constituents. We compare the performance of SAL and SRL on a small, synthetic dataset and a larger, natural dataset in Sec. 5.1.
Machine Translation results. We evaluate the baseline Transformer as well as SQ-Transformer on the CoGnition compositional generalization test set and WMT test sets and report their BLEU 4 (Papineni et al., 2002) scores. For CoGnition, we also report the novel compound translation error
400 300 200 100 0 100 200 300 400
400
300
200
100
0
100
200
300
400
left
right
thrice
twice
around
opposite
after and
look
run walk
turn
jump
(a) Source embeddings trained with no quantization.
400 300 200 100 0 100 200 300 400
400
300
200
100
0
100
200
300
400
left rightthrice twice
around opposite
afterand
look run walkturn
jump
(b) Source embeddings trained with SoVQ (6 classes).
Figure 3: T-SNE visualization of embeddings learned on SCAN ADDJUMP dataset (Lake and Baroni, 2018).
SQ-Transformer ADDJUMP CoGnition
w. SAL 99.42±0.98⋆ 59.85±0.49⋆ w. SRL 47.36±20.83† 62.35±0.52† None 53.79±18.36 61.11±0.34
- SoVQ 78.44±34.01⋆ 60.93±0.13† - Brown 97.75±3.93⋆ 61.52±0.22†
Table 4: Ablation: test accuracy on SCAN ADDJUMP and BLEU on CoGnition (averaged over 5 runs). ‘- SoVQ’ uses the original Vector Quantization. ‘- Brown’ uses the MMI objective from Stratos (2019). Results with ⋆ use SAL while results with † use SRL.
(CTER) (Li et al., 2021). It examines whether all of the atoms (tokens) in the novel compound are correctly translated in the generated Chinese sentence. Specifically, instance-level CTER denotes the percentage of the test instances in which one or more atoms in the novel compound are translated incorrectly. Aggregate-level CTER denotes the percentage of novel compounds that are translated wrong in at least one instance. Compared to the Transformer baseline, SQ-Transformer obtains significantly higher BLEU scores on CoGnition En→Zh (Table 2), WMT17 En↔De, and WMT14 En↔Fr tasks (Table 3). On CoGnition, SQ-Transformer also achieves substantially lower instance and aggregate compound error rate in its Chinese translation. This improvement shows that SoVQ and SRL enable the model to correctly translate more novel compounds: for example, translating “the dog he liked” by generalizing from training expressions of the same structure like “the dress she liked” and “an animal she liked”.

heading: 5 Analysis
text: 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).

heading: 5.1 Ablation Study
text: We conduct an ablation study on the components of SQ-Transformer in Table 4. Most notably, we find that SAL only works on small, synthetic datasets like SCAN ADDJUMP but not on CoGnition; while SRL works on larger, more natural datasets like CoGnition and COGS but not on SCAN tasks. This finding corroborates other structure-based attention (Li et al., 2019; Russin et al., 2020) that only achieve strong performance on small, synthetic datasets. It suggests that generalizing on SCAN requires learning its grammar from an extreme data setting, which might only be possible via strictly structural inductive biases like SAL’s hard invariance to primitive substitution. Natural languages, on the other hand, require learning both structural and non-structural relations (discussed in Sec. 6), which is not possible using SAL. Later (in Sec. 5.3), we will show how SRL injects a similar, but soft invariance into attention to perform strongly on more natural tasks. Moreover, we show that using the original attention layer (‘None’) or original VQ (‘-SoVQ’) results in a significant drop in performance on SCAN ADDJUMP and CoGnition. Finally, optimizing our Variational Brown Clustering loss brings extra benefits compared to the MMI (‘-Brown’) objective Stratos (2019).

heading: 5.2 Analyzing the Embedding Space
text: Visualizing SCAN embedding space. We visualize the embedding matrices using 2-d tdistributed Stochastic Neighbor Embedding (tSNE) , which projects each embedding into a 2- dimensional coordinate (Hinton and Roweis, 2002). In Fig. 3b, we show that the source embeddings learned by SoVQ on SCAN ADDJUMP are clustered based on their syntactic functions in a sentence structure: the conjunction words (‘and’, ‘af-
ter’), direction adverbs (‘left’, ‘right’), prepositions (‘around’, ‘opposite’), and adverbs (‘twice’, ‘thrice’) are clustered together respectively. Most importantly, the rare primitive ‘jump’ is clustered together with other primitives. This enables the SQ-Transformer with SAL to generalize to unseen expressions like “jump twice” by reusing the same attention pattern as it has learned for “walk twice”. On the contrary, in Fig. 3a, words of the same syntactic functions (e.g., ‘jump’ and ‘walk’) are distant apart in the t-SNE space. This prevents the Transformer from generalizing to novel expressions like “jump twice”. Finally, SoVQ also learns structurebased word clusters on COGS that has a larger vocab (discussed in Appendix C.2).
Case Study: Learning the syntactic equivalence of ‘jump’ and ‘walk’. The major challenge in SCAN ADDJUMP is to recognize the equivalent syntactic function of the rare primitive ‘jump’ and other common primitives like ‘walk’ based on the only syntactic structure10 that has both ‘jump’ and ‘walk’ as a constituent. Next, we present a case study of how SoVQ manages to learn this equivalence and demonstrate three necessary preconditions. First, it is important to choose the proper number of clusters so that the model cannot afford to reserve a cluster for ‘jump’ only. For example, if we initialize 12 classes instead of 6 classes for the vector quantization, the model will put ‘jump’ into a separate class from ‘walk’ and ‘run’. Second, the model must be exposed to a sufficient number of examples that put ‘jump’ and ‘walk’ in the same syntactic structure (context). In Fig. 4, we show the t-SNE visualization of the source embeddings and generalization accuracy when the model is trained with different (na) repe-
10Atomic expressions like “jump7→JUMP”, “walk 7→WALK.”
titions of atomic expressions for each primitive.11 In all four cases, SoVQ can roughly cluster words based on their syntactic functions (e.g., ‘twice’ and ‘thrice’ are always together). However, there are some subtle differences regarding the clustering of verb primitives (shown in red). When na equals 1 or 10 (Fig. 4a, Fig. 4b), SQ-Transformer can hardly generalize (accuracy < 5%) and the ‘jump’ is located distantly with other verbs in red. With 100/1000 atomic expressions per primitive (Fig. 4c, Fig. 4d), SQ-Transformer achieves significantly improved generalization and learns more compact clusters of verbs. Finally, our variational Brown Clustering objective is indispensable for clustering words based on their syntactic functions. This can be seen in both final results (Table 4) and the distribution of the embedding space (Fig. 6).

heading: 5.3 Analyzing the Attention Patterns
text: In this part, we reveal that SQ-Transformer indeed learns attention patterns that systematically generalize to novel expressions. As we have shown above, SoVQ can cluster words in the SCAN vocab into multiple structural equivalence classes {Ci}. Therefore, by computing the queries and keys using the quantized embeddings, Systematic Attention Layers (SAL) are guaranteed to produce the same attention maps given any examples of a common structure (e.g., $x around left ∀$x ∈ Ci).
Next, we focus on analyzing the attention maps learned by the Systematically Regularized Layer (SRL). Unlike SAL, SRL still computes the attention weights using word embeddings as the queries and keys, but is regularized to keep its outputs close to the code representations from SAL. To evaluate the effectiveness of regularization, we collect 48 pairs of source sentences from the CoGnition test set. The sentences within each pair are quantized into the same sequence of codes (e.g., “he stopped every girl.” and “he found each child.”) by SoQV.12 We then count the percentage of attention heads that assign the highest weight to the same token when processing two examples in a pair. In this metric, SQ-Transformer achieves 79.8% compared to the baseline with 72.8%. This demonstrates that SRL learns systematic attention patterns in representing structures, while maintaining generalizability to natural data like CoGnition and WMT.
11“na=1000” means there are 1000 “jump 7→JUMP”, 1000 “walk 7→WALK”, etc. in the training set.
12We show all 48 pairs of sentences in Table 6.

heading: 6 Discussion
text: 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 positions13 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 (Levesque et al., 2012) 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 doesn’t fit in the brown suitcase because it’s too big. What is too big? Answer: The trophy.
• The trophy doesn’t fit in the brown suitcase because it’s too small. What is too small? Answer: The suitcase.
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
13E.g., “The cat is awake” and “The dog is asleep”.
designed for capturing local dependency within a sentence. It cannot scale to represent the complex relationships between words across sentences. Previously, Sartran et al. (2022) 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 (Table 4), which include a much larger vocab plus more natural and longer expressions. Therefore, as we motivated in Sec. 3.3, we instead use the more flexible SRL to achieve strong performance on these tasks.

heading: 7 Other Related Work
text: Compositional generalization. Earlier works investigated compositionality of neural networks in language learning (Wong and Wang, 2007; Brakel and Frank, 2009), compositional counting (Wiles, 1998; Weiss et al., 2018), and syntax learning (Linzen et al., 2016). Recent works (Lake and Baroni, 2018; Kim and Linzen, 2020; Loula et al., 2018; Bastings et al., 2018; Keysers et al., 2020; Tsarkov et al., 2021; Hupkes et al., 2020) embed the compositional challenge into semantic parsing tasks and directly evaluate seq2seq models on an out-of-distribution test set.
Previous works have also proposed many novel methods to improve the compositional generalization of neural models. Such methods include novel architectures (Li et al., 2019; Russin et al., 2020; Dessì and Baroni, 2019; Gordon et al., 2020; Oren et al., 2020; Zheng and Lapata, 2021), grammarbased approaches (Shaw et al., 2021; Kim, 2021), task decomposition (Herzig et al., 2021), data augmentation (Andreas, 2020; Akyürek et al., 2021; Akyürek and Andreas, 2022), careful architecture selection (Csordás et al., 2021), and novel learning methods (Lake, 2019; Conklin et al., 2021; Jiang et al., 2022; Xu et al., 2022).
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 (Hewitt and Manning, 2019; Phang et al., 2019; Clark et al., 2019; Limisiewicz et al., 2020; Murty et al., 2023a). 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 (Devlin et al., 2019) 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.
Incorporating linguistic knowledge into models. Many previous works have tried to incorporate linguistically-informed labels, especially syntactic labels, into neural networks (Sennrich and Haddow, 2016; Strubell et al., 2018; Sachan et al., 2021; Qian et al., 2021; Sartran et al., 2022). Following this idea, later works explored the syntactic equivalence of ‘jump’ and other verbs to induce compositionality (Akyurek and Andreas, 2021; Jiang and Bansal, 2021; White and Cotterell, 2022) for SCAN. Most relevantly, Chakravarthy et al. (2022) manually implemented several abstract “roles” for tokens in SCAN and computed the attention using the role embeddings. Our work differs from most of these works in that we do not require any external labels (e.g., parse tree) of sentences and instead automatically infer the syntactic “roles” for each token using Structure-oriented Vector Quantization and leverage them in Systematic Attention Layers and Systematically Regularized Layers.

heading: 8 Conclusion
text: 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.

heading: 9 Limitations
text: In this work, the proposed Structure-oriented Vector Quantization (SoVQ) mainly clusters the uncontextualized, lexical constituents based on their syntactic roles. It does not explicitly encourage phrasal constituents with the same syntactic role to be close together in the representation space. Therefore, we find that SQ-Transformer does not achieve better performance than the vanilla Transformer on COGS test examples with “deeper recursion depth” or “novel combination modified phrases and grammatical roles”, both of which require generalizing to novel combinations of grammatical structures and phrasal constituents.
SoVQ on uncontextualized embeddings also does not consider polysemy, which requires assigning a token to potentially different syntactic classes based on its context (discussed in Appendix D.1). However, our work is the first to apply vector quantization to improve compositional generalization of Transformers. We showed that clustering the lexical embeddings can already improve the Transformer in learning a compositional representation of the sentences. Therefore, we believe releasing our promising results by structurally clustering lexical components can potentially inspire more researchers in the community to take on this challenging but rewarding task of structurally clustering contextualized, phrasal representations.
Same as other Vector Quantization methods, we need to manually set the number of classes for clustering/quantizing embeddings (hyperparameter search details discussed in Appendix B.2). To prove the effectiveness of SoVQ, we initialize 4 clusters for SCAN tasks (whose vocabulary size is smaller than 20) and 16/32 clusters for tasks like CoGnition and WMT. This is because having too many clusters loosens the information bottleneck and fails to induce generalization. Future work incorporating SoVQ into the pretraining stage of foundation models does not need to tune the number of classes for every new task, but only needs to set it once given the size of its universal vocabulary.

heading: 10 Ethics Statement
text: Our SQ-Transformer architecture is designed to build models with a stronger ability to generalize to unseen compositions. It can be used as the backbone of a large foundation model with better data efficiency in acquiring a generalizable understanding of natural languages. However, we note the
models to be built with our architecture should be used with careful consideration. Since this is only a study aiming to improve the generalization ability of Transformer models and we do not release or plan to release a model pretrained for general use, we thus do not study the undesired behavior of the proposed model. Therefore, future works, which intend to use our SQ-Transformer or any proposed components (SoVQ, Systematic Attention Layer, and Systematically Regularized Layer) in training a large language model for real-world usage, need to conduct a thorough study on its safety and potential risks (similar to usage of any other deep learning models), including but not limited to honesty, truthfulness, fairness, toxicity, etc.

heading: Acknowledgements
text: We thank Elias Stengel-Eskin, Peter Hase, Archiki Prasad, Swarnadeep Saha, and Yi-Lin Sung for their useful feedback. This work was supported by NSF-CAREER Award 1846185, DARPA MCS Grant N66001-19-2-4031, and an Apple Scholars in AI/ML PhD Fellowship. The views are those of the authors and not of the funding agency.

heading: Appendix

heading: A Structure-oriented Vector Quantization
text: Notations. We denote the source sequence as [x1...xN ] and the target sequence as [y1...yM ]. The framework consists of an encoder with word embeddings Ex ∈ RNx×D and an autoregressive decoder with word embeddings Ey ∈ RNy×D, where Nx and Ny are the number of tokens in vocabulary and D is the dimension of the embeddings. For quantizing Ex and Ey, we define two codebooks Zx ∈ RNx×Dz and Zy ∈ RNy×Dz , where Nx and Ny are the number of codes and Dz is the dimension of the code embedding.

heading: A.1 Generalized Brown Clustering Objective.
text: Here, we present the details of the Brown Clustering (Brown et al., 1992)., the generalized Brown Clustering objective (Stratos, 2019), and its Evidence Lower Bound (ELBO) Eqn. 3 (with proofs) discussed in Sec. 3.1.
Brown Clustering. Brown Clustering is an unsupervised word clustering algorithm that was proposed and popularized before the trend of neural networks. Brown clustering divides a vocabulary V into m clusters by maximizing the mutual information (MMI) between the clusters (ZX , ZY ) of a random bigram (X,Y ) in a corpus of N words (x1...xN ). The algorithm assumes a uniform distribution over consecutive word pairs (xi−1, xi) and optimizes the MMI objective:
max Z:V→[m]
= ∑
z,z′
#(z, z′) N log( #(z, z′)N #(z)#(z) ) (5)
where #(z, z′) denotes the number of occurrences of the cluster pair (z, z′) for any bigram in (x1...xN ). This objective is intractable, so (Brown et al., 1992) proposed a greedy algorithm that (1) initializes the clusters by assigning each word to a distinct class and then (2) merges the pair of classes that leads to the least loss in the average mutual information for a total of |V | −m times.
This original algorithm failed on vocabularies larger than 5000 words. To the remedy, the authors instead (1) initialize m classes that each contain one of the m most frequent words and (2)
repeatedly merge a pair of clusters that yields the smallest decrease in mutual information. However, this heuristic requires nontrivial combinatorial optimization and is difficult to scale and generalize.
The Generalized MMI Objective. Stratos (2019) generalized the Brown Clustering to the setting by maximizing the mutual information between the posterior clustering probability q(z|x) and prior p(z|x̂), where x is a random token from a sentence and x̂ is its surrounding context. Since p and q are conditionally independent given (x, x̂), we have p(z, z′|x, x̂) = q(z|x)p(z′|x̂). Thus, the mutual information between p and q given a single sentence (x, x̂) is:
Jx,x̂ = ∑ z,z′ q(z|x)p(z′|x̂)logq(z|x)p(z ′|x̂) q(z)p(z′) (6)
The author then maximizes Ex,x̂∼D[Jx,x̂]. This objective becomes the exact Brown Clustering in Eqn. 5 if (1) x̂ is a random context token; (2) p and q are tied and are hard clustering instead of probabilistic soft clustering.

heading: A.2 A Variational Lower Bound of MMI.
text: Stratos (2019) further improved the generalized Brown Clustering and derived a novel objective as the difference between the entropy of the cluster distribution H(Z) and the cross entropy between q and p:
Jvar = H(Z)−H(q, p) (7)
where H(q, p) is the cross entropy between q and p over samples:
H(q, p) = E x,x̂∼D
[ − ∑
z
q(z|x)logp(z|x̂) ] (8)
Intuitively, minimizing the cross entropy between q and p can improve their mutual information. Maximizing H(Z) encourages the equal occurrence of each cluster and can prevent the trivial solution that assigns all tokens x in the corpus into the same cluster.
The objective in Eqn. 7 can also be seen as the lower bound of the mutual information between two random variables I(X̂, Z), where X̂ is the context of a token X and Z is its cluster inferred from q(z|x). This lower bound is shown in McAllester (2018) and we replicate it below.
First, since I(X̂, Z) = H(Z) − H(Z|X̂) and by definition we have the conditional entropy:
H(Z|X̂)
=− ∑
x̂,z
p(x̂, z)log p(x̂, z)
p(x̂)
= ∑
x̂,z
p(x̂, z)log p(x̂)
p(x̂, z)
= ∑
x̂,z
p(x̂, z)log 1
π(z|x̂)
= E (x̂,x)∼D z∼q(·|x)
[log 1
π(z|x̂) ]
(9)
where π(z|x̂) is the ground-truth prior probability of the cluster of a token given its context. If we introduce a variational distribution p(z|x̂) to approximate π(z|x̂) and further expand Eqn. 8, we have
H(q, p)
= E (x,x̂)∼D
[ − ∑
z
q(z|x)logp(z|x̂) ]
= E (x,x̂)∼D
[∑
z
q(z|x)log 1 p(z|x̂)
]
= E (x̂,x)∼D z∼q(·|x)
[ log 1
p(z|x̂)
]
= E (x̂,x)∼D z∼q(·|x)
[ log
π(z|x̂) π(z|x̂)p(z|x̂)
]
= E (x̂,x)∼D z∼q(·|x)
[ log 1
π(z|x̂)
] + E
(x̂,x)∼D z∼q(·|x)
[ log
π(z|x̂) p(z|x̂)
]
=H(Z|X̂) +DKL(π||p) (10)
where DKL is Kullback–Leibler divergence. Therefore, we can see that H(p, q) is an upper bound of H(Z|X̂), and hence H(z)−H(q, p) is the lower bound of I(X̂, Z). Stratos (2019) argues that maximizing I(X̂, Z) over the cluster inference distribution q enforces “predictive coding” because it pushes the cluster inferred from q(z|X) to be more informative of the context X̂ .
As we can see in Eqn. 3, maximizing this ELBo is equivalent to (1) minimizing cross-entropy between the cluster inference posterior q(z|x) and cluster prediction prior p(z|x̂) and (2) maximizing the entropy H(Z) of the cluster distribution.
First, minimizing the cross-entropy H(q, p) enforces “predictive clustering”: the class of x must be predictable from its context x̂. We introduce a theorem to show how this leads to assigning words appearing in the same context to the same class.
Theorem 1. Let xa and xb be two tokens that only appear in the same sets of context X̂ . Let p′, q′ =
argmin p,q
H(q(z|xa), p(z|x̂))+H(q(z|xb), p(z|x̂))
Then, we have: q′(z|xa) = q′(z|xb) ∀z ∈ Z, which means xa and xb are clustered into the same class in the optimal solution.
The proof is straightforward: the minimum value of the cross-entropy loss is 0, and this can only be achieved when q(z|xa) = q(z|xb) = p(z|x̂). To better understand the mathematical intuition of this theorem, consider the case where all adjectives are categorized into the same cluster. Then, we can confidently predict the class of $x=‘amazing’ based on the context “The food tastes $x.”, thus achieving the low cross-entropy with the posterior.
Second, maximizing the entropy of the cluster distribution pushes the model to utilize every cluster z in the latent space with (almost) equal probability. It thus prevents the trivial solution that assigns all tokens to only one random cluster k: p(zk|x) = q(zk|x) = 1 to minimize the first crossentropy term (H(p, q) = 0). Empirically, this variational MMI objective achieves strong unsupervised POS tagging performance (Stratos, 2019).

heading: A.3 The Variational Brown Clustering
text: In this work, we propose another MMI objective that marries the original Brown Clustering objective (Eqn. 5) and the alternative objective (Eqn. 3). First, we redefine the cluster prediction distribution as p(z|ẑ), where ẑ are the quantized codes of all context tokens x̂ inferred from argmax(q(z|x̂)). This differs from the p(z|x̂) that predicts the cluster of x directly from its context x̂. Then, instead of maximizing the ELBO of I(X̂, Z), we maximize the ELBO of I(Ẑ, Z). Next, we derive the lower bound of I(Ẑ, Z) in a similar way to Eqn. 9 and Eqn. 10. First, by definition we have:
I(Ẑ, Z) = H(Z)−H(Z|Ẑ)
Then, we rewrite H(Z|Ẑ) as:
H(Z|Ẑ)
=− ∑
ẑ,z
p(ẑ, z)log p(ẑ, z)
p(ẑ)
= ∑
ẑ,z
p(ẑ, z)log p(ẑ)
p(ẑ, z)
= ∑
ẑ,z
p(ẑ, z)log 1
π(z|ẑ)
= ∑
ẑ,z
∑
x̂,x
p(x̂, x)p(ẑ, z|x̂, x)log 1 π(z|ẑ)
= E (x̂,x)∼D z∼q(·|x) ẑ∼q(·|x̂)
[log 1
π(z|ẑ) ]
(11)
where π(z|ẑ) is the ground-truth code prediction distribution given its context’s codes. If we introduce a variational distribution p(z|ẑ) and rewrite Eqn. 8 with the newly defined p(z|ẑ), we have:
H(q, p)
= E (x̂,x)∼D ẑ∼q(·|x̂)
[ − ∑
z
q(z|x)logp(z|ẑ) ]
= E (x̂,x)∼D ẑ∼q(·|x̂)
[∑
z
q(z|x)log 1 p(z|x̂)
]
= E (x̂,x)∼D ẑ∼q(·|x̂) z∼q(·|x)
[ log 1
p(z|ẑ)
]
= E (x̂,x)∼D ẑ∼q(·|x̂) z∼q(·|x)
[ log
π(z|ẑ) π(z|x̂)p(z|ẑ)
]
= E (x̂,x)∼D ẑ∼q(·|x̂) z∼q(·|x)
[ log 1
π(z|ẑ)
] + E
(x̂,x)∼D ẑ∼q(·|x̂) z∼q(·|x)
[ log
π(z|ẑ) p(z|ẑ)
]
=H(Z|Ẑ) +DKL(π||p)
Therefore, H(q, p) is the upper bound of H(Z|Ẑ), and hence H(Z)−H(q, p) is the lower bound of I(Ẑ, Z).
A.4 Implementation of MMI objectives Finally, we show our empirical implementation of the MMI objective. First, we implement q(z|x) as the cosine similarity between the word embedding
ex and code embedding z. We implement p(z|x̂) as a separate Transformer network that takes the context14 of x as the input and predicts the class of x. Following Stratos (2019), we estimate these terms from the training data:
q′(z) = 1
N
N∑
i=1
q(z|xi)
H ′(Z) = − ∑
z
q′(z)logq′(z)
H ′(p, q) = 1
N
N∑
i=1
(− ∑
z
q(z|xi)logp(z|ẑi))
(12)
where N is the number of tokens. We then add α(H ′(p, q)−H ′(Z)) to the overall loss, where α is a tunable coefficient.

heading: A.5 Systematic Attention Layer for Decoder
text: In Sec. 3.2, we introduce the Systematic Attention Layer (SAL) for a Transformer encoder. Here, we introduce SAL for a Transformer decoder with encoder-decoder cross attention. we can use the quantized embedding z as the queries and keys in computing the self-attention weights but use the word embedding x as the queries and keys for cross attention. This is because we find that long-term, cross-sentence relationships cannot be determined by words’ syntactic functions only (Sec. 6):
zl+ = SelfAttn(q = zl−1, k = zl−1, v = zl−1) zl+ = CrossAttn(q = zl, k = z x L, v = z x L)
yl+ = SelfAttn(q = zl−1, k = zl−1, v = yl−1)
yl+ = CrossAttn(q = yl, k = xL, v = xL)
where y0 and z0 are the non-contextualized word embeddings and their quantized code embeddings respectively, xL and zxL are the encoder’s final layer outputs (as shown in Fig. 2a). In the equation above, the cross-attention that computes yl (last row) is the same as a regular Transformer layer. After the final layer, we project zL and supervise it to predict the code of the next token, same as how we project xL to predict the next token.
A.6 Information Bottleneck Interpretation Tishby and Zaslavsky (2015) state that to have a generalizable deep neural network, it is necessary
14The context can be either the whole sentence with x masked (bidirectional context) or the preceding words of x only (left context).
to optimize the Information Bottleneck (IB) tradeoff between compression and prediction for each layer. It is equivalent to minimizing the Lagrangian I(Xl−1, Xl)− βI(Xl, Y ), where Xl is a mapping (e.g., representation produced by the l-th layer in a neural net) of X and I(Xl, Y ) is the mutual information between Xl and the label Y . This objective suggests that we need to find the most concise representation Xl that is also sufficient to encode I(Xl, Y ) at each layer l. On the one hand, minimizing I(Xl−1, Xl) can prevent redundant information from flowing to the next layer, which could be exploited to establish some spurious correlations between Xl and Y . On the other hand, maximizing I(Xl, Y ) ensures that sufficient information is encoded in layer l to enables the final prediction of Y .
We argue that SoVQ and SRL can implicitly minimize I(Xl−1, Xl) for l = 0, 1...L: (1) SoVQ clusters the N word embeddings X0 around K code embeddings (N > Z15), thus achieving a lower H(X0) compared to an unrestricted embedding space and minimizing I(X,X0); (2) SRL, on the other hand, clusters the layer’s outputs computed from word embeddings around the layer’s outputs computed from code embeddings. It thus reduces H(Xl) to minimize I(Xl−1, Xl) for l = 1...L. Therefore, from the standpoint of information theory, SoVQ and SRL impose information bottlenecks on the embedding layer and every attention layer to improve the generalization of the entire network.

heading: B Experiments

heading: B.1 Datasets
text: We use a series of semantic parsing and machine translation tasks requiring compositional generalization and the common WMT tasks. All semantic parsing datasets are in English. None of the datasets includes any information that can name or uniquely identify individual people or offensive content. Since SCAN, COGS, and CoGnition do not have a compositional generalization validation set, we randomly sample 20% of the test data and use it as the validation set.
SCAN ADDJUMP (Lake and Baroni, 2018) tests the models’ ability to generalize syntactic structures (e.g., “$x twice”) to a novel entity ($x =
15For example, we only use 16 codes (K=16) to quantize 40356 subword embeddings for the WMT17 En-De tasks.
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS] wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(a) Attention maps encoding “walk around left”, trained on original ADDJUMP.
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 4
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 4
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(b) Attention maps encoding “look around left”, trained on original ADDJUMP.
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 4
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 4
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(c) Attention maps encoding “run around left”, trained on original ADDJUMP. (d) Attention maps encoding “jump around left”, trained on original ADDJUMP.
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(e) Attention maps encoding “walk around left”, trained on 20x augmented ADDJUMP.
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 1, Head 4
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 2, Head 4
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 1
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 2
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 3
look around left [EOS]
lo ok ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(f) Attention maps encoding “look around left”, trained on 20x augmented ADDJUMP.
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 1, Head 4
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 2, Head 4
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 1
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 2
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 3
run around left [EOS]
ru n ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(g) Attention maps encoding “run around left”, trained on 20x augmented ADDJUMP.
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(h) Attention maps encoding “jump around left”, trained on 20x augmented ADDJUMP.
Figure 5: Attention maps encoding training examples “walk around left”, “look around left”, “run around left” and a test example “jump around left” from the Transformers trained on the original SCAN ADDJUMP training set (a, b, c, d), and 20x augmented training set (Zhou et al., 2023) (e, f, g, h). We highlight the attention patterns in (d) that differ from (a), (b), (c) in red boxes. When trained on 20x augmented training set, the model encodes the two examples (c and d) with highly similar attention maps across all layers and heads.
‘jump’). Here we use the augmented training set with 2 times more primitives (Jiang et al., 2022), which includes 97906 examples. The validation and test sets have 1541 and 7706 examples respectively.
SCAN AROUNDRIGHT (Loula et al., 2018) tests the models’ ability to generalize a common syntactic structure “$x1 around $x2” to an entity ($x2=‘right’) that is only associated with other structures during the training. The train, validation, and test sets have 15225, 895, and 4476 examples respectively. The SCAN ADDJUMP and AROUNDRIGHT dataset are released under BSD License.
COGS (Kim and Linzen, 2020) challenges models to parse a diverse set of natural language sentences into their corresponding logical forms based on lambda calculus to accurately reflect the semantic representation of the natural sentence. The train, validation, and test sets have 24155, 4200, and 21000 examples respectively. The COGS dataset is released under MIT License.
CoGnition (Li et al., 2021) is an English-toChinese translation dataset with a synthetic OOD test set, where each sentence contains novel compositions of seen structures and constituents. The
train, validation, and test sets have 196246, 2160, and 10800 examples respectively. The CoGnition dataset is released for public use and has no explicit license.
WMT. We use WMT17 English↔German (Bojar et al., 2017) and WMT14 English↔French (Bojar et al., 2014) translation tasks. WMT17 En↔De has 3961179 English-German sentence pairs in the training set, 40058 sentence pairs in the validation set, and 3003 sentence pairs in the test set. WMT17 En↔Fr has 35762532 English-German sentence pairs in the training set, 26854 sentence pairs in the validation set, and 3003 sentence pairs in the test set.

heading: B.2 Experimental Setup
text: Semantic parsing experiments. We use a 3- layer Transformer encoder and a 3-layer Transformer decoder with 4 heads per layer, a hidden size of 256, and a feedforward size of 512. We share input and output embeddings of the decoder. We optimize the model using Adam (Kingma and Ba, 2015), with β1 = 0.1, β2 = 0.98. All models are trained for 100,000 steps and we choose the best checkpoint on validation set for evaluation. On SCAN tasks where the original vocab size is small (17 for source and 10 for target, including special
tokens), we try source codebook sizes [4,6,8] and target codebook sizes [3,4,5], and end up using 6 codes for quantizing source tokens and 4 codes for quantizing target tokens. On COGS with 747 source tokens and 687 target tokens, we try source and target codebook sizes [8,16,32], and end up using 32 codes for quantizing source tokens and 16 codes for quantizing target tokens.
Machine translation experiments. We use a 6- layer Transformer encoder and a 6-layer Transformer decoder with 8 heads per layer, a hidden size of 512, and a feedforward size of 1024. It has the same size as the Transformer used in Yin et al. (2023). We share input and output embeddings of the decoder. The model parameters are optimized by Adam (Kingma and Ba, 2015), with β1 = 0.1, β2= 0.98. All models are trained for 1 million steps on CoGnition and 2 million steps on WMT training sets. We then choose the best checkpoint on the validation set for evaluation. During decoding, we use a beam size of 5 and a maximum generation length of ax + b where a=1.2 and b=10. On CoGnition with 2004 source English tokens and 5500 target Chinese tokens (including special tokens), we try source and target codebook sizes [16,32,64,128], and end up using 64 codes for quantizing source tokens and 32 codes for quantizing target tokens. On WMT tasks, we follow the tokenization and preprocessing steps in fairseq16, and use 16 codes each for quantizing source tokens and target tokens.

heading: C Additional Analyses

heading: C.1 Probing the Code Embeddings
text: In Sec. 3.1, we argue that word embeddings x are being quantized into codes z largely based on their syntactic function in a linguistic structure. To support this argument, we conduct a probing experiment on the CoGnition dataset by training a linear classifier on top of the codes z (frozen) to predict the Part-of-speech (POS) tag of the corresponding words x. We find that this classifier can correctly predict the POS tag in 40.8% of the cases. Given the fact that z contains no context information and thus is inherently limited in predicting the finer POS tags (whether “like” is a verb or a preposition), the result of this probing experiment supports our argument that the quantized embeddings (code) z
16https://github.com/facebookresearch/ fairseq
encode a large number of syntactic roles of lexical features.
C.2 Visualizing COGS Embedding Space In Fig. 7, we show the t-SNE of the source embedding matrices learned on COGS. Compared to SCAN, COGS has a much larger vocabulary (748 VS 13) and more diverse syntactic structures. Again, in Fig. 7c we show that SoVQ can cluster the word embeddings based on their syntactic functions. For example, the “red cluster” (in the middle of 2-d t-SNE space) is mostly made of names like “Andrew” and “Olivia”. The “yellow cluster” on the right side is comprised of verbs in their past participle forms (e.g., “smiled, rolled”). The “green” cluster at the bottom includes animals like “fish”, “bear”, and “giraffe”. In Table 5, we show 10 words sampled from the source vocabulary and their closest neighbors in the 2-d t-SNE space. One interesting finding is that the word ‘can’ is clustered with nouns like ‘block’, ‘bee’, and ‘ring’. This is because in the COGS training set, ‘can’ is only used as a noun and never used as a modal verb.
The embedding space learned by the Transformer baseline, on the other hand, does not demonstrate any patterns that can connect the distribution and the syntactic role of a word Fig. 7a. We can observe that sometimes words that are different tenses of the same verb (e.g., “given” and “gave” in the red circle) or have connections in their semantics (e.g., “sleep” and “bed” in the blue circle) are clustered together. The Transformer with vanilla Vector Quantization, although learns a more clusterseparated embedding space Fig. 7b, also does not demonstrate any noticeable similarities of words within a cluster.
Overall, based on the difference in how words are clustered in the embedding space, we can state that Structural-oriented Vector Quantization (SoVQ) can effectively cluster words based on their syntactic functions.
C.3 Visualizing the attention patterns of vanilla Transformer.
As is discussed in Sec. 2, we train a vanilla Transformer model with 3 layers and 4 heads on the SCAN ADDJUMP training set of different complexities (Zhou et al., 2023). For example, the 20x augmented training set has 20 times more primitives (excluding jump) that appear in a variety of syntax structures. In Fig. 8, we visualize its attention maps from processing two separate expres-
400 300 200 100 0 100 200 300 400
400
300
200
100
0
100
200
300
400
left
right
thrice
twice
around
opposite
after and
look
run walk
turn
jump
(a) Source embeddings trained with no quantization.
400 300 200 100 0 100 200 300 400
400
300
200
100
0
100
200
300
400
leftright
thrice
twice
around
opposite
after and
look
run
walk
turn
jump
(b) Source embeddings trained with vanilla VQ.
400 300 200 100 0 100 200 300 400
400
300
200
100
0
100
200
300
400
left rightthrice twice
around opposite
afterand
look run walkturn
jump
(c) Source embeddings trained with SoVQ.
Figure 7: T-SNE of embeddings learned on COGS (Kim and Linzen, 2020), with 32 clusters used in VQ.
Source Word Closest Words
William Sophia, Riley, Carter, Nora, Madison, John, Jack, Lillian, Sebastian, Christopher draw float, love, scream, double, see, shark, help, roll, frown, worship ate gave, wanted, noticed, laughed, smirked, hoped, talked, napped, doubled bear manager, writer, cow, warrior, governor, crocodile, bicycle, boulder, bag, cloud cookie piano, rug, car, banana, melon, bench, bottle, bible, storage, turtle improve redden, poke, pierce, discover, throw, toss, slide, freeze, disintegrate, Paula bed taxi, pot, sphere, cot, couch, bunker, backpack, glacier, vehicle, bin preferred jogged, smiled, sketched, craved, touched, gasped, yearned, supported, saw, crumple teacher strawberry, raisin, beast, soap, monster, clock, rose, child, lawyer, chief can block, bee, ring, blender, tripod, seat, jacket, dog, donkey, beer table stage, speaker, desk, barrel, boat, trunk, house, room, stand
Table 5: Words sampled from the source vocabulary of COGS and their closest words in the 2-d t-SNE space.
sions: “walk around left”, which is in the training set, and “jump around left”, which is excluded from the training set. We also show the generalization accuracy and Kullback–Leibler (KL) divergence between the two attention distributions, averaged over all heads and layers. In summary, we show that as the number of distinct primitives increases, the KL
divergence between the attention distributions of these two examples decreases. This improvement in attention similarity positively correlates with the improvement of the generalization accuracy. With 20x more lexical primitives than the original training set, the KL divergence reaches 0.001 and the accuracy reaches 100%. This observation provides
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS] wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(a) “walk around left”, trained on ADDJUMP. (b) “jump around left”, trained on ADDJUMP (KL: 0.337).
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(c) “walk around left”, trained on ADDJUMP 2x.
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(d) “jump around left”, trained on ADDJUMP 2x (KL: 0.080).
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(e) “walk around left”, trained on ADDJUMP 20x.
jump around left [EOS]
ju m p ar ou nd le ft [E OS ] Layer 1, Head 1 jump around left [EOS] ju m p ar ou nd le ft [E OS ] Layer 1, Head 2 jump around left [EOS] ju m p ar ou nd le ft [E OS ] Layer 1, Head 3 jump around left [EOS] ju m p ar ou nd le ft [E OS ] Layer 1, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ] Layer 2, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ] Layer 3, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(f) “jump around left”, trained on ADDJUMP 20x (KL:0.001).
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 1, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 2, Head 4
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 1
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 2
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 3
walk around left [EOS]
wa lk ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(g) “walk around left”, trained on ADDJUMP 200x.
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 1, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 2, Head 4
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 1
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 2
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 3
jump around left [EOS]
ju m p ar ou nd le ft [E OS ]
Layer 3, Head 4
0.2
0.4
0.6
0.8
(h) “jump around left”, trained on ADDJUMP 200x (KL:0.001).
Figure 8: Encoder’s attention maps and their average KL-divergence between the two examples, from the vanilla Transformer trained on the SCAN ADDJUMP datasets of different number of distinct primitives (Zhou et al., 2023) (original, 2x, 20x, 200x), where the model achieves 3.67%, 15.78%, 100%, and 100% accuracy on the entire test set respectively.
a mechanistic explanation of how a Transformer trained on more complex data acquires better outof-domain generalization to novel combinations of structure and lexical entities.

heading: C.4 Analyzing the Attention Pattern of SRL.
text: We show the example pairs used in Sec. 5.3. We collect a total of 48 pairs of source sentences from the test set. The source sentences within each pair are quantized into the same sequence of clusters (e.g., “he stopped every girl.” and “he found each child.”) by our SQ-Transformer.

heading: D More Discussion

heading: D.1 Polysemy
text: Vector Quantization on token embeddings assigns a token to one of several mutually exclusive classes (Eqn. 1). Therefore, it does not explicitly consider polysemy, which requires assigning a token to potentially multiple syntactic classes based on its context. For example, depending on the context, “sign” could belong to the verb class or the noun class. Empirically we find that words with multiple syntactic roles indeed have non-zero quantization probabilities to multiple classes: after being trained on the WMT En-De, the embedding of “sign” has a 0.3 probability of being quantized to a class mostly including verbs, and a 0.65 probability of being quantized to a class mostly including nouns. This
is because, while the cluster inference posterior q(z|x) is conditioned on the lexical embeddings only, the cluster prediction prior p(z|ẑ) is conditioned on the context and thus can predict the two distinct syntactic roles of the word “sign”. As a result of minimizing their cross-entropy, the posterior q(z|x) also learns a bi-modal class distribution for x =“sign”.
With that being said, since a token can only be assigned to the top class (noun class for “sign”) before any contextualization, the verb sense of “sign” is lost after SoVQ. Therefore, to accurately quantize polysemies like “sign”, future works can explore quantizing contextualized representation, like the outputs after the first 3 layers, of words. This can potentially lead to better-quality, contextaware code representation used in the downstream SAL/SRL.

heading: D.2 Scaling Up Model and Data
text: Scaling up the training data. In this work, we trained SQ-Transformer on datasets up to the scale of WMT machine translation datasets, which are the most popular benchmarks in evaluating seq2seq neural architectures (as used in Vaswani et al. (2017)). We also support our main arguments via proof-of-concept experiments based on semantic parsing. We do not train SQ-Transformer on larger web-scale, language-modeling datasets and leave
this to future work. We hope the promising results of SQ-Transformer can inspire more researchers to experiment with it on web-scale data.
Scaling up SQ-Transformer. In this work, we investigate SQ-Transformer up to 6 layers, 8 heads, and 512 hidden dim, same as the model size in Vaswani et al. (2017). To further scale the model
up, there are two main challenges. First, we use two separate Transformer decoders as the cluster prediction prior p(z|ẑ) (Eqn. 4), which brings memory overhead. However, the cluster predictors can be much smaller than the main Transformer network because the structure-oriented latent code space (16 for WMT En-De) is much smaller than the main network’s semantic-enriched word space (40356 for WMT En-De). It is also possible to share the parameters between the main network and the cluster predictors. Second, most LLMs use some types of byte pair encoding (BPE) tokenization to build a universal, sub-word level vocabulary. As a result, the structural role of a sub-word (e.g., ‘token-’) could depend on its suffix: “token-ize” is a verb while “token-ization” is a noun. This is similar to the polysemy problem at the word level (discussed in Appendix D.1), but is much more frequent at the sub-word level. So far, our proposed Structureoriented Vector Quantization (SoVQ) is only performed on the sub-word embeddings. Therefore, without any context information, SoVQ cannot accurately quantize ‘token’ to either the verb class or the noun class. Nonetheless, we still showed SoVQ’s effectiveness in WMT machine translation datasets (tokenized with BPE). This suggests that the SoVQ is robust to some ambiguous sub-words like the prefix ‘token’. However, further scaling SoVQ to web-scale data with a much larger subword vocabulary could potentially be challenging, since some of the special sub-words are not meaningful by themselves.17 To address this challenge, we believe future works can experiment with quantizing the contextualized representation (e.g., the output representations of the first k layers) rather than quantizing the sub-word embeddings.

heading: D.3 How does Transformer Generalize Compositionally?
text: It has long been argued that neural networks are associative devices that cannot capture systematic compositionality (Fodor and Pylyshyn, 1988; Marcus, 1998; Fodor and Lepore, 2002; Marcus, 2003; Calvo and Symons, 2014). Specifically, Fodor and Pylyshyn (1988) claimed that “in traditional Associationism, the probability that one Idea will elicit another is sensitive to the strength of the association between them. ... Associative strength was not, however, presumed to be sensitive to features of the
17For example, “<0xF0><0x9F><0xA6><0x99>” forms a “llama” emoji but each sub-word itself does not have a fixed role in linguistics.
content or the structure of representations per se. Similarly, in Connectionist models, the selection of an output corresponding to a given input is a function of properties of the paths that connect them.” As a result, they further stated that “The syntactic/semantic structure of the representation of an input is not presumed to be a factor in determining the selection of a corresponding output since, as we have seen, syntactic/semantic structure is not defined for the sorts of representations that Connectionist models acknowledge.” After we showcase the systematic behavior of SQ-Transformer, readers might ask “how does our method overcome the inherent limitation of Connectionist models elicited in Fodor and Pylyshyn (1988)?”
First, the statement above made an important assumption about neural networks (i.e., Connectionist models) that syntactic/semantic structure does not determine the strength of association between neurons (through the form of attention or full connection). We agree that the strength of the association is decided by the correlations a neural model observed in data. Thus, models like Transformers and RNNs fail to execute “jump twice” because where “jump” and “twice” are never seen together in training and models are insensitive to the syntactic structure.
However, we argue that with the proper regularization (e.g., SoVQ and SRL) to the intermediate representations (incl. embeddings and layer outputs), neural models can take the sentence structure into account when determining the strength of inter-neuron association via attention weights or MLP connection. This is because neural models are not inherently limited to capturing shallow, word-level correlations. As is shown in Hewitt and Manning (2019), the attention maps can also be sensitive to the common, although latent, dependency between words, which is simply a kind of statistical correlation between multiple latent and explicit factors (word, position, order, etc). Therefore, the resulting representations also encode rich structural information. The extent to which the model can be sensitive to the structure depends on the data complexity (Zhou et al., 2023), model architecture (Murty et al., 2023b), and regularization (Jiang and Bansal, 2021; Yin et al., 2023). For example, there is a statistical correlation between the input word “twice” and the latent output structure (always repeating the action preceding “twice” 2 times). We showed that SQ-Transformer can very well capture this association systematically in its
attention maps.

heading: Future Work
text: We do not train SQ-Transformer on larger web-scale, language-modeling datasets and leave
this to future work. We hope the promising results of SQ-Transformer can inspire more researchers to experiment with it on web-scale data. Scaling up SQ-Transformer. In this work, we investigate SQ-Transformer up to 6 layers, 8 heads, and 512 hidden dim, same as the model size in Vaswani et al. (2017). To further scale the model
up, there are two main challenges. First, we use two separate Transformer decoders as the cluster prediction prior p(z|ẑ) (Eqn. 4), which brings memory overhead. However, the cluster predictors can be much smaller than the main Transformer network because the structure-oriented latent code space (16 for WMT En-De) is much smaller than the main network’s semantic-enriched word space (40356 for WMT En-De). It is also possible to share the parameters between the main network and the cluster predictors. Second, most LLMs use some types of byte pair encoding (BPE) tokenization to build a universal, sub-word level vocabulary. As a result, the structural role of a sub-word (e. g. , ‘token-’) could depend on its suffix: “token-ize” is a verb while “token-ization” is a noun. This is similar to the polysemy problem at the word level (discussed in Appendix D. 1), but is much more frequent at the sub-word level. So far, our proposed Structureoriented Vector Quantization (SoVQ) is only performed on the sub-word embeddings. Therefore, without any context information, SoVQ cannot accurately quantize ‘token’ to either the verb class or the noun class. Nonetheless, we still showed SoVQ’s effectiveness in WMT machine translation datasets (tokenized with BPE). This suggests that the SoVQ is robust to some ambiguous sub-words like the prefix ‘token’. However, further scaling SoVQ to web-scale data with a much larger subword vocabulary could potentially be challenging, since some of the special sub-words are not meaningful by themselves. 17 To address this challenge, we believe future works can experiment with quantizing the contextualized representation (e. g. , the output representations of the first k layers) rather than quantizing the sub-word embeddings