Rule Extrapolation in Language Models: A Study of Compositional Generalization on OOD Prompts (2409.13728v2)

Abstract: LLMs show remarkable emergent abilities, such as inferring concepts from presumably out-of-distribution prompts, known as in-context learning. Though this success is often attributed to the Transformer architecture, our systematic understanding is limited. In complex real-world data sets, even defining what is out-of-distribution is not obvious. To better understand the OOD behaviour of autoregressive LLMs, we focus on formal languages, which are defined by the intersection of rules. We define a new scenario of OOD compositional generalization, termed rule extrapolation. Rule extrapolation describes OOD scenarios, where the prompt violates at least one rule. We evaluate rule extrapolation in formal languages with varying complexity in linear and recurrent architectures, the Transformer, and state space models to understand the architectures' influence on rule extrapolation. We also lay the first stones of a normative theory of rule extrapolation, inspired by the Solomonoff prior in algorithmic information theory.

Overview of "Rule Extrapolation in LLMs: A Study of Compositional Generalization on OOD Prompts"

The paper "Rule Extrapolation in LLMs" by Mészáros et al. explores an intricate facet of autoregressive LLMs (AR LLMs): their capacity for compositional generalization in out-of-distribution (OOD) contexts, termed as "rule extrapolation." This work explores how well these models can extrapolate known rules when confronted with OOD prompts—phrases or sequences that violate defined rules inherent within a formal language framework—by leveraging formal languages. The paper presents comprehensive empirical analyses across several architectural paradigms, forming a robust discussion on model behavior in OOD environments.

Rule Extrapolation: Theoretical Framework

The authors establish "rule extrapolation" as a subset of compositional generalization, where LLMs are trained on sequences adhering to multiple rule intersections. Rule extrapolation occurs when a model, encountering a prompt that violates certain rules, still manages to fulfill the unviolated rules in its completion. For instance, in a language defined by sequences where "a" precedes "b" with an equal number of each, a prompt beginning with "b" but maintaining the count equality, albeit with inversion, challenges the model to generalize beyond its training data.

Empirical Investigation and Architectures

This paper evaluates several models - Linear transformers, LSTMs, the state-space model "Mamba," and the xLSTM. The authors extensively experiment across datasets structured by the Chomsky hierarchy - spanning from regular to context-sensitive languages - to understand rule extrapolation capabilities across disparate linguistic complexities. Findings suggest that no single architecture universally excels; Transformers show proficiency in context-sensitive tasks but lag in regular language-based rule extrapolation, whereas LSTMs and Mambas provide robust performance on regular languages.

Data Utilization and Test Dynamics

The paper meticulously designs its datasets around Chomsky's linguistic categorization, allowing for clearly defined rule intersections for model training and well-structured OOD challenge sets for testing rule extrapolation. Utilizing these controlled environments, the authors discern distinct patterns in model adaptation and partial rule adherence.

Normative Theory and Algorithmic Insights

In addition to empirical analyses, the paper ventures into algorithmic ideation. Drawing inspiration from Solomonoff's algorithmic information theory, the authors propose a normative theory to rationalize OOD sequence completion. This theory compels models towards simplicity in their extrapolative reasoning, thereby aligning computational behaviors with fundamental algorithmic theories. This is reflective of broader computational trends and biases towards minimal complexity, potentially explaining observed behaviors in predictive tasks.

Implications and Future Prospects

The implications of this paper are manifold. Practically, it underlines the necessity of model-selective paradigms based on task-specific OOD challenges, urging further specificity in architecture development and deployment. Theoretically, it keenly intersects with algorithmic theory to offer novel paths for reasoning about OOD phenomena in learned models. Future developments might explore the nuanced dynamism in rule extrapolation more integrally with natural language datasets or complex real-world scenarios.

Conclusion

Mészáros et al.'s work provides a valuable lens to comprehend how LLMs engage with compositional generalization beyond IV distribution training, advancing both our empirical understanding and theoretical grounding of AI model behaviors in linguistically diverse environments. Through their methodological precision, the authors establish a benchmark for evaluating OOD compositional generalization while contributing a normative framework that complements existing algorithmic theories. This represents a step toward more nuanced, adaptable AI systems capable of reasoning within and beyond predefined linguistic constraints.