Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On Provable Length and Compositional Generalization

Kartik Ahuja, Amin Mansouri

TL;DR

The paper provides the first provable guarantees for length and compositional generalization in sequence-to-sequence models, covering deep sets, transformers, SSMs, and RNNs under structured (limited-capacity) and broader hypothesis classes. It formalizes length generalization and compositional generalization, introduces the RASP-L conjecture to guide when these guarantees hold, and demonstrates linear identification between learned and true representations under realizability. Chain-of-thought supervision is shown to enhance length generalization, especially for high-capacity models, by supplying privileged intermediate computations. The work blends rigorous theory with comprehensive experiments, and extends to finite and infinite hypothesis classes, offering new perspectives on learning under distribution shifts beyond the training support. These results provide a foundational step toward understanding when and how provable OOD generalization can emerge in practical seq-to-seq architectures.

Abstract

Out-of-distribution generalization capabilities of sequence-to-sequence models can be studied from the lens of two crucial forms of generalization: length generalization -- the ability to generalize to longer sequences than ones seen during training, and compositional generalization: the ability to generalize to token combinations not seen during training. In this work, we provide first provable guarantees on length and compositional generalization for common sequence-to-sequence models -- deep sets, transformers, state space models, and recurrent neural nets -- trained to minimize the prediction error. We show that \emph{limited capacity} versions of these different architectures achieve both length and compositional generalization provided the training distribution is sufficiently diverse. In the first part, we study structured limited capacity variants of different architectures and arrive at the generalization guarantees with limited diversity requirements on the training distribution. In the second part, we study limited capacity variants with less structural assumptions and arrive at generalization guarantees but with more diversity requirements on the training distribution. Further, we also show that chain-of-thought supervision enables length generalization in higher capacity counterparts of the different architectures we study.

On Provable Length and Compositional Generalization

TL;DR

The paper provides the first provable guarantees for length and compositional generalization in sequence-to-sequence models, covering deep sets, transformers, SSMs, and RNNs under structured (limited-capacity) and broader hypothesis classes. It formalizes length generalization and compositional generalization, introduces the RASP-L conjecture to guide when these guarantees hold, and demonstrates linear identification between learned and true representations under realizability. Chain-of-thought supervision is shown to enhance length generalization, especially for high-capacity models, by supplying privileged intermediate computations. The work blends rigorous theory with comprehensive experiments, and extends to finite and infinite hypothesis classes, offering new perspectives on learning under distribution shifts beyond the training support. These results provide a foundational step toward understanding when and how provable OOD generalization can emerge in practical seq-to-seq architectures.

Abstract

Out-of-distribution generalization capabilities of sequence-to-sequence models can be studied from the lens of two crucial forms of generalization: length generalization -- the ability to generalize to longer sequences than ones seen during training, and compositional generalization: the ability to generalize to token combinations not seen during training. In this work, we provide first provable guarantees on length and compositional generalization for common sequence-to-sequence models -- deep sets, transformers, state space models, and recurrent neural nets -- trained to minimize the prediction error. We show that \emph{limited capacity} versions of these different architectures achieve both length and compositional generalization provided the training distribution is sufficiently diverse. In the first part, we study structured limited capacity variants of different architectures and arrive at the generalization guarantees with limited diversity requirements on the training distribution. In the second part, we study limited capacity variants with less structural assumptions and arrive at generalization guarantees but with more diversity requirements on the training distribution. Further, we also show that chain-of-thought supervision enables length generalization in higher capacity counterparts of the different architectures we study.
Paper Structure (54 sections, 18 theorems, 118 equations, 4 figures, 3 tables)

This paper contains 54 sections, 18 theorems, 118 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Hypothesis classes with structural assumptions
  3. General hypothesis classes with less structural assumptions
  4. Related Works
  5. Length generalization
  6. Compositional generalization
  7. Length and Compositional Generalization: Definitions and Framework
  8. Learning via expected risk minimization
  9. Contrast with existing works on distribution shifts:
  10. RASP-L conjecture
  11. Structured Hypothesis Classes
  12. Deep sets
  13. Deep sets
  14. Linear identification
  15. High capacity deep sets
  16. ...and 39 more sections

Key Result

Lemma 3

Let $\mathcal{X}\subseteq \mathbb{R}^{n}$. If $f:\mathcal{X} \rightarrow \mathbb{R}^{m}$ and $g:\mathcal{X}\rightarrow \mathbb{R}^{m}$ are continuously differentiable functions that satisfy $f(x)=g(x)$ almost everywhere in $\mathcal{X}$, where $\mathcal{X}$ is a regular non-empty closed set, then $f

Figures (4)

  • Figure 1: Illustrating support of train vs test distribution for (a) compositional generalization and (b) length generalization.
  • Figure 2: Length generalization: Test $\ell_2$ loss on sequences of different lengths. The models are trained only on sequences of length up to $T=10$. All models achieve small error values $\approx 10^{-4}-10^{-7}$ at all sequence lengths and thus length generalize. Since the error values are already quite small, the increasing or decreasing trends are not numerically significant.
  • Figure 3: A transformer model with softmax attention with two hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ length generalizes to sequences of length up to $100$.
  • Figure 4: A failure case of length generalization under arbitrary expressive generative model with (a,b) Deep sets, (c,d) and Transformer. The generative function on both cases introduces an offset to sequences longer than some critical length ($T_0$). The learner is once trained on sequences longer than $T_0$ and successfully generalizes (a,c), and once is trained only on sequences shorter than $T_0$ where the offset never appears, and hence fails to generalize beyond that.

Theorems & Definitions (20)

  • Definition 1
  • Definition 2
  • Lemma 3
  • Theorem 4
  • Theorem 5
  • Theorem 6
  • Theorem 7
  • Theorem 8
  • Theorem 9
  • Theorem 10
  • ...and 10 more