On Provable Length and Compositional Generalization

Authors: Kartik Ahuja, Amin Mansouri
Primary area: learning theory
Keywords: sequence-to-sequence models, length generalization, compositional generalization, out-of-distribution generalization

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. Taking a first principles perspective, we study the realizable case, i.e., the labeling function is realizable on the architecture.  We show that simple limited capacity versions of these different architectures achieve  both length and compositional generalization. Across different architectures, we also find that a linear relationship between the learned representation and the representation in the labeling function is necessary for length and compositional generalization.

Decision: reject
Reviewer scores: 5, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730319755139

thread_type:
official_review

<thread_object>
creator: reviewer_mssc

<content>
summary:
The paper theoretically studies length generalization and compositional generalization in various simplified models: deep sets, simplified one-layer transformers and simple RNNs and state space models. All results are given in the realizable case, i.e. when the ground-truth labels are generated by a model from the hypothesis class, and assuming some capacity restriction on the class. The authors define length generalization to be a setting where the model minimizes the expected risk up to length T and achieves zero generalization error on longer lengths. Similarly, they define compositional generalization of a model minimizing the expected risk on some input distribution and generalizing to the Cartesian product of the support of the different variables. Then, the authors show that the simplified models achieve length and compositional generalization, under assumptions on the capacity of the models. Additionally, the authors complement their results with experiments.

strengths:
To my knowledge, this is the first work defining and analyzing length generalization, and one of the first works to theoretically study compositional generalization. The authors study several different architectures for learning over sequences and sets that are used in practice and show positive results on length and compositional generalization under reasonable assumptions. Additionally, the authors show that without restricting the capacity of the models, length generalization and compositional generalization are not possible.

weaknesses:
-	While the paper shows some novel theoretical results, the setting analyzed is somewhat idealized. Primarily, to my understanding, the authors assume that the model achieves zero error on the training distribution. Typically, however, theory of learning focuses on achieving (arbitrarily) small error and using a finite sample. It is not clear how these results extend to a setting where the learner achieves small, but non-zero, error on the training distribution.
-	The result on transformers extends only to positional encodings where tokens far enough from one another do not affect each other, which does not apply to positional encodings used in practice.
-	The authors can be a bit more rigorous in stating the definitions and setting. For example, in Definition 1, it is not clear what “zero generalization error” means. Does this mean zero error on the target distribution? Additionally, when assuming “a model trained on sequences of length up to train T”, does this mean that the model achieves zero error on the distribution P(T)?
-	Assumption 15, used in Theorem 4, is deferred to the appendix, which makes reading and understanding the result harder. This assumption should be stated in the main paper.
-	It is not clear what is the takeaway from the experimental section. For example, the authors plot the error of different sequence models displaying different behaviors when varying the sequence length, but then admit that all errors are very low so the conclusion from these plots is not clear. What are the questions or hypotheses that these experiments are suppose to prove or validate?

questions:
-	The models studied in the paper are simplified, one layer, models. Can the results shown in the paper be extended to deep, multi-layer, models?
-	The discussion on high-capacity models, specifically the “negative results” on length and compositional generalization, assumes no constraints at all on the model, which may be too permissive, in contrast to the “positive results” which analyze a specific family of functions. Can these results be extended to general classes of functions with some more generic capacity control (e.g., finite classes or classes with bounded VC dimension)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730560984653

thread_type:
official_review

<thread_object>
creator: reviewer_n34f

<content>
summary:
This paper considers the problems of length/compositional generalization from a functional expressivity perspective. It considers a variety of similar architectures under assumptions of realizability, some intricate assumptions about the support of continuous tokens, and strongly limited capacity (limiting the interaction between tokens). In this setting it considers the ERM solutions and shows that they must recover the ground truth features up to a linear transformation which allows them to generalize. There are negative results presented for all of the architectures considered if the capacity is increased.

strengths:
1. The paper proposes a (to my knowledge) novel method to study length/compositional generalization. 

2. The paper does a good job of extending results to a family of related toy models of different types of architectures. This generality is a strength of the approach.

3. The paper does a good job of presenting both negative and positive results. In particular the negative results actually seem to show that once the restrictive limits on the capacity in the architectures are weakened, then generalization can easily fail.

weaknesses:
1. The perspective of the paper on the prospects of length/compositional generalization seems overly optimistic. In particular, in practice achieving length/compositional has proven very difficult. The papers cited on length generalization struggle to even achieve 2x length generalization despite using all sorts of unnatural tricks. This paper seems to approach the problem as "length generalization happens and we want to explain it", where the real situation appears to be "length generalization is very difficult to achieve if even possible generally". Even in the context of the paper, perhaps the more realistic results are the negative results for each architecture showing that if you allow for (still restrictive) higher capacity, the architectures are able to represent solutions that achieve zero in-distribution error but fail to generalize.

2. It is not clear why we would expect the simplified models presented in the paper to bear any relation to things we may see in practice. Of course in theory we want to study simplified models, but the ultimate goal is for those models to bring insights to practical settings. The negative results indicate that as soon as we expand the very tight capacity restrictions, things can Simplified architectures seem to have no chance of applying to real settings.

3. There is no consideration for statistical issues. Essentially the ERM approach is assuming infinite data. This is combined with the strong assumptions of full support over continuous variables to pin down the learned function to be similar to the realizable ground-truth solution. It is not clear if the arguments presented can apply at all to the case of finite datasets, even as those datasets become very large. 

4. The coverage assumptions and continuous support seem to be quite strong. It seems that the tools are very reliant on having full support over continuous variables. For transformers the assumption is even stronger than having full support over the marginals to full support over the joint distributions of each pair of tokens which is even stronger. It is not clear how these assumptions could be relaxed.

5. The paper does not study standard attention, but instead a simplified version. This should be made more clear in the abstract/intro.

6. The numerical experiments basically only verify the theory. Moreover the plots showing large trends in increasing or decreasing loss are a somewhat confusing way to claim that length generalization is always happening (because the y-axes are very small). Perhaps adding some baselines/alternatives that do not length generalize could make it more clear whether the ability of these architectures to generalize is interesting.

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730571546285

thread_type:
official_review

<thread_object>
creator: reviewer_jxtz

<content>
summary:
This paper analyzes length generalization and compositional generalization in deep sets, transformers, SSMs, and RNNs, proving that these models achieve both types of generalization under certain assumptions. Additionally, the paper demonstrates that the representations learned by each model in these settings are linearly related to the true labeling function, indicating linear identification. Experimental results further confirm the occurrence of linear identification in real models.

strengths:
- This paper provides a clear definition of length generalization and compositional generalization in sequence-to-sequence models and demonstrates valuable guidance on how to analyze these properties. This serves as a strong foundation for future work in this area.
- The figures effectively illustrate the concepts of length generalization and compositional generalization, and the rationale behind focusing on simple limited capacity models is explained in a convincing manner.
- The paper not only examines a range of models, but also extends its theorems to cover various scenarios, such as $C^1$-diffeomorphisms and discrete tokens.

weaknesses:
- **Strong assumption in hypothesis selection**: As acknowledged by the authors in the discussion and limitations section, the assumption, stated as equation 1, that a hypothesis is chosen to minimize expected risk may be a bit strong.
- **Repetitive proofs**: Some sections of the proofs appear to rely heavily on copy-pasting, leading to redundancy and potential errors. For example:
  + On page 28, line 1483, the reference to “equation 31 of equation 30” seems incorrect and should likely be “equation 67 of equation 66.”
  + At the beginning of Appendix C.3.1, “Assumption 2” appears, which may be intended to refer to “Assumption 15.”
  
  If possible, reducing redundancy by grouping similar arguments into a single lemma could help minimize errors and improve readability.
- **Extension to discrete tokens**: While the extension of the theorems to discrete tokens is particularly intriguing, some aspects of the proof raise questions:
  + This overlaps slightly with the second point. The phrase “regular closedness of the support” appears in Theorems 6,9,11, and 12, but for discrete tokens, the support is not regular closed, and the “equality almost everywhere -> equality everywhere” arguments may not be necessary after all.
  + The reasoning from equation 30 to 33 is somewhat unclear; could you refine this explanation a bit? Is it not possible to derive equation 33 directly by setting $x_1 = x_2$ in equation 30?
- **Presentation issues**:
  + In the statements of Theorems 3 and 4, Assumptions 12 and 15 are referenced before being defined. It would be helpful to include these assumptions in the main text, or alternatively, to provide informal statements of Theorems 3 and 4 in the main text and formalize them in the Appendix.
  + In equation 66, $\psi$ and $\phi$ are swapped.
  + In the second line of equation 108, the left-hand side would be more readable if written as $\left.\frac{\partial^p \sigma(s)}{\partial s^p}\right|_{s=v^\top y}$.
  + The formula on page 37, line 1987, should have line breaks before and after it, with an added explanation to improve readability.

questions:
- This question applies to the proofs of nearly all the theorems; as an example, let’s consider the case of deep sets: if we are considering a hypothesis that minimizes the expected risk $R(h;T) = \sum_{t=1}^T \mathbb{E}\left[\ell(h(x_{\leq t},y_t)\right]$, why is it not possible to directly derive $A\psi(x) = B\phi(x)$ for all $x \in [0,1]^n$ from the equality $\omega\left(\psi(x)\right) = \rho\left(\phi(x)\right)$ for all $x \in [0,1]^n$ in the case of $t=1$?
- How realistic is Assumption 18? It would be helpful to see concrete examples that illustrate this assumption.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699923783

thread_type:
official_review

<thread_object>
creator: reviewer_4sdj

<content>
summary:
This paper addresses two key challenges in out-of-distribution generalization for seq-to-seq models: length generalization and compositional generalization. While previous works have primarily focused on empirical approaches to these challenges, this paper aims to provide a theoretical, provable guarantee for achieving both types of generalization. For various architectures (deep sets, Transformers, SSMs, RNNs), the authors provide conditions on models for achieving length and compositional generalization, under the expected risk minimization setup and the realizable assumption. The paper validates their theoretical findings by presenting the experimental results on toy examples.

strengths:
- The paper addresses important challenges in seq-to-seq generalization, length generalization and compositional generalization, across multiple architectures.
- The paper is well-structured and easy to follow.
- This paper is the first work in the literature that presents provable guarantees for length and compositional generalization for seq-to-seq models. Also, the theoretical framework introduced in the paper is a pioneering contribution to the ML community.
- I briefly read the proof, and the proof seems correct and solid.

weaknesses:
1. The models investigated in the paper are highly simplified. For example, the authors describe the decoder-only Transformer model as $w(\sum_{j=1}^i \frac{1}{i} \psi(x_i, x_j))$, where $w$ is a single-layer perception with a bijective activation. This model can not capture softmax operation (as explained in the paper), and apply only to single-layer Transformers. I understand that models with arbitrary capacity (no constraints on $w, \psi$ ) will not exhibit any generalization (as explained in the paper), but the current restrictions are too strong. 
- Continuing from the first point, although the theoretical framework is novel, the derived results are not particularly surprising. Given that $w$ is a bijective single-layer perceptron, one can immediately derive $A\sum_{j\le T} \psi(x_j) = B \sum_{j \le T} \phi(x_j)$ (where the right-hand side corresponds to the realizable solution). While the paper presents rigorous proofs involving mathematical analysis techniques to establish the length and compositional generalization from this equation, the results feel somewhat expected, and I think this predictability diminishes the impact and novelty of the findings.
2. The theorems are built under the expected risk minimization setup, which deviates from the practical scenario where we can only access a finite number of samples.
3. While the paper validates their findings by providing experimental results, they are conducted on simplified models. Including experiments on more realistic scenarios would significantly strengthen the paper’s applicability.

questions:
1. Would it be possible to incorporate the optimization process into the theoretical framework?
2. The key point of the proof relies on showing that the Jacobians of the realizable solution f and an arbitrary expected risk minimization solution g are identical. On a related note, it seems that the approach in [1] also involves the use of Jacobians. Could you briefly explain the main differences between your proof and that of [1]? I ask this as I am not familiar with theoretical works on compositional generalization.

[1] Wiedemer, Thaddäus, et al. "Compositional generalization from first principles." *Advances in Neural Information Processing Systems* 36 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731760589538

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall Response

comment:
We thank all the reviewers for their efforts and insightful feedbacks. We thank the reviewers for encouraging words such as "serves as a strong foundation for future work in this area", "first work in the literature", "pioneering contribution to the community". We have revised the manuscript to address the concerns that were raised (with major changes marked in teal color). Below we start with the description of a new result followed by a brief summary of responses to key concerns from each reviewer.

1. **New result (Theorem 6) on finite hypothesis classes**  We assume that the hypothesis class is finite and make no other constraints on the model. For a finite hypothesis class, we show that if the training length is sufficiently large, then the solutions to expected risk minimization achieve length generalization.  This captures as a special case settings where we have a transformer network with many layers with standard softmax attention, where each weight vector is constrained to only take a finite of values. 


1. **Reviewer 4sDj**:  Theorem 1 and 2 are indeed on the simpler side and are presented to build intuition. Theorem 3-6 and 8, which include the new result described above, apply to more complex model classes than the ones used in Theorem 1 and 2. We believe there is a fair bit of surprise in Theorem 3-6 and 8, and each proof involves several interesting elements. 

2. **Reviewer JxtZ:** In the Appendix, following the proofs to the Theorem 1 and 2, we have remarks that state that the theorems continue to hold even if we do not observe data at t=1. We have fixed all the presentation issues you raised. We have also corrected the facts you pointed to in the proofs of theorems concerning discrete tokens. 


3. **Reviewer n34f:** We strongly believe there is value in doing theoretical research even if the results do not immediately translate into practical insights. Seminal works on PAC learning theory or convergence analysis (mostly about shallow neural networks) are valuable theoretical contributions but have not found practical insights. Our work is the first to develop theoretical guarantees for the hard problem of length generalization.  We also present a new result (Theorem 6), which applies to arbitrary models (e.g., transformers with many layers) with controlled capacity.


4. **Reviewer MsSC:** We have included a new result on finite hypothesis classes that addresses your brilliant question. Finite hypothesis classes contain as a special case deeper transformers with discretized weights and hard-coded positional encodings. We corrected other presentation issues (e.g., moving assumption to the main body) and changed the definitions as you stated.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission aims to provide theoretical guarantees on length and compositional generalization for various simplified sequence-to-sequence architectures (deep sets, single-layer transformer variants, simplified RNNs, and SSMs) under the realizability assumption and certain restrictive conditions on capacity and data support. The authors’ motivation is commendable: length and compositional generalization are important challenges, and providing the first theoretical framework to study them is a noteworthy effort. The paper takes initial steps toward understanding these complex phenomena.