Sometimes I am a Tree: Data Drives Unstable Hierarchical Generalization

Authors: Tian Qin, Naomi Saphra, David Alvarez-Melis
Primary area: interpretability and explainable AI
Keywords: Language models, simplicity bias, random variations, OOD generalization

Abstract:
Neural networks often favor shortcut heuristics based on surface-level patterns. Language models (LMs), for example, behave like n-gram models early in training. However, to correctly apply grammatical rules, LMs must instead rely on hierarchical syntactic representations rather than on surface-level heuristics derived from n-grams. In this work, we use cases studies of English grammar to explore how latent structures in training data drives models toward improved out-of-distribution (OOD) generalization. We then investigate how data composition can lead to inconsistent behavior across random seeds. Our results show that models stabilize in their OOD behavior only when they commit to either a surface-level linear rule or a hierarchical rule. The hierarchical rule, furthermore, is induced by grammatically complex sequences with deep embedding structures, whereas the linear rule is induced by simpler sequences. When the data contains a mix of simple and complex examples, potential rules compete; each independent training run either stabilizes by committing to a single rule or remains unstable in its OOD behavior. We also identify an exception to the relationship between stability and generalization: Models which memorize patterns from homogeneous training data can overfit stably, with different rules for memorized and unmemorized patterns. While existing works have attributed similar generalization behavior to training objective and model architecture, our findings emphasize the critical role of training data in shaping generalization patterns and how competition between data subsets contributes to inconsistent generalization outcomes.

Decision: reject
Reviewer scores: 5, 8, 3

Reviews and discussion:
<thread>
creation date:
1730045463492

thread_type:
official_review

<thread_object>
creator: reviewer_sgqh

<content>
summary:
This paper investigate generalization of LMs.  Specilically, data composition significantly influences a model's out-of-distribution (OOD) generalization.  Insufficient data diversity can lead models to rely on memorization rather than achieving true generalization. They primarily validate these observations through learning on grammatical tasks.

model: 12M  decoder-only Transformer.
tasks: question formation, tense inflection.

strengths:
The paper is well-written, with comprehensive visuals and thorough experimental results.

weaknesses:
1. The conclusion in this paper have many overlaps with previous studies[1],[2],[3]; please specifiy your difference and contribution detailed.
2. The experiments only contain two dataset: question formation and tense inflection. This limited validation is insufficient to support the overall conclusions of the paper, generally speaking.  More tasks and models are needed.  Suggest tasks: language modeling, mathematics reasoning … …
3. The authors present only the experimental observation that "lower diversity data tends to promote memorization, while higher diversity data encourages generalization." These are  obvious conclusions compared with previous papers. Are there any new insights offering deeper theoretical reasoning or more controlled experiments to support these findings? or fuerther verfiication on LLMs, e.g. 7B ?

### Grammar

Line 113: “. . “

Line 119: multiple abbreviation definitions

Table 1: revise “Left” and “right” to “Top” and “Bottom”

[1] Liu Z, Michaud E J, Tegmark M. Omnigrok: Grokking beyond algorithmic data[C]//The Eleventh International Conference on Learning Representations. 2022.

[2] Zhu X, Fu Y, Zhou B, et al. Critical data size of language models from a grokking perspective[J]. arXiv preprint arXiv:2401.10463, 2024.

[3] Wang B, Yue X, Su Y, et al. Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization[J]. arXiv preprint arXiv:2405.15071, 2024.

questions:
See in Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730121170557

thread_type:
official_review

<thread_object>
creator: reviewer_4elk

<content>
summary:
This paper investigates the influence of training data composition on the generalization behaviors of language models, focusing on their ability to learn hierarchical syntactic representations versus relying on surface-level heuristics such as n-gram models. Through case studies involving English grammar tasks (specifically question formation and tense inflection) the authors show that the complexity and diversity of training data play pivotal roles in determining if models adopt hierarchical rules or simpler linear heuristics. 
Interesting findings of the paper: 
i) Adding sentences with deep syntactic trees in the training data encourages models to develop hierarchical syntactic representations  (and this seems to enable out of distribution generalization).
ii) When the training dataset have a mix of simple and complex grammatical structures, models show unstable training dynamics and inconsistent rule commitments across different random seeds (this aligns with findings by McCoy et al. (2018, 2020))
iii) In case of low data diversity, models tend to memorize patterns without learning robust hierarchical or linear rules, resulting in poor generalization.

I think the main message of this paper is the relevance of training data features in shaping the inductive biases of neural networks.

strengths:
- Interesting analysis of the mechanisms of rule commitment.
- The identification of a memorization regime, where models stabilize without learning either hierarchical or linear rules.
- The findings are replicated across two distinct grammatical tasks, and backed up by linguistics theroies.

weaknesses:
- The experiments utilize relatively small transformer models trained on synthetic datasets with 100K samples (it doesn't make the study less valid, but larger models may exhibit different inductive biases and learning dynamics that are not captured by smaller-scale experiments).
- The authors use a fixed set of hyperparameters (learning rate of 1e-4, Adam optimizer, specific layer configurations) across all experiments, without addressing how variations in hyperparameters might influence the model's ability to learn hierarchical rules or affect training stability.

questions:
- It might be worthy study how other objectives (like MLM) interact with data composition to affect rule learning. 
- Have you considered how the role of data composition might vary with languages that have different syntactic structures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730346626774

thread_type:
official_review

<thread_object>
creator: reviewer_tfb9

<content>
summary:
This study analyzes the role of the syntactic complexity in language models’ acquisition of hierarchical inductive biases. Specifically, by generating sentences with varying syntax tree depths and of varying structures (some with mixed scoping and others with forward scoping), the features that lead to linear or hierarchical generalization are isolated. Additionally, by mixing these types of data, the authors conduct a grokking-style analysis of when stable, unstable, hierarchical, linear, or mixed generalizations are likely to be learned. 

It is found that including sentences with a syntactic tree depth of at least 3 is required for hierarchical generalizations to be stably acquired, and that mixing depths in the training data can lead to unstable generalizations. It is also found that if the data is not sufficiently diverse, memorization (rather than any consistent generalization) is the preferred strategy.

strengths:
* The paper adds to a rich literature on directly evaluating the inductive biases of language models. This paper approaches it from a grokking perspective, which, to my knowledge, has not been done before.
* In Fig. 4, it is interesting that such a small variety of syntactic structures in the declaratives would lead to hierarchical generalization. I had assumed that a more naturalistic and varied distribution would be necessary to prevent memorization.
* Thorough analysis of variation across many random seeds.

weaknesses:
1. The findings could be much better contextualized with related work throughout the paper. For example, at L150: this is not the first work to analyze the effect of data on syntactic generalization. Mueller et al. (2023) (who are currently cited only in the Appendix) do a similar analysis, finding that simpler corpora lead to syntactic generalization with less data. Additionally, when referring to the syntactic transformations task setup at L144, Frank & Mathis (2007) [1] is the correct citation. At L154-161 (and App. F): it would be nice to contextualize this with analyses from McCoy et al. (2020) on instability across random seeds. In Sec. 4.3, these findings are very related to those of Papadimitriou & Jurafsky (2023) [2] (who are not cited) on training with recursive and/or cross-serial dependencies.
2. The findings are somewhat obvious in light of the above citations, and in light of grokking work (which, by the way, should also be better cited and discussed; see [3,4,5]). Data determines generalization because neural networks are a statistical approach to learning. There are probably new insights to be gleaned from this new task setting compared to past grokking work, and it could be nice to explicitly enumerate these in the paper.
3. L212-215: It is not clear why these hyperparameters were chosen, and why different hyperparameters were used for the two tasks.
4. Related to the first point (but more minor), the decision of what to include in the main paper vs. the appendix Related Work sections currently feels arbitrary.

questions:
Questions
===
1. Where did the hyperparameters in L212-215 come from? Was there some tuning involved, and if so, would it be possible to show results for other tested hyperparameters?
2. By syntactic complexity, do you generally mean syntactic tree depth? The rarity/difficulty for humans in processing of particular structures? Number of edges in the tree (which could correlate with sentence length, even if the maximum depth of the tree were low)? I assume the first given the analyses in Sec. 4, but would be nice to explicitly define the term.
3. L420: which studies?
4. At first glance, Fig. 4 and Sec. 6 feel contradictory. There is some interesting nuance here that could be interesting to explore: diversity is necessary to prevent memorization, but at the same time, including a diverse mixture of complex and simple sentences leads to *worse* generalization than simply using more complex sentences. Why is this?

Suggestions/Typos
===
* L182: there are *at least* two strategies. The model could rely on other heuristics, such as always moving the affirmative verb in QF as opposed to the main or first verb (as attested in Mueller et al., 2024).
* Table 1 caption: by “left” and “right”, do you mean “top” and “bottom”?
* L113: two periods

References
===
[1] Robert Frank & Donald Mathis (2007). “Transformational networks.” Models of Human Language Acquisition. https://bpb-us-e2.wpmucdn.com/websites.umass.edu/dist/a/27637/files/2017/06/cogsci-2007.pdf

[2] Isabel Papadimitriou & Dan Jurafsky (2023). “Injecting structural hints: Using language models to study inductive biases in language learning.” Findings of EMNLP. https://aclanthology.org/2023.findings-emnlp.563/

[3] Yifei Huang et al. (2024). “Unified View of Grokking, Double Descent and Emergent Abilities: A Perspective from Circuits Competition.” COLM. https://arxiv.org/abs/2402.15175

[4] Vikrant Varma et al. (2023). “Explaining grokking through circuit efficiency.” https://arxiv.org/abs/2309.02390

[5] Ziming Liu, Eric J. Michaud, & Max Tegmark (2023). “Omnigrok: Grokking Beyond Algorithmic Data.” ICLR. https://arxiv.org/abs/2210.01117
</content>
</thread_object>

</thread>
<thread>
creation date:
1731788974556

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Meta-Response

comment:
We want to thank all the reviewers for their time and valuable feedback! Because some reviewers were unclear on what our novel contributions are, we have realized that our presentation does not sufficiently clarify our new framework or results. We have revised the paper to explain this framework more clearly, and map it out below.

Previous work has investigated the role of diversity in promoting generalization over memorization. Other work has also investigated some data factors that influence the selection of one generalization rule over another. However, none of these works unify all our three themes: (1) complexity AND diversity, (2) memorization/generalization AND rule selection, (3) training stability AND consistency across seeds. This relationship can be summarized through a series of points:

* Diverse data promotes general rules over memorization. This is an extension of existing results in grokking and compositionality to a new setting (i.e., structural grokking), with new precision over the definition of diversity (in our case, latent structure diversity).
  * In our updated version, we have supported and refined this point by confirming that in the memorization regime, models still apply rules consistently to familiar syntactic structures. In the generalization regime, by contrast, models extrapolate those rules to unseen syntactic structures. 

* Complex data promotes complex rules (i.e., rules that rely on latent structure). This is an extension of existing results in OOD generalization from ambiguous training data, though again we have controlled the setting more precisely and found more specific complexity factors (center embeddings) driving hierarchical generalization.

  * The revised paper further narrows down the relevant complexity factor, from mixed scoping broadly to center embeddings specifically (i.e., the "degree 2" embedding structures that were previously only addressed in the discussion section). We have added experiments confirming that center embeddings explain rule selection better than other sources of syntactic complexity.

* Intermediate Diversity and Complexity: Intermediate diversity creates an unstable regime where memorization and rule-based strategies compete. This intermediate diversity regime is a novel discovery, and may be considered surprising, since it is not obvious that intermediate diversity will lead to more instability than either heterogeneous or homogeneous data. Similarly, intermediate complexity promotes instability due to competition between generalization rules. These findings support prior work on oscillatory dynamics and subset competition (e.g., Elan et al. [1]).

Furthermore, we have connected these unstable training dynamics and data properties to the much-overlooked role of random variation across seeds. While some existing work discusses random variation in generalization rules [2] or in the timing of grokking [3], to our knowledge, this is the first paper to make a direct connection between clustered OOD behavior and the likelihood of unstable training. This connection yields the horseshoe-shaped curves in Fig. 4. In particular, we show:

* Stable training runs are correlated, across random seeds, with either memorization or systematic application of one of the generalization rules. Models that do not commit to one of these strategies---memorization, linear rule, or hierarchical rule---tend to be unstable.

  * This might be an extension of existing findings: omnigrok [4] shows that in their specific setting, generalization is associated with more stability, but they don't show that memorization can also be stable. They also do not consider correlation across random seeds, meaning that this stability could be a direct effect of their regularizer rather than representing rule commitment more generally under natural variation.

* Data settings that create a higher likelihood of unstable training runs also allow models to learn either rule, determined only by the random initialization seed.

  * This finding addresses a missing gap in the literature. Given findings like Juneja et al. [2]  and Zhou et al. [5], it has thus far been unclear why some settings lead to inconsistent OOD behavior across random seeds. Our findings confirm that the inconsistency is a direct result of intermediate levels of data complexity or data diversity that promote different model strategies.

  * Our work also shines light on why clusters of OOD behavior can form across random seeds, per juneja et al [2]. Our work illustrates that these clusters have formed around stable model strategies and that training runs are unstable unless they find these basins.

  * Clustering effects tie to a number of key ideas in recent literature. Sharp transitions between generalization rules lead to phenomena like emergence, where sudden breakthroughs occur because there is no way to interpolate between rules.

We have emphasized the connections above in our updated presentation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies a particular aspect of how the composition of the training data influences the generalization (vs memorization) behavior of language models. Reviewers appreciated the study of inductive biases of language models. However, there was generally a lack of clarity on exactly what the contributions here are. In their response, the authors have attempted to clarify this. However, this seems to represent a major revision, requiring a resubmission of the paper.