On Provable Length and Compositional Generalization

## Abstract

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. In all our results across different architectures, we find that the learned representations are linearly related to the representations generated by the true labeling function.

## 1 Introduction

Large language models (LLMs), such as the GPT models (achiam2023gpt) and the Llama models (touvron2023llama), have led to a paradigm shift in the development of future artificial intelligence (AI) systems. The accounts of their successes (bubeck2023sparks; gunasekar2023textbooks) as well as their failures, particularly in reasoning and planning (bubeck2023sparks; stechly2023gpt; valmeekam2023can), continue to rise. The successes and failures of these models have sparked a debate about whether they actually learn general algorithms or if their success is primarily due to memorization and a superficial form of generalization (dziri2024faith).

A model’s ability to perform well across different distribution shifts highlights its ability to learn general algorithms. For models with fixed-dimensional inputs, considerable efforts have led to methods with provable out-of-distribution (OOD) generalization guarantees (rojas2018invariant; rame2022diverse; chaudhuri2023does; wiedemer2023compositional; eastwood2024spuriosity). For sequence-to-sequence models, a large body of empirical works have investigated OOD generalization (anil2022exploring; jelassi2023length) but we lack efforts that study provable OOD generalization guarantees for these models. These provable guarantees provide a stepping stone towards explaining the success of the existing paradigm and also shine a light on where the existing paradigm fails.

OOD generalization capabilities of sequence-to-sequence models can be studied from the lens of two 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. While transformers (vaswani2017attention) are the go-to sequence-to-sequence models for many applications, recently, alternative architectures based on state-space models, as noted by gu2021efficiently, orvieto2023resurrecting, and gu2023mamba, have shown a lot of promise. This motivates us to study a range of natural sequence-to-sequence architectures, including deep sets (zaheer2017deep), transformers, state space models (SSMs), and recurrent neural networks (RNNs). We focus on the realizable case, i.e., the labeling function is in the hypothesis class of the architecture. Further, in our theoretical analysis, we make certain simplifications to permit tractable analysis, for instance, we study RNNs with a limit on hidden state dimension. Our key contributions and insights are summarized below.

$\bullet$ Simple limited capacity versions of the different architectures namely deep sets, transformers, SSMs, and RNNs, provably achieve length and compositional generalization.

$\bullet$ In all our results across different architectures, we find that the learned representations are linearly related to the representations generated by the true labeling function, which is also termed *linear identification* (khemakhem2020variational; roeder2021linear).

$\bullet$ Through a range of experiments, we show the success in both forms of generalization, matching the predictions of the theory and even going beyond.

To the best of our knowledge, our provable guarantees for length and compositional generalization for sequence-to-sequence models are the first in the literature.

## 2 Related Works

##### Length generalization

In the field of length generalization, many important empirical insights have been synthesized over the last few years. shaw2018self discovered the drawbacks of absolute positional embeddings and suggested relative positional embeddings as an alternative. Subsequent empirical analyses, notably by anil2022exploring and jelassi2023length, explored length generalization in different settings for transformer-based models. Key findings revealed that larger model sizes don’t necessarily enhance generalization, the utility of scratchpads varies, and the effectiveness of relative positional embeddings appeared task-dependent. In kazemnejad2024impact, the authors did a comprehensive study of different positional embeddings and provided evidence to show that explicit use of positional encodings is perhaps not essential. In deletang2022neural, the authors conducted experiments on tasks divided based on their placement in the Chomsky hierarchy and showed the importance of structured memory (stack, tape) in length generalization. In a recent work, zhou2023algorithms proposed RASP conjecture, which delineates the tasks where transformers excel or fall short in length generalization, emphasizing the necessity of task simplicity for the transformer and data diversity. Our work is inspired by the experimental findings of their work. While zhou2023algorithms provide empirical evidence for the conjecture, our work formalizes and proves simpler versions of the conjecture for a range of architectures.

On the theoretical side of length generalization, in abbe2023generalization, the authors showed an implicit bias of neural network training towards min-degree interpolators. This bias was used to explain the failures of length generalization on the parity task from anil2022exploring. In xiao2023conditions, the authors leverage directed acyclic graphs (DAGs) to formulate the computation in reasoning tasks and characterize conditions under which there exist functions that permit length generalization. Our results crucially differ, we show a range of conditions under which models learned via standard expected risk minimization achieve length and compositional generalization.

##### Compositional generalization

Compositionality has long been seen as a key piece to the puzzle of human-level intelligence (fodor1988connectionism; hinton1990mapping; plate1991holographic; montague1970pragmatics).
Compositionality is a large umbrella term associated with several aspects (hupkes2020compositionality). In this work, we focus on systematicity, which evaluates a model’s capability to understand known parts and combine them in new contexts. The breadth of research on compositional generalization, encompassing studies like lake2018generalization; loula2018rearranging; gordon2019permutation; hupkes2020compositionality; kim2020cogs; xu2022compositional; arora2023theory; zhang2024can, is too expansive to address comprehensively here, refer to these surveys (lin2023survey; sinha2024survey) for more detail.

In recent years, several works have taken first steps towards theoretical foundations of compositionality.
We leverage the mathematical definition of compositionality from wiedemer2023compositional, which focuses on generalization to the Cartesian product of the support of individual features. In dong2022first, the authors analyze the conditions that provably guarantee generalization to the Cartesian product of the support of individual training features. dong2022first studied additive models, i.e., labeling function is additive over individual features. In (wiedemer2023provable), the authors focus on a more general model class than dong2022first, where the labeling function is of the form $f(x_{1},\cdots,x_{n})=C(\psi_{1}(x_{1}),\cdots,\psi_{n}(x_{n}))$. However, to guarantee compositional generalization, wiedemer2023compositional require that the learner needs to know the exact function $C$ that is used to generate the data. In our work, we do not make such an assumption, our data generation is dictated by the architecture in question, e.g., RNN, and we constrain the dimension of its hidden state. lachapelle2023additive; brady2023provably extend these precursor results from dong2022first from the supervised setting to the unsupervised setting. In particular, lachapelle2023additive; brady2023provably are inspired by the success of object-centric models and show additive decoder based autoencoders achieve compositional generalization.

## 3 Provable Length and Compositional Generalization

We are given a dataset comprising of a sequence of inputs $\{x_{1},\cdots,x_{t}\}$ and a corresponding sequence of labels $\{y_{1},\cdots,y_{t}\}$, where each $x_{i}\in\mathbb{R}^{n}$ and $y_{i}\in\mathbb{R}^{m}$. Observe that this formulation includes both standard downstream tasks such as arithmetic tasks, e.g., $y_{i}=\sum_{j=1}^{i}x_{j}$, $y_{i}=\Pi_{j=1}^{i}x_{j}$ etc., as well as next-token prediction task, where $\{y_{1},\cdots,y_{t}\}=\{x_{2},\cdots,x_{t+1}\}$. We denote a sequence $\{s_{1},\cdots,s_{t}\}$ as $s_{\leq t}$, $X_{k}$ is random variable for token at $k^{th}$ position and its realization is $x_{k}$. Consider a sequence $\{x_{j}\}_{j=1}^{\infty}$, which is sampled from $\mathbb{P}_{X}$, and a subsequence of this sequence $x_{\leq t}=\{x_{j}\}_{j=1}^{t}$, whose distribution is denoted as $\mathbb{P}_{X_{\leq t}}$. The label $y_{t}=f(x_{\leq t})$, where $f$ is the labeling function. The tuple of base distribution and the labeling function is denoted as $\mathcal{P}=\Big{\{}\mathbb{P}_{X},f\Big{\}}$ and the tuple of base distribution up to length $t$ is denoted as $\mathcal{P}(t)=\Big{\{}\mathbb{P}_{X_{\leq t}},f\Big{\}}$. The support of $k^{th}$ token $X_{k}$ in the sequence sampled from $\mathbb{P}_{X}$ is denoted $\mathsf{supp}(X_{k})$. Given training sequences of length $T$ from $\mathcal{P}(T)$, we are tasked to learn a model from the dataset that takes a sequence $x_{\leq t}$ as input and predicts the true label $y_{t}$ as well as possible. If the model succeeds to predict well on sequences that are longer than maximum training length $T$, then it is said to achieve length generalization (a more formal definition follows later). Further, if the model succeeds to predict well on sequences comprising of combination of tokens that are never seen under training distribution, then it is said to achieve compositional generalization (a more formal definition follows later.). We study both these generalization forms next.

##### Learning via expected risk minimization

Consider a map $h$ that accepts sequences of $n$-dimensional inputs to generate a $m$-dimensional output. We measure the loss of predictions of $h$, i.e., $h(x_{\leq t})$, against true labels as $\ell\big{(}h(x_{\leq t}),y_{t}\big{)}$, where $y_{t}$ is the true label for sequence $x_{\leq t}$. In what follows, we use the $\ell_{2}$ loss. Given sequences sampled from $\mathcal{P}(T)$, the expected risk across all time instances up to maximum length $T$ is defined as $R(h;T)\coloneq\sum_{t=1}^{T}\mathbb{E}\big{[}\ell(h(x_{\leq t}),y_{t})\big{]}$.
The learner aims to find an $h^{*}$ that solves

$$ $h^{*}\in\operatorname*{arg\,min}_{h\in\mathcal{H}}R(h;T),$ (1) $$

where $\mathcal{H}$ is the hypothesis class of models. We seek to understand the properties of solutions to equation 1 through the lens of following questions.
{mdframed}
When do common sequence-to-sequence models $\mathcal{H}$ succeed at length & compositional generalization and when do they fail?

###### Definition 1 .

Consider the setting where a model is trained on sequences $(x_{\leq t},y_{\leq t})$ of length up to $T$ drawn from $\mathcal{P}(T)$. If the model achieves zero generalization error on sequences $(x_{\leq t},y_{\leq t})$ of length up to $\tilde{T}$ drawn from $\mathcal{P}(\tilde{T}),\forall\;\tilde{T}\geq T$, then it achieves length generalization w.r.t. $\mathcal{P}$.

In the above definition of length generalization, we simply ask if the model generalizes to longer sequences.
We drop the phrase w.r.t $\mathcal{P}$ hereafter to avoid repetition. We now define a test distribution that evaluates compositional generalization capabilities. We consider sequences of fixed length $T$. Define a uniform distribution $\mathbb{Q}_{X_{\leq T}}$ such that the support of $\mathbb{Q}_{X_{\leq T}}$ equals the Cartesian product of the support of each token $X_{k}$ from $\mathbb{P}_{X}$, we write this joint support as $\Pi_{j=1}^{T}\mathsf{supp}(X_{j})$. In this case as well, the labeling function continues to be $f$. Hence, we obtain the tuple $\mathcal{Q}(T)=\{\mathbb{Q}_{X_{\leq T}},f\}$.

Figure: Figure 1: Illustrating support of train vs test distribution for (a) compositional generalization and (b) length generalization.
Refer to caption: NeurIPS/FIg1_new.pdf

###### Definition 2 .

Consider the setting where a model is trained on sequences $(x_{\leq t},y_{\leq t})$ of length up to $T$ drawn from $\mathcal{P}(T)$. If the model achieves zero generalization error on sequences $(x_{\leq t},y_{\leq t})$ of length up to $T$ drawn from $\mathcal{Q}(T)$, then it achieves compositional generalization.

This definition of compositionality above is based on wiedemer2023compositional; brady2023provably. In this definition, we ask if the model generalizes to new combinations of seen tokens.

Illustrative example We teach the model multiplication on sequences of length $2$ , where each $x_{j}$ is a scalar, $y_{i}=\Pi_{j=1}^{i}x_{j}$. Say the support of the entire sequence drawn from $\mathbb{P}_{X}$ is $\big{\{}x\;\big{|}\;\|(x_{1},x_{2})-\frac{1}{2}\bm{1}\|_{1}\leq\frac{1}{2},\;x
_{k}\in[0,1],\;\forall k\geq 3\big{\}}$. The support of training distribution $\mathbb{P}_{X_{\leq 2}}$ is $\big{\{}x\;\big{|}\|(x_{1},x_{2})-\frac{1}{2}\bm{1}\|_{1}\leq\frac{1}{2}\big{\}}$ shown in the pink region in Figure 1a. In Figure 1a, we illustrate compositional generalization, the model is trained on pink region and asked to generalize to the yellow region. Further, if the model continues to correctly multiply on longer sequence lengths in $\mathbb{P}_{X_{\leq\tilde{T}}}$ for $\tilde{T}\geq T$, then it achieves length generalization shown in Figure 1b.

###### Definition 1 .

###### Definition 2 .

##### A preview of the technical challenges

Both notions of compositional generalization and length generalization introduced above involve testing on distributions whose support is not contained in the training distributions. The long line of work on distribution shifts (sugiyama2007covariate; david2010impossibility; ben2014domain; rojas2018invariant; arjovsky2019invariant; ahuja2021invariance) assume the support of test is contained in the support of train distribution.
In recent years, there has been development of theory for distribution shifts under support mismatch dong2022first; abbe2023generalization; wiedemer2023compositional; netanyahu2023learning; shen2023engression. Our work is closer to the latter line of work but it comes with its own *technical challenges*, which involve building new proofs different from the above line of work, as we study a new family of models, i.e., sequence-to-sequence models, and a new form of generalization, i.e., length generalization.

##### RASP conjecture

zhou2023algorithms propose a conjecture backed by empirical evidence, which delineates the conditions that suffice for length generalization for transformers.
The conjecture places three requirements – a) realizability: the task of interest is realizable on the transformer, b) simplicity: the task can be expressed as a short program in RASP-L language, c) diversity: the training data is sufficiently diverse such that there is no shorter program that achieves in-distribution generalization but not OOD generalization. We leverage assumptions similar to a) and b). We assume realizability, which means labeling function $f$ is in the hypothesis class $\mathcal{H}$. As to simplicity, we consider hypothesis class $\mathcal{H}$ with limited capacity, e.g., we study one block transformer, or RNNs with a limit on hidden state dimension. We emphasize that the third assumption c) on diversity from zhou2023algorithms is quite strong. In our setting, we do not invoke it and instead, we require that the support of test distribution is not larger than the Cartesian product of the marginal distribution of the tokens. We now move to proving simplified versions of this conjecture for different architectures.

### 3.1 Deep sets

Deep sets are a natural first choice of architecture to study here. These take sets as inputs and thus handle inputs of arbitrary lengths. These were introduced in zaheer2017deep. Informally stated, zaheer2017deep show that a large family of permutation-invariant functions can be decomposed as $\rho\big{(}\sum_{x\in\mathcal{X}}\phi(x)\big{)}$. Consider the examples of the sum operator or the multiplication operator, which take $\{x_{1},x_{2},\cdots,x_{k}\}$ as input, and return the sum $y=\sum_{j=1}^{k}x_{j}$ or the product $y=\Pi_{j=1}^{k}x_{j}$. These operations are permutation invariant and can be expressed using the decomposition above. For the sum operator $\rho$ and $\phi$ are identity and for the multiplication operator $\rho=\exp$ and $\phi=\log$. Consider another example from language. We construct a bag of words sentiment classifier, where $\{x_{1},x_{2},\cdots,x_{i}\}$ is the set of words that appear in the sentence, $\phi(x_{j})$ is the feature embedding for word $j$. $\sum_{j\leq i}\phi(x_{j})$ is the representation of the entire sentence which is passed to the final layer $\rho$ that generates the sentiment label. In what follows, we aim to understand when such a classifier generalizes to sentences beyond training lengths and to new sentences comprised of unseen word combinations.

###### Assumption 1 .

Each function in the hypothesis class $\mathcal{H}$ takes a sequence $\{x_{1},\cdots,x_{i}\}$ as input and outputs $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j\leq i}\psi(x_{j})\Big{)}$, where $\omega$ is a single layer perceptron with a continuously differentiable bijective activation (e.g., sigmoid) and $\psi$ is a map that is differentiable.

A simple mathematical example of a function from the above family when $\psi(x_{j})=[x_{j},x_{j}^{2},x_{j}^{3}]$ is a polynomial map of degree $3$ and each $x_{j}$ is a scalar – $\sigma\big{(}a\sum_{j\leq i}x_{j}+b\sum_{j\leq i}x_{j}^{2}+c\sum_{j\leq i}x_{j
}^{3}\big{)}$. In the assumption that follows, we assume that the support of the sequences is regular closed in the standard topology, i.e., the set is equal to the closure of its interior.

###### Assumption 2 .

The joint support $\mathsf{supp}(X_{\leq i})$ is a regular closed set for all $i\leq T$.

In most of our results in the main body, we invoke Assumption 2. This assumption is satisfied in many cases if the tokens are continuous random variables but it is not satisfied for discrete random variables. In the Appendix, we extend several of our key results to discrete tokens.

###### Assumption 1 .

###### Assumption 2 .

##### Linear identification

Each architecture that we study in this work relies on a hidden representation that is passed on to a non-linearity to generate the label. Under the realizability condition for deep sets, the labeling function takes the form $f(\mathcal{X})=\rho(\sum_{x\in\mathcal{X}}\phi(x))$, where $\phi(x)$ is the hidden representation. If the learned deep set is denoted by $\omega(\sum_{x\in\mathcal{X}}\psi(x))$, then the learned hidden representation is $\psi(x)$. If $\psi(x)=A\phi(x)$, then the learned representation is said to *linearly identify* the true data generating representation $\phi(x)$. We borrow this definition from the representation identification literature (khemakhem2020variational; roeder2021linear).

###### Theorem 1 .

If $\mathcal{H}$ follows Assumption 1, the realizability condition holds, i.e., $f\in\mathcal{H}$, $\mathsf{supp}(X_{j})=[0,1]^{n},\;\forall j\geq 1$, and the regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss generalizes to all sequences in the hypercube $[0,1]^{nt},\;\forall t\geq 1$ and thus achieves length and compositional generalization.

Due to space constraints, we only give proof sketch for some of the results in the main body.

Proof sketch. The detailed proof is in Section C.1. Since we use $\ell_{2}$ loss, every solution $h^{*}$ to equation 1 equals the true labeling function $f$ on the training support almost everywhere. We invoke the regular closedness assumption and continuity of $f$ to show that $h^{*}=f$ everywhere on the training support. We then show that the learned representation of a sequence in training is a linear transform of its true representation. Owing to the regular closedness assumption, we show that the Jacobian of the learned representation of a token is a linear transform of the Jacobian of the true representation. This leads to the result that the learned representation of a token is a linear transform of the true representation of the token, which implies linear identification.
We use linear identification over sequences of arbitrary length and combine it with the fact that the learned non-linearity undoes the linear transformation of the true representation to show length generalization.

In the above result, we require the support of the marginal distribution of each token to be $[0,1]^{n}$. The support of $T$ token length sequence under the joint training distribution can still be a much smaller subset of $[0,1]^{nT}$, as illustrated in Figure 1a (and Figure 4 in the Appendix). Despite this the model generalizes to all sequences in $[0,1]^{nt}$ for all $t$. An important insight from the proof is that the hidden representation learned by the model is a linear transform of the true hidden representation, i.e., it achieves linear identification $\psi=A\phi$ (Further details are in Corollary 1).

###### Theorem 1 .

##### Extensions of Theorem 1

In Theorem 7, we extend Theorem 1 to $\omega$ from $C^{1}$-diffeomorphisms(^1^11$C^{1}$-diffeomorphism - a continuously differentiable map that has a continuously differentiable inverse.). As a by product, we obtain length & compositional generalization for multiplication operator. In Theorem 6, we extend the above result to discrete tokens. Further, most results in this work translate to settings where we do not observe labels at all lengths from $1$ to $T$ (further discussion in Appendix).

##### High capacity deep sets

In the above results, we operated with some constraints on the deep sets. In Theorem 1, we used limited capacity $\omega$ that are represented via a single layer perceptron. In Theorem 7, we used $\omega$ that are represented via $C^{1}$-diffeomorphisms, which implies the output dimension of $\psi$ equals label dimension $m$ and cannot be larger. What happens when we work with deep sets with arbitrary capacity, i.e., no constraints on $\omega$ and $\psi$? These models then express a large family of permutation invariant maps (zaheer2017deep). Suppose $\mathcal{H}$ is the class of all permutation invariant maps and the labeling function $f\in\mathcal{H}$. Consider a map $h$ such that $h=f$ for all sequences of length up to $T$, and $h=f+c$ otherwise. Observe that $h$ is permutation invariant and also belongs to $\mathcal{H}$. $h$ achieves zero generalization error on training sequences of length $T$ but a non-zero error on longer sequences. Thus in the setting of high capacity deep sets, there exist solutions to equation 1, which do not achieve length generalization. We can construct the same argument for compositional generalization as well and say $h=f$ on the training distribution (pink region) in Figure 1a and $h=f+c$ on the testing distribution (yellow region) in Figure 1a. In order to show successful generalization (length or compositional) in Theorem 1, we require all solutions to risk minimization in equation 1 to match the predictions of true labeling function on data beyond the support of the training distribution. In order to show that high capacity models are not guaranteed to succeed, we focused on showing that there exists a solution to equation 1 that does not generalize beyond the support of training distribution. A more nuanced argument for failure should show that there exist solutions reachable via gradient descent that do not generalize. We leave a rigorous theoretical exploration of this to future work. However, we conduct experiments with high capacity models in the Appendix (Section D.3) to illustrate failures in high capacity regime.

### 3.2 Transformers

Ever since their introduction in vaswani2017attention, transformers have revolutionized all domains of AI. In this section, we seek to understand length generalization for these models. Transformer architectures are represented as alternating layers of attention and position-wise non-linearity. We drop layer norms for tractability. Following similar notation as previous section, we denote position-wise non-linearity as $\rho$ and attention layer as $\phi$. We obtain the simplest form of causal transformer model as $\rho\Big{(}\sum_{j=1}^{i}\frac{1}{i}\cdot\phi(x_{i},x_{j})\Big{)}$. This decomposition captures linear attention, ReLU attention, sigmoid attention, ReLU squared attention, which were studied previously in wortsman2023replacing; hua2022transformer; shen2023study and found to be quite effective in several settings. This decomposition does not capture softmax-based attention and developing provable length generalization guarantees for the same is an exciting future work. Other works (bai2023transformers) also replaced softmax with other non-linear attention for a more tractable analysis. We illustrate the sigmoid-based transformer from wortsman2023replacing below. Let $W_{q}\in\mathbb{R}^{k\times n}$, $W_{k}\in\mathbb{R}^{k\times n}$, and $W_{v}\in\mathbb{R}^{k\times n}$ be the query, key and value matrices. $\rho$ is parametrized via a multi-layer perceptron denoted as $\mathsf{MLP}$.

$$ $\begin{split}&q_{i}=W_{q}x_{i},\;k_{j}=W_{k}x_{j},v_{j}=W_{v}x_{j},\phi(x_{i}, x_{j})=\sigma\bigg{(}\frac{q_{i}^{\top}k_{j}}{\sqrt{d}}\bigg{)}v_{j},\;\mathsf {MLP}\Big{(}\sum_{j=1}^{i}\frac{1}{i}\cdot\phi(x_{i},x_{j})\Big{)}.\end{split}$ (2) $$

In the above feedforward computation, the output of attention for the current query is computed and sent to the MLP to generate the label.

###### Assumption 3 .

Each function in the hypothesis class $\mathcal{H}$ takes a sequence $\{x_{1},\cdots,x_{i}\}$ as input and outputs $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j\leq i}\frac{1}{i}\cdot\psi(x_{i},x_
{j})\Big{)}$, where $\omega$ is a single layer perceptron with continuously differentiable bijective activation (e.g., sigmoid) and $\psi$ is a map that is differentiable.

We denote the joint support of two tokens $X_{i},X_{j}$ as $\mathsf{supp}(X_{i},X_{j})$.

###### Theorem 2 .

If $\mathcal{H}$ follows Assumption 3, the realizability condition holds, i.e., $f\in\mathcal{H}$, $\mathsf{supp}(X_{i},X_{j})=[0,1]^{2n},\;\forall i\not=j$ and the regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 (with $T\geq 2$) with $\ell_{2}$ loss generalizes to all sequences in the hypercube $[0,1]^{nt},\;\forall t\geq 1$ and thus achieves length and compositional generalization.

Similar to Theorem 1, we observe linear identification here too, i.e., learned attention representation denoted $\psi$ is a linear transform of the true attention representation denoted $\phi$, i.e., $\psi(x_{i},x_{j})=C\phi(x_{i},x_{j})$, (details in Section C.2, see Corollary 2). We now extend Theorem 2 from single layer perceptron $\omega$ to $C^{1}$-diffeomorphism. We also extend Theorem 2 to discrete tokens in Theorem 9.

###### Assumption 4 .

Each function in $\mathcal{H}$ takes $\{x_{1},\cdots,x_{i}\}$ as input and outputs $h(x_{1},\cdots,x_{i})=\omega(\sum_{j=1}^{i-1}\frac{1}{i-1}\cdot\psi(x_{i},x_{j
}))$, where $\omega$ is a $C^{1}$-diffeomorphism, $\omega(0)=0$.

The reader would notice that the summation is up to $i-1$ and hence it computes attention scores w.r.t all other terms in the context except $x_{i}$. We conjecture that the theorem that we present next extends to the more general case where summation includes the $i^{th}$ term.

###### Theorem 3 .

If $\mathcal{H}$ follows Assumption 4, the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further assumption on the support (Assumption 12) holds, then the model trained to minimize the risk in equation 1 (with $T\geq 3$) with $\ell_{2}$ loss generalizes to all sequences in $[0,1]^{nt},\forall t\geq 1$ and thus achieves length and compositional generalization.

Proof sketch. The complete proof is provided in Section C.2. Following earlier proof sketch, we see that $h^{*}=f$ everywhere on the support of the training distribution. We use the property of diffeomorphisms to show that learned attention representation of a query, key combination is a bijective transform of the true representation generated by the labeling function $f$. We then arrive at the condition that the bijective transformation applied to the sum of representations of query, key combinations is equal to the sum of bijective transformation applied to each query, key combination separately. This does not yield linearity immediately but points to it. We then exploit the fact that the bijective transformation is continuous to show that this transform is indeed linear. Similar to the earlier proof sketch, we use this linearity to achieve the generalization guarantees.

Observe that similar to Theorem 2, the above proof sketch argues that the learned attention representation is a linear transform of the true attention representation (See Corollary 3).

###### Assumption 3 .

###### Theorem 2 .

###### Assumption 4 .

###### Theorem 3 .

##### Multiple attention heads and positional encoding

While the discussion in this section used a single attention head $\phi$, the results extend to multiple attention heads as shown in Section C.2. The model of transformers discussed so far uses the current query and compares it to keys from the past, it does not distinguish the keys based on their positions. For many arithmetic tasks such as computing the median, maximum etc., the positions of keys do not matter but for other downstream tasks such as sentiment classification, the position of the words can be important. In Section C.2, we adapt the architecture to incorporate relative positional encodings and show how some of the results extend. We modify the model as $\rho(\sum_{j=1}^{i}\frac{1}{i}\cdot\phi_{i-j}(x_{i},x_{j}))$, where $\phi_{i-j}(x_{i},x_{j})$ computes the query key inner product while taking the relative position $i-j$ into account. We show that if $\phi_{i-j}=0$ for $i-j>T_{\max}$, i.e., two tokens sufficiently far apart do not impact the data generation, then length generalization and compositional generalization are achieved.

##### High capacity transformers

In the above results, we operated with constraints on transformers, which limit their capacity. Similar to the setting of deep sets, observe that Assumption 3 constrains $\omega$ to single layer perceptron, Assumption 4 constraints $\omega$ to $C^{1}$-diffeomorphisms. What happens if we work with transformers with no constraint on $\omega$ and $\psi$? If $\psi(x,y)=\psi(\tilde{x},y),\forall x\not=\tilde{x}$, then the decomposition for the causal transformer $\omega\Big{(}\sum_{j=1}^{i}\frac{1}{i}\cdot\psi(x_{i},x_{j})\Big{)}$ becomes $\omega\Big{(}\sum_{j=1}^{i}\frac{1}{i}\cdot\psi(x_{j})\Big{)}$, which is very similar to deep sets. In such a case, we can adapt arguments similar to that of arbitrary capacity deep sets and argue that there exist solutions to equation 1 that do not achieve length and compositional generalization. We now move to state-space models and RNNs.

### 3.3 State space models

In recent years, state space models gu2021efficiently; orvieto2023resurrecting have emerged as a promising competitor to transformers. In (orvieto2023universality; orvieto2023resurrecting), the authors used the lens of linear recurrent layer followed by position-wise non-linearities as the main building block to understand these models. We illustrate the dynamics of these models to show the generation of $x_{\leq t}$ and $y_{\leq t}$ next. Given the current input $x_{t}$, we combine it linearly with the hidden state from the past to obtain the current hidden state. The hidden state is input to $\rho$, which generates the label as follows

$$ $\begin{split}&h_{1}=Bx_{1};\;\;\;h_{2}=\Lambda h_{1}+Bx_{2};\;\cdots,h_{t}= \Lambda h_{t-1}+Bx_{t},\;\\ &y_{1}=\rho(h_{1});\;y_{2}=\rho(h_{2});\;\cdots\cdots,\;\;\;\;\;y_{t}=\rho(h_{ t}),\end{split}$ (3) $$

where $h_{t}\in\mathbb{R}^{k}$ is hidden state at point $t$, $\Lambda\in\mathbb{R}^{k\times k},B\in\mathbb{R}^{k\times n}$ and $\rho:\mathbb{R}^{k}\rightarrow\mathbb{R}^{m}$. We can succinctly write $h_{t}=\sum_{j=0}^{t-1}\Lambda^{j}Bx_{t-j}$.

###### Assumption 5 .

Each function in the hypothesis class $\mathcal{H}$ takes a sequence $\{x_{1},\cdots,x_{i}\}$ as input and outputs $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j}\Big{)}$, where $\omega:\mathbb{R}^{k}\rightarrow\mathbb{R}^{m}$ is a $C^{1}$-diffeomorphism, $B$ and $\Lambda$ are square invertible. As a result, $m=k=n$.

###### Theorem 4 .

If $\mathcal{H}$ follows Assumption 5, and the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further condition on the support, i.e., Assumption 15, holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss $(T\geq 2)$ achieves length and compositional generalization.

The proof is provided in Section C.3. Similar to previous theorems, the hidden state estimated by the learned model, $\tilde{h}_{t}$, and the true hidden state, $h_{t}$, bear a linear relationship (Corollary 4), i.e., linear identification is achieved. We extend Theorem 4 to discrete tokens in Theorem 11.

###### Assumption 5 .

###### Theorem 4 .

##### High capacity SSMs

In the above result, we operated with certain constraints on SSMs, i.e., the input dimension, output dimension, and the hidden state dimension are equal. These constraints limit their capacity. What happens if we put no constraints on $\Lambda$, $B$ and $\omega$? orvieto2023universality showed that SSMs with appropriately large $\Lambda$ and $B$ matrices can approximate a sequence-to-sequence mapping up to some length with arbitrary precision. Consider the true labeling function $f$ and another function $h$, which is equal to $f$ for all sequences of length up to $T$ and $f+c$ for larger lengths. If we use such arbitrary capacity SSMs as our hypothesis class, then this hypothesis class contains both $f$ and $h$. As a result, $h$ is a solution to equation 1 and it does not achieve length generalization. We can extend the same argument to compositional generalization as well.

### 3.4 Vanilla recurrent neural networks

Standard RNNs have a non-linear recurrence unlike the linear recurrence studied in the previous section. We use the same notation as the previous section and only add an activation for non-linear recurrence. We illustrate the dynamics to show the generation of $x_{\leq t}$ and $y_{\leq t}$ below.

$$ $\begin{split}&h_{1}=\sigma(Bx_{1});\;\;\;h_{2}=\sigma(\Lambda h_{1}+Bx_{2});\; \cdots,h_{T}=\sigma(\Lambda h_{T-1}+Bx_{T})\;\\ &y_{1}=\rho(h_{1});\;\;\;\;\;\;y_{2}=\rho(h_{2});\;\cdots\cdots\cdots,\;\;\;\; \;y_{T}=\rho(h_{T}),\end{split}$ (4) $$

###### Assumption 6 .

Each function in the hypothesis class $\mathcal{H}$ is a vanilla RNN of the form equation 4, where the position-wise non-linearity is a single layer perceptron $\sigma\circ A$, and $\Lambda,B$ govern the hidden state dynamics (equation 4). $A,\Lambda,B$ are square invertible matrices, and $\sigma$ is the sigmoid activation.

###### Theorem 5 .

If $\mathcal{H}$ follows Assumption 6, and the realizability condition holds, i.e., $f\in\mathcal{H}$ and regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss (with $T\geq 2$) achieves length and compositional generalization.

Proof sketch. We first show that the parameters of learned and true RNN are related via a permutation matrix $\Pi$. The learned matrices $\tilde{B},\tilde{A},\tilde{\Lambda}$ relate to the true matrices as follows, i.e., $\tilde{B}=\Pi^{\top}B$, $\tilde{A}=A\Pi$, and $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi$. We first establish $\tilde{B}=\Pi^{\top}B$ and $\tilde{A}=A\Pi$. To establish this, we first show that the higher order derivatives of sigmoid activation can be expressed as a polynomial of sigmoids, and we also rely on properties of the measure of zeros of an analytic function (mityagin2015zero). We then substitute $\tilde{B}=\Pi^{\top}B$ and $\tilde{A}=A\Pi$ in the identity equating true label and predicted label at subsequent length to show $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi$. Finally, we invoke induction principle to show that learned hidden state $\tilde{h}_{t}=\Pi h_{t},\forall t$, which is then used to show length and compositional generalization.

From the proof sketch above, it follows that linear identification is achieved. The hidden state estimated by the learned model, i.e., $\tilde{h}_{t}$, and the true hidden state $h_{t}$, bear a linear relationship (See Corollary 5 in Section C.4 for details), where the linear relationship is a permutation map. We extend Theorem 5 to discrete tokens in Theorem 12.

###### Assumption 6 .

###### Theorem 5 .

##### High capacity RNNs

In our result above, similar to previous sections we showed that limited capacity RNNs can achieve length and compositional generalization. How about RNNs with arbitrary capacity, i.e., no constraint on $\Lambda,B$ and $\rho$? These systems can approximate sequence-to-sequence models to arbitrary precision (sontag1992neural; guhring2020expressivity). Hence, we can use the same argument as previous sections to argue that if $\mathcal{H}$ corresponds to RNNs with arbitrary capacity, then there exist solutions to equation 1 that do not achieve length and compositional generalization.

##### Remark on proofs

Finally, we would like the reader to appreciate that our proofs follow different strategies in comparison to wiedemer2023compositional; dong2022first, due to the fact that we cater to sequence-to-sequence models. Consider the proofs in wiedemer2023provable, which reduce the solutions of equation 1 to solutions of set of ordinary differential equations, which under their assumptions are unique. That leads to exact identification in contrast to linear identification.

## 4 Experiments

We present the empirical evaluation of compositional and length generalization capabilities of the architectures from the previous section. All the experiments are carried out in the realizable case where $f\in\mathcal{H}$, i.e., depending on the architecture in question, we use a random instance of the architecture to generate the labels. We train a model $h$ from the same architecture class to minimize the $\ell_{2}$ loss between $h$ and $f$. Under different scenarios, we ask if $h$ achieves length generalization and compositional generalization. We also seek to understand the relationship between the hidden representations of $h$ and hidden representations of $f$.

### 4.1 Length generalization

We sample sequences $x_{\leq t}$ of varying length with a maximum length of $T=10$. Each token $x_{i}\sim\text{Uniform}[0,1]^{n}$, where $n=20$. The sequences are then fed to the labeling $f$, which comes from the hypothesis class of the architecture, to generate the labels. We minimize the empirical risk version of equation 1 over the same hypothesis class with $\ell_{2}$ loss. For evaluation, we present the $\ell_{2}$ loss on the test datasets. We also evaluate $R^{2}$ of linear regression between the learned hidden representations denoted $\psi(x_{i})$ and true hidden representations $\phi(x_{i})$ for all $x_{i}\in x_{\leq t}$ from the test dataset sequences. This metric is often used to evaluate the claims of linear identification (khemakhem2020variational), i.e., the higher this value, the closer the linear relationship. We present results averaged over five seeds for models with two hidden layer MLPs for $\rho$ ($\phi$ is two hidden layer MLP for deep sets). Figure 2 shows a very small test loss of models on increasing sequence lengths when only trained with sequences of up to length $T=10$, which is in agreement with Theorem 1-5. Further, in Figure 3, we show an exemplar sequence from test set and how the trained transformer tracks it. Table 2 shows the average of $R^{2}$ score of $\psi(x_{i}),\phi(x_{i})$ across different positions $i$ at test time. These results demonstrate a linear relationship between learned and true hidden representations, which agrees with our theoretical claims. In Section D, we show that when realizability condition does not hold, i.e., $f\not\in\mathcal{H}$, then length generalization is not achieved. We also present *additional experiments with discrete tokens, failures in the high capacity settings*, and other experimental details in Section D.

Figure: 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.
Refer to caption: x1.png

**Table 1: Average test $R^{2}$ of true and learned hidden representations $\psi(x_{i}),\phi(x_{i})$ across all positions $i$ at various lengths unseen during training. A strong linear relationship is observed for all models across lengths.**
| Model | $R^{2}$ ($t=20$) | $R^{2}$ ($t=100$) |
| --- | --- | --- |
| Deep set | $0.97\pm 0.01$ | $0.97\pm 0.01$ |
| Transformer | $0.99\pm 0.01$ | $0.99\pm 0.01$ |
| SSM | $0.99\pm 0.01$ | $0.99\pm 0.01$ |
| RNN | $0.99\pm 0.01$ | $0.99\pm 0.01$ |

### 4.2 Compositional generalization

For compositional generalization, we generate data following the illustration in Figure 1a. During training, we sample each component $k$ of a token from $\text{Uniform}[0,1]$ and accept the sampled sequences that satisfy the following for all components $i$:
$-0.5\leq\sum_{j=1}^{T}(x_{j}^{k}-0.5)\leq 0.5\hskip 5.69054pt\forall k$, where $x_{j}^{k}$ is the $k^{th}$ component of token $j$. During testing, we sample $x_{\leq t}$ from the complementary set of the training set, i.e., corners of hypercube $[0,1]^{nt}$. We present the $\ell_{2}$ loss on the test dataset, as well as the mean $R^{2}$, where the results are averaged over 5 seeds. The rest of the details are the same as the previous section, i.e., $T=10$, $n=20$, $\rho$ is a two hidden layer MLP ($\phi$ is also a two hidden layer MLP for deep sets). Table 2 shows the test $\ell_{2}$ loss and $R^{2}$ scores for linear identification.

Figure: 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$.
Refer to caption: x2.png

## 5 Discussion and limitations

Our work is a step towards theoretical foundations of successes and failures of length and compositional generalization in sequence-to-sequence models. We prove simplified versions of the recently proposed RASP conjecture under weaker data diversity assumptions. In our analysis, we make certain simplifications, e.g., on the architectures considered, which motivates some of the important conjectures for future work. The main conjectures go as follows –
a) Conjecture 1: Theorem 2 and 3 currently incorporate different non-linear attentions but not the softmax attention. We believe these guarantees on transformers extend to softmax attention given the experimental evidence in Section 4.
b) Conjecture 2: Theorem 2 and 3 use one block of attention and one block of non-linearity. We believe that it is possible to extend these results to more expressive $\mathcal{H}$, e.g., with more alternating blocks.
c) Conjecture 3: Our results focus on the generalization properties of all the possible solutions to risk minimization equation 1. However, in practice the optimization procedure may be biased towards a subset of those. Does accounting for the bias of optimization procedure give way to explaining the success of generalization in even higher capacity architectures?

## Appendix

## Appendix A Illustration of the test support for compositional generalization

The notion of compositional generalization we study requires us to evaluate the model on the Cartesian product of the support of individual token distributions.
In Figure 4, we give some additional examples besides the ones shown in Figure 1a to illustrate the difference between the Cartesian product set and the observed support. These examples illustrate the support of the observed training data distribution can be much smaller than the Cartesian product of the support of the individual tokens.

Figure: Figure 4: Illustration of observed support and its Cartesian product. These examples illustrate the support of the observed training data distribution can be much smaller than the Cartesian product of the support of the individual tokens.
Refer to caption: ICLR/Fig_supp.pdf

## Appendix B Supplement on Related Works

We briefly discuss some other relevant works here, which could not be mentioned in the main body due to space constraints. In schug2023discovering, the authors exploit compositionality in the context of meta learning, where each task parameter is specified via a linear combination of some basis module parameters. They construct an approach that achieves provable compositional guarantees and outperforms meta-learning approaches such as MAML and ANIL. In a concurrent work hou2024universal propose an interesting scratch pad strategy inspired from the operation of Turing machines. They call this strategy Turing programs. The scratch pad emulates the operation of a Turing machine. The authors argue that there exist short RASP program ($O(n)$ length) that can simulate the operation of a Turing machine for sufficiently long number of steps ($O(\exp(n))$). Our current framework does not incorporate scratchpad strategies into it, and it is a promising future work to investigate provable length generalization guarantees with scratchpad.

## Appendix C Proofs

In all the results that follow, we work with standard topology in $\mathbb{R}^{nt}$, where $n$ is dimension of each token and $t$ is the sequence length. We remind the reader of the definition of a regular closed set – if a set is equal to the closure of its interior, then it is said to be a regular closed set. In all the results that follow, we either work with continuous random variables for which the Radon-Nikodym derivative of $X_{\leq t}$ is absolutely continuous w.r.t Lebesgue measure $\forall t$ or we work with discrete random variables for which the Radon-Nikodym derivative of $X_{\leq t}$ is absolutely continuous w.r.t counting measure $\forall t$.

###### Lemma 1 .

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(x)=g(x),\forall x\in\mathcal{X}$ and $\nabla f(x)=\nabla g(x),\forall x\in\mathcal{X}$, where $\nabla$ is the Jacobian w.r.t $x$.

###### Proof.

Let us consider the interior of $\mathcal{X}$ and denote it as $\mathcal{X}^{\mathsf{int}}$. We first argue that the two functions $f$ and $g$ are equal at all points in the interior. Suppose there exists a point $x\in\mathcal{X}^{\mathsf{int}}$ at which $f(x)\not=g(x)$.
Consider a ball centered at $x$ of radius $r$ denoted as $B(x,r)\subset\mathcal{X}^{\mathsf{int}}$ (such a ball exists as this point is in the interior of $\mathcal{X}$.). We argue that there exists at least one point $x_{1}$ in this ball at which $f(x_{1})=g(x_{1})$. If this were not the case, then the equality will not hold on the entire ball, which would contradict the condition that the equality $f(x)=g(x)$ can only be violated on a set of measure zero. Note this condition holds true for all $r>0$. Suppose the distance of $x_{1}$ from $x$ is $r_{1}\leq r$. Consider another ball with radius $r_{2}<r_{1}$ and let $x_{2}\in B(x,r_{2})$ where the equality holds. By repeating this argument, we can construct a sequence $\{x_{k}\}_{k\in\mathbb{N}}$ that converges to $x$, where $\mathbb{N}$ is the set of natural numbers. On this sequence, the following conditions hold.

$$ $\begin{split}&f(x_{k})=g(x_{k}),\forall k\in\mathbb{N}\end{split}$ (5) $$

Further, from the continuity of $f$ and $g$ it follows that

$$ $\begin{split}\lim_{k\rightarrow\infty}f(x_{k})=f(x),\lim_{k\rightarrow\infty}g (x_{k})=g(x)\\ \end{split}$ (6) $$

Combining the above two conditions, we get that $f(x)=g(x)$. This leads to a contradiction since we assumed that $f(x)\not=g(x)$. Thus there can be no such $x$ in the interior at which $f(x)\not=g(x)$. From this it follows that $f(x)=g(x)$ for all $x\in\mathcal{X}^{\mathsf{int}}$. Now let us consider the closure of $\mathcal{X}^{\mathsf{int}}$, which is $\mathcal{X}$ itself since it is a regular closed set. Every point $x\in\mathcal{X}$ in the closure can be expressed as limit of points in $\mathcal{X}^{\mathsf{int}}$.
Consider an $x\in\mathcal{X}$ and from the definition of regular closed set it follows that $\lim_{k\rightarrow\infty}x_{k}=x$, where $x_{k}\in\mathcal{X}^{\mathsf{int}}$. We already know from the fact that $f$ and $g$ are equal in the interior

$$ $f(x_{k})=g(x_{k}),\forall k\in\mathbb{N}$ (7) $$

From the continuity of $f$ and $g$ it follows

$$ $\begin{split}\lim_{k\rightarrow\infty}f(x_{k})=f(x),\lim_{k\rightarrow\infty}g (x_{k})=g(x)\\ \end{split}$ (8) $$

Combining the above two we get that $f(x)=g(x)$ for all $x\in\mathcal{X}$.
After this we can use Lemma 6 from (lachapelle2023additive) to conclude that
$\nabla f(x)=\nabla g(x),\forall x\in\mathcal{X}$.
We repeat their proof here for completeness. For all points in the interior of $\mathcal{X}$, it follows that $\nabla f(x)=\nabla g(x),\forall x\in\mathcal{X}^{\mathsf{int}}$.

Now consider any point $x\in\mathcal{X}$. Since $\mathcal{X}$ is a regular closed set, $\lim_{k\rightarrow\infty}x_{k}=x$.
Since each $x_{k}$ is in the interior of $\mathcal{X}$ it follows that

$$ $\begin{split}\nabla f(x_{k})=\nabla g(x_{k}),\forall k\in\mathbb{N}\end{split}$ (9) $$

From the continuity of $\nabla f$ and $\nabla g$ it follows that

$$ $\lim_{k\rightarrow\infty}\nabla f(x_{k})=\nabla f(x),\lim_{k\rightarrow\infty} \nabla g(x_{k})=\nabla g(x)$ (10) $$

Combining the above conditions, we get that $\nabla f(x)=\nabla g(x)$. This completes the proof.

∎

###### Lemma 1 .

###### Proof.

### C.1 Deep sets

In this section, we provide the proofs for length and compositional generalization for deep sets. We first provide the proof for Theorem 1, followed by Corollary 1, where we establish linear identification. We then present the discrete tokens counterpart to Theorem 1 in Theorem 6.
In the next part of this section, we extend Theorem 1 with $\omega$ from $C^{1}$-diffeomorphism in Theorem 7.

We restate the theorems from the main body for convenience of the reader. In what follows, we remind the reader that we denote the labeling function $f(\mathcal{X})=\rho(\sum_{x\in\mathcal{X}}\phi(x))$ and the function learned is denoted as $h(\mathcal{X})=\omega(\sum_{x\in\mathcal{X}}\psi(x))$.

See 1

###### Proof.

Consider any $h$ that solves equation 1. Since $\ell$ is $\ell_{2}$ loss and realizability condition holds, $f$ is one of the optimal solutions to equation 1. For all $x_{\leq T}\in\mathsf{supp}(X_{\leq T})$ except over a set of measure zero the following condition holds

$$ $h(x_{\leq T})=f(x_{\leq T}).$ (11) $$

The above follows from the fact that $h$ solves equation 1, i.e., $\mathbb{E}[\|h-f\|^{2}]=0$ and from Theorem 1.6.6. (ash2000probability). Since $\mathsf{supp}(X_{\leq T})$ is regular closed, $f,h$ are both continuously differentiable, we can use Lemma 1, it follows that the above equality holds for all $x_{\leq T}\in\mathsf{supp}(X_{\leq T})$.
From realizability condition it follows that true $f(x_{\leq T})=\rho\Big{(}\sum_{j\leq T}\phi(x_{j})\Big{)}$. We substitute the functional decomposition from Assumption 1 to get

$$ $\begin{split}\omega\Big{(}\sum_{j\leq T}\psi(x_{j})\Big{)}=\rho\Big{(}\sum_{j \leq T}\phi(x_{j})\Big{)}.\\ \end{split}$ (12) $$

$\omega$ and $\rho$ are both single layer perceptron with a bijective activation $\sigma$. We substitute the parametric form of $\omega$ and $\rho$ to obtain

$$ $\begin{split}&\sigma\Big{(}A\sum_{j\leq T}\psi(x_{j})\Big{)}=\sigma\Big{(}B \sum_{j\leq T}\phi(x_{j})\Big{)}\implies A\sum_{j\leq T}\psi(x_{j})=B\sum_{j \leq T}\phi(x_{j}).\\ \end{split}$ (13) $$

The second equality in the above simplification follows from the fact that the activation $\sigma$ is bijective, the inputs to $\sigma$ are equal. We take the derivative of the expressions above w.r.t $x_{r}$ to get the following condition and equate them (follows from Lemma 1). For all $x_{r}\in\mathsf{supp}(X_{r})$, i.e., $x_{r}\in[0,1]^{n},$

$$ $\begin{split}\nabla_{x_{r}}\Big{(}A\sum_{j\leq T}\psi(x_{j})\Big{)}=\nabla_{x_ {r}}\Big{(}B\sum_{j\leq T}\phi(x_{j})\Big{)}.\\ \end{split}$ (14) $$

We drop the subscript $r$ to simplify the notation. Therefore, for all $x\in[0,1]^{n}$

$$ $\begin{split}A\nabla_{x}\psi(x)=B\nabla_{x}\phi(x),\end{split}$ (15) $$

where $\nabla_{x}\psi(x)$ is the Jacobian of $\psi(x)$ w.r.t $x$ and $\nabla_{x}\phi(x)$ is the Jacobian of $\phi(x)$ w.r.t $x$.
We now take the derivative w.r.t some component $x^{k}$ of vector $x=[x^{1},\cdots,x^{n}]$. Denote the components other than $k$ as $x^{-k}=x\setminus x^{k}$. From the above condition, it follows that for all $x\in[0,1]^{n}$

$$ $\begin{split}A\frac{\partial\psi(x)}{\partial x^{k}}=B\frac{\partial\phi(x)}{ \partial x^{k}}.\end{split}$ (16) $$

Using fundamental theorem of calculus, we can integrate both sides for fixed $x^{-k}$ and obtain the following for all $x^{k}\in[0,1],$

$$ $\begin{split}&A\psi(x^{k},x^{-k})=B\phi(x^{k},x^{-k})+C_{k}(x^{-k})\implies A \psi(x)-B\phi(x)=C_{k}(x^{-k}).\end{split}$ (17) $$

The above condition is true of all $k\in\{1,\cdots,n\}$. Hence, we can deduce that for all $x\in[0,1]^{n}$ and for $k\not=j$, where $j,k\in\{1,\cdots,d\}$,

$$ $A\psi(x)-B\phi(x)=C_{k}(x^{-k})=C_{j}(x^{-j}).$ (18) $$

Take the partial derivative of $C_{k}(x^{-k})$ and $C_{j}(x^{-j})$ w.r.t $x^{j}$ to obtain, for all $x^{j}\in[0,1]$,

$$ $\begin{split}\frac{\partial C_{k}(x^{-k})}{\partial x^{j}}=\frac{\partial C_{j }(x^{-j})}{\partial x^{j}}=0.\end{split}$ (19) $$

In the above simplification, we use the fact that $\forall x^{j}\in[0,1],\frac{\partial C_{j}(x^{-j})}{\partial x^{j}}=0$. Therefore, $C_{k}(x^{-k})$ cannot depend on $x^{j}$. We can apply the same condition on all $j\not=k$. As a result, $C_{k}(x^{-k})$ is a fixed constant vector denoted as $C$. We write this as

$$ $A\psi(x)=B\phi(x)+C.$ (20) $$

Substitute the above into $A\sum_{j\leq T}\psi(x_{j})=B\sum_{j\leq T}\phi(x_{j})$ to obtain

$$ $B\sum_{j\leq T}\phi(x_{j})+CT=B\sum_{j\leq T}\phi(x_{j})\implies C=0.$ (21) $$

Therefore, we get

$$ $\forall x\in[0,1]^{n},A\psi(x)=B\phi(x).$ (22) $$

We now consider any sequence $x_{\leq\tilde{T}}$ from $[0,1]^{n\tilde{T}}$. The prediction made by $h$ is

$$ $\begin{split}h(x_{\leq\tilde{T}})=\sigma\Big{(}A\sum_{j\leq\tilde{T}}\psi(x_{j })\Big{)}=\sigma\Big{(}B\sum_{j\leq\tilde{T}}\phi(x_{j})\Big{)}=f(x_{\leq \tilde{T}}).\end{split}$ (23) $$

We use equation 22 in the simplification above. From the above, we can conclude that $h$ continues to be optimal for distribution $\mathbb{P}_{X_{\leq\tilde{T}}}$.

∎

###### Corollary 1 .

If $\mathcal{H}$ follows Assumption 1 with the condition that the output layer weight matrix is left invertible, the realizability condition holds, i.e., $f\in\mathcal{H}$, $\mathsf{supp}(X_{j})=[0,1]^{n},\;\forall j\geq 1$, and the regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss achieves linear identification. Further, under the stated conditions linear identification is necessary for compositional and length generalization.

###### Proof.

We follow the exact same steps as in the previous proof of Theorem 1 up to equation 22. We restate equation 22 below.

$$ $\begin{split}\forall x\in[0,1]^{n},\;&A\psi(x)=B\phi(x)\\ &\psi(x)=A^{-1}B\phi(x)\end{split}$ (24) $$

The above condition establishes linear identification, i.e., the learned model’s representation for a token is a linear transform of the true model’s representation. From the above, we can write that $x_{\leq\tilde{T}}$ from $[0,1]^{n\tilde{T}}$

$$ $\sum_{j\leq\tilde{T}}\psi(x_{j})=A^{-1}B\sum_{j\leq\tilde{T}}\phi(x_{j})$ (25) $$

The above shows linear relationship holds for the entire sequence as well.

Now let us turn to the part on necessity. From the proof of previous theorem, we know that

$$ $\forall x_{\leq T}\in\mathsf{supp}(X_{\leq T}),\;\;\sigma\Big{(}A\sum_{j\leq T }\psi(x_{j})\Big{)}=\sigma\Big{(}B\sum_{j\leq T}\phi(x_{j})\Big{)}\implies \forall x\in[0,1]^{n},A\psi(x)=B\phi(x)$ (26) $$

Thus if $\forall x\in[0,1]^{n},A\psi(x)=B\phi(x)$ is not true, then $\sigma\Big{(}A\sum_{j\leq T}\psi(x_{j})\Big{)}=\sigma\Big{(}B\sum_{j\leq T}
\phi(x_{j})\Big{)}$ cannot be true either. Therefore, in the absence of linear identification neither length nor compositional generalization are achievable.
∎

###### Proof.

###### Corollary 1 .

###### Proof.

##### Remarks

A few remarks and observations from the proof are in order. Firstly, observe that we do not require $\phi$ and $\psi$ to have the same output dimension for the above proof to go through. Secondly, in Theorem 1, we observe all the labels from $t=1$ to $T$, i.e., $y_{1}$ to $y_{T}$. The result continues to hold if we only observe label at length $T$, i.e., $y_{T}$. Finally, we make an observation in this result, which would apply to all the subsequent theorems. The definition of compositional generalization requires generalization to the Cartesian product over sequences of length $T$, where $T$ is the training length. Since our model generalizes to the hypercube $[0,1]^{nt},\forall t$, we achieve compositional generalization even beyond the training lengths.

#### C.1.1 Extending Theorem 1 to discrete tokens

In our discussion, we have focused on settings where the support of each token has a non-empty interior (Assumption 2). In practice of language modeling, we use discrete tokens and hence Assumption 2 does not hold anymore. In this section, we discuss the adaptation of results for deepsets to setting when the the support of tokens is a finite set.

###### Assumption 7 .

The marginal support of token for all positions is the same and denoted as $\mathcal{X}$. The joint support of first and second token is $\mathcal{X}\times\mathcal{X}$.

###### Theorem 6 .

If $\mathcal{H}$ follows Assumption 1, the realizability condition holds, i.e., $f\in\mathcal{H}$, and Assumption  7 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss generalizes to all sequences in the hypercube $[0,1]^{nt},\;\forall t\geq 1$ and thus achieves length and compositional generalization.

###### Proof.

Consider any $h$ that solves equation 1. Since $\ell$ is $\ell_{2}$ loss and realizability condition holds, $f$ is one of the optimal solutions to equation 1. For all $x_{\leq T}\in\mathsf{supp}(X_{\leq T})$ except over a set of measure zero the following condition holds

$$ $h(x_{\leq T})=f(x_{\leq T}).$ (27) $$

The above follows from the fact that $h$ solves equation 1, i.e., $\mathbb{E}[\|h-f\|^{2}]=0$ and from Theorem 1.6.6. (ash2000probability). Since $\mathsf{supp}(X_{\leq T})$ is regular closed, $f,h$ are both continuously differentiable, we can use Lemma 1, it follows that the above equality holds for all $x_{\leq T}\in\mathsf{supp}(X_{\leq T})$.
From realizability condition it follows that true $f(x_{\leq T})=\rho\Big{(}\sum_{j\leq T}\phi(x_{j})\Big{)}$. We substitute the functional decomposition from Assumption 1 to get

$$ $\begin{split}\omega\Big{(}\sum_{j\leq T}\psi(x_{j})\Big{)}=\rho\Big{(}\sum_{j \leq T}\phi(x_{j})\Big{)}.\\ \end{split}$ (28) $$

$\omega$ and $\rho$ are both single layer perceptron with a bijective activation $\sigma$. We substitute the parametric form of $\omega$ and $\rho$ to obtain

$$ $\begin{split}&\sigma\Big{(}A\sum_{j\leq T}\psi(x_{j})\Big{)}=\sigma\Big{(}B \sum_{j\leq T}\phi(x_{j})\Big{)}\implies A\sum_{j\leq T}\psi(x_{j})=B\sum_{j \leq T}\phi(x_{j}).\\ \end{split}$ (29) $$

The second equality in the above simplification follows from the fact that the activation $\sigma$ is bijective, the inputs to $\sigma$ are equal.

From Assumption 7, it follows that for all $x_{1},x_{2}\in\mathcal{X}\times\mathcal{X}$

$$ $A\phi(x_{1})+A\phi(x_{2})=B\psi(x_{1})+B\psi(x_{2})$ (30) $$

Consider a reference point $x^{r}\in\mathcal{X}$. Consider the point $(x^{r},x_{2})$ in the support (this point exists in the support due to Assumption 7).

$$ $A\phi(x^{r})+A\phi(x_{2})=B\psi(x^{r})+B\psi(x_{2})$ (31) $$

Take the difference of equation 31 of equation 30 to obtain

$$ $A\phi(x_{1})=B\psi(x_{1})+C,\forall x_{1}\in\mathcal{X}.$ (32) $$

where $C=\Big{(}B\psi(x^{r})-A\phi(x^{r})\Big{)}$. Substitute equation 32 into equation 30 we obtain that $C=0$.

Therefore, we get

$$ $\forall x\in\mathcal{X},A\psi(x)=B\phi(x).$ (33) $$

We now consider any sequence $x_{\leq\tilde{T}}$ from $\mathcal{X}^{\tilde{T}}$. The prediction made by $h$ is

$$ $\begin{split}h(x_{\leq\tilde{T}})=\sigma\Big{(}A\sum_{j\leq\tilde{T}}\psi(x_{j })\Big{)}=\sigma\Big{(}B\sum_{j\leq\tilde{T}}\phi(x_{j})\Big{)}=f(x_{\leq \tilde{T}}).\end{split}$ (34) $$

We use equation 22 in the simplification above. From the above, we can conclude that $h$ continues to be optimal for distribution $\mathbb{P}_{X_{\leq\tilde{T}}}$.

∎

###### Assumption 7 .

###### Theorem 6 .

###### Proof.

#### C.1.2 Extending Theorem 1 to ω 𝜔 \omega italic_ω from C 1 superscript 𝐶 1 C^{1} italic_C start_POSTSUPERSCRIPT 1 end_POSTSUPERSCRIPT -diffeomorphisms class

###### Assumption 8 .

Each function in $\mathcal{H}$ is expressed as $h(x_{1},\cdots,x_{i})=\omega(\sum_{j=1}^{i}\psi(x_{j}))$, where $\omega$ is a $C^{1}$-diffeomorphism.

###### Assumption 9 .

The joint support $\mathsf{supp}(X_{\leq i})$ is a regular closed set for all $i\leq T$. The support of all tokens is equal, i.e., $\mathsf{supp}(X_{j})=[0,1]^{n}$, where $j\geq 1$. The support of $[\phi(X_{1}),\phi(X_{2})]$ is $\mathbb{R}^{2m}$, where $\phi$ is the embedding function for the labeling function $f(\mathcal{X})=\rho(\sum_{x\in\mathcal{X}}\phi(x))$.

We provide a remark on the assumption and where it is used following the proof of the next theorem.

###### Theorem 7 .

If $\mathcal{H}$ follows Assumption 8, the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further assumption on the support (Assumption 9) holds, then the model trained to minimize the risk in equation 1 (with $T\geq 2$) with $\ell_{2}$ loss generalizes to all sequences in $[0,1]^{nt},\forall t\geq 1$ and thus achieves length and compositional generalization.

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$. We first use the fact $h(x_{\leq i})=f(x_{\leq i})$ almost everywhere in the support. We can use the continuity of $h,f$ and regular closedness of the support to extend the equality to all points in the support (follows from the first part of Lemma 1) to obtain the following. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$

$$ $\begin{split}&\omega\Big{(}\sum_{j\leq i}\psi(x_{j})\Big{)}=\rho\Big{(}\sum_{j \leq i}\phi(x_{j})\Big{)}\implies\sum_{j\leq i}\psi(x_{j})=\omega^{-1}\circ \rho\Big{(}\sum_{j\leq i}\phi(x_{j})\Big{)}\implies\\ &\sum_{j\leq i}\psi(x_{j})=a\Big{(}\sum_{j\leq i}\phi(x_{j})\Big{)},\end{split}$ (35) $$

where $a=\omega^{-1}\circ\rho$. In the above simplification, we used the parametric form for the true labeling function and the learned labeling function and use the invertibility of $\omega$. Let us consider the setting when $i=1$. In that case summation involves only one term. Substitute $x_{1}=x$. We obtain $\forall x\in[0,1]^{n},$

$$ $\psi(x)=a(\phi(x)).$ (36) $$

The above expression implies that $\psi$ bijectively identifies $\phi$. Let us consider the setting when $i=2$. Substitute $x_{1}=x$ and $x_{2}=y$. We obtain

$$ $a(\phi(x))+a(\phi(y))=a\big{(}\phi(x)+\phi(y)\big{)}.$ (37) $$

We now use the that assumption $[\phi(x),\phi(y)]$ spans $\mathbb{R}^{2m}$, where $\phi(x)$ and $\phi(y)$ individually span $\mathbb{R}^{m}$. Substitute $\phi(x)=\alpha$ and $\phi(y)=\beta$. We obtain $\forall\alpha\in\mathbb{R}^{m}$, $\forall\beta\in\mathbb{R}^{m}$

$$ $a(\alpha)+a(\beta)=a\big{(}\alpha+\beta\big{)}.$ (38) $$

Observe that $a(0)=0$ (substitute $\alpha=\beta=0$ in the above).

We use equation 38 to show that $a$ is linear. To show that, we need to argue that $a(c\alpha)=ca(\alpha)$ as we already know $a$ satisfies additivity condition.

From the identity above, we want to show that equation 74 $a(p\alpha)=pa(\alpha)$, where $p$ is some integer.

Substitute $\beta=-\alpha$ in $a(\alpha+\beta)=a(\alpha)+a(\beta)$. We obtain $a(0)=a(\alpha)+a(-\alpha)\implies a(-\alpha)=-a(\alpha)$. Suppose $p$ is a positive integer. We simplify $a(p\alpha)$ as follows $a(\alpha+(p-1)\alpha)=a(\alpha)+a((p-1)\alpha)$. Repeating this simplification, we get $a(p\alpha)=pa(\alpha)$. Suppose $p$ is a negative integer. We can write $a(p\alpha)=a(-p\times-\alpha)=-pa(-\alpha)$. Since $a(-\alpha)=-a(\alpha)$, we get $a(p\alpha)=pa(\alpha)$.

Suppose $c$ is some rational number, i.e., $c=p/q$, where $p$ and $q$ are non-zero integers. We already know $a(p\alpha)=pa(\alpha)$. Further, we obtain

$a(q\frac{1}{q}\alpha)=qa(\frac{1}{q}\alpha)\implies a(\frac{1}{q}\alpha)=\frac
{1}{q}a(\alpha)$, where $q$ is some integer.

Now combine these $a(p/q\alpha)=pa(1/q\alpha)=\frac{p}{q}a(\alpha)$. We have established the homogeneity condition for rationals.

We will now use the continuity of the function $a$ and density of rationals to extend the claim for irrationals. Suppose $c$ is some irrational. Define a sequence of rationals that approach $c$ (this follows from the fact that rationals are dense in $\mathbb{R}$).

$a(c\alpha)=a(\lim_{n\rightarrow\infty}q_{n}\alpha)=\lim_{n\rightarrow\infty}a(
q_{n}\alpha).$

In the second equality above, we use the definition of continuity ($a$ is continuous since composition of continuous functions is continuous). We can also use the property that we already showed for rationals to further simplify

$\lim_{n\rightarrow\infty}a(q_{n}\alpha)=a(\alpha)\lim_{n\rightarrow\infty}q_{n
}=ca(\alpha).$

Observe that $a:\mathbb{R}^{m}\rightarrow\mathbb{R}^{m}$ and for any $\alpha,\beta\in\mathbb{R}^{m}$ $a(\alpha+\beta)=a(\alpha)+a(\beta)$ and $a(c\alpha)=ca(\alpha)$. From the definition of a linear map it follows that $a$ is linear. As a result, we can write $\forall x\in[0,1]^{n}$

$$ $\psi(x)=A(\phi(x))$ (39) $$

Observe that $a$ is invertible because both $\rho$ and $\omega$ are invertible. As a result, we know that $A$ is an invertible matrix. From this we get

$$ $\phi(x)=A^{-1}\psi(x)=C(\psi(x))$ (40) $$

For all $z\in\mathbb{R}^{m}$, we obtain

$$ $a(z)=\rho^{-1}\circ\omega(z)=Cz\implies\omega(z)=\rho(Cz)$ $$

Let us consider any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$.
We use the above conditions

$$ $\omega\big{(}\sum_{j\leq\tilde{T}}\psi(x_{j})\big{)}=\rho(C\sum_{j\leq\tilde{T }}\psi(x_{j}))=\rho\big{(}\sum_{j\leq\tilde{T}}\phi(x_{j})\big{)}$ $$

Thus we obtain length and compositional generalization.

∎

###### Assumption 8 .

###### Assumption 9 .

###### Theorem 7 .

###### Proof.

##### Remark on Assumption 9

In Assumption 9, we require that the support of $[\phi(X_{1}),\phi(X_{2})]$ is $\mathbb{R}^{2m}$. This assumption is used in the proof in equation equation 38. We used this assumption to arrive at $a(\alpha+\beta)=a(\alpha)+a(\beta),\forall\alpha,\beta\in\mathbb{R}^{m}$. We then used continuity of $a$ to conclude $a$ is linear. Now suppose $[\phi(X_{1}),\phi(X_{2})]$ is some subset $\mathcal{Z}\subseteq\mathbb{R}^{2m}$. We believe that it is possible to extend the result to more general $\mathcal{Z}$, it might still be possible to arrive at $a$ is linear. We leave this investigation to future work.

##### Remark on expressivity under Assumption 8 and Assumption 9

Assumption 9 requires $\omega$ is a $C^{1}$-diffeomorphism.
Suppose the label is one dimensional, i.e., $m=1$. From Assumption 9 output dimension of $\phi$ is restricted to be one dimensional. Consider the map $h(x_{1},\cdots,x_{i})=\rho(\sum_{j\leq i}\phi(x_{j}))$. The output dimension of $\phi$ is required to grow with sequence length to express all permutation invariant maps (See Theorem 7 in (zaheer2017deep)). Thus by restricting the output dimension of $\phi$ to one, we cannot express all the permutation invariant maps.

##### Multiplication operator

Consider the multiplication operator $y_{i}=\Pi_{j=1}^{i}x_{i}$, where each $x_{i}>0$. Observe that we can rewrite this as $y_{i}=\exp(\sum_{j=1}^{i}\log(x_{j}))$. This operator is realizable on deep sets from hypothesis class described by Assumption 8 with $\omega=\exp$ and $\psi=\log$. In Assumption 9, we require the support of $[\phi(X_{1}),\phi(X_{2})]$ to be $\mathbb{R}^{2}$. We let the support of $X_{1}$ and $X_{2}$ be $(0,\infty)$. In Assumption 9 we require that the support of each token was equal to $[0,1]$. However, the proof of Theorem 7 still goes through even if support is $(0,\infty)$. Hence, we can use Theorem 7 to conclude that deep sets trained to predict the output of multiplication can multiply longer sequences and also multiply new token combinations.

### C.2 Transformers

In this section, we provide the proofs for length and compositional generalization for transformers.
We first provide the proof for Theorem 2, followed by Corollary 2, where we establish linear identification. We present an extention of Theorem 2 to incorporate positional encoding in Theorem 8. We then present the discrete tokens counterpart to Theorem 2 in Theorem 9. In the next part of this section, we extend Theorem 2 with $\omega$ from $C^{1}$-diffeomorphism in Theorem 3. Theorem 10 adapts Theorem 3 to incorporate positional encodings.

We restate the theorems from the main body for convenience of the reader. In what follows, we want to remind the reader we denote the labeling function $f(x_{1},\cdots,x_{i})=\rho(\sum_{j\leq i}\phi(x_{i},x_{j}))$ and the function learned is denoted as $h(x_{1},\cdots,x_{i})=\omega(\sum_{j\leq i}\psi(x_{i},x_{j}))$.

See 2

###### Proof.

Consider any $h$ that solves equation 1. Since $\ell$ is $\ell_{2}$ loss and realizability condition holds, $f$ is one of the optimal solutions to equation 1. For all $i\leq T,x_{\leq i}\in\mathsf{supp}(X_{\leq i})$ except over a set of measure zero the following condition holds

$$ $h(x_{\leq i})=f(x_{\leq i}).$ (41) $$

The above follows from the fact that $h$ solves equation 1, i.e., $\mathbb{E}[\|h-f\|^{2}]=0$ and from Theorem 1.6.6. (ash2000probability). Since $\mathsf{supp}(X_{\leq i})$ is regular closed, $f,h$ are both continuously differentiable, we can use Lemma 1, it follows that the above equality holds for all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$.
From realizability condition it follows that true $f(x_{\leq i})=\rho\Big{(}\sum_{k\leq i}\phi(x_{i},x_{k})\Big{)}$. We substitute the parametric forms from Assumption 3 to get

$$ $\begin{split}\omega\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\psi(x_{i},x_{k})\Big{ )}=\rho\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\phi(x_{i},x_{k})\Big{)}.\\ \end{split}$ (42) $$

Since $\omega$ and $\rho$ are single layer perceptron with bijective activation $\sigma$. We substitute the parametric form of $\omega$ and $\rho$ to obtain the following condition. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$,

$$ $\begin{split}&\sigma\Big{(}A\sum_{k\leq i}\frac{1}{i}\cdot\psi(x_{i},x_{k}) \Big{)}=\sigma\Big{(}B\sum_{k\leq i}\frac{1}{i}\cdot\phi(x_{i},x_{k})\Big{)} \implies A\sum_{k\leq i}\psi(x_{i},x_{k})=B\sum_{k\leq i}\phi(x_{i},x_{k}).\\ \end{split}$ (43) $$

The second equality follows from the fact that the activation $\sigma$ is bijective and hence the inputs to $\sigma$ are equal. We take the derivative of the expressions above w.r.t $x_{j}$ to get the following (follows from Lemma 1). For $j<i$ (there exists a $j<i$ as $T\geq 2$ and we can set $i\geq 2$) and for all $x_{j}\in\mathsf{supp}(X_{j})$, i.e., $x_{j}\in[0,1]^{n},$

$$ $\begin{split}&\nabla_{x_{j}}\Big{(}A\sum_{k\leq i}\psi(x_{i},x_{k})\Big{)}= \nabla_{x_{j}}\Big{(}B\sum_{k\leq i}\phi(x_{i},x_{k})\Big{)}\implies\\ &A\nabla_{x_{j}}\psi(x_{i},x_{j})=B\nabla_{x_{j}}\phi(x_{i},x_{j}),\end{split}$ (44) $$

where $\nabla_{x_{j}}\psi(x_{i},x_{j}),\nabla_{x_{j}}\phi(x_{i},x_{j})$ are the Jacobians of $\psi$ and $\phi$ w.r.t $x_{j}$ for a fixed $x_{i}$. Note that $A\nabla_{x_{j}}\psi(x_{i},x_{j})=B\nabla_{x_{j}}\phi(x_{i},x_{j})$ holds for all $x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n}$ (here we use the fact that joint support of every pair of tokens spans $2n$ dimensional unit hypercube assumed in the Theorem 8). In this equality, we now consider the derivative w.r.t some component $x_{j}^{k}$ of $x_{j}$. Denote the remaining components as $x_{j}^{-k}$. From the above condition it follows that for all $x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n},$

$$ $\begin{split}A\frac{\partial\psi(x_{i},x_{j})}{\partial x_{j}^{k}}=B\frac{ \partial\phi(x_{i},x_{j})}{\partial x_{j}^{k}}.\end{split}$ (45) $$

Using fundamental theorem of calculus, we can integrate both sides for fixed $x_{j}^{-k}$ and obtain the following for all $x_{j}^{k}\in[0,1],$

$$ $\begin{split}&A\psi\big{(}x_{i},[x_{j}^{k},x_{j}^{-k}]\big{)}=B\phi\big{(}x_{i },[x_{j}^{k},x_{j}^{-k}]\big{)}+C_{k}\big{(}x_{i},x_{j}^{-k}\big{)}=\\ &A\psi(x_{i},x_{j})=B\phi(x_{i},x_{j})+C_{k}(x_{i},x_{j}^{-k}).\end{split}$ (46) $$

The same condition is true of all $k$. Hence, $\forall x_{i}\in[0,1]^{d},\forall x_{j}\in[0,1]^{d}$ and for $k\not=q$, where $q,k\in\{1,\cdots,d\},$

$$ $A\psi(x_{i},x_{j})-B\phi(x_{i},x_{j})=C_{k}(x_{i},x_{j}^{-k})=C_{q}(x_{i},x_{j }^{-q}).$ (47) $$

Take the partial derivative of both sides w.r.t $x_{j}^{q}$ to obtain, $\forall x_{j}^{q}\in[0,1],$

$$ $\begin{split}\frac{\partial C_{k}(x_{i},x_{j}^{-k})}{\partial x_{j}^{q}}=\frac {\partial C_{q}(x_{i},x_{j}^{-q})}{\partial x_{j}^{q}}=0.\end{split}$ (48) $$

Therefore, $C_{k}(x_{i},x_{j}^{-k})$ cannot depend on $x_{j}^{q}$. We can apply the same condition on all $q\not=k$. As a result, $C_{k}(x_{i},x_{j}^{-k})$ is only a function of $x_{i}$ denoted as $C(x_{i})$. Therefore, for $j<i$ and for all $x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n}$

$$ $A\psi(x_{i},x_{j})=B\phi(x_{i},x_{j})+C(x_{i}).$ (49) $$

If we substitute $x_{i}=x_{j}=x$, then the above equality extends for $i=j$ and thus we get

$$ $A\psi(x_{i},x_{i})=B\phi(x_{i},x_{i})+C(x_{i}).$ (50) $$

Substitute the above equation 49, equation 50 into $A\sum_{k\leq i}\psi(x_{i},x_{k})=B\sum_{k\leq i}\phi(x_{i},x_{k})$ to obtain

$$ $\begin{split}&B\sum_{k\leq i}\phi(x_{i},x_{k})+(i)C(x_{i})=B\sum_{k\leq i}\phi (x_{i},x_{k})\implies C(x_{i})=0.\end{split}$ (51) $$

Thus we obtain

$$ $\forall x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n}\hskip 5.69054ptA\psi(x_{i},x_{j})= B\phi(x_{i},x_{j}).$ (52) $$

We now consider any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$. The prediction made by $h$ is

$$ $\begin{split}h(x_{\leq\tilde{T}})=\sigma\Big{(}A\sum_{j\leq\tilde{T}}\psi(x_{ \tilde{T}},x_{j})\Big{)}=\sigma\Big{(}B\sum_{j\leq\tilde{T}}\phi(x_{\tilde{T}} ,x_{j})\Big{)}=f(x_{\leq\tilde{T}})\end{split}$ (53) $$

We use equation 52 in the simplification above. From the above, we can conclude that $h$ continues to be optimal for all sequences in
$[0,1]^{n\tilde{T}}$.

∎

###### Corollary 2 .

If $\mathcal{H}$ follows Assumption 3 with the condition that the output layer weight matrix is left invertible, the realizability condition holds, i.e., $f\in\mathcal{H}$, $\mathsf{supp}(X_{i},X_{j})=[0,1]^{2n},\;\forall i\not=j$ and the regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 (with $T\geq 2$) with $\ell_{2}$ loss achieves linear identification. Further, linear identification is necessary for both length and compositional generalization.

###### Proof.

We follow the exact same steps as in the previous proof of Theorem 2 up to equation 52. We restate equation 52 below.

$$ $\begin{split}\forall x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n}\hskip 5.69054pt&A\psi (x_{i},x_{j})=B\phi(x_{i},x_{j})\\ &\psi(x_{i},x_{j})=A^{-1}B\phi(x_{i},x_{j})\end{split}$ (54) $$

In the second step above, we use left invertibility of $A$. The above condition establishes linear identification, i.e., the learned model’s representation is a linear transform of the true model’s representation. From this we obtain that for any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$

$$ $\sum_{j\leq\tilde{T}}\psi(x_{\tilde{T}},x_{j})=A^{-1}B\Big{(}\sum_{j\leq\tilde {T}}\psi(x_{\tilde{T}},x_{j})\Big{)}$ (55) $$

The above establishes a linear relationship between the learned representation of the sequence and the representation of the sequence under the true model. Now let us turn to the part on necessity. From the proof of previous theorem, we know that

$$ $\omega\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\psi(x_{i},x_{k})\Big{)}=\rho\Big{( }\sum_{k\leq i}\frac{1}{i}\cdot\phi(x_{i},x_{k})\Big{)}\implies\forall x_{i} \in[0,1]^{n},x_{j}\in[0,1]^{n}\hskip 5.69054ptA\psi(x_{i},x_{j})=B\phi(x_{i},x _{j})$ (56) $$

Thus from the above it follows that in the absence of linear identification neither length nor compositional generalization are achievable.
∎

###### Proof.

###### Corollary 2 .

###### Proof.

##### On the absence of labels at all lengths from t = 1 𝑡 1 t=1 italic_t = 1 to t = T 𝑡 𝑇 t=T italic_t = italic_T

A few important remarks are to follow. In the proof above, we do not require to observe all the labels from $t=1$ to $t=T$, where $T\geq 2$. The proof goes through provided we observe data at two different lengths.

#### C.2.1 Extension of Theorem 2 to incorporate positional encodings

In what follows, we extend the above result (Theorem 2) to incorporate positional encoding. We start with extension of the hypothesis class to incorporate positional encoding.

###### Assumption 10 .

Each function in the hypothesis class $\mathcal{H}$ used by the learner is given as $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j\leq i}\frac{1}{i}\psi_{i-j}(x_{i},x
_{j})\Big{)}$, where $\omega$ is a single layer perceptron with continuously differentiable bijective activation (e.g., sigmoid) and each $\psi_{k}$ is a map that is differentiable. Also, $\psi_{k}=0$ for $k\geq T_{\mathsf{max}}$, i.e., two tokens that are sufficiently far apart do not interact.

In the above assumption, we incorporate relative positional encodings by making the function $\psi_{i-j}$ depend on the relative positional difference between token $x_{i}$ and token $x_{j}$. We would like to emphasize the reasons why we assume that the tokens that are sufficiently far apart do not interact. Suppose $T_{\mathsf{max}}=\infty$, which implies tokens at all positions interact. As a result, during training since we only see sequences of finite length $T$, we will not see the effect of interactions of tokens that are separated at a distance larger than $T$ on the data generation, which makes it impossible to learn anything about $\phi_{i-j}$, where $i-j\geq T-1$.

In the theorem that follows, we show that we can achieve length and compositional generalization for the above hypothesis class.

###### Theorem 8 .

If $\mathcal{H}$ follows Assumption 10, the realizability condition holds, i.e., $f\in\mathcal{H}$, $\mathsf{supp}(X_{i},X_{j})=[0,1]^{2n},\;\forall i\not=j\in\{1,\cdots,\infty\}$, the regular closedness condition in Assumption 2 holds and $T\geq T_{\mathsf{max}}\geq 2$, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss generalizes to all sequences in the hypercube $[0,1]^{nt},\;\forall t$ and thus achieves length and compositional generalization.

###### Proof.

Consider any $h$ that solves equation 1. Since $\ell$ is $\ell_{2}$ loss and realizability condition holds, $f$ is one of the optimal solutions to equation 1. For all $i\leq T$ and for all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$ except over a set of measure zero the following condition holds

$$ $h(x_{\leq i})=f(x_{\leq i}).$ (57) $$

The above follows from the fact that $h$ solves equation 1, i.e., $\mathbb{E}[\|h-f\|^{2}]=0$ and from Theorem 1.6.6. (ash2000probability). Since $\mathsf{supp}(X_{\leq i})$ is regular closed, $f,h$ are both continuously differentiable, we can use Lemma 1, it follows that the above equality holds for all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$.
From realizability condition it follows that true $f(x_{\leq i})=\rho\Big{(}\sum_{k\leq i}\phi_{i-k}(x_{i},x_{k})\Big{)}$. We substitute the parametric forms from Assumption 3 to get

$$ $\begin{split}\omega\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\psi_{i-k}(x_{i},x_{k} )\Big{)}=\rho\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\phi_{i-k}(x_{i},x_{k})\Big{ )}.\\ \end{split}$ (58) $$

Since $\omega$ and $\rho$ are single layer perceptron with bijective activation $\sigma$. We substitute the parametric form of $\omega$ and $\rho$ to obtain the following condition. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$,

$$ $\begin{split}&\sigma\Big{(}A\sum_{k\leq i}\frac{1}{i}\cdot\psi_{i-k}(x_{i},x_{ k})\Big{)}=\sigma\Big{(}B\sum_{k\leq i}\frac{1}{i}\cdot\phi_{i-k}(x_{i},x_{k}) \Big{)}\implies\\ &A\sum_{k\leq i}\psi_{i-k}(x_{i},x_{k})=B\sum_{k\leq i}\phi_{i-k}(x_{i},x_{k}) .\\ \end{split}$ (59) $$

The second equality follows from the fact that the activation $\sigma$ is bijective and hence the inputs to $\sigma$ are equal. We take the derivative of the expressions above w.r.t $x_{j}$ to get the following (follows from Lemma 1). The equality holds true for all $i\leq T$.

From the above, we can use $i=1$ and obtain

$$ $A\psi_{0}(x_{1},x_{1})=B\phi_{0}(x_{1},x_{1}),\forall x_{1}\in[0,1]^{n}.$ $$

From $i=2$, we obtain

$$ $A\psi_{0}(x_{2},x_{2})+A\psi_{1}(x_{2},x_{1})=B\phi_{0}(x_{2},x_{2})+B\phi_{1} (x_{2},x_{1}),\forall x_{1}\in[0,1]^{n},x_{2}\in[0,1]^{n}$ $$

Combining the two conditions we get

$$ $A\psi_{1}(x_{2},x_{1})=B\phi_{1}(x_{2},x_{1}),\forall x_{1}\in[0,1]^{n},x_{2} \in[0,1]^{n}.$ $$

We can use this argument and arrive at

$$ $A\psi_{i-1}(x_{i},x_{1})=B\phi_{i-1}(x_{i},x_{1}),\forall x_{i}\in[0,1]^{n},x_ {1}\in[0,1]^{n},\forall i\leq T.$ $$

Thus we obtain

$$ $\forall i-j\leq T-1,\forall x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n},\hskip 5.69054 ptA\psi_{i-j}(x_{i},x_{j})=B\phi_{i-j}(x_{i},x_{j}).$ (60) $$

From Assumption 10 and $T\geq T_{\mathsf{max}}$, we already know that

$$ $\forall i-j\geq T,\forall x_{i}\in[0,1]^{n},x_{j}\in[0,1]^{n},\hskip 5.69054 ptA\psi_{i-j}(x_{i},x_{j})=B\phi_{i-j}(x_{i},x_{j})=0.$ (61) $$

If $A$ is left invertible, then the above condition implies that linear representation identification is necessary for both compositional and length generalization.

We now consider any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$. The prediction made by $h$ is

$$ $\begin{split}h(x_{\leq\tilde{T}})=\sigma\Big{(}A\sum_{j\leq\tilde{T}}\psi_{ \tilde{T}-j}(x_{\tilde{T}},x_{j})\Big{)}=\sigma\Big{(}B\sum_{j\leq\tilde{T}} \phi_{\tilde{T}-j}(x_{\tilde{T}},x_{j})\Big{)}=f(x_{\leq\tilde{T}})\end{split}$ (62) $$

We use equation 52 in the simplification above. From the above, we can conclude that $h$ continues to be optimal for all sequences in
$[0,1]^{n\tilde{T}}$.

∎

###### Assumption 10 .

###### Theorem 8 .

###### Proof.

#### C.2.2 Extending Theorem 2 to discrete tokens

In the above result we used Assumption 2. In practice of language modeling, we use discrete tokens and hence Assumption 2 does not hold anymore. In this section, we discuss the adaptation of results for transformers to setting when the the support of tokens is a finite set.

###### Assumption 11 .

The marginal support of token for all positions is the same and denoted as $\mathcal{X}$. The joint support of first three tokens is $\mathcal{X}\times\mathcal{X}\times\mathcal{X}$.

###### Theorem 9 .

If $\mathcal{H}$ follows Assumption 3, the realizability condition holds, i.e., $f\in\mathcal{H}$, and Assumption  11 holds, then the model trained to minimize the risk in equation 1 (with $T\geq 2$) with $\ell_{2}$ loss generalizes to all sequences in the hypercube $[0,1]^{nt},\;\forall t\geq 1$ and thus achieves length and compositional generalization.

###### Proof.

Consider any $h$ that solves equation 1. Since $\ell$ is $\ell_{2}$ loss and realizability condition holds, $f$ is one of the optimal solutions to equation 1. For all $i\leq T,x_{\leq i}\in\mathsf{supp}(X_{\leq i})$ except over a set of measure zero the following condition holds

$$ $h(x_{\leq i})=f(x_{\leq i}).$ (63) $$

The above follows from the fact that $h$ solves equation 1, i.e., $\mathbb{E}[\|h-f\|^{2}]=0$ and from Theorem 1.6.6. (ash2000probability). Since $\mathsf{supp}(X_{\leq i})$ is regular closed, $f,h$ are both continuously differentiable, we can use Lemma 1, it follows that the above equality holds for all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$.
From realizability condition it follows that true $f(x_{\leq i})=\rho\Big{(}\sum_{k\leq i}\phi(x_{i},x_{k})\Big{)}$. We substitute the parametric forms from Assumption 3 to get

$$ $\begin{split}\omega\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\psi(x_{i},x_{k})\Big{ )}=\rho\Big{(}\sum_{k\leq i}\frac{1}{i}\cdot\phi(x_{i},x_{k})\Big{)}.\\ \end{split}$ (64) $$

Since $\omega$ and $\rho$ are single layer perceptron with bijective activation $\sigma$. We substitute the parametric form of $\omega$ and $\rho$ to obtain the following condition. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$,

$$ $\begin{split}&\sigma\Big{(}A\sum_{k\leq i}\frac{1}{i}\cdot\psi(x_{i},x_{k}) \Big{)}=\sigma\Big{(}B\sum_{k\leq i}\frac{1}{i}\cdot\phi(x_{i},x_{k})\Big{)} \implies A\sum_{k\leq i}\psi(x_{i},x_{k})=B\sum_{k\leq i}\phi(x_{i},x_{k}).\\ \end{split}$ (65) $$

The second equality follows from the fact that the activation $\sigma$ is bijective and hence the inputs to $\sigma$ are equal.

From Assumption 11, it follows that for all $x_{1},x_{2},x_{3}\in\mathcal{X}\times\mathcal{X}\times\mathcal{X}$

$$ $A\phi(x_{3},x_{1})+A\phi(x_{3},x_{2})=B\psi(x_{3},x_{1})+B\psi(x_{3},x_{2})$ (66) $$

Consider a reference point $x^{r}\in\mathcal{X}$. Consider the point $(x_{1}^{r},x_{2},x_{3})$ in the support (this point exists in the support due to Assumption 11).

$$ $A\phi(x_{3},x_{1}^{r})+A\phi(x_{3},x_{2})=B\psi(x_{3},x_{1}^{r})+B\psi(x_{3},x _{2})$ (67) $$

Take the difference of equation 31 of equation 30 to obtain

$$ $A\phi(x_{3},x_{1})=B\psi(x_{3},x_{1})+C(x_{3}),\forall x_{1}\in\mathcal{X}.$ (68) $$

where $C(x_{3})=\Big{(}B\psi(x_{3},x^{r})-A\phi(x_{3},x^{r})\Big{)}$. Substitute equation 68 into equation 66 we obtain that $C(x_{3})=0$ for all $x_{3}\in\mathcal{X}$.

Thus we obtain

$$ $\forall x_{i}\in\mathcal{X},x_{j}\in\mathcal{X}\hskip 5.69054ptA\psi(x_{i},x_{ j})=B\phi(x_{i},x_{j}).$ (69) $$

We now consider any sequence $x_{\leq\tilde{T}}\in\mathcal{X}^{\tilde{T}}$. The prediction made by $h$ is

$$ $\begin{split}h(x_{\leq\tilde{T}})=\sigma\Big{(}A\sum_{j\leq\tilde{T}}\psi(x_{ \tilde{T}},x_{j})\Big{)}=\sigma\Big{(}B\sum_{j\leq\tilde{T}}\phi(x_{\tilde{T}} ,x_{j})\Big{)}=f(x_{\leq\tilde{T}})\end{split}$ (70) $$

We use equation 69 in the simplification above. From the above, we can conclude that $h$ continues to be optimal for all sequences in
$[0,1]^{n\tilde{T}}$.

∎

###### Assumption 11 .

###### Theorem 9 .

###### Proof.

#### C.2.3 Extending Theorem 2 to ω 𝜔 \omega italic_ω from C 1 superscript 𝐶 1 C^{1} italic_C start_POSTSUPERSCRIPT 1 end_POSTSUPERSCRIPT -diffeomorphisms

###### Assumption 12 .

The joint support $\mathsf{supp}(X_{\leq i})$ is a regular closed set for all $i\leq T$. The support of all pairs of tokens is equal, i.e., $\mathsf{supp}(X_{i},X_{j})=[0,1]^{2n}$, where $i\not=j$, $i\geq 1,j\geq 1$. The support of $[\phi(X_{1},X_{2}),\phi(X_{1},X_{3})]$ is $\mathbb{R}^{2m}$, where $\phi$ is the embedding function for the labeling function $\rho(\sum_{j\leq i}\phi(x_{i},x_{j}))$.

See 3

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$. We first use the fact $h(x_{\leq i})=f(x_{\leq i}),\forall i\leq T$ almost everywhere in the support. We can use the continuity of $h,f$ and regular closedness of the support to extend the equality to all points in the support (follows from the first part of Lemma 1) to obtain the following. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$

$$ $\begin{split}&\omega\Big{(}\sum_{j<i}\frac{1}{i-1}\cdot\psi(x_{i},x_{j})\Big{) }=\rho\Big{(}\sum_{j<i}\frac{1}{i-1}\cdot\phi(x_{i},x_{j})\Big{)}\implies\\ &\sum_{j<i}\frac{1}{i-1}\psi(x_{i},x_{j})=\omega^{-1}\circ\rho\Big{(}\sum_{j<i }\frac{1}{i-1}\cdot\phi(x_{i},x_{j})\Big{)}\implies\\ &\sum_{j<i}\frac{1}{i-1}\psi(x_{i},x_{j})=a\Big{(}\sum_{j<i}\frac{1}{i-1}\phi( x_{i},x_{j})\Big{)},\end{split}$ (71) $$

where $a=\omega^{-1}\circ\rho$. In the above simplification, we used the parametric form for the true labeling function and the learned labeling function and use the invertibility of $\omega$. Let us consider the setting when $i=2$. In that case summation involves only one term. Substitute $x_{1}=y$ and $x_{2}=x$. We obtain $\forall x\in[0,1]^{n},y\in[0,1]^{n},$

$$ $\psi(x,y)=a(\phi(x,y)).$ (72) $$

The above expression implies that $\psi$ bijectively identifies $\phi$. Let us consider the setting when $i=3$ (this is possible since $T\geq 3$). We substitute $x_{3}=x$, $x_{2}=y$, $x_{1}=z$ and obtain

$$ $\frac{1}{2}\Big{[}a(\phi(x,y))+a(\phi(x,z))\Big{]}=a\big{(}\frac{1}{2}\big{(} \phi(x,y)+\phi(x,z)\big{)}\big{)}.$ (73) $$

Substitute $\phi(x,y)=\alpha$ and $\phi(x,z)=\beta$. In the simplifcation that follows, we use the that assumption $[\phi(x,y),\phi(x,z)]$ spans $\mathbb{R}^{2m}$, where $\phi(x,y)$ and $\phi(x,z)$ individually span $\mathbb{R}^{m}$.

$$ $\frac{1}{2}(a(\alpha)+a(\beta))=a\big{(}\frac{1}{2}(\alpha+\beta)\big{)}.$ (74) $$

Observe that $a(0)=0$ because $\omega^{-1}\circ\rho(0)=0$ because $\omega^{-1}(0)=\rho(0)=0$.

$$ $\begin{split}&\frac{1}{2}(a(2\alpha)+a(0))=a\big{(}\frac{1}{2}(2\alpha+0)\big{ )}\\ &a(2\alpha)=2a(\alpha)\end{split}$ (75) $$

Next, substitute $\alpha$ with $2\alpha$ and $\beta$ with $2\beta$ in equation 74 to obtain

$$ $\begin{split}&\frac{1}{2}(a(2\alpha)+a(2\beta))=a\big{(}\frac{1}{2}(2\alpha+2 \beta)\big{)}\\ &a(\alpha+\beta)=a(\alpha)+a(\beta)\end{split}$ (76) $$

We use equation 76 to show that $a$ is linear. To show that, we need to argue that $a(c\alpha)=ca(\alpha)$ as we already know $a$ satisfies additivity condition.

Suppose $c$ is some rational number, i.e., $c=p/q$, where $p$ and $q$ are non-zero integers.

From the identity it is clear that $a(p\alpha)=pa(\alpha)$, where $p$ is some integer.

$a(q\frac{1}{q}\alpha)=qa(\frac{1}{q}\alpha)\implies a(\frac{1}{q}\alpha)=\frac
{1}{q}a(\alpha)$, where $q$ is some integer.

Now combine these $a(p/q\alpha)=pa(1/q\alpha)=\frac{p}{q}a(\alpha)$. We have established the homogeneity condition for rationals.

We will now use the continuity of the function $a$ and density of rationals to extend the claim for irrationals. Suppose $c$ is some irrational. Define a sequence of rationals that approach $c$ (this follows from the fact that rationals are dense in $\mathbb{R}$).

$a(c\alpha)=a(\lim_{n\rightarrow\infty}q_{n}\alpha)=\lim_{n\rightarrow\infty}a(
q_{n}\alpha).$

In the second equality above, we use the definition of continuity ($a$ is continuous since composition of continuous functions is continuous). We can also use the property that we already showed for rationals to further simplify

$\lim_{n\rightarrow\infty}a(q_{n}\alpha)=a(\alpha)\lim_{n\rightarrow\infty}q_{n
}=ca(\alpha).$

Observe that $a:\mathbb{R}^{m}\rightarrow\mathbb{R}^{m}$ and for any $\alpha,\beta\in\mathbb{R}^{m}$ $a(\alpha+\beta)=a(\alpha)+a(\beta)$ and $a(c\alpha)=ca(\alpha)$. From the definition of a linear map it follows that $a$ is linear. As a result, we can write $\forall x\in[0,1]^{n},y\in[0,1]^{n}$

$$ $\psi(x,y)=A(\phi(x,y))$ (77) $$

Observe that $a$ is invertible because both $\rho$ and $\omega$ are invertible. As a result, we know that $A$ is an invertible matrix. From this we get

$$ $\forall x\in[0,1]^{n},y\in[0,1]^{n},\phi(x,y)=A^{-1}\psi(x,y)=C(\psi(x,y))$ (78) $$

For all $z\in\mathbb{R}^{m}$, we obtain

$$ $a(z)=\rho^{-1}\circ\omega(z)=Cz\implies\omega(z)=\rho(Cz)$ $$

Let us consider any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$.
We use the above conditions

$$ $\omega\big{(}\sum_{j<\tilde{T}}\psi(x_{\tilde{T}},x_{j})\big{)}=\rho(C\sum_{j< \tilde{T}}\psi(x_{\tilde{T}},x_{j}))=\rho\big{(}\sum_{j<\tilde{T}}\phi(x_{ \tilde{T}},x_{j})\big{)}.$ $$

Thus we obtain length and compositional generalization.

∎

###### Corollary 3 .

If $\mathcal{H}$ follows Assumption 4, the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further assumption on the support (Assumption 12) holds, then the model trained to minimize the risk in equation 1 (with $T\geq 3$) with $\ell_{2}$ loss achieves linear identification. Further, under the stated conditions linear identification is necessary for both length and compositional generalization.

###### Proof.

We follow the exact same steps as in the previous proof of Theorem 3 up to equation 78. We restate equation 78 below.

$$ $\forall x\in[0,1]^{n},y\in[0,1]^{n},\phi(x,y)=C(\psi(x,y))$ (79) $$

The above condition directly implies linear identification. We can use this to obtain that for any sequence $x_{\leq\tilde{T}}\in[0,1]^{n\tilde{T}}$

$$ $\sum_{j<\tilde{T}}\psi(x_{\tilde{T}},x_{j})=\sum_{j<\tilde{T}}\phi(x_{\tilde{T }},x_{j})$ (80) $$

To show necessity of linear identification, from the proof of Theorem 3 observe that

$$ $\begin{split}&\forall i\leq T,\forall x_{\leq i}\in\mathsf{supp}(X_{\leq i}), \;\omega\Big{(}\sum_{j<i}\frac{1}{i-1}\cdot\psi(x_{i},x_{j})\Big{)}=\rho\Big{( }\sum_{j<i}\frac{1}{i-1}\cdot\phi(x_{i},x_{j})\Big{)}\implies\\ &\forall x\in[0,1]^{n},y\in[0,1]^{n},\phi(x,y)=C(\psi(x,y))\end{split}$ (81) $$

Thus from the above it follows that in the absence of linear identification neither length nor compositional generalization are achievable.
∎

###### Assumption 12 .

###### Proof.

###### Corollary 3 .

###### Proof.

##### On absence of labels at all lengths from 1 1 1 1 to T 𝑇 T italic_T

We argue that the above proof can be adapted to the setting where we do not observe labels at all lengths from $1$ to $T$. Suppose we only observe label at length $T$. Take equation equation 71 and substitute $x_{i}=x$ and $x_{j}=y$ for all $j<i$ to obtain the same condition as equation equation 72.
Suppose $T$ is odd and larger than or equal to $3$. Fix $x_{i}=x$, $x_{2j-1}=y,\forall j\in\{1,\cdots,(T-1)/2\}$, $x_{2j}=z,\forall j\in\{1,\cdots,(T-1)/2\}$. We obtain the same condition as equation equation 73. Rest of the proof can be adapted using a similar line of reasoning.

##### Remark on Assumption 4

We require that the support of $[\phi(X_{1},X_{2}),\phi(X_{1},X_{3})]$ is $\mathbb{R}^{2m}$. This assumption is used in the proof in equation equation 76. We used this assumption to arrive at $a(\alpha+\beta)=a(\alpha)+a(\beta),\forall\alpha,\beta\in\mathbb{R}^{m}$. We then used continuity of $a$ to conclude $a$ is linear. Now suppose $[\phi(X_{1},X_{2}),\phi(X_{1},X_{3})]$ is some subset $\mathcal{Z}\subseteq\mathbb{R}^{2m}$. We believe that it is possible to extend the result to more general $\mathcal{Z}$, it might still be possible to arrive at $a$ is linear. We leave this investigation to future work.

#### C.2.4 Extending Theorem 3 to incorporate positional encodings

We next present the result when $\omega$ is continuously differentiable and invertible.

###### Assumption 13 .

Each function in the hypothesis class $\mathcal{H}$ used by the learner is given as $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j\leq i}\psi_{i-j}(x_{i},x_{j})\Big{)}$, where $\omega$ is a $C^{1}$-diffeomorphism. Also, $\psi_{i-j}=0$ for $i-j>T_{\mathsf{max}}-1$, i.e., two tokens that are sufficiently far apart do not interact. For all $k\leq T_{\mathsf{max}}-1$ each $x\in[0,1]^{n},$ $\exists\;y\in[0,1]^{n}$ we $\psi_{k}(x,y)=0$.

In the theorem that follows, we require the support of training distribution under consideration is already sufficiently diverse and hence we only seek to prove length generalization guarantees.

###### Assumption 14 .

The joint support $\mathsf{supp}(X_{\leq T})=[0,1]^{T}$. The support of $[\phi_{1}(X_{1},X_{2}),\phi_{2}(X_{1},X_{3})]$ is $\mathbb{R}^{2k}$, where $\phi_{i-j}$ is the embedding function for the labeling function $\rho(\sum_{j\leq i}\phi_{i-j}(x_{i},x_{j}))$.

###### Theorem 10 .

If $\mathcal{H}$ follows Assumption 13, the realizability condition holds, i.e., $f\in\mathcal{H}$, Assumption 14 holds and $T\geq T_{\mathsf{max}}$, then the model trained to minimize the risk in equation 1 (with $T\geq 2$) with $\ell_{2}$ loss achieves length generalization.

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$. We first use the fact $h(x_{\leq i})=f(x_{\leq i})$ almost everywhere in the support. We can use the continuity of $h,f$ and regular closedness of the support to extend the equality to all points in the support (follows from the first part of Lemma 1) to obtain the following. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$

$$ $\begin{split}&\omega\Big{(}\sum_{j<i}\frac{1}{i-1}\psi_{i-j}(x_{i},x_{j})\Big{ )}=\rho\Big{(}\sum_{j<i}\frac{1}{i-1}\phi_{i-j}(x_{i},x_{j})\Big{)},\\ &\sum_{j<i}\frac{1}{i-1}\psi_{i-j}(x_{i},x_{j})=\omega^{-1}\circ\rho\Big{(} \sum_{j<i}\frac{1}{i-1}\phi_{i-j}(x_{i},x_{j})\Big{)},\\ &\sum_{j<i}\frac{1}{i-1}\psi_{i-j}(x_{i},x_{j})=a\Big{(}\sum_{j<i}\frac{1}{i-1 }\phi_{i-j}(x_{i},x_{j})\Big{)},\end{split}$ (82) $$

where $a=\omega^{-1}\circ\rho$. In the above simplification, we used the parametric form for the true labeling function and the learned labeling function. We also used the invertibility of $\rho$. Let us consider the setting when $i=2$. In that case summation involves only one term. Substitute $x_{1}=y$ and $x_{2}=x$. We obtain $\forall x\in[0,1]^{n},y\in[0,1]^{n},$

$$ $\psi_{1}(x,y)=a(\phi_{1}(x,y)).$ (83) $$

For $i=3$, substitute $x_{1}=x$, $x_{3}=z$ and set $x_{2}=y$ in such a way that $\phi_{1}(x,y)=0$ (follows from Assumption 13). Thus we obtain

$$ $\psi_{2}(x,y)=a(\phi_{2}(x,y)).$ (84) $$

Similarly, we can obtain the following. For all $k\leq T_{\mathsf{max}}$

$$ $\psi_{k}(x,y)=a(\phi_{k}(x,y)).$ (85) $$

The above expression implies that $\psi$ bijectively identifies $\phi$. Let us consider the setting when $i=3$ (this is possible since $T\geq 3$). We substitute $x_{3}=x$, $x_{2}=y$, $x_{1}=z$ to give

$$ $\frac{1}{2}\big{(}a(\phi_{1}(x,y))+a(\phi_{2}(x,z))\big{)}=a\big{(}\frac{1}{2} (\phi_{1}(x,y)+\phi_{2}(x,z))\big{)}.$ (86) $$

We now use the that assumption $[\phi_{1}(x,y),\phi_{2}(x,z)]$ spans $\mathbb{R}^{2k}$ and substitute $\phi_{1}(x,y)=\alpha$ and $\phi_{2}(x,z)=\beta$

$$ $\frac{1}{2}(a(\alpha)+a(\beta))=a\big{(}\frac{1}{2}(\alpha+\beta)\big{)}.$ (87) $$

Rest of the proof follows the same strategy as proof of Theorem 3.
∎

###### Assumption 13 .

###### Assumption 14 .

###### Theorem 10 .

###### Proof.

#### C.2.5 Extending Theorem 3 to incorporate multiple attention heads

Our choice of the archictecture did not invoke multiple attention heads. If we include multiple attention heads, then also we can arrive at the same length generalization guarantees. The model class with two attention heads $\psi_{1},\psi_{2}$ can be stated as follows $\omega\Big{(}\sum_{j<i}A[\psi_{1}(x_{i},x_{j}),\psi_{2}(x_{i},x_{j})]^{\top}
\Big{)}$, where $A$ combines the outputs of the attention heads linearly. Following the same steps of proof of Theorem 3, we obtain the following.

$$ $\begin{split}&\omega\Big{(}\sum_{j<i}A[\psi_{1}(x_{i},x_{j}),\psi_{2}(x_{i},x_ {j})]^{\top}\Big{)}=\rho\Big{(}\sum_{j<i}B[\phi_{1}(x_{i},x_{j}),\phi_{2}(x_{i },x_{j})]^{\top}\Big{)},\\ &\omega\Big{(}\sum_{j<i}\tilde{\psi}(x_{i},x_{j})\Big{)}=\rho\Big{(}\sum_{j<i} \tilde{\phi}(x_{i},x_{j})\Big{)},\\ &\sum_{j<i}\tilde{\psi}(x_{i},x_{j})=a\Big{(}\sum_{j<i}\tilde{\phi}(x_{i},x_{j })\Big{)},\end{split}$ (88) $$

where $a=\omega^{-1}\circ\rho$. In the above simplification, the RHS shows the labeling function and the RHS is the function that is learned. We can follow the same strategy as the proof of Theorem 3 for the rest of the proof. We set $i=2$ and obtain a condition similar to equation 72 and for $i=3$ we obtain a condition similar to equation 73. Following a similar proof technique, we obtain $a$ is linear and the proof extends to multiple attention heads.

### C.3 State space models

In this section, we first provide the proof to Theorem 4. We then provide Corollary 4, where we describe how the learned representations linearly identify the true representations. In Theorem 11, we present the discrete tokens counterpart to Theorem 4.

###### Assumption 15 .

The joint support $\mathsf{supp}(X_{\leq i})$ is a regular closed set for all $i\leq T$. The support of $X_{1}$ is $\mathbb{R}^{n}$. For some length $2\leq i\leq T$ an there exists $in$ sequences $x_{\leq i}$ such that their concatenation forms a $in\times in$ matrix of rank $in$.

See 4

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$. We first use the fact $h(x_{\leq i})=f(x_{\leq i}),\forall i\leq T$ almost everywhere in the support. We can use the continuity of $h,f$ and regular closedness of the support to extend the equality to all points in the support (from first part of Lemma 1) to obtain the following. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$.

$$ $\begin{split}&f(x_{\leq i})=h(x_{\leq i})=\\ &\rho(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\omega(\sum_{j=0}^{i-1}\tilde{ \Lambda}^{j}\tilde{B}x_{i-j})\implies\\ &\omega^{-1}\circ\rho(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1} \tilde{\Lambda}^{j}\tilde{B}x_{i-j}=\\ &c(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1}\tilde{\Lambda}^{j} \tilde{B}x_{i-j}\end{split}$ (89) $$

For $i=1$, $\forall x_{1}\in\mathbb{R}^{n},c(Bx_{1})=\tilde{B}x_{1}$. Substitute $Bx_{1}=x$, we obtain $\forall x\in\mathbb{R}^{n},c(x)=\tilde{B}B^{-1}x=Cx$, where we use the fact that $Bx_{1}$ spans $\mathbb{R}^{n}$ as $B$ is invertible.

From linearity of $c$, we obtain

$$ $\omega^{-1}\circ\rho(z)=Cz\implies\rho(z)=\omega(Cz),\forall z\in\mathbb{R}^{n}$ (90) $$

We use this linearity of $c$ to simplify

$$ $\begin{split}&c(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1}\tilde{ \Lambda}^{j}\tilde{B}x_{i-j}\implies\\ &C(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1}\tilde{\Lambda}^{j} \tilde{B}x_{i-j}\implies\\ &[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]\begin{bmatrix}x_{i}\\ x_{i-2}\\ \vdots\\ x_{1}\end{bmatrix}-[\tilde{B},\tilde{\Lambda}\tilde{B},\tilde{\Lambda}^{2} \tilde{B},\cdots,\tilde{\Lambda}^{i-1}\tilde{B}]\begin{bmatrix}x_{i}\\ x_{i-2}\\ \vdots\\ x_{1}\end{bmatrix}=0\implies\\ &\Big{[}[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]-[\tilde{B},\tilde {\Lambda}\tilde{B},\tilde{\Lambda}^{2}\tilde{B},\cdots,\tilde{\Lambda}^{i-1} \tilde{B}]\Big{]}\bm{X}=0,\\ \end{split}$ (91) $$

where $\bm{X}=\begin{bmatrix}x_{i}\\
x_{i-2}\\
\vdots\\
x_{1}\end{bmatrix}$.

Denote $R=\Big{[}[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]-[\tilde{B},
\tilde{\Lambda}\tilde{B},\tilde{\Lambda}^{2}\tilde{B},\cdots,\tilde{\Lambda}^{
i-1}\tilde{B}]\Big{]}$.
We collect a set of points $\bm{X}^{+}=[\bm{X}^{(1)},\cdots,\bm{X}^{(l)}]$ where $l\geq ni$ and rank of $\bm{X}^{+}=ni$ (from Assumption 15). Since the matrix $\bm{X}^{+}$ is full rank, we have

$$ $R\bm{X}^{+}=0\implies R=0.$ $$

This yields

$$ $CB=\tilde{B},C\Lambda B=\tilde{\Lambda}\tilde{B},\cdots,C\Lambda^{i}B=\tilde{ \Lambda}^{i}\tilde{B}.$ (92) $$

Observe that from the second equality, we get $\tilde{\Lambda}=C\Lambda C^{-1}$. Given the parameters $(\Lambda,B)$, the set of parameters $(\tilde{\Lambda},\tilde{B})$ that solve the first two equalities are – $\{\tilde{B}\text{ is an arbitrary invertible matrix,}\;\tilde{\Lambda}=C
\Lambda C^{-1},\;\;\text{where }C=\tilde{B}B^{-1}\}$.

Take any solution of the first two equalities and compute

$$ $\tilde{\Lambda}^{i}\tilde{B}=C\Lambda^{i}C^{-1}\tilde{B}=C\Lambda^{i}B,\forall i\geq 1$ (93) $$

From equation 93 and equation 90, we obtain that for all $x_{\leq i}\in\mathbb{R}^{ni}$

$$ $h(x_{\leq i})=\omega(\sum_{j=0}^{i-1}\tilde{\Lambda}^{j}\tilde{B}x_{i-j})= \omega(C\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\rho(\sum_{j=0}^{i-1}\Lambda^{j} Bx_{i-j})=f(x_{\leq i})$ (94) $$

This establishes both compositional and length generalization.

∎

###### Corollary 4 .

If $\mathcal{H}$ follows Assumption 5, and the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further condition on the support, i.e., Assumption 15, holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss $(T\geq 2)$ achieves linear identification. Further, under the stated conditions linear identification is necessary for both length and compositional generalization.

###### Proof.

We follow the same steps as proof of Theorem 4 up to equation 93. From that we obtain that for all $x_{\leq i}\in\mathbb{R}^{ni}$

$$ $\sum_{j=0}^{i-1}\tilde{\Lambda}^{j}\tilde{B}x_{i-j}=C(\sum_{j=0}^{i-1}\Lambda^ {j}Bx_{i-j})$ (95) $$

Recall that $\sum_{j=0}^{i-1}\tilde{\Lambda}^{j}\tilde{B}x_{i-j}=\tilde{h}_{j}$ and $\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j}=h_{j}$. From this it follows that $\tilde{h}_{j}=Ch_{j}$, which proves that learned hidden state are a linear transform of the hidden state underlying the labeling function. This establishes linear identification.

To show the necessity of linear identification, from the proof of Theorem 4 it follows that

$$ $\begin{split}&\forall i\leq T,\forall x_{\leq i}\in\mathsf{supp}(X_{\leq i}), \rho(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\omega(\sum_{j=0}^{i-1}\tilde{ \Lambda}^{j}\tilde{B}x_{i-j})\implies\\ &\forall i\geq 1,\forall x_{\leq i}\in\mathbb{R}^{ni},\tilde{h}_{j}=Ch_{j}\end {split}$ (96) $$

If the latter conditon in the above implication does not hold, then the former condition cannot hold. Hence, linear identification is necessary.

∎

###### Assumption 15 .

###### Proof.

###### Corollary 4 .

###### Proof.

#### C.3.1 Extending Theorem 4 to discrete tokens

In our discussion, we have focused on settings where the support of each token has a non-empty interior (Assumption 2). In practice of language modeling, we use discrete tokens and hence Assumption 2 does not hold anymore. In this section, we discuss the adaptation of results for SSMs to setting when the the support of tokens is a finite set.

###### Assumption 16 .

Each function in the hypothesis class $\mathcal{H}$ takes a sequence $\{x_{1},\cdots,x_{i}\}$ as input and outputs $h(x_{1},\cdots,x_{i})=\omega\Big{(}\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j}\Big{)}$, where $\omega:\mathbb{R}^{k}\rightarrow\mathbb{R}^{m}$ is a single layer perceptron denoted as $\sigma\circ A$. $A$, $B$ and $\Lambda$ are square invertible. As a result, $k=m=n$.

###### Assumption 17 .

For some length $2\leq i\leq T$ an there exists $in$ sequences $x_{\leq i}$ such that their concatenation forms a $in\times in$ matrix of rank $in$.

###### Theorem 11 .

If $\mathcal{H}$ follows Assumption 16, and the realizability condition holds, i.e., $f\in\mathcal{H}$, and a further condition on the support, i.e., Assumption 17, holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss $(T\geq 2)$ achieves length and compositional generalization.

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$. We first use the fact $h(x_{\leq i})=f(x_{\leq i}),\forall i\leq T$ almost everywhere in the support. We can use the continuity of $h,f$ and regular closedness of the support to extend the equality to all points in the support (from first part of Lemma 1) to obtain the following. For all $x_{\leq i}\in\mathsf{supp}(X_{\leq i})$.

$$ $\begin{split}&f(x_{\leq i})=h(x_{\leq i})=\\ &\rho(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\omega(\sum_{j=0}^{i-1}\tilde{ \Lambda}^{j}\tilde{B}x_{i-j})\implies\\ &\sigma(A\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sigma(\tilde{A}\sum_{j=0}^{i-1} \tilde{\Lambda}^{j}\tilde{B}x_{i-j})=\\ &C(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1}\tilde{\Lambda}^{j} \tilde{B}x_{i-j},\end{split}$ (97) $$

where $C=\tilde{A}^{-1}A$.

We simplify the last identity in the above further.

$$ $\begin{split}&C(\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\sum_{j=0}^{i-1}\tilde{ \Lambda}^{j}\tilde{B}x_{i-j}\implies\\ &[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]\begin{bmatrix}x_{i}\\ x_{i-2}\\ \vdots\\ x_{1}\end{bmatrix}-[\tilde{B},\tilde{\Lambda}\tilde{B},\tilde{\Lambda}^{2} \tilde{B},\cdots,\tilde{\Lambda}^{i-1}\tilde{B}]\begin{bmatrix}x_{i}\\ x_{i-2}\\ \vdots\\ x_{1}\end{bmatrix}=0\implies\\ &\Big{[}[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]-[\tilde{B},\tilde {\Lambda}\tilde{B},\tilde{\Lambda}^{2}\tilde{B},\cdots,\tilde{\Lambda}^{i-1} \tilde{B}]\Big{]}\bm{X}=0,\\ \end{split}$ (98) $$

where $\bm{X}=\begin{bmatrix}x_{i}\\
x_{i-2}\\
\vdots\\
x_{1}\end{bmatrix}$.

Denote $R=\Big{[}[CB,C\Lambda B,C\Lambda^{2}B,\cdots,C\Lambda^{i-1}B]-[\tilde{B},
\tilde{\Lambda}\tilde{B},\tilde{\Lambda}^{2}\tilde{B},\cdots,\tilde{\Lambda}^{
i-1}\tilde{B}]\Big{]}$.
We collect a set of points $\bm{X}^{+}=[\bm{X}^{(1)},\cdots,\bm{X}^{(l)}]$ where $l\geq ni$ and rank of $\bm{X}^{+}=ni$ (from Assumption 15). Since the matrix $\bm{X}^{+}$ is full rank, we have

$$ $R\bm{X}^{+}=0\implies R=0.$ $$

This yields

$$ $CB=\tilde{B},C\Lambda B=\tilde{\Lambda}\tilde{B},\cdots,C\Lambda^{i}B=\tilde{ \Lambda}^{i}\tilde{B}.$ (99) $$

Observe that from the second equality, we get $\tilde{\Lambda}=C\Lambda C^{-1}$. Given the parameters $(\Lambda,B)$, the set of parameters $(\tilde{\Lambda},\tilde{B})$ that solve the first two equalities are – $\{\tilde{B}\text{ is an arbitrary invertible matrix,}\;\tilde{\Lambda}=C
\Lambda C^{-1},\;\;\text{where }C=\tilde{B}B^{-1}\}$.

Take any solution of the first two equalities and compute

$$ $\tilde{\Lambda}^{i}\tilde{B}=C\Lambda^{i}C^{-1}\tilde{B}=C\Lambda^{i}B,\forall i\geq 1$ (100) $$

From equation 100, we obtain that for all $x_{\leq i}\in\mathbb{R}^{ni}$

$$ $h(x_{\leq i})=\omega(\sum_{j=0}^{i-1}\tilde{\Lambda}^{j}\tilde{B}x_{i-j})= \omega(C\sum_{j=0}^{i-1}\Lambda^{j}Bx_{i-j})=\rho(\sum_{j=0}^{i-1}\Lambda^{j} Bx_{i-j})=f(x_{\leq i})$ (101) $$

This establishes both compositional and length generalization.

∎

###### Assumption 16 .

###### Assumption 17 .

###### Theorem 11 .

###### Proof.

### C.4 Vanilla RNNs

In this section, we discuss RNNs and present the proof of Theorem 5. We first build some lemmas in the form of Lemma 2 and 3 that are used to prove Theorem 5. In Corollary 5, we explain the learned hidden state are a permutation transform of the true hidden state and also show that its a necessary condition for length and compositional generalization. Finally, in Theorem 12, we present the discrete token counterpart to Theorem 5.

###### Lemma 2 .

The $k^{th}$ derivative of sigmoid function denoted $\frac{\partial^{k}\sigma(s)}{\partial s^{k}}$ is not zero identically.

###### Proof.

The first derivative of the sigmoid function $\frac{\partial\sigma(s)}{\partial s}=\sigma(s)(1-\sigma(s))$.
We argue that the $\frac{\partial^{k}\sigma(s)}{\partial s^{k}}$ is a polynomial in $\sigma(s)$ with degree $k+1$. Consider the base case of $k=1$. This condition is true as $\frac{\partial\sigma(s)}{\partial s}=\sigma(s)(1-\sigma(s))$.
Now let us assume that $\frac{\partial^{k}\sigma(s)}{\partial s^{k}}$ is a polynomial of degree at most $k+1$ denoted as $P_{k+1}(\sigma(s))$. We simplify

$$ $\frac{\partial^{k}\sigma(s)}{\partial s^{k}}=P_{k+1}(\sigma(s))=\sum_{j=1}^{k+ 1}a_{j}(\sigma(s))^{j}$ $$

We take another derivative of the term above as follows.

$$ $\frac{\partial^{k+1}\sigma(s)}{\partial s^{k+1}}=\frac{\partial P_{k+1}(\sigma (s))}{\partial s}=\sum_{j=1}^{k+1}a_{j}\frac{\partial(\sigma(s))^{j}}{\partial s }=\sum_{j=1}^{k+1}a_{j}j\sigma(s)^{j-1}(\sigma(s)(1-\sigma(s)))$ $$

Observe that the $\frac{\partial^{k+1}\sigma(s)}{\partial s^{k+1}}$ is also a polynomial in $\sigma(s)$. Observe that the degree $k+2$ term has one term with coefficient $-a_{k+1}\cdot(k+1)$. Since $a_{k+1}\not=0$, the coefficient of degree $k+2$, $-a_{k+1}\cdot(k+1)$, is also non-zero. Since $\frac{\partial^{k}\sigma(s)}{\partial s^{k}}$ is a polynomial in $\sigma(s)$ with degree $k+1$ and hence, it cannot be zero identically.

∎

###### Lemma 3 .

Let $x\in\mathbb{R}^{n}$ and $A\in\mathbb{R}^{n\times n}$. Suppose $Ax=0,\forall x\in\mathcal{X}$, where $\mathcal{X}$ has a non-empty interior. Under these conditions $A=0$.

###### Proof.

Since $\mathcal{X}$ has a non-empty interior, we can construct a $\ell_{\infty}$ ball centered on $\theta$, defined as follows – $\tilde{\mathcal{X}}=\{\theta+\sum_{j=1}^{n}\alpha_{j}e_{j}\;|\|\alpha\|_{
\infty}\leq\alpha_{\mathsf{max}}\;\}$, where $e_{j}$ is a vector that is zero in all components and one on the $j^{th}$ component.
Suppose $A$ was non-zero. One of the columns say $a_{j}$ is non-zero. Consider two points in the ball $\tilde{\mathcal{X}}$ such that
$j^{th}$ coefficients are non-zero but rest of the coefficients are zero. We denote the $j^{th}$ components for the two components as $\alpha_{j}$ and $\tilde{\alpha}_{j}$, where $\alpha_{j}\not=\tilde{\alpha}_{j}$.
We now plug these two points into the condition that $Ax=0$

$$ $\begin{split}&A(\theta+\alpha_{j}e_{j})=0\implies A\theta=\alpha_{j}a_{j},\\ &A(\theta+\tilde{\alpha}_{j}e_{j})=0\implies A\theta=\tilde{\alpha}_{j}a_{j}, \\ \end{split}$ (102) $$

We take a difference of the two steps above and obtain

$$ $(\alpha_{j}-\tilde{\alpha_{j}})a_{j}=0\implies a_{j}=0$ $$

This is a contradiction. Hence, $A=0$.
∎

See 5

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$ everywhere in the support of the training distribution (using first part of Lemma 1).
We start with equating label at length 1, i.e., $y_{1}$. For all $x_{1}\in\mathsf{supp}(X_{1})$

$$ $\begin{split}&\sigma(A\sigma(Bx_{1}))=\sigma(\tilde{A}\sigma(\tilde{B}x_{1})) \implies A\sigma(Bx_{1})=\tilde{A}\sigma(\tilde{B}x_{1})\implies\\ &\sigma(B\tilde{B}^{-1}\tilde{B}x_{1})=A^{-1}\tilde{A}\sigma(\tilde{B}x_{1}) \end{split}$ (103) $$

Say $y=\tilde{B}x_{1}$, $A^{-1}\tilde{A}=U$, $B\tilde{B}^{-1}=V$. We substitute these expressions in the simplificaction below. We pick a $y$ in the interior of $\tilde{B}\cdot\mathsf{supp}(X_{1})$.

$$ $\sigma(Vy)=U\sigma(y)$ (104) $$

Take the first row of $V$ and $U$ as $v^{\top}$ and $u^{\top}$ to obtain

$$ $\sigma(v^{\top}y)=u^{\top}\sigma(y)$ (105) $$

Suppose there is some non-zero component of $v$ say $i$ but the corresponding component is zero in $u$.

$$ $\begin{split}\frac{\partial\sigma(v_{i}y_{i}+v_{-i}y_{-i})}{\partial y_{i}}= \sigma^{{}^{\prime}}(v_{i}y_{i}+v_{-i}y_{-i})v_{i}=\frac{\partial u_{-i}^{\top }\sigma(y_{-i})}{\partial y_{i}}=0\end{split}$ (106) $$

From the above we get $\sigma^{{}^{\prime}}(v^{\top}y)=0$. But sigmoid is strictly monotonic on $\mathbb{R}$, $\sigma^{{}^{\prime}}(x)>0,\forall x\in\mathbb{R}$ and $v^{\top}y\in\mathbb{R}$. Hence, $\sigma^{{}^{\prime}}(v^{\top}y)=0$ is not possible.
Similarly, suppose some component is non-zero in $u$ and zero in $v$.

$$ $\begin{split}\frac{\partial\sigma(v_{-i}^{\top}y_{-i})}{\partial y_{i}}=0= \frac{\partial(u_{i}\sigma(y_{i})+u_{-i}^{\top}\sigma(y_{-i}))}{\partial y_{i} }=u_{i}\sigma^{{}^{\prime}}(y_{i})\end{split}$ (107) $$

Since the derivative of $\sigma$ cannot be zero, the above condition cannot be true.

From the above, we can deduce that both $u$ and $v$ have same non-zero components.

Let us start with the case where $p\geq 2$ components of $u,v$ are non-zero. Below we equate the partial derivative w.r.t all components of $y$ that have non-zero component in $u$ (since $y$ is in the interior of the image of $\tilde{B}x_{1}$, we can equate these derivatives).

$$ $\begin{split}\sigma(v^{\top}y)=u^{\top}\sigma(y),\\ \frac{\partial^{p}\sigma(s)}{\partial s^{p}}\Pi_{v_{i}\not=0}v_{i}=0\implies \frac{\partial^{p}\sigma(s)}{\partial s^{p}}=0.\\ \end{split}$ (108) $$

Since support $X_{1}$ has a non-empty interior, the set of values $v^{\top}y$ takes also has a non-empty interior in $\mathbb{R}$. Hence, the above equality is true over a set of values $s$, which have a non-empty interior. Since $\sigma(s)$
is analytic, $\frac{\partial^{p}\sigma(s)}{\partial s^{p}}$ is also analytic. From (mityagin2015zero), it follows that $\frac{\partial^{p}\sigma(s)}{\partial s^{p}}=0$ everywhere. From Lemma 2, we know this condition cannot be true.

We are left with the case where $u$ and $v$ have one non-zero component each.

$\frac{1}{1+e^{-vy}}=\frac{u}{1+e^{-y}}\implies 1+e^{-y}=u+ue^{-vy}$
In the simplification above, we take derivative w.r.t $y$ to obtain $e^{-(v-1)y}=1/uv$. We now again take derivative again w.r.t $y$ to get $v=1$ and substitute it back to get $u=1$. Note that no other row of $U$ or $V$ can have same non-zero element because that would make matrix non invertible. From this we deduce that $U$ and $V$ are permutation matrices. From $\sigma(Vy)=U\sigma(y)$ it follows that $U=V=\Pi$. Thus $B=\Pi\tilde{B}$ and $\tilde{A}=A\Pi$.

Next, we equate predictions for $y_{2}$ to the ground truth (label $y_{2}$ exists as $T\geq 2$). For all $x_{1}\in\mathsf{supp}(X_{1})$

$$ $\begin{split}&\sigma(A\sigma(\Lambda\sigma(Bx_{1})+Bx_{2}))=\sigma(\tilde{A} \sigma(\tilde{\Lambda}\sigma(\tilde{B}x_{1})+\tilde{B}x_{2}))\implies\\ &A\sigma(\Lambda\sigma(Bx_{1})+Bx_{2})=\tilde{A}\sigma(\tilde{\Lambda}\sigma( \tilde{B}x_{1})+\tilde{B}x_{2})\implies\\ &\tilde{A}\sigma(\tilde{\Lambda}\sigma(\tilde{B}x_{1})+\tilde{B}x_{2})=A\Pi \sigma(\tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+\Pi^{\top}Bx_{2})=A\sigma(\Pi \tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+Bx_{2}).\\ \end{split}$ (109) $$

We use the simplification in the second step to equate to LHS in the first step as follows.

$$ $\begin{split}&A\sigma(\Pi\tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+Bx_{2})=A \sigma(\Lambda\sigma(Bx_{1})+Bx_{2})\\ &\implies(\Pi\tilde{\Lambda}\Pi^{\top}-\Lambda)\sigma(Bx_{1})=0.\\ \end{split}$ (110) $$

Since $\sigma(Bx_{1})$ spans a set that has a non-empty interior, we get that $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi$ (from Lemma 3).

From the above conditions, we have arrived at $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi,\tilde{B}=\Pi^{\top}B,\tilde{A}=A\Pi$.

We want to show that for all $k\geq 1$

$$ $h_{k}=\Pi\tilde{h}_{k},$ (111) $$

where $h_{k}=\sigma(\Lambda h_{k-1}+Bx_{k})$ and $\tilde{h}_{k}=\sigma(\tilde{\Lambda}\tilde{h}_{k-1}+\tilde{B}x_{k})$ and $h_{0}=\tilde{h}_{0}=0$. In other words, we define $T_{k}$ as a mapping that takes $x_{\leq k}$ as input and outputs $h_{k}$, i.e., $T_{k}(x_{\leq k})=h_{k}$. Similarly, we write $\tilde{T}_{k}(x_{\leq k})=\tilde{h}_{k}$. We want to show

$$ $T_{k}=\Pi\tilde{T}_{k},\forall k$ (112) $$

We show the above by principle of induction. Let us consider the base case below. For all $x_{1}\in\mathbb{R}^{n}$

$$ $\tilde{A}\sigma(\tilde{B}x_{1})=A\Pi\sigma(\Pi^{\top}Bx_{1})=A\sigma(Bx_{1})= Ah_{1}\implies h_{1}=\Pi\tilde{h}_{1}\implies T_{1}(x_{1})=\Pi\tilde{T}_{1}(x_ {1})$ (113) $$

Suppose $\forall j\leq k,T_{j}=\Pi\tilde{T}_{j}$.

Having shown the base case and assumed the condition for $j\leq k$, we now consider the mapping $\tilde{T}_{k+1}$

$$ $\Pi\tilde{T}_{k+1}(x_{\leq k+1})=\Pi\sigma(\tilde{\Lambda}\tilde{h}_{k}+\tilde {B}x_{k+1})=\Pi\sigma(\Pi^{\top}\Lambda\Pi\tilde{h}_{k}+\Pi^{\top}Bx_{k})= \sigma(\Lambda h_{k}+Bx_{k})=T_{k+1}(x_{\leq k+1}).$ (114) $$

The prediction from the model $(\tilde{A},\tilde{\Lambda},\tilde{B})$ at a time step $k$ is denoted as $\tilde{y}_{k}$ and it relates to $\tilde{h}_{k}$ as follows $\tilde{y}_{k}=\sigma(\tilde{A}\tilde{h}_{k})$. We use the above condition in equation equation 130 to arrive at the following result. For all $x_{\leq k}\in\mathbb{R}^{nk}$

$\tilde{y}_{k}=\sigma(\tilde{A}\tilde{h}_{k})=\sigma(\tilde{A}\tilde{T}(x_{\leq
k
}))=\sigma(A\Pi\tilde{T}(x_{\leq k}))=\sigma(AT(x_{\leq k}))=y_{k}$

This completes the proof.
∎

###### Corollary 5 .

If $\mathcal{H}$ follows Assumption 6, and the realizability condition holds, i.e., $f\in\mathcal{H}$ and regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss (with $T\geq 2$) achieves permutation identification. Further, under the stated conditions permutation identification is necessary for both length and compositional generalization.

###### Proof.

We follow the exact steps from the proof of Theorem 5 up to equation 132. From equation 132 it follows that for all $x_{\leq k}\in\mathbb{R}^{nk}$

$$ $T_{k}(x_{\leq k})=\Pi\tilde{T}(x_{\leq k})\implies h_{k}=\Pi\tilde{h}_{k}$ (115) $$

The above implies permutation identification. To show the necessity of permutation identification, from the proof of Theorem 5 observe that

$$ $\forall i\leq T,\forall x_{\leq i}\in\mathsf{supp}(X_{\leq i}),\;\sigma(A \sigma(Bx_{i}+\Lambda h_{i-1}))=\sigma(\tilde{A}\sigma(\tilde{B}x_{i}+\Lambda \tilde{h}_{i-1}))\implies T_{k}(x_{\leq k})=\Pi\tilde{T}(x_{\leq k})$ (116) $$

The latter condition implies permutation identification. If it does not hold, then the condition in LHS cannot hold and hence neither length nor compositional generalization can be achieved.
∎

###### Lemma 2 .

###### Proof.

###### Lemma 3 .

###### Proof.

###### Proof.

###### Corollary 5 .

###### Proof.

#### C.4.1 Extending Theorem 5 to discrete tokens

In our discussion, we have focused on settings where the support of each token has a non-empty interior (Assumption 2). In practice of language modeling, we use discrete tokens and hence Assumption 2 does not hold anymore. In this section, we discuss the adaptation of results for vanilla RNNs to setting when the the support of tokens is a finite set.

Define $\mathcal{S}=\{y=Bx\;|\;x\in\mathcal{X}\}$, where $\mathcal{X}$ is the marginal support of each token.

###### Assumption 18 .

a) For each component $i$ of $y$, $\mathcal{S}$ contains two pairs where the first coordinate differs by the same amount. Mathematically stated, the two pairs are $\Big{(}(y_{i},y_{-i}),(y_{i}+\delta,y_{-i}))\Big{)}$ and $\Big{(}(y_{i}^{{}^{\prime}},y_{-i}^{{}^{\prime}}),(y_{i}^{{}^{\prime}}+\delta,
y_{-i}^{{}^{\prime}}))\Big{)}$.

b)For every pair of components $i,j$ of $y$, $\mathcal{S}$ contains a point $y$ that satisfies the following. There exists three points in $\mathcal{S}$ such that they only differ in $y_{i},y_{j}$, and form a rectangle, $(y_{i},y_{j}),$ $(y_{i}^{{}^{\prime}},y_{j}),$ $(y_{i},y_{j}^{{}^{\prime}}),$ $(y_{i}^{{}^{\prime}},y_{j}^{{}^{\prime}})$. Similarly, there exists another set of points where $y_{i}^{{}^{\prime}}<y_{i}$ and $y_{j}^{{}^{\prime}}<y_{j}$.

###### Theorem 12 .

If $\mathcal{H}$ follows Assumption 6, and the realizability condition holds, i.e., $f\in\mathcal{H}$ and regular closedness condition in Assumption 2 holds, then the model trained to minimize the risk in equation 1 with $\ell_{2}$ loss (with $T\geq 2$) achieves length and compositional generalization.

###### Proof.

We start with the same steps as earlier proofs and equate the prediction of $h$ and $f$ everywhere in the support of the training distribution (using first part of Lemma 1).
We start with equating label at length 1, i.e., $y_{1}$. For all $x_{1}\in\mathsf{supp}(X_{1})$

$$ $\begin{split}&\sigma(A\sigma(Bx_{1}))=\sigma(\tilde{A}\sigma(\tilde{B}x_{1})) \implies A\sigma(Bx_{1})=\tilde{A}\sigma(\tilde{B}x_{1})\implies\\ &\tilde{A}^{-1}A\sigma(Bx_{1})=\sigma(\tilde{B}B^{-1}Bx_{1})\end{split}$ (117) $$

Say $y=Bx_{1}$, $\tilde{A}^{-1}A=U$, $\tilde{B}B^{-1}=V$. We substitute these expressions in the simplificaction below. We pick a $y$ in the interior of $\tilde{B}\cdot\mathsf{supp}(X_{1})$.

$$ $\sigma(Vy)=U\sigma(y)$ (118) $$

Take the first row of $V$ and $U$ as $v^{\top}$ and $u^{\top}$ to obtain

$$ $\sigma(v^{\top}y)=u^{\top}\sigma(y)$ (119) $$

Say $v_{i}\not=0$ and $u_{i}=0$. We consider a $(y_{i},y_{-i})$ and $(y_{i}^{{}^{\prime}},y_{-i})$ satisfying Assumption 18 a. We substitute these points in equation 119 and take the difference of the LHS and RHS in equation 119 to obtain.

$$ $\sigma(v_{i}y_{i}^{{}^{\prime}}+v_{-i}y_{-i})-\sigma(v_{i}y_{i}+v_{-i}y_{-i})=0$ (120) $$

$\sigma$ is strictly monotonic and thus the above cannot be true. Similarly, we can rule out the case when $u_{i}\not=0$ and $v_{i}=0$. Thus we can deduce that both $u$ and $v$ have same non-zero components.

Let us start with the case where $p\geq 2$ components of $u,v$ are non-zero. Without loss of generality say the first two components are among coordinates that are non-zero. Pick a $y\in\mathcal{S}$ that satisfies Assumption 18 b. Suppose $v^{\top}y\geq 0$. We select the neighbors of $y$ that form the rectangle such that each coordinate is greater than $y$. We substitute these points in equation 119 and the simplification procedure works as follows.
Let

$$ $\begin{split}&s_{1}=v_{1}y_{1}^{{}^{\prime}}+v_{2}y_{2}^{{}^{\prime}}+\cdots v _{n}y_{n},\;\;s_{3}=v_{1}y_{1}^{{}^{\prime}}+v_{2}y_{2}+\cdots v_{n}y_{n}\\ &s_{2}=v_{1}y_{1}+v_{2}y_{2}^{{}^{\prime}}+\cdots v_{n}y_{n},\;\;s_{4}=v_{1}y_ {1}+v_{2}y_{2}+\cdots v_{n}y_{n}\end{split}$ (121) $$

Observe that $s_{1}>s_{2}>s_{4}$ and $s_{1}>s_{3}>s_{4}$. It is possible that $s_{2}\geq s_{3}$ or $s_{3}>s_{2}$. Suppose $s_{2}\geq s_{3}$.

We can write

$$ $\begin{split}&\sigma(s_{1})=u_{1}\sigma(y_{1}^{{}^{\prime}})+u_{2}\sigma(y_{2} ^{{}^{\prime}})+\cdots+u_{n}\sigma(y_{n}),\sigma(s_{2})=u_{1}\sigma(y_{1})+u_{ 2}\sigma(y_{2}^{{}^{\prime}})+\cdots+u_{n}\sigma(y_{n})\\ &\sigma(s_{3})=u_{1}\sigma(y_{1}^{{}^{\prime}})+u_{2}\sigma(y_{2})+\cdots+u_{n }\sigma(y_{n}),\sigma(s_{4})=u_{1}\sigma(y_{1})+u_{2}\sigma(y_{2})+\cdots+u_{n }\sigma(y_{n})\end{split}$ (122) $$

We take a difference of the first two and the latter two, and subtract these differences to get

$$ $\Big{(}\sigma(s_{1})-\sigma(s_{2})\Big{)}-\Big{(}\sigma(s_{3})-\sigma(s_{4}) \Big{)}=0$ (123) $$

From mean value theorem, we get that $\sigma^{{}^{\prime}}(\tilde{s})=\sigma^{{}^{\prime}}(s^{\dagger})$, where $\sigma^{{}^{\prime}}$ is the derivative of $\sigma$, $\tilde{s}$ is a value between $s_{1}$ and $s_{2}$, and $s^{\dagger}$ is a value between $s_{3}$ and $s_{4}$. Since $s_{1}>s_{2}>s_{3}>s_{4}>0$, $\tilde{s}>s^{\dagger}>0$. Since $\sigma^{{}^{\prime}}$ strictly decreases on positive values, the above equality $\sigma^{{}^{\prime}}(\tilde{s})=\sigma^{{}^{\prime}}(s^{\dagger})$ is not possible. Similarly, we can tackle the case $v^{\top}y<0$.

We are left with the case where $u$ and $v$ have one non-zero component each. From Assumption 18a, we select two pairs that differe exactly in the non-zero component. We can resort to dealing with scalars as follows. We start with first pair $(y,y+\delta)$.

$$ $\begin{split}&\sigma(vy)=u\sigma(y)\implies\frac{1}{1+e^{-vy}}=\frac{u}{1+e^{- y}}\implies 1-u=ue^{-vy}-e^{-y}\\ &\sigma(v(y+\delta))=u\sigma(y+\delta)\implies 1-u=ue^{-v(y+\delta)}-e^{-(y+ \delta)}\end{split}$ (124) $$

By equating the RHS in the above, we obtain

$$ $\frac{1-e^{-\delta}}{1-e^{-v\delta}}=ue^{-(v-1)y}$ (125) $$

For the second pair $(y^{{}^{\prime}},y^{{}^{\prime}}+\delta)$, we obtain

$$ $\frac{1-e^{-\delta}}{1-e^{-v\delta}}=ue^{-(v-1)y^{{}^{\prime}}}$ (126) $$

If we compare the RHS of equation 125 and equation 126, we obtain $ue^{-(v-1)y}=ue^{-(v-1)y^{{}^{\prime}}}$. Since $u$ is non-zero, we obtain that $v=1$. Substituting this into $\sigma(vy)=u\sigma(y)$, we also obtain $u=1$.

Note that no other row of $U$ or $V$ can have same non-zero element because that would make matrix non invertible. From this we deduce that $U$ and $V$ are permutation matrices. From $\sigma(Vy)=U\sigma(y)$ it follows that $U=V=\Pi$. Thus $B=\Pi\tilde{B}$ and $\tilde{A}=A\Pi$.

Next, we equate predictions for $y_{2}$ to the ground truth (label $y_{2}$ exists as $T\geq 2$). For all $x_{1}\in\mathsf{supp}(X_{1})$

$$ $\begin{split}&\sigma(A\sigma(\Lambda\sigma(Bx_{1})+Bx_{2}))=\sigma(\tilde{A} \sigma(\tilde{\Lambda}\sigma(\tilde{B}x_{1})+\tilde{B}x_{2}))\implies\\ &A\sigma(\Lambda\sigma(Bx_{1})+Bx_{2})=\tilde{A}\sigma(\tilde{\Lambda}\sigma( \tilde{B}x_{1})+\tilde{B}x_{2})\implies\\ &\tilde{A}\sigma(\tilde{\Lambda}\sigma(\tilde{B}x_{1})+\tilde{B}x_{2})=A\Pi \sigma(\tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+\Pi^{\top}Bx_{2})=A\sigma(\Pi \tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+Bx_{2}).\\ \end{split}$ (127) $$

We use the simplification in the second step to equate to LHS in the first step as follows.

$$ $\begin{split}&A\sigma(\Pi\tilde{\Lambda}\Pi^{\top}\sigma(Bx_{1})+Bx_{2})=A \sigma(\Lambda\sigma(Bx_{1})+Bx_{2})\\ &\implies(\Pi\tilde{\Lambda}\Pi^{\top}-\Lambda)\sigma(Bx_{1})=0.\\ \end{split}$ (128) $$

Since $\sigma(Bx_{1})$ spans a set that has a non-empty interior, we get that $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi$ (from Lemma 3).

From the above conditions, we have arrived at $\tilde{\Lambda}=\Pi^{\top}\Lambda\Pi,\tilde{B}=\Pi^{\top}B,\tilde{A}=A\Pi$.

We want to show that for all $k\geq 1$

$$ $h_{k}=\Pi\tilde{h}_{k},$ (129) $$

where $h_{k}=\sigma(\Lambda h_{k-1}+Bx_{k})$ and $\tilde{h}_{k}=\sigma(\tilde{\Lambda}\tilde{h}_{k-1}+\tilde{B}x_{k})$ and $h_{0}=\tilde{h}_{0}=0$. In other words, we define $T_{k}$ as a mapping that takes $x_{\leq k}$ as input and outputs $h_{k}$, i.e., $T_{k}(x_{\leq k})=h_{k}$. Similarly, we write $\tilde{T}_{k}(x_{\leq k})=\tilde{h}_{k}$. We want to show

$$ $T_{k}=\Pi\tilde{T}_{k},\forall k$ (130) $$

We show the above by principle of induction. Let us consider the base case below. For all $x_{1}\in\mathbb{R}^{n}$

$$ $\tilde{A}\sigma(\tilde{B}x_{1})=A\Pi\sigma(\Pi^{\top}Bx_{1})=A\sigma(Bx_{1})= Ah_{1}\implies h_{1}=\Pi\tilde{h}_{1}\implies T_{1}(x_{1})=\Pi\tilde{T}_{1}(x_ {1})$ (131) $$

Suppose $\forall j\leq k,T_{j}=\Pi\tilde{T}_{j}$.

Having shown the base case and assumed the condition for $j\leq k$, we now consider the mapping $\tilde{T}_{k+1}$

$$ $\Pi\tilde{T}_{k+1}(x_{\leq k+1})=\Pi\sigma(\tilde{\Lambda}\tilde{h}_{k}+\tilde {B}x_{k+1})=\Pi\sigma(\Pi^{\top}\Lambda\Pi\tilde{h}_{k}+\Pi^{\top}Bx_{k})= \sigma(\Lambda h_{k}+Bx_{k})=T_{k+1}(x_{\leq k+1}).$ (132) $$

The prediction from the model $(\tilde{A},\tilde{\Lambda},\tilde{B})$ at a time step $k$ is denoted as $\tilde{y}_{k}$ and it relates to $\tilde{h}_{k}$ as follows $\tilde{y}_{k}=\sigma(\tilde{A}\tilde{h}_{k})$. We use the above condition in equation equation 130 to arrive at the following result. For all $x_{\leq k}\in\mathcal{X}^{k}$

$\tilde{y}_{k}=\sigma(\tilde{A}\tilde{h}_{k})=\sigma(\tilde{A}\tilde{T}(x_{\leq
k
}))=\sigma(A\Pi\tilde{T}(x_{\leq k}))=\sigma(AT(x_{\leq k}))=y_{k}$

This completes the proof.
∎

###### Assumption 18 .

###### Theorem 12 .

###### Proof.

## Appendix D Experiments

Here we provide additional experimental results as well as the training details.

##### Model Architecture

In all the architectures, there are two types of non-linearities, $\omega$ that generates the target label, $\psi$ that operates on inputs (used in deep sets and transformers). We use MLPs to implement these non-linearities. We instantiate MLPs with $l$ hidden layers, and the input, output, and hidden dimensions are all the same $m=n=k$. Recall that under the realizability assumption $f\in\mathcal{H}$. Therefore, we need to select the labeling function from $\mathcal{H}$. To do so, the weights of MLP are initialized according to $\mathcal{N}(\mu,\sigma^{2})$, where $\mu=0.0,\sigma=0.6$. For RNNs and SSMs, $A,B,\Lambda$ are initialized separately for the learner and true generating process as orthogonal matrices. All hidden layers, as well as the output layer are followed by a sigmoidal activation function.

##### Training Details and Hyperparameter Selection

We train all models with AdamW optimizer (loshchilov2018decoupled) with a learning rate of $10^{-3}$, weight decay of $0.01$, $\epsilon=10^{-8},\beta_{1}=0.9,\beta_{2}=0.95$. We reduce the learning rate by a factor of $0.8$ if the validation loss is not improved more than $10^{-6}$ for 1 epoch. This drop is followed by a cool-down period of 1 epoch, and the learning rate cannot decrease to lower than $10^{-7}$. For all datasets we use a streaming dataset where each epoch contains 100 batches of size 256 sampled online from the specified training and test distributions, and we train all models for 100 epochs. Therefore, the size of the training dataset is $256\times 10^{4}$ and the size of the testing dataset is $256\times 10^{2}$. Since our models are generally small, running the experiments is rather inexpensive, and we carried out each experiment on 4 CPU cores using 20 GB of RAM. For inference, specially for SSM and RNN with very long sequences, we use RTX8000 GPUs.

### D.1 Length Generalization

In Figure 5, we present additional findings for length generalization capability of all architectures when both the learner and the generating process MLPs all consist of one hidden layers with input, output, and hidden size matching $n=m=k=20$.

Figure: Figure 5: 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^{-5}-10^{-6}$ 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.
Refer to caption: x3.png

To complement Figure 3, in Figures 6, 7 we present the prediction behaviour of SSM and RNN architectures with two hidden layer MLPs for $\omega$ trained on sequences output by two hidden layer MLPs for $\rho$.

Figures 8, 9, 10, 11 present the prediction behaviour of deep set, Transformer with softmax attention, SSM, and RNN architectures with one hidden layer in $\rho$ (and one hidden layer MLPs for the learner $\omega$). Training procedure remains the same. We can observe that all models length generalize.

Additionally, to support the theory on other types of attention, Figures 12, 13 demonstrate the loss and prediction of a Transformer with ReLU attention and one hidden layer MLPs for $\omega,\psi$ trained on output sequences of a Transformer with ReLU attention and one hidden layer MLP for $\rho,\phi$. Similarly, all these models were trained to predict sequences of length up to $T=10$ output by a true labeling function $f$ in their respective hypothesis classes $\mathcal{H}$, and were tested with sequences of length up to 100. As a reminder, the output tokens $y_{i}\in\mathbb{R}^{m}$, where $m=20$, and the figures below show only one representative dimension for illustration. All models demonstrate strong length generalization capacity.

##### Discrete Tokens

In Table 3 we present the results for successful length generalization of the different architectures when the inputs are discrete. We sample all components from $[0,1]$ interval and discretize the values to one of the 5 levels in
$[0.0,0.2,0.4,0.6,0.8]$. Note that the small scale of values of loss at longer lengths indicate successful generalization. For a visual depiction of results, please see Fig. 14. Also note that in all architectures $\rho$ in $f$ and $\omega$ in $h$ are comprised of two hidden layers.

**Table 3: Length generalization of different architectures when the input tokens are discrete. Models are trained in sequences of length up to $T=10$ and show successful generalization on much longer sequences.**
| Model | Test Loss $\times 10^{6}$ ($t=10$) | Test Loss $\times 10^{6}$ ($t=90$) |
| --- | --- | --- |
| Deep set | $3.48\pm 0.15$ | $52.7\pm 0.88$ |
| Transformer | $1.72\pm 0.27$ | $48.8\pm 2.44$ |
| SSM | $0.2\pm 0.0$ | $4.06\pm 0.0$ |
| RNN | $0.22\pm 0.0$ | $1.3\pm 0.0$ |

Figure: Figure 6: A SSM model 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$.
Refer to caption: x4.png

Figure: Figure 7: A RNN model 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$.
Refer to caption: x5.png

Figure: Figure 8: A deep set model with one hidden layer MLP for $\psi,\omega$ trained on sequences of length up to $T=10$ shows perfect generalization to sequences of length up to $100$.
Refer to caption: x6.png

Figure: Figure 9: A Transformer model with softmax attention and one hidden layer MLP trained on sequences of length up to $T=10$ shows perfect generalization to sequences of length up to 100.
Refer to caption: x7.png

Figure: Figure 10: A SSM model with one hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ length generalizes to sequences of length up to $100$.
Refer to caption: x8.png

Figure: Figure 11: A RNN model with one hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ length generalizes to sequences of length up to $100$.
Refer to caption: x9.png

Figure: Figure 12: Test loss of a transformer model with ReLU attention and one hidden layer MLP for $\omega,\psi$ trained on sequences of length up to $T=10$ length generalizes to sequences of length up to $100$. The results are averaged over five seeds.
Refer to caption: x10.png

Figure: Figure 13: A transformer model with ReLU attention and one hidden layer MLP for $\omega,\psi$ trained on sequences of length up to $T=10$ length generalizes to sequences of length up to $100$.
Refer to caption: x11.png

Figure: (a)
Refer to caption: x12.png

### D.2 Compositional Generalization

Here we present the prediction behavior of different architectures on the test sequences that consist of unseen token combinations during training. This helps us better interpret qualitatively how the model actually performs in following the true labels. Figures 16- 19 show the prediction trajectories for different architectures. We can observe that not only do these models perform quite well on unseen sequences of length up to $T=10$, but they also length generalize and continue to remain consistent with the true labels on unseen combinations at longer lengths than the training. Table 4 presents the test loss and $R^{2}$ on the test set when the model is only trained on the red region in Figure 15. All models generalize to unseen combination of tokens and the learned representations linearly identify the true hidden representations.

Figures 20, 21, 22, 23 present the prediction behaviour of deep set, transformer with softmax attention, SSM, and RNN architectures with two hidden layers in $\rho$ (and two hidden layer MLPs for the learner $\omega$) when trained on sequences of length up to $T=10$ sampled from the red region in Figure 15. We can observe that all models continue to generalize to unseen combinations beyond their training length.

##### Discrete Tokens

Evaluating compositional generalization with discrete tokens introduces additional challenges. This is because we have to sample the training and test distribution according to Fig. 1 (and 15). There are multiple ways to achieve this but they become infeasible with long sequences of interest in practice:

- •
We could continuously sample from the regions in Fig. 1 and 15 and then round up or down the components to one of the predefined set of values. However, with longer sequences this translates to sampling and then rounding values in a high-dimensional hyper-diamond where points are increasingly spread out toward the boundaries. Rounding up results in corners becoming part of the training samples, corrupting the test set. Rounding down will result in a training set in which the support of tokens no longer follows Fig. 1 (i.e., does not cover the discrete set of values predefined in $[0,1]$).
- •
We could instead sample continuously and then discretize based on finding the nearest neighbour of each point to the points in a discrete grid of values in $\mathbb{R}^{T}$. Having as few as 5 discrete levels renders this sampling procedure impossible for long sequences due to the complexity of finding nearest neighbours.
- •
Lastly, one could construct the set of discrete points in $\mathbb{R}^{T}$ that satisfy the constraints in Fig. 1 and then sample from this set, however, this enumeration also proves infeasible as the search space grows exponentially.

Therefore, evaluating compositional generalization in the discrete case is not straightforward beyond very short sequences.

**Table 4: Compositional generalization: Test $\ell_{2}$ loss and $R^{2}$ score for models with one hidden layers on sequences of length $T=10$. A strong linear relationship is observed for all models for new sequences made of unseen token combinations.**
| Model | Test Loss $\times 10^{6}$ | $R^{2}$ |
| --- | --- | --- |
| Deep set | $1.27\pm 0.24$ | $0.96\pm 0.01$ |
| Transformer | $4.50\pm 3.28$ | $1.00\pm 0.00$ |
| SSM | $11.00\pm 10.92$ | $1.00\pm 0.00$ |
| RNN | $1.22\pm 0.12$ | $0.99\pm 0.00$ |

##### Practical Considerations

For training and evaluating compositional aspect of generalization, we follow the sampling procedure described in Figure 1 with a slight modification that allows for faster sampling and easier training. This procedure is illustrated in Figure 15, and results in a more difficult testing strategy, as the test set spans a smaller area than the complement of the training set.

We opted for such a procedure because rejection sampling from the complement of the training set given in Figure 1 is extremely slow. In particular, given our batch size of 256, token dimension $n=20$, and having 100 batches per epoch, constructing the full test set requires finding $256\times 100\times 20$ sequences of length $t\leq T$ that are rejected by the original constraints. This becomes quite inefficient and slow specially in higher dimensions as the sum of the sequence along each component tends to concentrate more around $t/2$, therefore it becomes harder to find such sequences (the sum follows Irwin-Hall distribution since the components come from the Uniform distribution).
To improve the speed of sampling the test dataset, we sample token dimensions $x_{i}^{k}$ from the smaller corners shown in Figure 15 which allows for parallel sampling. These corners correspond to sampling $x_{i}^{k}\sim\text{Uniform}[0,1/2T]$ or $x_{i}^{k}\sim\text{Uniform}[1/2+1/2T,1]$. This way we can sample token components independently and in parallel without having to reject any samples, since by construction no test sequence coincides with the training set. This procedure leaves a gap (see Figure 15) that will not be sampled neither during training nor testing.

Figure: Figure 15: Illustrating the modified support of train vs test distribution for compositional generalization. This enables speed up in the sampling procedure, while keeping the challenging aspect of generalization to the corners.
Refer to caption: NeurIPS/comp_gen_practice.pdf

Figure: Figure 16: A deep set model with one hidden layer MLP for $\omega,\psi$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences (Figure 15). Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x15.png

Figure: Figure 17: A Transformer model with softmax attention and one hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences (Figure 15). Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x16.png

Figure: Figure 18: A SSM model with one hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences (Figure 15). Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x17.png

Figure: Figure 19: A RNN model with one hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences (Figure 15). Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x18.png

Figure: Figure 20: A deep set model with two hidden layer MLP for $\omega,\psi$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences. Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x19.png

Figure: Figure 21: A Transformer model with softmax attention and two hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences. Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x20.png

Figure: Figure 22: A SSM model with two hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences. Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x21.png

Figure: Figure 23: A RNN model with two hidden layer MLP for $\omega$ trained on sequences of length up to $T=10$ sampled according to Figure 15 can generalize to unseen test sequences. Additionally, the compositional generalization holds even beyond the training length.
Refer to caption: x22.png

### D.3 Failure Cases

Although most of our focus has been on the success scenarios for length and compositional generalization, here we provide examples to show how a model might fail.

##### f 𝑓 f italic_f is not realizable in ℋ ℋ \mathcal{H} caligraphic_H

In Figures 25, 26, 27, 28, we present the predictions of learned models from different architectures initialized in their respective $\mathcal{H}$ that does not contain the true $f$. In particular, we have the following for the different architectures:

- •
Deep set: The true labeling function $f$ is a deep set with one hidden layer MLPs for $\rho,\phi$, but the learner uses $h$, a deep set model for which the MLPs $\psi,\omega$ have no hidden layers.
- •
Transformer: The true labeling function $f$ is a Transformer with 1 hidden layer in $\rho$, but the learner uses $h$, an RNN with 1 hidden layer in $\omega$.
- •
SSM: The true labeling function $f$ is an SSM with 1 hidden layer in $\rho$, but the learner uses $h$, an RNN without any hidden layers in $\omega$.
- •
RNN: The true labeling function $f$ is an RNN with 1 hidden layer in $\rho$, but the learner uses $h$, a Transformer with 1 hidden layer in $\omega$.

In each case, the learner is trained on sequences of length up to $T=10$ and its performance on the test set at longer lengths indicates whether generalization is possible or not. For a visual illustration of such failures beyond the training length, see Figures 25, 26, 27, 28. We can observe that the model can predict the test sequence well up to the length it has learned during training, but starts to diverge from the true labels beyond that. This demonstrates a failure case in which the realizability condition is violated.

##### f 𝑓 f italic_f is realizable in a high capacity ℋ ℋ \mathcal{H} caligraphic_H

For a given $\mathcal{H}$, if all solutions to 1 achieve length generalization or compositional generalization, then we can guarantee length or compositional generalization regardless of the training procedure. When the capacity of $\mathcal{H}$ becomes very large, it continues to contain the right solutions but it starts to contain many incorrect solutions that match the true solution only on the support of training distribution. In such a case, there is no reason to presume that our learning procedure picks the right solution to 1 that also achieves length and compositional generalization. Figure 24 show experiments illustrating the above. We experiment with the following scenarios for deep sets and transformers:

- •
Deep set: We use the labeling function that takes the following form $f=\rho(\sum_{i\leq t}\phi(x_{i}))$ for $t\leq T$ and $f=\rho(\sum_{i\leq t}\phi(x_{i}))+c$ for $t>T$ with $c=0.2,T=5$. We use 1 hidden layer MLPs for $\rho,\phi$ (with no activation on the output of $\rho$). We use 2 hidden layer MLPs for $\omega,\psi$ for $h$ so that it can express the above labeling function. The input, hidden, and output dimensions are all equal $m=n=k=20$ for $f,h$. We train on sequences of length longer than $T$ to demonstrate this expressivity claim. When the model is trained on sequences of length less than $T$, due to the simplicity bias of the training procedure model learns $\rho(\sum_{i\leq t}\phi(x_{i}))$ and uses it on longer sequences and hence fails.
- •
Transformer: We use the labeling function that takes the following form $f=\rho(\sum_{j=1}^{i}\frac{1}{i}\phi(x_{i},x_{j}))$ for $t\leq T$ and $f=\rho(\sum_{j=1}^{i}\frac{1}{i}\phi(x_{i},x_{j}))+c$ for $t>T$ with $c=0.1,T=10$. We use 1 hidden layer MLPs for $\rho$ (with no activation on the output of $\rho$). We use a Transformer with 3 hidden layer MLPs for $\omega$ so that it can express the above labeling function. The input, hidden, and output dimensions are all equal $m=n=k=20$ for $f,h$. We train on sequences of length longer than $T$ to demonstrate this expressivity claim. When the model is trained on sequences of length less than $T$, due to the simplicity bias of the training procedure model learns $f=\rho(\sum_{j=1}^{i}\frac{1}{i}\phi(x_{i},x_{j}))$ and uses it on longer sequences and hence fails.

The failures of such degenerate solutions can be visualized in Figure 24 (right), where the predictions diverge from the true values when the model is only trained on sequences shorter than $T_{0}$. Figure 24 (left) shows that when the model is trained on sequences longer than $T_{0}$, it can successfully generalize to longer lengths. Table 5 further validates this observation numerically. It presents the test loss of each model at lengths shorter and longer than $T_{0}$ under the two training schemes: a) When trained only on sequences of length shorter than $T_{0}$ (rows corresponding to Fig 2-b and 2-d which result in failure due to degenerate solution), b) when trained on sequences of length longer than $T_{0}$ (rows corresponding to Fig 2-a and 2-c which result in successful generalization).

**Table 5: Length generalization of different architectures when the hypothesis class $\mathcal{H}$ is highly expressive. For further details see Fig. 24**
| Deep set | Loss ($t<T_{0}$) | Loss ($t\geq T_{0}$) |
| --- | --- | --- |
| Fig 2-a | $0.001\pm 10^{-4}$ | $0.002\pm 3\times 10^{-4}$ |
| Fig 2-b | $0.0007\pm 10^{-4}$ | $0.007\pm 0.001$ |
| Transformer | Loss ($t<T_{0}$) | Loss ($t\geq T_{0}$) |
| Fig 2-c | $0.0006\pm 10^{-4}$ | $0.006\pm 3\times 10^{-3}$ |
| Fig 2-d | $10^{-5}\pm 10^{-6}$ | $0.01\pm 0.003$ |

Figure: (a)
Refer to caption: x23.png

Figure: Figure 25: A failure case of length generalization in the unrealizble setting: The predictions come from a deep set with linear layers for $\psi,\omega$ trained to predict the sequences (of length up to $T$) output by a deep set with 1 hidden layer MLPs for $\phi,\rho$. In this case the realizability condition does not hold, and the learner fails to length generalize.
Refer to caption: x27.png

Figure: Figure 26: A failure case of length generalization in the unrealizble setting: The predictions come from an RNN with 1 hidden layer in $\omega$ trained to predict the sequences (of length up to $T=10$) output by a Transformer with 1 hidden layer in $\rho$. In this case the realizability condition does not hold, and the learner fails to length generalize.
Refer to caption: x28.png

Figure: Figure 27: A failure case of length generalization: The predictions come from an RNN without any hidden layers in $\omega$ trained to predict the sequences (of length up to $T=10$) output by an SSM with 1 hidden layer in $\rho$. In this case the realizability condition does not hold, and the learner fails to length generalize.
Refer to caption: x29.png

Figure: Figure 28: A failure case of length generalization: The predictions come from a Transformer with 1 hidden layer in $\omega$ trained to predict the sequences (of length up to $T=10$) output by an RNN with 1 hidden layer in $\rho$. In this case the realizability condition does not hold, and the learner fails to length generalize.
Refer to caption: x30.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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).
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
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.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See weaknesses
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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)?
soundness: 3
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.
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.
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?

Meta-review: 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.

Decision: Reject