Sneaking Syntax into Transformer Language Models with Tree Regularization (2025)

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 research.

## introduction a substantial body of research @xcite @xcite suggests that human language processing is inherently hierarchical: syntactic rules combine word-level meanings to give rise to semantics for larger constituents, which are 1 https://github.com/ananjan-nandi-9/tree_ regularization then combined into sentences. in contrast, transformers @xcite , 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 @xcite , but recent studies find that state-of-the-art llms still struggle with compositional generalization @xcite @xcite , the ability to understand familiar words in novel contexts. prior work @xcite 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 @xcite 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. possible bracketings of "is very happy" are ("is very", "happy") with score scin(2, 3) + scin(4, 4) and ("is", "very happy") with score scin(2, 2) + scin(3, 4). loss for this span (l (2,4) ) encourages the second bracketing. l @xmath0,5) + l (1,4) + l (2,4) includes analogous losses for spans "he is very happy" and "he is very happy now". the use of treereg improves syntactic generalization across model and data scales. on diagnostic sentence transformation tasks such as tense inflection and question formation @xcite , it significantly enhances grokking behavior ( § 4). it also improves syntactic generalization on the blimp @xcite and syntaxgym @xcite test suites by up to 3.2 and 9.5 points, both when used during pre-training from scratch ( § 5), and during continued pre-training ( § 6.1). lms with treereg surpass the syntactic generalization of baseline lms with less than half the training data ( § 5.2). most importantly, treereg improves outof-distribution language understanding. when used during pre-training of lms on fully parsed ( § 5.1) as well as a mixture of parsed and unparsed ( § 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 @xcite , it mitigates performance drops by up to 41.2 points on adversarial nli benchmarks ( § 6.2). we release all code at https://github. com/ananjan-nandi-9/tree_regularization . ## background and notation tree structures in language processing. consider a sentence @xmath1, x 2 , . . . , x n } where x j is the j th token in the sentence. let tree t (s) over sentence s be a set of spans s a;@xmath2, . . . , x b } such that if s a;b and s b+1;@xmath3, then s a;@xmath4. 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 @xcite @xcite suggests that human language processing is hierarchicalsentences 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 t (s) is the tree form by all constituents of s. treestructured models @xcite 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. ## our approach treereg is a differentiable loss function that injects the hierarchical structure of input constituency parses into transformer circuits. inspired by @xcite , 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 (scin, § 3.1). treereg maximizes scin for spans corresponding to input constituents while simultaneously minimizing it for other spans ( § 3.2). in an optimally trained model, the tree with the highest cumulative 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. ## 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 @xcite . for treereg lm, we autoparse both datasets with the berkeley neural parser @xcite , and use l tr on 2 out of 8 attention heads using hidden states from layer 2 of the lm, once every 20 steps of l lm . to account for variance across training runs, we report averages as well as best performance at the end of training across 4 seeds ( 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. ## sentence-level language modeling can treereg provide gains when introduced in pre-training? we investigate this under settings where data passed to l lm and l tr is the same ( § 5.1) and different ( § 5.3). ## continued pre-training and fine-tuning next, we apply treereg to llms. in particular, we use the sheared llama-1.3b model of @xcite and explore treereg under continued pretraining and fine-tuning settings. model syntactic generalization (↑) ppl (↓) blimp sg bllip ptb base lm 73.6 80.7 23.8 42.5 cpt lm 80.7 83.9 9.24 15.2 treereg cpt lm 81.9 85.5 9.41 14.3 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 § 5.1, which we conjecture is due to the greater challenge in rewiring the unstructured inductive biases learned by the llm during largescale pre-training. ## 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 § 3.2 to recover induced parses from treereg lm ( § 5.1) and treereg cpt lm ( § 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 @xcite . since treereg induces unlabeled binarized parses, we report unlabeled f1 scores against binarized parses in 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-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 @xcite 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 § 5.1 (trained with l 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 l 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, built 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 § 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 topdown 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 ( ## 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 finetuning. 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 @xcite ) 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.