On Provable Length and Compositional Generalization

We thank the reviewer for their encouraging review and brilliant questions. Please see our responses below. We look forward to hearing back from you.

1. Regarding finite sample issues: We want to emphasize that the challenge of out-of-distribution generalization persists even when the training data consists of infinite samples from the training distribution. Due to this most works that study out-of-distribution generalization (such as Wiedemer et al., Lachapelle et al., Arjovsky et al., Dong et al.), assume that there is infinite data from training distribution. In the experiments, we operate with finite data.

2. About experiments: Our theory predicts that solutions to expected risk minimization under stated conditions achieve length and compositional generalization. Our theory also predicts that learned representations are linear transforms of true representations. When working with experiments, we do not have access to infinite data and the training procedure relies on gradient descent that is not guaranteed to reach the exact minima. Hence, the experiments section is meant to validate the theory under non-ideal conditions. In Figure 2, we observe small numbers for errors, which confirms that the model achieves length generalization and validates the theory. In Table 1, we verify the linearity of representations and in Table 2, we verify the claim of compositionality.

3. Regarding finite hypothesis classes: You asked a brilliant question, i.e., if the theory extends to finite hypothesis classes or bounded VC dimensions. We have added a new result (Theorem 6), which proves this claim for finite hypothesis classes. 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. Finite hypothesis classes contain arbitrarily deep transformers with weights constrained to a finite set of values as a special case.

4. The assumption we make on positional encoding is similar to the Hard-Alibi positional encoding scheme used in Hou et al., Jelassi et al.. Jelassi et al. also constructed the positional encoding for theoretical purposes but found it to be effective as well.

5. We have made changes to the definition of length generalization to remove any confusions you raised. We have moved the Assumption 15 to the main body.

References

[Wiedemer et al.] "Compositional generalization from first principles." Advances in Neural Information Processing Systems 36 (2024)

[Lachapelle et al.] "Additive decoders for latent variables identification and cartesian-product extrapolation." Advances in Neural Information Processing Systems 36 (2024).

[Arjovsky et al.] "Invariant risk minimization." arXiv preprint arXiv:1907.02893 (2019).

[Dong et al.] "First steps toward understanding the extrapolation of nonlinear models to unseen domains." arXiv preprint arXiv:2211.11719 (2022).

[Jelassi et al]. "Repeat after me: Transformers are better than state space models at copying." arXiv preprint arXiv:2402.01032 (2024).

[Hou et al.] "Universal length generalization with turing programs." arXiv preprint arXiv:2407.03310 (2024).

We thank the reviewer for their time and effort in their useful review. Please see our responses below. We look forward to hearing back from you.

1. We agree that length generalization has been quite difficult to achieve in general. However, there have been empirical works that have put forth a conjecture -- RASP-L conjecture that states when such generalization is achievable https://arxiv.org/pdf/2310.16028. Even if the empirical evidence indicates generalization up to 2x length, taking first steps towards theoretical explanations is important and our work is an endeavor in that direction. For the models we considered in our experiments, we already show length generalization up to 10x training length.

2. On insights from simplified models: a) Firstly, we believe that our results are the first on the theory of length generalization and we should be credited for that. Secondly, we also add a new result on finite hypothesis classes (described below), which we hope helps you see our work in a more positive light. Thirdly, we believe that there are many theoretical papers that make simplifying assumptions and it is not always possible to have immediate practical insights from them. Take for example the whole line of work on PAC learning theory, or analysis of convergence in neural networks (often done with very shallow depth two (Du et al.), or the recent analysis of in-context learning (Oswald et al.) or analysis of out-of-distribution generalization (Dong et al., Wiedemer et al., Lachapelle et al.). We believe all these works are very valuable but may not give any immediate insights that change practice. However, in the big picture they contribute valuable knowledgable which slowly turns into mental models and useful practical insights.

   b) We introduce a new result (Theorem 6) that assumes the hypothesis class is finite and impose 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. To consider example of finite hypothesis classes, imagine we have a transformer network with many layers with softmax attention or an RNN. Instead of operating with the property that each scalar value of the weight can take all possible real values, we constrain the weights to only take discrete set of values.

3. On statistical issues: We want to emphasize that the challenge of out-of-distribution generalization persists even when the training data consists of infinite samples from the training distribution. Due to this most works that study out-of-distribution generalization (such as Wiedemer et al., Lachapelle et al., Arjovsky et al.), assume that there is infinite data from training distribution. In the experiments, we operate with finite data.

4. Full support over marginals: Full support over marginals (or pairs) is a far weaker assumption than any assumption on the joint support. Look at Figure 1 and Figure 4 for illustration to see that the joint support occupies a much smaller portion of the region. Such assumptions have also been made in other works such as (Wiedemer et al., Dong et al., and Lachapelle et al.). Also, observe that our results for RNN and our new result (Theorem 6) above with finite hypothesis classes does not invoke such an assumption.

5. To supplement Figure 2, we also have Figures 6-11, 13,14, 16-23 in the Appendix that track the predicted sequence against the true sequence for lengths beyond the training length.

References

[Wiedemer et al.] "Compositional generalization from first principles." Advances in Neural Information Processing Systems 36 (2024)

[Lachapelle et al.] "Additive decoders for latent variables identification and cartesian-product extrapolation." Advances in Neural Information Processing Systems 36 (2024).

[Arjovsky et al.] "Invariant risk minimization." arXiv preprint arXiv:1907.02893 (2019).

[Dong et al.] "First steps toward understanding the extrapolation of nonlinear models to unseen domains." arXiv preprint arXiv:2211.11719 (2022).

[Du et al.] "Gradient descent provably optimizes over-parameterized neural networks." arXiv preprint arXiv:1810.02054 (2018).

[Oswald et al.] "Transformers learn in-context by gradient descent." International Conference on Machine Learning. PMLR, 2023.

We thank the reviewer for their response and raising the score. Please take our responses below into consideration.

1. Coverage-based assumptions: Our coverage-based assumptions are inspired by Wiedemer et al. These assumptions seem necessary to even be able to meaningfully define the notion of compositionality.

2. Abstract and introduction: Our abstract and introduction explicitly state that our results are developed for simple, limited-capacity models. We would appreciate it if you could specify which part of the abstract and introduction you found overly optimistic.

3. Simplifying assumptions in theoretical frameworks: Most theoretical frameworks in deep learning often make simplifying assumptions. For instance, analysis of shallow networks is often used to draw stylized insights for the general phenomenon at hand. Our theory is inspired by the RASP-L conjecture proposed in Zhou et al., which identifies conditions under which transformers empirically generalize on real algorithmic tasks, such as copy, mode, addition, and sort. The RASP-L conjecture identifies three conditions - realizability, simplicity, and data diversity - but did not provide a formal proof. Our theory captures these three conditions and proves results for a range of architectures and a range of tasks (e.g., multiplication, addition, tasks expressed by finite automata through our result on finite hypothesis classes).

We hope our above response changes your perspective on our work.

References

Wiedemer, et al. "Compositional generalization from first principles." Advances in Neural Information Processing Systems 36 (2024).

Zhou et al. "What algorithms can transformers learn? a study in length generalization." arXiv preprint arXiv:2310.16028 (2023).

We thank the reviewer for their efforts in a dilligent review and constructive feedback. We especially appreciate reviewer's eye for detail. Please see our responses below. We look forward to hearing back from you.

1. Regarding equation 1: It is not clear what part of equation 1 you are objecting to. We presume you are objecting to the realizability condition, i.e., the labeling function is in the hypothesis class. Firstly, the realizability assumption is a part of the RASP-L conjecture (Zhou et al.) that motivated our work. Secondly, the assumption of realizability is made in most works studying different forms of out-of-distribution generalization (Wiedemer et al., Lachapelle et al., Arjovsky et al., Dong et al.). It is tricky to study generalization beyond the training distribution in the absence of such an assumption.

2. Repetitive proofs and other fixes: We kept some of the proofs separate as some steps change between similar proofs and may not be obvious to the reader. We have fixed the equation referencing errors. We have removed the statements on regular closedness from the discrete token proofs. We have shortened the description from equation 30-33 and used your suggestion. We have included the assumptions that were referred to from th Appendix in the main body. We have corrected the psi and phi notation swaping error. We have made other presentation changes that you pointed to.

3. On equating predictions at $t=1$: Your comment intends to show that we can obtain $A\psi = B\phi$ from $t=1$. Firstly, your comment is applicable to only Theorem 1 and Theorem 2 and there is an important extension that I discuss in a bit. The comment is not applicable to more general extensions for deep sets and transformers presented in Theorem 8 and Theorem 3, and also not to any of the results in SSMs and RNNs. Further, in the Appendix, below Theorem 1 and 2, we explain that our results can also be extended to setting where we do not observe data for t=1 and only observe data at some length T>1. In this case, when we do not observe data at t=1 we cannot apply your observation.

4. Regarding Assumption 18: Assumption 18 is a form of diversity assumption. Suppose change along each dimension of the token captures some concept (e.g., gender). The first part of the assumption requires that we observe at least two instances of the concept change appear -- ("Man", "Woman"), ("Boy", "Girl"). The second part of the assumption requires that we observe concept change independently for each dimension -- ("Man, Boy"), ("Woman, Boy"), ("Man, Girl"), ("Woman, Girl"). In this example, along each of the two dimensions, we observe gender change.

References

[Orvieto et al.] "On the universality of linear recurrences followed by nonlinear projections." arXiv preprint arXiv:2307.11888 (2023).

[Wiedemer et al.] "Compositional generalization from first principles." Advances in Neural Information Processing Systems 36 (2024)

[Lachapelle et al.] "Additive decoders for latent variables identification and cartesian-product extrapolation." Advances in Neural Information Processing Systems 36 (2024).

[Arjovsky et al.] "Invariant risk minimization." arXiv preprint arXiv:1907.02893 (2019).

[Dong et al.] "First steps toward understanding the extrapolation of nonlinear models to unseen domains." arXiv preprint arXiv:2211.11719 (2022).

[Zhou et al.] "What algorithms can transformers learn? a study in length generalization." arXiv preprint arXiv:2310.16028 (2023).

We thank the reviewer for their responses.

1. Adding further clarification on true risk minimizer: Yes, we would be happy to add that clarification. In particular, we can refer to other works on compositionality that analyze true risk minimizers like us.

2. Regarding linear identification from T=1: Yes, you are right in Theorem 7 and Theorem 10, we could have used T=1 but there is an important point that we mention below. We were trying to keep the setting similar to Theorem 1 and Theorem 2 as they are the extensions of Theorem 1 and Theorem 2 respectively. Similar to Theorem 1 and Theorem 2, we can show that there is an extension of Theorem 7 and Theorem 10 to settings where we do not observe data at T=1. Suppose we observe data at some length T>1. In that case, we would require that the support for that length T to be the product $\mathcal{X} \times \cdots \text{T times} \; \mathcal{X}$ so that we can set $x_1=x_2...x_T=x$. We agree with your observations for the other Theorems.

We hope that our explanations are satisfactory and help you view our paper more positively.

We thank the reviewer for their efforts in giving an insightful review. Please see our responses below. We look forward to hearing back from you.

1. Regarding simplified models

   a) The simplicity of Theorems 1 and 2 is a feature and not a bug. We present these primarily to build intuition and progressively go towards more complex results in subsequent theorems (Theorem 3-6 and Theorem 8). In Theorem 3 for instance, we move beyond single layer perceptron to C1-diffeomorphisms. In Theorem 4, we consider the model for SSMs used in Orvieto et al. In Theorem 5, we use standard RNN except we only impose restriction on the dimension of the matrices parametrizing the model.

   b) We also introduce a new result (Theorem 6) on finite hypothesis classes 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. To illustrate finite hypothesis classes, consider transformer network with many layers with softmax attention. Instead of operating with the property that each scalar value of the weight can take all possible real values, we constrain the weights to only take discrete set of values.

2. Regarding expected risk minimization We want to emphasize that the challenge of out-of-distribution generalization persists even when the training data consists of infinite samples from the training distribution. Due to this most works that study out-of-distribution generalization (such as Wiedemer et al., Lachapelle et al., Arjovsky et al.), assume that there is infinite data from training distribution. In the experiments, we operate with finite data.

3. On incorporating optimization process: There are many important and interesting challenges to overcome before a general framework can be built to incorporate optimization process. Due to this we mentioned this problem in the discussions and limitations section as a conjecture. Even in many non-sequence to sequence models studying out-of-distribution generalization, it is common to not include the optimization process (Wiedemer et al., Lachapelle et al., Arjovsky et al, Dong et al.).

4. Comparisons with Wiedemer et al.: Our proof of Theorem 1 has similarity in the Jacobian step but rest of the proof is quite different from Wiedemer et al.. Other theorems -- Theorem 3 for transformers, Theorem 4 for SSMs, Theorem 5 RNNs and Theorem 6 for finite hypothesis classes -- each use different proof technique . We elaborate more below but before that we want to say that Wiedemer is a very nice work. We view our contributions as complementary to them. They make a strong assumption, which is not practical for our case and hence we have different proofs. They assume that part of the mixing function that generates the data is completely known.

   i) The proofs in [Wiedemer et al.] use the fact that function $C$ is known. The remaining unknowns are $\phi$, which are solutions to a set of PDEs. They show that with the help of the knowledge of value of true $\phi$ at one point, one can reduce the solution space of PDEs to that of ODEs. These ODEs under standard regularity assumptions have a unique solution, which leads to exact identification of true $\phi$. In contrast, we do not need to exactly identify and linear identification suffices.

   ii) The proof for transformers carefully uses properties of diffeomorphisms and spanning conditions to show that attention heads are linearly identified. We first arrive at the condition $a(\frac{\alpha + \beta}{2}) = \frac{1}{2}(a(\alpha) + a(\beta))$, where $a$ dictates the relationship between true and learned attenteion head, i.e., $\phi = a \circ \psi$. Note that this condition immediately does not imply that $a$ is linear. Standard definition of a linear transform requires that $a(\alpha+\beta) = a(\alpha) + a(\beta)$ and that $a$ is also homogeneous. We use the fact that $a$ is continuous to prove homogeneity.

   iii) For RNNs, we show that the set of all the solutions to ERM in (1) are all permutations of one another in the weight space. We first show a permutation relationship between the matrices that act on the input directly $B$ and $\tilde{B}$. We use this at the subsequent length to show a permutation like relationship between the matrices that operate on the hidden state. To arrive at these results, we prove some important properties of higher order derivatives of sigmoid activation (how it can be viewed as a polynomial) and use properties of zeros of an analytic function. Once this is established we use principle of induction to show that these permutation relationships suffice to match the true labels at all the lengths.