Venue: NAACL 2025

# Sneaking Syntax into Transformer Language Models

with Tree Regularization

Ananjan Nandi, Christopher D. Manning & Shikhar Murty

Department of Computer Science

Stanford University

Stanford, CA 94305, USA

{ananjan,smurty}@stanford.edu

###### Abstract

While compositional accounts of human language understanding are based on a hierarchical tree-like process, neural models like transformers lack a direct inductive bias for such tree structures. Introducing syntactic inductive biases could unlock more robust and data-efficient learning in transformer language models (LMs), but existing methods for incorporating such structure greatly restrict models, either limiting their expressivity or increasing inference complexity. This work instead aims to _softly_ inject syntactic inductive biases into given transformer _circuits_, through a structured regularizer. We introduce TreeReg, an auxiliary loss function that converts bracketing decisions from silver parses into a set of differentiable orthogonality constraints on vector hidden states. TreeReg integrates seamlessly with the standard LM objective, requiring no architectural changes. LMs pre-trained with TreeReg on natural language corpora such as WikiText-103 achieve up to 10% lower perplexities on out-of-distribution data and up to 9.5 point improvements in syntactic generalization, requiring less than half the training data to outperform standard LMs. TreeReg still provides gains for pre-trained LLMs: Continued pre-training of Sheared Llama with TreeReg results in improved syntactic generalization, and fine-tuning on MultiNLI with TreeReg mitigates degradation of performance on adversarial NLI benchmarks by 41.2 points. We release all code to guide future research1.

Footnote 1: [https://github.com/ananjan-nandi-9/tree_regularization](https://github.com/ananjan-nandi-9/tree_regularization)

## 1 Introduction

A substantial body of research [6, 8, 7] suggests that human language processing is inherently _hierarchical_: syntactic rules combine word-level meanings to give rise to semantics for larger constituents, which are then combined into sentences. In contrast, transformers [26], the architecture behind large language models (LLMs), allow for arbitrary, non-hierarchical routing of information while processing an input sentence.

LLMs are pre-trained on massive amounts of data to achieve reasonable generalization [1, 12], but recent studies find that state-of-the-art LLMs still struggle with compositional generalization [1, 13, 14], the ability to understand familiar words in novel contexts. Prior work [27, 28] proposed explicitly tree-structured models capable of data-efficient compositional generalization, but they are too constrained and not amenable to pre-training at scale.

Consequently, people have proposed inductive biases that encourage hierarchical computation in transformers, including Syntactic Language Models (SLMs, 25, 26) that model a joint distribution over words and syntax trees. While SLMs often improve data efficiency and syntactic generalization, they have several drawbacks: They often involve additional parameters to model syntax, constrain the attention mechanisms of the underlying model, or involve more complex and slower inference methodologies.

In this work, we instead devise a new differentiable loss function that _softly_ injects syntactic inductive biases into a given _circuit_ of the transformer: TreeReg. These biases are soft structural constraints that aim to ensure that the hidden state representation from the transformer circuit for each constituent in an input sentence is _maximally orthogonal_ to that of its surrounding context (Figure 1). TreeReg is simply added as a regularizer to the LM loss during training. Crucially, an LM trained with TreeReg is _completely indistinguishable_ from a standard LM in both architecture and inference mechanism.



The use of TreeReg improves syntactic generalization across model and data scales. On diagnostic sentence transformation tasks such as tense inflection and question formation (McCoy et al., 2020), it significantly enhances grokking behavior (SS 4). It also improves syntactic generalization on the BLiMP Warstadt et al. (2020) and SyntaxGym Gauthier et al. (2020) test suites by up to 3.2 and 9.5 points, both when used during pre-training from scratch (SS 5), and during continued pre-training (SS 6.1). LMs with TreeReg surpass the syntactic generalization of baseline LMs with less than half the training data (SS 5.2). Most importantly, TreeReg improves out-of-distribution language understanding. When used during pre-training of LMs on fully parsed (SS 5.1) as well as a mixture of parsed and unparsed (SS 5.3) data, TreeReg reduces out-of-distribution perplexities by up to 9.2%. When TreeReg is used for fine-tuning an LLM on MultiNLI Williams et al. (2018), it mitigates performance drops by up to 41.2 points on adversarial NLI benchmarks (SS 6.2). We release all code at [https://github.com/ananjan-nandi-9/tree_regularization](https://github.com/ananjan-nandi-9/tree_regularization).

## 2 Background and Notation

Tree Structures in Language Processing.Consider a sentence \(S=\{x_{1},x_{2},\ldots,x_{n}\}\) where \(x_{j}\) is the \(j^{\text{th}}\) token in the sentence. Let tree \(T(S)\) over sentence \(S\) be a set of spans \(S_{a;b}=\{x_{a},\ldots,x_{b}\}\) such that if \(S_{a;b}\) and \(S_{b+1;c}\in S\), then \(S_{a;c}\in S\). We say \(S_{a;c}\) is _split_ at index \(b\) to get spans \(S_{a;b}\) and \(S_{b+1;c}\) of \(T(S)\). Prior work Pallier et al. (2011); van Schijndel et al. (2013); Hale et al. (2018) suggests that human language processing is hierarchical--sentences contain "constituents" or spans of words that act as semantic units, with smaller constituents ("is", "very happy") recursively combining to form larger ones ("is very happy"), building up meaning incrementally. The constituency parse \(\hat{T}(S)\) is the tree form by all constituents of \(S\). Treestructured models Socher et al. (2013); Tai et al. (2015) build sentence representations through recursively bottom-up computation over constituency parses, whereas for transformers, any hierarchically structured computation is at most implicitly learned in the attention layers.

Hierarchical Structure in Transformers.Recent evidence suggests that pre-trained transformers behave similarly to tree-structured models on certain tasks Murty et al. (2023c, a). The _tree projection_ metric from Murty et al. (2023c) quantifies how well a transformer's computation can be approximated by a tree-structured model. It is based on the idea that bottom-up computations on trees create context-invariant representations at intermediate nodes. Therefore, the tree built from maximally context-invariant spans aligns most closely with the transformer's computation. Another line of work builds Syntactic Language Models Suratran et al. (2022); Murty et al. (2023b); Hu et al. (2024) that learn joint distributions over inputs and their parses to introduce syntactic inductive biases into transformer language models. These ap 

proaches involve substantial changes to the transformer architecture--explicitly modifying attention patterns, as well as changing the output space or inference mechanisms. TreeReg provides an alternative, biasing transformer circuits towards hierarchical structure through regularization.

Evaluating Syntactic Generalization.Syntactic generalization refers to an LM's ability to apply syntactic rules implicitly learned during training to new, unseen data. We use the BLiMP Warstadt et al. (2020) and SyntaxGym SG Gauthier et al. (2020) test suites to evaluate LMs for syntactic generalization. In BLiMP, models are presented with pairs of grammatical and ungrammatical sentences that differ minimally, and are evaluated on assigning lower perplexity to the grammatical sentences. The SG test suites evaluate six distinct syntactic phenomena, with each suite involving a particular inequality constraint related to the surprisal of different continuations of input sentences based on their prefixes. These constraints are rooted in theories of syntax; for example, assigning higher surprisal to the third word in the sentence "I know who you introduced to them." compared to "I know that you introduced to them."

## 3 Our Approach

TreeReg is a differentiable loss function that projects the hierarchical structure of input constituency parses into transformer circuits. Inspired by Murty et al. (2023c), TreeReg specifies a set of soft constraints to ensure that constituent representations are maximally independent of surrounding context. For any span, we quantify this independence with the Span Contextual Independence Score (\({\rm SCIN}\), SS 3.1). TreeReg maximizes \({\rm SCIN}\) for spans corresponding to input constituents while simultaneously minimizing it for other spans (SS 3.2). In an optimally trained model, the tree with the highest cumulative \({\rm SCIN}\) thus recovers the constituency parse. We then provide a greedy algorithm to extract the parse tree induced by TreeReg from a given model for any input sentence. A deeper discussion of the design choices underlying our method can be found in Appendix A.

### Span Contextual Independence Score

Let \(\mathbf{h}^{l}_{a,1},\mathbf{h}^{l}_{a,2},\ldots,\mathbf{h}^{l}_{a,n}\) be \(L_{2}\)-normalized vector hidden states for each token in \(S\) produced by attention head \(a\) at layer \(l\) of a causal auto-regressive transformer LM with multi-headed self-attention. Since most tasks benefit from a blend of hierarchical and unstructured computation, we only apply TreeReg on the circuit formed by a given subset of attention heads \(A\) at a given layer \(l\).

Span vectors.For any \(j\leq n\), \(\mathbf{h}^{l}_{A,j}\) is the concatenation (\(\#\)) of vector representations from the attention heads in \(A\) at layer \(l\) for prefix \(x_{\leq j}\):

\[\mathbf{h}^{l}_{A,j}=\#_{a\in A}\mathbf{h}^{l}_{a,j} \tag{1}\]

Then, for span \(S_{i;j}\), \(\mathbf{h}^{l}_{A,i-1}\) contains information before the start of \(S_{i;j}\), or the _context_ for the span.

Contextual Independence.We operationalize contextual independence in terms of orthogonality constraints between span representations. In particular, for \(S_{i;j}\) to be contextually invariant, we expect \(\mathbf{h}^{l}_{A,j}\) to be independent of, or _orthogonal to_\(\mathbf{h}^{l}_{A,i-1}\). We also constrain \(\mathbf{h}^{l}_{A,j+1}\) to be orthogonal to \(\mathbf{h}^{l}_{A,j}\), as the information in \(S_{i;j}\) is expected to be independent of what comes after it in the sentence. Therefore, we define \({\rm SCIN}^{l}_{A}\) for \(S_{i;j}\) at layer \(l\) for attention heads in \(A\) as

\[{\rm SCIN}^{l}_{A}(i,j) =\|{\rm orth}(\mathbf{h}^{l}_{A,j},\mathbf{h}^{l}_{A,i-1})\|\] \[+\|{\rm orth}(\mathbf{h}^{l}_{A,j+1},\mathbf{h}^{l}_{A,j})\|, \tag{2}\]

where \(\|\|\) is the \(L_{2}\) norm, and \({\rm orth}(x,y)=x-(x^{\top}y)y\) for any two normalized vectors \(x\) and \(y\). For the sake of completeness, \(\|{\rm orth}(h_{j},h_{-1})\|~{}\forall j\) and \(\|{\rm ordt}(h_{k+1},h_{k})\|\) are assumed to be 0. Since \(l\) and \(A\) are fixed during training, we simplify \({\rm SCIN}^{l}_{A}\) to \({\rm SCIN}\) from here on.

### TreeReg Loss

Given sentence \(S\), the TreeReg loss (\(\mathcal{L}_{\text{TR}}\)) biases the model towards driving up \({\rm SCIN}\) for all spans in the constituency parse \(\hat{T}(S)\) while driving down \({\rm SCIN}\) for other spans. Specifically, for constituent \(S_{i;j}\in\hat{T}(S)\) split at index \(p\), we first compute split scores \(s(i,q,j)\) for \(i\leq q\leq j-1\) as

\[s(i,q,j) ={\rm SCIN}(i,q)\] \[+{\rm SCIN}(q+1,j). \tag{3}\]

We then use these scores to compute a span-level log loss \(l_{(i,j)}\). These losses are computed as

\[l_{(i,j)}=\log\left[\sum_{q=i}^{j-1}\exp s(i,q,j)\right]-s(i,p,j).\] (4) 

\({\mathcal{L}_{\text{TR}}}\) is then the sum of span-level losses \(l_{(i,j)}\) for all constituents in \(\hat{T}(S)\),

\[{\mathcal{L}_{\text{TR}}}=\sum_{S_{ij}\in\hat{T}(S)}l_{(i,j)}. \tag{5}\]

In practice, \({\mathcal{L}_{\text{TR}}}\) is computed recursively top-down on \(\hat{T}(S)\) following Alg. 1. \({\mathcal{L}_{\text{TR}}}\) is added as an auxiliary loss to the LM loss during training, resulting in the training objective \({\mathcal{L}_{\text{LM}}}+\alpha{\mathcal{L}_{\text{TR}}}\) where \({\mathcal{L}_{\text{LM}}}\) is the LM loss. Since both \({\mathcal{L}_{\text{TR}}}\) and \({\mathcal{L}_{\text{LM}}}\) are cross-entropy losses, \(\alpha\) can generally be set to 1. \({\mathcal{L}_{\text{TR}}}\) and \({\mathcal{L}_{\text{LM}}}\) can also consume data from different datasets. For example, we can perform LM on a large pre-training dataset while passing batches from a smaller, parsed dataset to \({\mathcal{L}_{\text{TR}}}\).

Recovering parses during inference.During inference, we can use a top-down greedy decoding algorithm to recover the unique parse tree encoded in the TreeReg constraints for any given \(S\) from the hidden states of a circuit. Given \(\mathrm{SCIN}\) scores for all spans in \(S\), we recover the tree that maximizes the sum of \(\mathrm{SCIN}\) scores across its spans. Specifically, the split for any span \(S_{i;j}\) happens at the index \(\hat{p}\) that maximizes \(s(i,q,j)\),

\[\hat{p}=\operatorname*{arg\,max}_{q\in[i,j-1]}s(i,q,j). \tag{6}\]

\(S_{i;\hat{p}}\) and \(S_{\hat{p};j}\) can then be recursed on to obtain more constituents of this _induced_ parse tree. Details of this method can be found in Alg. 2.

### Implementation Details

\(\mathrm{SCIN}\) computation does not require additional model calls and can be performed during the forward pass that computes \({\mathcal{L}_{\text{LM}}}\). \(\mathrm{SCIN}\) scores for all spans in a sentence can be calculated simultaneously and efficiently using vectorization, as detailed in Appendix C. Additionally, updates from \({\mathcal{L}_{\text{LM}}}\) and \({\mathcal{L}_{\text{TR}}}\) can occur at different frequencies during training, for example, \({\mathcal{L}_{\text{TR}}}\) may be applied once every \(k\) steps of \({\mathcal{L}_{\text{LM}}}\). Despite the quadratic time complexity of computing \({\mathcal{L}_{\text{TR}}}\) (which scales with the square of input token count), these factors collectively ensure that TreeReg remains efficient in practice. As per our implementation, there is an increase of around \(25\%\) in training time if \({\mathcal{L}_{\text{TR}}}\) is applied once every 10 \({\mathcal{L}_{\text{LM}}}\) steps, on \(25\%\) of attention heads. This excludes the time required for parsing the data used in TreeReg. The parses are either included with the dataset (BLLIP-LG, MultiNLI) or can be computed once and reused for every training run.

## 4 Warm-up: Improving Grokking on Sentence Transformation Tasks

We start by training transformer LMs on two diagnostic tasks derived from PCFGs (See Appendix D for examples and data statistics). In Tense Inflection (TI), the model is provided with an input in the past tense, and required to generate the same input in the present tense. In Question Formation (QF), the model is required to transform a declarative sentence into a question. For QF, we report first word accuracy of the decoded question, and for TI, we report the fraction of test inputs for which the target verb is correctly inflected.

Setup.We train 4-layer transformer LMs for 500k steps (Base LM), performing extended training to enable grokking as reported in prior work Murty et al. (2023a). For TreeReg LM, we auto-parse both datasets with the Berkeley Neural Parser Benepar et al. (2022), and use \({\mathcal{L}_{\text{TR}}}\) on 2 out of 8 attention heads using hidden states from layer 2 of the LM, once every 20 steps of \({\mathcal{L}_{\text{LM}}}\). To account for variance across training runs, we report averages as well as best performance at the end of training across 4 seeds (Table 1). We also report the training iteration (averaged over all runs) after which the performance on the test set converges.

Results.LMs trained with TreeReg grok faster and achieve higher performance than Base LMs on average, across datasets. Notably, on QF, TreeReg LM obtains a 57.5 pt gain over a standard LM (generalizing perfectly) while also grokking over 10 times faster on average.

\begin{table}
\begin{tabular}{|c c c c|} \hline \hline
**Model** & **Avg. Acc. (\(\uparrow\))** & **Best Acc. (\(\uparrow\))** & **itr. (\(\downarrow\))** \\ \hline  & \multicolumn{3}{c|}{_Dense Inflection (TI)_} \\ Base LM & 47.2 \(\pm\) 16.7 & 71.1 & 427k \(\pm\) 41k \\ TreeReg LM & **90.4 \(\pm\) 6.3** & **98.3** & **391k \(\pm\) 35k** \\  & \multicolumn{3}{c|}{_Question Formation (QF)_} \\ Base LM & 42.1 \(\pm\) 15.4 & 66.9 & 460k \(\pm\) 7k \\ TreeReg LM & **99.6 \(\pm\) 0.7** & **100.0** & **43k \(\pm\) 26k** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Evaluations on TI and QF tasks. We report averaged test accuracy across seeds (**Avg. Acc.**), the best tester accuracy (**Best Acc.**), and the average iteration after which test performance converges (**itr.**). TreeReg LMs grok faster and achieve better performance.

 



## 5 Sentence-Level Language Modeling

Can TreeReg provide gains when introduced in pre-training? We investigate this under settings where data passed to \({\cal L}_{\text{LM}}\) and \({\cal L}_{\text{TR}}\) is the same (SS 5.1) and different (SS 5.3).

### Language Modeling on BLLIP-LG

Setup.We train 16-layer LMs on the BLLIP-LG dataset Hu et al. (2020) (Base LM, detailed hyperparameters in Appendix F). For TreeReg LM, we use \({\cal L}_{\text{TR}}\) at layer 12 of the LM, on 2 out of 8 attention heads, once every 10 steps of \({\cal L}_{\text{LM}}\), on the already-parsed BLLIP-LG dataset. A discussion of the layer and attention heads used for \({\cal L}_{\text{TR}}\) can be found in Appendix G. We benchmark syntactic generalization using the BLiMP and SG test suites (See SS 2). To evaluate out-of-distribution generalization, we report perplexity on the Penn TreeBank PTB; Marcus et al. (1993) test set.

To the best of our knowledge, all prior approaches to induce hierarchical computations in autoregressive causally-trained transformer language models do not leave the underlying transformer architecture untouched, and thus are not directly comparable with TreeReg. Nevertheless, we also report results for Parsing as Language Model PLM; Qian et al. (2021), Transformer Grammars TG; Sartran et al. (2022), and Pushdown Layers Pushdown LM; Murty et al. (2023b) in Table 2.

Results.TreeReg achieves a 2.6 point (pt) gain on BLiMP and 8.1 pt overall gain on SG test suites. Figure 2 shows TreeReg improves over the baseline on 4 out of 6 SG test suites, with a notable 17 pt gain on Licensing. TreeReg LM also generalizes better to PTB with 9.2% lower perplexity, with a marginal 0.7 pt perplexity increase on BLLIP-LG.

While PLM, TG and Pushdown LM outperform TreeReg at syntactic generalization on BLiMP and SG test suites, they require additional parameters or incur inference overheads. Except for Pushdown LM, these models also have substantially higher perplexities on the BLLIP-LG test set compared to TreeReg, showing the benefits of the unchanged transformer architecture.

Sample Efficiency Setup.To measure sample efficiency gains from TreeReg, we train 16-layer LMs with the same configuration as SS 5.1 on [10, 50, 100]% of BLLIP-LG. Only parses of the data used for LM is used for TreeReg in each case. We present syntactic generalization results on SG in Figure 3.

Results.While syntactic generalization improves with more training data for both the Base and

Figure 2: Comparing TreeReg LM with Base LM from Table 2 on SG test suites. TreeReg LM outperforms the Base LM on 4 out of 6 test suites, with 1 tie.



TreeReg LMs, TreeReg LMs outperform the baseline across all settings. Moreover, the performance gap widens as the amount of training data increases. Notably, TreeReg surpasses the baseline's maximum syntactic generalization performance on SG by 4.4 points when trained on only 50% of the data, demonstrating substantially more sample-efficient syntactic generalization. Therefore, even though TreeReg increases training time by 25% compared to a baseline trained for the same number of iterations (see SS 3.3), it reaches better syntactic generalization in only 62.5% of the baseline's training time.

Language Modeling on WikiText Setup.Auto-parsing text corpora may not always be practical for large-scale language modeling, particularly when the corpus is very large or when the text lacks clear syntactic structure, as is the case for web data. Therefore, we experiment with a setting where the pre-training and parsed data come from different data sources. We train a 12-layer GPT-2-small model (Base LM) on chunks of 1024 tokens from WikiText-103 [2] with the exact hyper-parameters and tokenization of GPT-2 small (with the addition of dropout (\(=0.1\)) to prevent overfitting). For TreeReg LM, we reuse the BLLIP-LG parses and use \({\mathcal{L}_{\text{TR}}}\) at layer 4, on 3 out of 12 attention heads, once every 10 steps of \({\mathcal{L}_{\text{LM}}}\). We also test a setting where LM training is additionally performed on one batch from BLLIP-LG for every 20 batches from WikiText-103 (Table 3) for both Base LM and TreeReg LM. PTB is again used to evaluate out-of-distribution generalization.

Results.We observe improvements in syntactic generalization with TreeReg in both settings, obtaining up to 2.6 pt gain on BlLiMP and 9.5 pt overall gain on SG test suites. The gains are relatively smaller when LM training is not performed on BLLIP-LG, and perplexities on PTB increase by 24% over the baseline. In this setting, we find that perplexities are also high on BLLIP-LG, averaging around 300. As a result, we conjecture that the hidden states used in TreeReg are not adequately trained in this setting, leading to worse out-of-distribution perplexities.

In the other setting, Base LM struggles to utilize the LM training on BLLIP-LG, resulting in a 2.2 pt overall drop on SG. In contrast, TreeReg shows an overall increase of 3 pt on SG test suites. Perplexities on PTB are again 2.4 pts lower with TreeReg in this setting compared to Base LM, showing better generalization from the BLLIP-LG LM training.

## 6 Continued Pre-training and Fine-tuning

Next, we apply TreeReg to LLMs. In particular, we use the Sheared Llama-1.3B model of Xia et al. (2024) and explore TreeReg under continued pre-training and fine-tuning settings.

### Continued Pre-training on BLLIP-LG Setup.

We perform continued pre-training of Sheared Llama-1.3B (Base LM) on BLLIP-LG, with a context window of 512 tokens, to get CPT LM. For TreeReg LM, BLLIP-LG parses are used, and \({\mathcal{L}_{\text{TR}}}\) is applied once every 5 steps of \({\mathcal{L}_{\text{LM}}}\), on 6 out of 24 attention heads at layer 16 of the LLM. We report scores on the BlLiMP and SG test suites and perplexity on PTB (Table 4).

Figure 3: Plot of Syntactic Generalization on SG test suites vs Percentage of BLLIP-LG data used to train LMs from scratch. TreeReg LM exceeds the maximum syntactic generalization performance of Base LM with less than 50% of the data.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline \multirow{2}{*}{**Model**} & \multicolumn{2}{c}{**Syntactic Generalization (\(\uparrow\))**} & \multicolumn{2}{c}{**PPL (\(\downarrow\))**} \\ \cline{2-5}  & **BLIMP** & **SG** & **WikiText** & **PTB** \\ \hline \multicolumn{5}{r}{_No LM on BLLIP-LG_} \\ \hline Base LM & 71.0 & 69.4 & **17.5** & **331.5** \\ TreeReg LM & **72.5** & **73.7** & 17.8 & 411.1 \\ \hline \multicolumn{5}{r}{_LM on BLLIP-LG_} \\ \hline Base LM & 71.5 & 67.2 & **18.3** & 53.2 \\ TreeReg LM & **74.1** & **76.7** & 19.7 & **50.8** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Syntactic Generalization on BlLiMP and SG test suites, and perplexities on the WikiText and PTB test sets for GPT-2-small models trained on WikiText, with and without LM on BLLIP-LG. TreeReg LMs show better syntactic generalization in both settings.

 

Results.TreeReg results in a 1.2 pt gain on BLiMP and 1.6 pt overall gain on SG test suites, along with a 0.9 pt decrease in PTB perplexity. These improvements are less significant than those reported in SS 5.1, which we conjecture is due to the greater challenge in rewiring the unstructured inductive biases learned by the LLM during large-scale pre-training.

### Fine-tuning on MultiNLI

Can TreeReg prevent LLMs fine-tuned on classification tasks from learning spurious shortcuts? To answer this, we consider Natural Language Inference (NLI), a task that involves classifying the logical relationship between a _premise_ and a _hypothesis_ sentence as _entailment_, _contradiction_, or _neutral_. Prior work McCoy et al. (2019); Geiger et al. (2020) suggests that models often learn spurious shortcuts when trained on broad coverage NLI datasets, causing them to fail on diagnostic out-of-distribution datasets.

Setup.We finetune Sheared Llama-1.3B (Base LM) on MultiNLI Williams et al. (2018) (training details and evaluation methodology in Appendix E) to get FT LM. For TreeReg LM, we binarize the parses of the premise and hypothesis sentences provided in the dataset and apply \({\mathcal{L}}_{\text{TR}}\) at layer 16 on 6 out of 24 attention heads, every 10 steps of \({\mathcal{L}}_{\text{LM}}\). We report accuracy on the MultiNLI test set, which includes a _matched_ split from the same sources as the training data, and a _mismatched_ split from different sources. We also present results (Table 5) on two diagnostic datasets: MoNLI Geiger et al. (2020) and MED Yanaka et al. (2019).

Results.We note improvements of up to 2.2 pt on MultiNLI from TreeReg. Next, we find a decrease of 48.4 pt on MoNLI and 6.9 pt on MED when Base LM is finetuned on MultiNLI. With TreeReg, the finetuned model is substantially more robust, with a moderate decrease of 6.8 pt on MoNLI and 4.2 pt on MED compared to Base LM. We conclude that TreeReg discourages the LM from learning spurious shortcuts for this task.

## 7 Analysis

Parsing.TreeReg induces a tree-structured inductive bias in transformer LMs. We explore how well these induced trees align with given constituency parses. In particular, we use the procedure detailed in SS 3.2 to recover induced parses from TreeReg LM (SS 5.1) and TreeReg CPT LM (SS 6.1) trained on BLLIP-LG. We evaluate parsing on the auto-parsed BLLIP-LG and the manually annotated PTB test sets. Additionally, to estimate parsing on domains other than Newswire, we evaluate parsing on the 4000 Questions dataset Judge et al. (2006). Since TreeReg induces unlabeled binarized parses, we report unlabeled F1 scores against binarized parses in Table 6.

Although TreeReg only implicitly encodes constituency parses in transformer hidden states, we observe that it demonstrates remarkable parsing capabilities. Parses induced by TreeReg align closely with given parses on all datasets, with F1 scores near or above 90. Parsing per

\begin{table}
\begin{tabular}{l c c c} \hline \hline
* & **94.7** & 87.7 \\ \hline \hline \end{tabular}
\end{table}
Table 6: Unlabeled F1 scores against parses from the BLLIP-LG, PTB and 4000 Questions (4kQ) test sets for TreeReg LMs trained from scratch (**TreeReg LM**) and continued pre-trained from Sheared Llama-1.3B (**TreeReg CPT LM**) on BLLIP-LG. F1 scores are near or above 90 for all datasets. We also present unlabeled F1 scores from Benepar Kitaev et al. (2022) for comparison.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline  & \multicolumn{2}{c}{**Static Generalization** (\(\uparrow\))} & \multicolumn{2}{c}{**PPL** (\(\downarrow\))} \\ \cline{2-5} Model & **BLiMP** & **SG** & **BLILP** & **PTB** \\ \hline Base LM & 73.6 & 80.7 & 23.8 & 42.5 \\ \hline CPT LM & 80.7 & 83.9 & **9.24** & 15.2 \\ TreeReg CPT LM & **81.9** & **85.5** & 9.41 & **14.3** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Syntactic Generalization on BLiMP and SG test suites, and perplexities on BLLIP-LG and PTB test sets for Sheared Llama-1.3B (**LM**) continued pre-trained (**CPT**) on BLLip-LG. With TreeReg, CPT LM shows better syntactic generalization and PTB perplexities.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline \multirow{2}{*}{**Model**} & \multicolumn{2}{c}{**MultiNLI** (\(\uparrow\))} & \multicolumn{2}{c}{**Adversarial** (\(\uparrow\))} \\ \cline{2-5}  & **Matched** & **Mismatched** & **MoNLI** & **MED** \\ \hline Base LM & 35.5 & 35.1 & 50.3 & 50 \\ \hline FT LM & 65.9 & 66.3 & 1.9 & 43.1 \\ TreeReg FT LM & **68.1** & **68.0** & **43.5** & **45.8** \\ \hline \hline \end{tabular}
\end{table}
Table 5: NLI accuracies on MultiNLI test splits and two **Adversarial** NLI evaluation datasets for Sheared Llama-1.3B (**LM**) finetuned (**FT**) on MultiNLI. TreeReg results in more gains on MultiNLI and lower decreases on the adversarial benchmarks.

 

formance on PTB and 4000 Questions is slightly better with TreeReg CPT LM, likely due to improved representations from Sheared Llama-1.3B (which benefits from large-scale pre-training) on these datasets, that are relatively out-of-distribution from the BLLIP-LG training data. Notably, induced parses from TreeReg align better with silver parses than Benepar Kitaev et al. (2022) on the 4000 Questions datasets, by 6.8 F1 pts. We provide some example parses on each dataset in Appendix H.

Emergent structure across layers.We investigate the evolution of tree structure across layers in models trained with TreeReg. To do this, we first infer trees for the BLLIP-LG test set from the circuit formed by the attention heads used to train TreeReg, at each layer of the 16-layer TreeReg LM from SS 5.1 (trained with \({\mathcal{L}_{\text{TR}}}\) at layer 12). We report unlabeled F1 scores of parses recovered at every layer against given parses in Figure 4.

The F1 scores reveal that at layer 1, the parses are nearly random. As we progress to layer 12, they increasingly align with given parses. Beyond layer 12, there is a sharp drop in F1 scores back to near-random by layer 16. This suggests that the circuit involved in \({\mathcal{L}_{\text{TR}}}\) becomes increasingly aligned with given parses up to layer 12, and then rapidly loses this structure. We conclude that to perform next-word prediction, TreeReg LMs implement a non-syntactic function over the syntactic representations built at layer 12. These syntactic representations are, in turn, build gradually over the preceding layers of the transformer, even in the absence of explicit supervision from TreeReg at these layers.

Training on Randomized Parses.To assess if the gains from TreeReg are specifically due to hierarchical computation that align closely with constituency parses, we conduct a controlled experiment. We train a 16-layer LM using the exact setup described in SS 5.1, but instead of using gold parses, we provide TreeReg with random parses of BLLIP-LG sentences (Randomized TreeReg LM). These randomized parses are generated top-down recursively, by choosing random split points for each sentence span. We then compare performance with the TreeReg LM and Base LM on the BLiMP and SG test suites (Table 7).

As expected, overall performance on the SG test suites drops marginally compared to Base LM when randomized parses are fed to TreeReg. On the other hand, BLiMP accuracies increase by 1.2 pts even in this setting. We hypothesize that this discrepancy arises from BLiMP's less granular sentence-level perplexity-based evaluation, in contrast to the surprisal-based approach used in SG.

Dependence on amount of parsed data.We train 16-layer LMs on BLLIP-LG using TreeReg, with the same setup as SS 5.1, but vary the percentage of parsed BCLlip-LG used for TreeReg across [1; 5; 10; 20; 40; 60; 80; 100]%. We report the overall SG test suite performance for these settings in Figure 5.

As the amount of parsed data provided to TreeReg increases, we see a general trend of improvement in syntactic generalization on the SG test suites. Remarkably, we see improvements over Base LM from TreeReg even when it is provided with parses for only 1% of the data used for LM.

Figure 4: Unlabeled F1 scores on the BLLIP-LG test set for parse trees induced from every layer of a 16-layer TreeReg LM trained on BLLIP-LG with \({\mathcal{L}_{\text{TR}}}\) at layer 12. Circuits become increasingly tree-structured till layer 12, then rapidly become unstructured.

 



## 8 Other Related Work

Generalization Failures of Transformer LMs.A substantial body of research has pointed out the limitations of transformer-based LMs in achieving robust compositional generalization. They have been shown to struggle with simple tasks, such as determining the parity of binary strings Bhattamishra et al. (2020); Chiang and Cholak (2022), balancing bracket sequences Hahn (2020), as well as more complex diagnostic natural language tasks McCoy et al. (2020); Geiger et al. (2020); Dziri et al. (2023). Large-scale pre-training does not fully resolve these issues. For instance, Berglund et al. (2024) shows that GPT-4 Achiam et al. (2023) fails to generalize on bijective relationships (for example, answering "Who is Tom Cruise's mother?" correctly but not "Who is Mary Lee South's son?"). Hale and Stanojevic (2024) also find that Gemini Pro Team et al. (2023) does not learn syntactic universals such as the Final-over-Final condition Sheehan et al. (2017) for low-resource languages such as Basque.

Tree-structure in Neural Networks.Building on theories of tree-structured human language processing, prior work explores various tree-structured model architectures Socher et al. (2013); Tai et al. (2015); Le and Zuidema (2015); Dyer et al. (2016); Shen et al. (2019); Hu et al. (2021). These models often excel at data-efficient compositional generalization due to their inherently hierarchical computation. However, they have been largely overshadowed by non-hierarchical pre-trained transformer-based LLMs Brown et al. (2020); Achiam et al. (2023).

Linguistic Structure in Transformers.One promising alternative to tree-structured models is to inject syntactic inductive biases into transformers either by jointly modeling sentences and parse structure Qian et al. (2021); Sartran et al. (2022); Murty et al. (2023b), among others), or via constraints on self-attention Strubell et al. (2018); Wang et al. (2019); Deshpande and Narasimhan (2020); Sartran et al. (2022). Although these models bring data efficiency and generalization improvements, they often require additional parameters, impose rigid constraints on attention mechanisms, or complicate inference. To the best of our knowledge, TreeReg is the first approach to introduce explicit, albeit soft, syntactic biases into the transformer without modifying the architecture. This is achievable because standard transformers can encode hierarchical languages of bounded depth Yao et al. (2021) when trained appropriately.

## 9 Conclusion

We propose TreeReg, a structured regularizer that injects soft syntactic inductive biases into given circuits of a transformer LM, converting constituency parses of input sentences into differentiable orthogonality constraints on vector hidden states. Without requiring any architectural changes to the transformer, models pre-trained with TreeReg achieve up to 10% lower perplexities on out-of-distribution data and enhance syntactic generalization by up to 9.5 points on standard test suites. When applied to pre-trained LLMs, TreeReg improves syntactic generalization during continued pre-training and mitigates degradation on adversarial NLI benchmarks by up to 41.2 points when used during fine-tuning. TreeReg more than doubles the sample efficiency of syntactic generalization and does not require the entire training dataset to be parsed, making it practical and efficient. To the best of our knowledge, this work is among the first to translate insights from a mechanistic interpretability-related work Murty et al. (2023c) into an actionable modeling innovation. We leave unsupervised alternatives to TreeReg, pre-training at scale with TreeReg, and application of TreeReg to languages other than English as future work.

Figure 5: Syntactic Generalization on SG test suites vs Percentage of parsed BLILIP-LG data provided to TreeReg, for LMs trained from scratch on BLILIP-LG. Even with 1% of the data, TreeReg LMs have better syntactic generalization than baseline LMs.





### Limitations

TreeReg requires constituency-parsed datasets, which might be difficult to obtain for languages other English due to the lack of easily available constituency parsers. This also makes TreeReg inapplicable to languages that do not have constituency structure, and we limit the experiments presented here to English. TreeReg also introduces some additional hyperparameters that require tuning, such as the layer and subset of attention heads on which TreeReg is be applied. As a heuristic, applying TreeReg at the middle layer of the model, on 25% of the attention heads, tends to work well across settings (see our ablations in Appendix G for more details). Finally, while we optimize TreeReg computation through vectorization, it still adds some computational overhead during training.

## Acknowledgments

We thank Ananth Agarwal, Jasper Jian, Derek Chong, Harshit Joshi, Martijn Bartelds and other members of the Stanford NLP Group, and the reviewers for discussions and feedback.

## References

* Achiam et al. (2023) Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. GPT-4 technical report. arXiv preprint arXiv:2303.08774. Cited by: SS1.
* Z. Allen-Zhu and Y. Li (2023)Physics of language models: part 1, learning hierarchical language structures. arXiv preprint arXiv:2305.13673. Cited by: SS1.
* L. Berglund, M. Tong, M. Kaufmann, M. Balseni, A. C. Stickland, T. Korbak, and O. Evans (2024)The reversal curse: LLMs trained on "A is B" fail to learn "B is A". In Proceedings of the International Conference on Learning Representations, External Links: Link Cited by: SS1.
* S. Bhattacharya, K. Ahuja, and N. Goyal (2020)On the Ability and Limitations of Transformers to Recognize Formal Languages. In Proceedings of the Conference on Empirical Methods in Natural Language Processing, External Links: Link Cited by: SS1.
* T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, Arvind Neelakantan, P. Shyram, G. Sastry, A. Asskell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei (2020)Language models are few-shot learners. In Proceedings of the International Conference on Neural Information Processing Systems, External Links: Link Cited by: SS1.
* D. Chiang and P. Cholak (2022)Overcoming a theoretical limitation of self-attention. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. . Cited by: SS1.
* J. Cocke (1969)Programming languages and their compilers: preliminary notes. New York University. Cited by: SS1.
* S. Crain and M. Nakayama (1987)Structure dependence in grammar formation. Language522-543. Cited by: SS1.
* A. Deshpande and K. Narasimhan (2020)Guiding attention for self-supervised learning with transformers. In Findings of the Association for Computational Linguistics: EMNLP, Cited by: SS1.
* A. Dubey, A. Jauhr, A. Pandey, A. Kadian, A. Al-Dahale, A. Letman, A. Mathur, A. Schelten, A. Yang, A. Fan, et al. (2024)The llama 3 herd of models. arXiv preprint arXiv:2407.21783. Cited by: SS1.
* C. Dyer, A. Konucoro, M. Ballesteros, and N. A. Smith (2016)Recurrent neural network grammars. In Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, External Links: Link Cited by: SS1.
* N. Dziri, X. Lu, M. Solar, X. Lorraine) Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi (2023)Faith and fate: limits of transformers on compositionality. In Proceedings of the International Conference in Neural Information Processing Systems, External Links: Link Cited by: SS1.
* A. Eriguchi, K. Hashimoto, and Y. Tsuruoka (2016)Tree-to-sequence attentional neural machine translation. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), External Links: Link Cited by: SS1.
* J. Gauthier, J. Hu, E. Wilcox, P. Qian, and R. Levy (2020)Syntaxgym: an online platform for targeted evaluation of language models. In Proceedings of the Annual Meeting of the Association for Computational Linguistics: System Demonstrations, External Links: Link Cited by: SS1.



Yinuo Guo, Zeqi Lin, Jian-Guang Lou, and Dongmei Zhang. 2020. Hierarchical poset decoding for compositional generalization in language. In _Proceedings of the International Conference on Neural Information Processing Systems_.
* Hahn (2020) Michael Hahn. 2020. Theoretical limitations of self-attention in neural sequence models. _Transactions of the Association for Computational Linguistics_, 8:156-171.
* Hale et al. (2018) John Hale, Chris Dyer, Adhiguna Kuncoro, and Jonathan Brennan. 2018. Finding syntax in human encephalopathy with beam search. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_.
* Hale and Stanojevic (2024) John T. Hale and Milos Stanojevic. 2024. Do LLMs learn a true syntactic universal? In _Proceedings of the Conference on Empirical Methods in Natural Language Processing_.
* Hu et al. (2020) Jennifer Hu, Jon Gauthier, Peng Qian, Ethan Wilcox, and Roger Levy. 2020. A systematic assessment of syntactic generalization in neural language models. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics_.
* Hu et al. (2024) Xiang Hu, Pengyu Ji, Qingyang Zhu, Wei Wu, and Kewei Tu. 2024. Generative pretrained structured transformers: Unsupervised syntactic language models at scale. _arXiv preprint arXiv:2403.08293_.
* Hu et al. (2021) Xiang Hu, Haitao Mi, Zuijie Wen, Yafang Wang, Yi Su, Jing Zheng, and Gerard de Melo. 2021. R2D2: Recursive transformer based on differentiable tree for interpretable hierarchical language modeling. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics and the International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_.
* Judge et al. (2006) John Judge, Aoife Cahill, and Josef van Genabith. 2006. QuestionBank: Creating a corpus of parse-annotated questions. In _Proceedings of the International Conference on Computational Linguistics and Annual Meeting of the Association for Computational Linguistics_.
* Kasami (1966) Tadao Kasami. 1966. An efficient recognition and syntax-analysis algorithm for context-free languages. _Coordinated Science Laboratory Report no. R-257_.
* Kitaev et al. (2022) Nikita Kitaev, Thomas Lu, and Dan Klein. 2022. Learned incremental representations for parsing. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_.
* Le and Zuidema (2015) Phong Le and Willem Zuidema. 2015. Compositional distributional semantics with long short term memory. In _Proceedings of the Joint Conference on Lexical and Computational Semantics_.
* Marcus et al. (1993) Mitch Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. 1993. Building a large annotated corpus of english: The penn treebank. _Computational linguistics_, 19(2):313-330.
* McCoy et al. (2020) R Thomas McCoy, Robert Frank, and Tal Linzen. 2020. Does syntax need to grow on trees? sources of hierarchical inductive bias in sequence-to-sequence networks. _Transactions of the Association for Computational Linguistics_, 8:125-140.
* McCoy et al. (2019) Tom McCoy, Ellie Pavlick, and Tal Linzen. 2019. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics_.
* Murty et al. (2023a) Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. 2023a. Grokking of hierarchical structure in vanilla transformers. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_.
* Murty et al. (2023b) Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. 2023b. Pushdown layers: Encoding recursive structure in transformer language models. In _Proceedings of the Conference on Empirical Methods in Natural Language Processing_.
* Murty et al. (2023c) Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher D Manning. 2023c. Characterizing intrinsic compositionality in transformers with tree projections. In _Proceedings of the International Conference on Learning Representations_.
* Pallier et al. (2011) Christophe Pallier, Anne-Dominique Devauchelle, and Stanislas Dehaene. 2011. Cortical representation of the constituent structure of sentences. _Proceedings of the National Academy of Sciences_, 108(6):2522-2527.
* Qian et al. (2021) Peng Qian, Tahira Naseem, Roger Levy, and Ramon Fernandez Astudillo. 2021. Structural guidance for transformer language models. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics and the International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_.
* Sartran et al. (2022) Laurent Sartran, Samuel Barrett, Adhiguna Kuncoro, Milos Stanojevic, Phil Blunsom, and Chris Dyer. 2022. Transformer grammars: Augmenting transformer language models with syntactic inductive biases at scale. _Transactions of the Association for Computational Linguistics_, 10:1423-1439.
* Sheehan et al. (2017) Michelle Sheehan, Theresa Biberauer, Ian Roberts, and Anders Holmberg. 2017. _The Final-Over-Final Condition: A Syntactic Universal_, volume 76. MIT Press.
* Shen et al. (2019) Yikang Shen, Shawn Tan, Alessandro Sordoni, and Aaron Courville. 2019. Ordered neurons: Integrating tree structures into recurrent neural networks. In _Proceedings of the International Conference on Learning Representations_ 

Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. 2013. Recursive deep models for semantic compositionality over a sentiment treebank. In _Proceedings of the Conference on Empirical Methods in Natural Language Processing_.
* Strubell et al. (2018) Emma Strubell, Patrick Verga, Daniel Andor, David Weiss, and Andrew McCallum. 2018. Linguistically-informed self-attention for semantic role labeling. In _Proceedings of the Conference on Empirical Methods in Natural Language Processing_.
* Tai et al. (2015) Kai Sheng Tai, Richard Socher, and Christopher D Manning. 2015. Improved semantic representations from tree-structured long short-term memory networks. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics and the International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_.
* Team et al. (2023) Gemini Team, Rohan Anil, Sebastian Borgeaud, Jeannigste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. 2023. Gemini: a family of highly capable multimodal models. _arXiv preprint arXiv:2312.11805_.
* van Schijndel et al. (2013) Marten van Schijndel, Andy Exley, and William Schuler. 2013. A model of language processing as hierarchic sequential prediction. _Topics in cognitive science_, 5(3):522-540.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In _Proceedings of the International Conference on Neural Information Processing Systems_.
* Wang et al. (2019) Yashian Wang, Hung-Yi Lee, and Yun-Nung Chen. 2019. Tree transformer: Integrating tree structures into self-attention. In _Proceedings of the Conference on Empirical Methods in Natural Language Processing and the International Joint Conference on Natural Language Processing_.
* Warstadt et al. (2020) Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R Bowman. 2020. Blimp: The benchmark of linguistic minimal pairs for english. _Transactions of the Association for Computational Linguistics_, 8:377-392.
* Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel Bowman. 2018. A broad-coverage challenge corpus for sentence understanding through inference. In _Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers)_.
* Xia et al. (2024) Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. 2024. Sheared llama: Accelerating language model pre-training via structured pruning. In _Proceedings of the International Conference on Learning Representations_.
* Yanaka et al. (2019) Hitomi Yanaka, Koji Mineshima, Daisuke Beki, Kentaro Inui, Satoshi Sekine, Lasha Abizaniadze, and Johan Bos. 2019. Can neural networks understand monotonicity reasoning? In _Proceedings of the ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_.
* Yao et al. (2021) Shunyau Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. 2021. Self-attention networks can process bounded hierarchical languages. In _Proceedings of the Annual Meeting of the Association for Computational Linguistics and the International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_.
* Younger (1967) Daniel H Younger. 1967. Recognition and parsing of context-free languages in time n3. _Information and control_, 10(2):189-208.

## Appendix A Discussion of Design Decisions

### Last Hidden State as Representation for Prefix

We use the hidden state for the last token in a span as an easily computable representation of the span. There is some evidence suggesting that in the middle layers of an autoregressive transformer language model, the burden of computation of representations tends towards the right, indicating that intermediate hidden states summarize corresponding prefixes to some extent. For instance, Allen-Zhu and Li, 2023 find that when a transformer is trained on generations from a Probabilistic Context-Free Grammar (PCFG), ancestors of non-terminals are encoded at the hidden state of the last token in the span produced from the non-terminal. Therefore, we leverage the prefix information summarized in the hidden state of the last token of a span and rely on the LM's substantial expressive capabilities to embed TreeReg's orthogonality constraints into its hidden states.

### Orthogonality for Operationalizing Contextual Independence

We choose to operationalize contextual independence between spans as orthogonality of correspoding hidden states, to design an efficiently computable loss function. There exist other ways of operationalizing contextual independence but these are inefficient and often require multiple forward passes on the transformer. For instance, Murty et al., 2023c defines contextual independence as the similarity between the sum of hidden states for a span with unconstrained attention, compared to those from a transformer with masked attention for 

the prefix of the span. Another efficient alternative is using the negative absolute value of cosine similarity between span vector representations to represent contextual independence. This does not work as well, resulting in a syntactic generalization of \(76.4\) on SG test suites (a drop of \(3.6\) points) and \(70.6\) on BLiMP (a drop of \(4.2\) points) when used during the pretraining of a 16-layer LM from scratch on BLLIP-LG (SS 5.1).

TreeReg softly encourages higher \({\rm SCIN}\) for constituents of the parse tree, and lower \({\rm SCIN}\) for other constituents. However, a low \({\rm SCIN}\) score indicates greater similarity between span vectors, due to a smaller orthogonal component between them, which can result in degenerate hidden states. In practice, this is mitigated by opposing pressure from \({\cal L}_{\text{LM}}\), that encourages divergent hidden states for different tokens in the sentence. To confirm this, we compute the average \({\rm SCIN}\) for constituent and non-constituent spans in the PennTreeBank test dataset, for the Base LM and TreeReg LM trained in SS 5.1. For the TreeReg LM, the average \({\rm SCIN}\) for constituents is \(1.29\), and that for non-constituents is \(0.98\). In contrast, for the Base LM, the average \({\rm SCIN}\) for constituents is \(0.66\), and that for non-constituents is \(0.65\). We therefore conclude that the average \({\rm SCIN}\) for non-constituents is substantially positive for a model trained with TreeReg, and in fact is even higher than that for a model trained without TreeReg.

### Computation of \({\rm SCIN}\)

Our method for computation of \({\rm SCIN}\) involves the addition of two terms, the first one being \(\|{\rm orth}(\boldsymbol{h}_{j},\boldsymbol{h}_{i-1})\|\), which captures the contextual independence between the information in a span and its prefix. The second term is \(\|{\rm orth}(\boldsymbol{h}_{j+1},\boldsymbol{h}_{j})\|\), which is a measure of how much the suffix of a span is dependent on the information contained in the span. The information in constituent spans should be mostly self-contained in a tree-structured computation, and therefore the hidden state after the span is expected to be independent of it. While this could result in enforced orthogonalities between consecutive words that may belong to the same constituent, TreeReg only enforces soft constraints, and there is opposing pressure from \({\cal L}_{\text{LM}}\) mitigating any undesirable behavior. We also found that TreeReg improves syntactic generalization even without this future-masking constraint, reaching \(77.4\) on SG test suites (a drop of \(2.6\) points) and \(72.8\) on BLiMP (a drop of \(2.0\) points) when used during pre-training of a 16-layer LM from scratch on BLLIP-LG (SS 5.1).

### Greedy Decoding for Recovering Parses

An alternative to the proposed top-down greedy decoding algorithm for recovering induced parse trees (SS3.2) is a dynamic programming approach like the CKY algorithm [11, 12]. However, while computing \({\cal L}_{\text{TR}}\) during training, we only apply supervision on spans that are possible splits of constituents in the silver parse tree. For example, when applying TreeReg to "the quick brown fox jumped over the river", there is no supervision on the \({\rm SCIN}\) for the span "box jumped over", as "fox" is a part of the constituent "the quick brown fox", and "jumped over" is a part of the constituent "jumped over the river". A globally optimal parsing algorithm such as CKY will also consider \({\rm SCIN}\) for spans that had no supervision on similar spans during training, and these can be arbitrarily large, resulting in worse parsing performance.

## Appendix B Algorithms for Computing TreeReg Loss and Induced Parse Trees

Alg. 1 details the procedure used to compute TreeReg loss for a given sentence. Alg. 2 details the procedure used to recover induced parsed trees from TreeReg-trained models.

``` functionloss(\(SCIN,P,i,j\));//Start:loss(\(SCIN,P,0,n\)) Input : \(SCIN\) : Dictionary of \({\rm SCIN}\) for all spans \(P\): Dictionary of split points of constituents in \(\hat{T}(S)\) \(i\): Starting index of span \(S_{i;j}\) \(j\): Ending index of span \(S_{i;j}\) Output :\(l_{i,j}\), contribution to \({\cal L}_{\text{TR}}\) from \(S_{i;j}\) if\(j-i<1\)then /*Ifspanhaslength1or2,onlyonepossiblesplit*/ return0; scores=\([SCIN(i,q)+SCIN(q+1,j)\)for\(q\in[i,j-1]\) p=\(P(S_{i;j})\) span_loss=\(cross\_entropy(scores,p)\) left_loss=\({\rm loss}(SCIN,P,i,p)\) right_loss=\({\rm loss}(SCIN,P,p+1,j)\)returnleft_loss+span_loss+right_loss 



## Appendix C Efficient Computation of \({\rm SCIN}\)

In this section, we detail our approach to efficiently compute \({\rm SCIN}\) for all spans in a sentence simultaneously through vectorization. For sentence \(S\), let **H** be the tensor containing all hidden states \(\mathbf{h}^{l}_{A,j}\):

\[\textbf{H}=[\mathbf{h}^{l}_{A,1},\ldots,\mathbf{h}^{l}_{A,n}] \tag{7}\]

Note that **H** has dimension \((n,kd)\), where \(d\) is the dimensionality of the hidden states from a single attention head of the model. We create a tensor **O** whose elements correspond to orthogonals dropped from vectors in **H** to every other vector in **H**, with dimensionality \((n,n,kd)\). In PyTorch, this can be implemented as follows:

\[\textbf{H'}=\textbf{H}.{\rm unsqueeze}(\dim=1) \tag{8}\] \[\textbf{D}={\rm torch.sum}(\textbf{H*H'},\dim=-1)\] (9) \[\textbf{O}=\textbf{H}-\textbf{D}.{\rm unsqueeze}(\dim=-1)* \textbf{H'} \tag{10}\]

We then have:

\[{\rm orth}(\mathbf{h}^{l}_{A,j},\mathbf{h}^{l}_{A,i})=\textbf{O}[i,j,:] \tag{11}\]

and:

\[{\rm SCIN}(i,j) =\|\textbf{O}[i-1,j,:]\|\] \[+\|\textbf{O}[j,j+1,:]\| \tag{12}\]

where \(\|\|\) is the \(L_{2}\) norm as before.

## Appendix D Dataset Statistics

In this section, we provide statistics for all datasets used in our experiments in Table 8. We also provide some examples from the Tense Inflection, Question Formation, MultiNLI, MoNLI and MED datasets.

### Tense Inflection Examples

* **Input.** my quail upon my peacocks **remembered** her unicorns.
* **Output.** my quail upon my peacocks **remembers** her unicorns.
* **Input.** the vultures upon my unicorn **waited**.
* **Output.** the vultures upon my unicorn **wait**.

### Question Formation Examples

* **Input.** my raven that doesn't sleep **does** change.
* **Output. does** my raven that doesn't sleep change?
* **Input.** our orangutans who do swim **don't** eat.
* **Output. don't** our orangutans who do swim cat?

### MultiNLI Examples

* product and geography.
* **Hypothesis.** Product and geography are what make cream skimming work.
* **Label. Neutral**
* **Premise.** How do you know? All this is their information again.
* **Hypothesis.** This information belongs to them.

\begin{table}
\begin{tabular}{p{56.9pt} p{34.1pt} p{34.1pt}} \hline \hline  & function parse\((SCIN,i,j)\); // Start: parse\((SCIN,0,n)\) \\ Input : & \(SCIN\) : Dictionary of \({\rm SCIN}\) for all spans, \(i\): Starting index of span \(S_{i;j}\) \\ \(j\): Ending index of span \(S_{i;j}\) \\ Output : Constituents in parse tree induced by \({\rm SCIN}\) for \(S_{i;j}\) if\(i==j\)then /* If span has length 1, terminate recursion */ return \(S_{st;st}\); split_point = \(argmax_{q\in[i,j-1]}\) \(SCIN(i,q)+SCIN(q+1,j)\) left_constituents = parse\((SCIN,i,split\text{$Split\_point$})\) right_constituents = parse\((SCIN,split\text{$Split\_point$}+1,j)\) return \(S_{i;j}\) + left_constituents + right_constituents \\ \hline \end{tabular}
\end{table}
Table 8: Statistics for each dataset used in our paper. \({}^{\dagger}\): Number of data points. *: Number of tokens when tokenized using the GPT-2 tokenizer.



* **Label. Entailment**
* **Premise.** Fun for adults and children.
* **Hypothesis.** Fun for only children.
* **Label. Contradiction**

### MoNLI Examples

* **Premise.** The man does not own a dog.
* **Hypothesis.** The man does not own a mammal.
* **Label. Neutral**
* **Premise.** The man does not own a mammal.
* **Hypothesis.** The man does not own a dog.
* **Label. Entailment**

### MED Examples

* **Premise.** No delegate finished the report on time.
* **Hypothesis.** No delegate finished the report.
* **Label. Neutral**
* **Premise.** Some delegates finished the survey on time.
* **Hypothesis.** Some delegates finished the survey.
* **Label. Entailment**

## Appendix E NLI Setup

We now discuss the training and evaluation setup used for our NLI experiments. We format the training examples using the following prompt:

Determine if the hypothesis is an entailment, contradiction or neutral in relation to the premise. If the hypothesis directly follows from the premise, it is an entailment. If the hypothesis directly contradicts the premise, it is a contradiction. Any relationship that does not fit in with the above definitions is considered neutral. Return A if the second sentence is an entailment, B if it is a contradiction, and C if it is neutral.

#### Premise

{premise}

##### Hypothesis

{hypothesis}

##### Answer

We frame training as a classification task on the final logit from the LLM when the prompt above is passed as input. Instead of using a classification head, we directly apply cross-entropy loss over the log-probabilities of 'A', 'B', and 'C' as possible continuations of the prompt, with the groundtruth class label as the target. During evaluation, we select the continuation with the highest log-probability at the final logit, with 'A', 'B' and 'C' as choices. Note that we train on MultiNLI, which includes examples of the entailment, contradiction and neutral labels. On the other hand, MoNLI and MED datasets treat NLI as a two-class problem (entailment or neutral). However, we retain all three classes as options during inference on these datasets, for consistency with training.

## Appendix F Experiment Hyperparameters

In this section, we detail the hyperparameters used for all of our main experiments. All experiments with LMs trained from scratch use tied input and output weight matrices. For all experiments, linear warmup is used for learning rates for the first 10% of training steps, after which cosine decay is used to reach a learning rate of 0 at the end of training. We use the AdamW optimizer with epsilon of 1e-7 and clip gradients to a maximum \(L_{2}\) norm of 1.0.

### SS 4

* Number of layers: 4
* Number of attention heads: 8
* Hidden dimension: 512
* Training steps: 500,000
* Batch size: 8
* Learning rate: 1e-4
* Weight decay: 0.01
* Dropout: 0.1

### SS 5.1

* Number of layers: 16
* Number of attention heads: 8
* Hidden dimension: 512
* Training steps: 100,000
* Batch size: 160
* Learning rate: 1e-4
* Weight decay: 0.01
* Dropout: 0.1 



### $ 5.3
* Number of layers: 12
* Number of attention heads: 12
* Hidden dimension: 768
* Training steps: 40,000
* Batch size: 480
* Learning rate: 6e-4
* Weight decay: 0.1
* Dropout: 0.1

### $ 6.1
* Training steps: 10,000
* Batch size: 32
* Learning rate: 2e-5
* Weight decay: 0.01
* Dropout: 0.1

### $ 6.2
* Training steps: 20,000
* Batch size: 32
* Learning rate: 1e-5
* Weight decay: 0.01
* Dropout: 0.1

## Appendix G Selection of Layer and Number of Attention Heads for TreeReg

What is the ideal layer and number of attention heads for the \({\cal L}_{\text{TR}}\) computation? To explore, we vary the layer used in TreeReg as [2,4,6,8,10,12,14] and the number of attention heads used as [2,4,8] out of 8 in the experiment setup described in Section 5.1 for the training of a 16-layer LM from scratch on the BLLIP-LG dataset. For the resulting models, we report the overall SyntaxGym performance in Figure 6 and PTB perplexities in Figure 7.

We find that applying TreeReg at layer 12 on 2 attention heads results in the highest overall SyntaxGym performance as well as lowest PTB perplexity. With a couple of exceptions (4 attention heads at layer 4 and 6), TreeReg consistently outperforms the Base LM in both SyntaxGym performance and PTB perplexities across hyperparameter settings. Both syntactic generalization and out-of-distribution perplexities show generally improving trends till layer 12, after which they start degrading. This finding suggests that TreeReg LMs represent syntax most naturally at intermediate layers, on a small subset of attention heads.

## Appendix H Examples of Parses Induced by TreeReg

In this section, we provide some examples of parses induced from the TreeRegLM trained in SS5.1 on the BLLIP-LG, PTB and 4000 Questions test sets.

### Bllip-lg

Figure 9 is an example where the induced parse from TreeReg matches the silver parse. Figure 9 and Figure 10 represents a case where the induced parse differs from the silver parse



### Ptb

Figure 11 is an example where the induced parse from TreeReg matches the silver parse. Figure 12 and Figure 13 represents a case where the induced parse differs from the silver parse.

### 4000 Questions

Figure 14 is an example where the induced parse from TreeReg matches the silver parse. Figure 15 and Figure 16 represents a case where the induced parse differs from the silver parse.

## Appendix H

Figure 8: Silver and induced parse for “The company declined to elaborate on the disclosure.”

Figure 10: Induced parse for “Mr. Helguera is a lawyer, economist and rancher from Buenos Aires.”

Figure 12: Silver parse for ‘Four million common shares, via Salomon Brothers Inc.”

Figure 13: Induced parse for “Four million common shares, via Salomon Brothers Inc.”



Figure 16: Induced parse for “What debts did Qintex group leave?”