L EARNING L OW D IMENSIONAL S TATE S PACES WITH O VERPARAMETERIZED R ECURRENT N EURAL N ET - WORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Overparameterization in deep learning typically refers to settings where a trained
Neural Network (NN) has representational capacity to ﬁt the training data in many ways, some of which generalize well, while others do not.
In the case of Recurrent
Neural Networks (RNNs), there exists an additional layer of overparameterization, in the sense that a model may exhibit many solutions that generalize well for sequence lengths seen in training, some of which extrapolate to longer sequences, while others do not.
Numerous works studied the tendency of Gradient Descent (GD) to ﬁt overparameterized NNs with solutions that generalize well.
On the other hand, its tendency to ﬁt overparameterized RNNs with solutions that extrapolate has been discovered only lately, and is far less understood.
In this paper, we analyze the extrapolation properties of GD when applied to overparameterized linear RNNs.
In contrast to recent arguments suggesting an implicit bias towards short-term memory, we provide theoretical evidence for learning low dimensional state spaces, which can also model long-term memory.
Our result relies on a dynamical characterization which shows that GD (with small step size and nearzero initialization) strives to maintain a certain form of balancedness, as well as on tools developed in the context of the moment problem from statistics (recovery of a probability distribution from its moments).
Experiments corroborate our theory, demonstrating extrapolation via learning low dimensional state spaces with both linear and non-linear RNNs.
I NTRODUCTION
Neural Networks (NNs) are often overparameterized , in the sense that their representational capacity far exceeds what is necessary for ﬁtting training data.
Surprisingly, training overparameterized NNs via (variants of) Gradient Descent (GD) tends to produce solutions that generalize well, despite existence of many solutions that do not.
This implicit generalization phenomenon attracted considerable scientiﬁc interest, resulting in various theoretical explanations (see, e.g., Woodworth et al.
Yun et al.
; Zhang et al.
; Li et al.
; Ji & Telgarsky (2018); Lyu & Li (2019)).
Recent studies have surfaced a new form of implicit bias which arises in Recurrent Neural Networks (RNNs) and their variants (e.g., Long Short-Term Memory Hochreiter & Schmidhuber (1997) and
Gated Recurrent Units Chung et al.
For such models, the length of sequences in training is often shorter than in testing, and it is not clear to what extent a learned solution will be able to extrapolate beyond the sequence lengths seen in training.
In the overparameterized regime, where the representational capacity of the learned model exceeds what is necessary for ﬁtting short sequences, there may exist solutions which generalize but do not extrapolate, meaning that their accuracy is high over short sequences but arbitrarily poor over long ones (see Cohen-Karlik et al.
In practice however, when training RNNs using GD, accurate extrapolation is often observed.
We refer to this phenomenon as the implicit extrapolation of GD.
As opposed to the implicit generalization of GD, little is formally known about its implicit extrapolation.
Existing theoretical analyses of the latter focus on linear RNNs — also known as Linear
Dynamical Systems (LDS) — and either treat inﬁnitely wide models (Emami et al., 2021), or models of ﬁnite width that learn from a memoryless teacher (Cohen-Karlik et al., 2022).
In these regimes,
GD has been argued to exhibit an implicit bias towards short-term memory.
While such results are informative, their generality remains in question, particularly since inﬁnitely wide NNs are known to substantially differ from their ﬁnite-width counterparts, and since a memoryless teacher essentially neglects the main characteristic of RNNs (memory).
In this paper, we theoretically investigate the implicit extrapolation of GD when applied to overparameterized ﬁnite-width linear RNNs learning from a teacher with memory.
We consider models with symmetric transition matrices, in the case where a student (learned model) with state space dimension d is trained on sequences of length k generated by a teacher with state space dimension ˆ d .
Our interest lies in the overparameterized regime, where d is greater than both k and ˆ d , meaning that the student has state space dimension large enough to fully agree with the teacher on sequences of length k , while potentially disagreeing with it on longer sequences.
As a necessary assumption on initialization, we follow prior work and focus on a certain balancedness condition, which is known (see experiments in Cohen-Karlik et al.
(2022), as well as our theoretical analysis) to capture near-zero initialization as commonly employed in practice.
Our main theoretical result states that GD originating from a balanced initialization leads the student to extrapolate, irrespective of how large its state space dimension is .
Key to the result is a surprising connection to a moment matching theorem from Cohen & Yeredor (2011), whose proof relies on ideas from compressed sensing (Elad, 2010; Eldar & Kutyniok, 2012) and neighborly polytopes (Gale, 1963).
This connection may be of independent interest, and in particular may prove useful in deriving other results concerning implicit properties of GD.
We corroborate our theory with experiments, which demonstrate extrapolation via learning low dimensional state spaces in both the analyzed setting and ones involving non-linear RNNs.
The implicit extrapolation of GD is a new and exciting area of inquiry.
Our results suggest that short-term memory is not enough for explaining it as previously believed.
We hope the techniques developed in this paper will contribute to a further understanding of this phenomenon.
R ELATED W ORK
The study of linear RNNs, or LDS, has a rich history dating back to at least the early works of Kalman (Kalman, 1960; 1963).
An extensively studied question relevant to extrapolation is that of system identiﬁcation , which explores when the parameters of a teacher LDS can be recovered (see Ljung (1999)).
Another related topic concerns ﬁnding compact realizations of systems, i.e. realizations of the same input-output mapping as a given LDS, with a state space dimension that is lower (see
Antoulas (2005)).
Despite the relation, our focus is fundamentally different from the above — we ask what happens when one learns an LDS using GD.
Since GD is not explicitly designed to ﬁnd a low dimensional state space, it is not clear that the application of GD to an overparameterized student allows system identiﬁcation through a compact realization.
The fact that it does relates to the implicit properties of GD, and to our knowledge has not been investigated in the classic LDS literature.
The implicit generalization of GD in training RNNs has been a subject of theoretical study for at least several years (see, e.g., Hardt et al.
(2016); Allen-Zhu et al.
; Lim et al.
In contrast, works analyzing the implicit extrapolation of GD have surfaced only recently, speciﬁcally in Emami et al.
and Cohen-Karlik et al.
Emami et al.
(2021) analyzes linear RNNs in the inﬁnite width regime, suggesting that in this case GD is implicitly biased towards impulse responses corresponding to short-term memory.
Cohen-Karlik et al.
(2022) studies ﬁnite-width linear
RNNs (as we do), showing that when the teacher is memoryless (has state space dimension zero),
GD emanating from a balanced initialization successfully extrapolates.
Our work tackles an arguably more realistic and challenging setting — we analyze the regime in which the teacher has memory.
Our results suggest that the implicit extrapolation of GD does not originate from a bias towards short-term memory, but rather a tendency to learn low dimensional state spaces.
1 We note that there have been works studying extrapolation in the context of non-recurrent NNs, e.g. Xu et al.
, this et al teacher stable
Our work on the other hand is not limited by these assumptions, and we show that in cases where they are violated, learning yields solutions with low dimensional state spaces that do not result in short-term memory.
type of extrapolation deals with the behavior of learned functions outside the support of the training distribution, and is therefore fundamentally different from the type of extrapolation we consider which deals with the behavior of learned functions over sequences longer than those seen in training. Linear RNNs fundamentally differ from the commonly studied model of linear (feed-forward) NNs (see, e.g., Arora et al. ; Ji & Telgarsky (2018); Arora et al. (2019a;b);
Razin & Cohen (2020)).
One of the key differences is that a linear RNN entails different powers of a parameter (transition) matrix, leading to a loss function which roughly corresponds to a sum of losses for multiple linear NNs having different architectures and shared weights.
This precludes the use of a vast array of theoretical tools tailored for linear NNs, rendering the analysis of linear RNNs technically challenging.
On the empirical side, extrapolation of NNs to sequence lengths beyond those seen in training has been experimentally demonstrated in numerous recent works, covering both modern language and attention models (Press et al., 2022; Anil et al., 2022; Zhang et al., 2022), and RNNs with transition matrices of particular forms (Gu et al., 2022; 2021; 2020; Gupta, 2022).
The current paper is motivated by these ﬁndings, and takes a step towards theoretically explaining them.
L INEAR R ECURRENT N EURAL N ETWORKS
Our theoretical analysis applies to single-input single-output (SISO) linear RNNs with symmetric transition matrices.
Given a state space dimension d P N , this model is deﬁned by the update rule:
[Equation]
[Equation]
[Equation]
) where A P R d ˆ d , B P R d ˆ 1 and C P R 1 ˆ d are conﬁgurable parameters, with the transition matrix A satisfying A “ A J ; x 0 , x 1 , . . .
P R form an input sequence; y 0 , y 1 , . . .
P R form the corresponding output sequence; and s t P R d
ˆ 1 represents the internal state at time t , assumed to be equal to zero at the outset (i.e. it is assumed that s 0 “ 0 ).
As with any linear time-invariant system (Porat, 1996), the input-output mapping realized by the RNN is determined by its impulse response.
Deﬁnition 1.
The impulse response of the RNN is the output sequence corresponding to the input sequence p x 0 , x 1 , x 2 , . . .
q “ p 1 , 0 , 0 , . . .
q .
Namely, it is the sequence p CB , CAB , CA 2 B , . . .
q .
For brevity, we employ the shorthand ⇥ : “ p A , B , C q .
The d ˆ d symmetric transition matrix A is parameterized through a d p d ` 1 q{ 2 -dimensional vector holding its upper triangular elements, and with a slight overloading of notation, the symbol A is also used to refer to this parameterization.
We note that our theory readily extends to multiple-input multiple-output (MIMO) networks, and the focus on the SISO case is merely for simplicity of presentation.
Note also that the restriction to symmetric transition matrices is customary in both theory (Hazan et al., 2018) and practice (Gupta, 2022), and represents a generalization of the canonical modal form , which under mild non-degeneracy conditions does not limit generality (Boyd & Lessard, 2006).
Given a length k input sequence, x “ p x 0 , . . .
, x k ´ 1 q P R k , consider the output at the last time step, i.e. y : “ y k P R , and denote it by RNN p x q .
Using this output as a label, we deﬁne an empirical loss induced by a training set S “ ` x p 1 q , y p 1 q [Equation]
[Equation]
where ` p y, ˆ y q “ p y ´ ˆ y q 2 is the square loss.
By the update rule of the RNN (Equation 3.1), we have:
[Equation]
Suppose that ground truth labels are generated by an RNN as deﬁned in Equation 3.1, and denote the state space dimension and parameters of this teacher network by ˆ d and ˆ ⇥ “ p ˆ A , ˆ B , ˆ C q respectively.
We employ the common assumption (e.g., see Hardt et al.
(2016)) by which input sequences are drawn from a whitened distribution, i.e. a distribution where E r x j x j 1 s equals 1 if j “ j 1 and 0 otherwise.
The population loss over length k sequences can then be written as (see Lemma D.1):
[Equation]
Equation 3.4 implies that a solution ⇥ “ p A , B , C q achieves zero population loss over length k sequences if and only if CA j B “ ˆ C ˆ A j ˆ B for j “ 0 , . . .
, k ´ 1 .
To what extent does such a solution imply that the student (i.e., the learned RNN) extrapolates to longer sequences?
This depends on how close CA j B is to ˆ C ˆ
A j ˆ B for j • k .
Deﬁnition 2.
For ✏ • 0 and q P N , we say that the student ✏ -extrapolates with horizon q with respect to (w.r.t) the teacher if:
[Equation]
If the above holds for all q P N then the student is said to ✏ -extrapolate w.r.t the teacher, and if it holds for all q P N with ✏ “ 0 then the student is simply said to extrapolate w.r.t the teacher.
Per Deﬁnition 2, ✏ -extrapolation with horizon q is equivalent to the ﬁrst q elements of the student’s impulse response being ✏ -close to those of the teacher’s, whereas extrapolation means that the student’s impulse response fully coincides with the teacher’s.
The latter condition implies that the student realizes the same input-output mapping as the teacher, for any sequence length (this corresponds to the notion of system identiﬁcation; see Section 2).
When the student is overparameterized , in the sense that d is greater than k and ˆ d , it may perfectly generalize, i.e. lead the population loss over length k sequences (Equation 3.4) to equal zero, and yet fail to extrapolate as stated in the following proposition.
Proposition 3.
Assume d ° k , and let ✏ • 0 and q P t k ` 1 , k ` 2 , . . .
u .
Then, for any teacher parameters ˆ ⇥ , there exist student parameters ⇥ with which the population loss in Equation 3.4 equals zero, and yet the student does not ✏ -extrapolate with horizon q .
Proof sketch (for complete proof see Appendix D.1.2).
The result follows from the fact that the ﬁrst d elements of the student’s impulse response can be assigned freely via a proper choice of ⇥ .
We are interested in the extent to which student parameters learned by GD extrapolate in the overparameterized regime.
Proposition 3 implies that, regardless of how many (length k ) sequences are used in training, if GD leads to any form of extrapolation, it must be a result of some implicit bias induced by the algorithm.
Note that in our setting, extrapolation cannot be explained via classic tools from statistical learning theory, as evaluation over sequences longer than those seen in training violates the standard assumption of train and test data originating from the same distribution.
To decouple the question of extrapolation from that of generalization, we consider the case where the training set S is large, or more formally, where the empirical loss L S p¨q (Equation 3.3) is well represented by the population loss L p¨q (Equation 3.4).
We model GD with small step size via
Gradient Flow (GF), as customary in the theory of NNs — see Saxe et al.
; Gunasekar et al.
(2017); Arora et al.
(2018; 2019b)
Lyu & Li (2019); Li et al.
; Azulay et al.
(2021) for examples
and Elkabetz & Cohen (2021) for theoretical justiﬁcation.
That is, we analyze the following dynamics:
[Equation]
where ↵ P t A , B , C u .
If no assumption on initialization is made, no form of extrapolation can be established (indeed, the initial point may be a global minimizer of L p¨q that fails to extrapolate, and GF will stay there).
Following prior work (see Cohen-Karlik et al.
, we assume that initialization adheres to the following balancedness condition:
Deﬁnition 4.
An RNN with parameters ⇥ “ p A , B , C q is said to be balanced if B “ C J .
It was shown empirically in Cohen-Karlik et al.
(2022) that the balancedness condition captures near-zero initialization as commonly employed in practice.
We support this ﬁnding theoretically in
Section 4.3.
Aside from the initialization of the student, we will also assume that the teacher adheres to the balancedness condition.
T HEORETICAL A NALYSIS
We turn to our theoretical analysis.
Section 4.1 proves that in the setting of Section 3, convergence of
GF to a zero loss solution leads the student to extrapolate, irrespective of how large its state space dimension is .
Section 4.2 extends this result by establishing that, under mild conditions, approximate convergence leads to approximate extrapolation.
The results of sections 4.1 and 4.2 assume that GF emanates from a balanced initialization, which empirically is known to capture near-zero initialization as commonly employed in practice (see Section 3).
Section 4.3 theoretically supports this empirical premise, by showing that with high probability, random near-zero initialization leads to balancedness.
We introduce notations that will be used throughout the analysis.
For a matrix Q P R m ˆ n , we let } Q } F , } Q } 8 and } Q } 2 denote the Frobenius, ` 8 and ` 2 (spectral) norms, respectively.
For a vector v P R m , we use } v } to denote the Euclidean norm and v i to denote its i th entry.
4. C ONVERGENCE L EADS TO E XTRAPOLATION
Theoretical analyses of implicit generalization often assume convergence to a solution attaining zero loss (see, e.g., Azulay et al.
; Gunasekar et al.
; Lyu & Li (2019); Woodworth et al.
Under such an assumption, Theorem 5 below establishes implicit extrapolation, i.e. that the solution to which GF converges extrapolates, irrespective of how large the student’s state space dimension d is .
A condition posed by the theorem is that the training sequence length k is greater than two times the teacher’s state space dimension ˆ d .
We empirically demonstrate the necessity of this condition in Section 5, by showing that smaller values for k lead to failure in extrapolation.
Theorem 5.
Assume that d ° k ° 2 ˆ d , the teacher parameters ˆ ⇥ are balanced (Deﬁnition 4), and the student parameters ⇥ are learned by applying GF to the loss L p¨q (Equation 3.4) starting from a balanced initialization.
Then, if GF converges to a point ⇥ ˚ satisfying L p ⇥ ˚ q “ 0 , this point extrapolates (Deﬁnition 2).
In order to prove Theorem 5, we introduce two lemmas:
Lemma 6, which shows that balancedness is preserved under GF; and Lemma 7, which (through a surprising connection to a moment problem from statistics) establishes that a balanced solution attaining zero loss necessarily extrapolates.
With
Lemmas 6 and 7 in place, the proof of Theorem 5 readily follows.
Lemma 6.
Let ⇥ p ⌧ q , with ⌧ • 0 , be a curve brought forth by applying GF to the loss L p¨q starting from a balanced initialization.
Then, ⇥ p ⌧ q is balanced for every ⌧ • 0 .
Proof sketch (for complete proof see Appendix D.1.5).
The result follows from the symmetric role of B and C in the loss L p¨q .
Lemma 7.
Suppose that d ° k ° 2 ˆ d , the teacher is balanced, and that the student parameters ⇥ are balanced and satisfy L p ⇥ q “ 0 .
Then ⇥ extrapolates.
Proof sketch (for complete proof see Appendix D.2).
The proof is based on a surprising connection to the moment problem from statistics (recovery of a probability distribution from its moments), which has been studied for decades (see, e.g., Schmüdgen (2017)).
Without loss of generality, we may assume that ˆ A is diagonal (if this is not the case then we apply an orthogonal eigendecomposition to ˆ A and subsequently absorb eigenvectors into ˆ B and ˆ C ).
We may also assume that ˆ C ˆ B “ 1 (otherwise we absorb a scaling factor into B and/or C ).
Since ˆ C J “ ˆ B (teacher is balanced), we may deﬁne a probability vector (i.e. a vector with non-negative entries summing up to one) ˆ p P R ˆ d via ˆ p
i “ ˆ C
i ˆ B i , i “ 1 , . . .
, ˆ d .
We let ˆ Z denote the random variable whose support is t ˆ A 11 , . . .
, ˆ A ˆ d ˆ d u and which assumes the value ˆ A ii with probability ˆ p
i , i “ 1 , . . .
, ˆ d .
Notice that for every j P N :
[Equation]
meaning that the elements of the teacher’s impulse response are precisely the moments of ˆ Z .
Similarly to above we may assume A is diagonal, and since L p ⇥ q “ 0 it holds that CB “ ˆ C ˆ B “ 1 .
We may thus deﬁne a probability vector p P R d via p i “ C i B i , i “ 1 , . . .
, d , and a random variable Z which assumes the value A ii with probability p
i , i “ 1 , . . .
, d .
For every j P N :
[Equation]
and so the elements of the student’s impulse response are precisely the moments of Z .
The probabilistic formulation we set forth admits an interpretation of extrapolation as a moment problem.
Namely, since L p ⇥ q “ 0 (i.e. CA j B “ ˆ C ˆ A j ˆ B for j “ 0 , . . .
, k ´ 1 ) the random variables Z and ˆ Z agree on their ﬁrst k ´ 1 moments, and the question is whether they agree on all higher moments as well.
We note that this question is somewhat more challenging than that tackled in classic instances of the moment problem, since the support of the random variable whose moments we match ( Z ) is not known to coincide with the support of the random variable we seek to recover ( ˆ Z ).
Luckily, a recent powerful result allow addressing the question we face — Cohen & Yeredor (2011) showed that the ﬁrst 2 n moments of a discrete random variable X taking at most n P N values uniquely deﬁne X , in the sense that any discrete random variable agreeing with these 2 n moments must be identical to X .
Translating this result to our setting, we have that if Z agrees with ˆ Z on its ﬁrst 2 ˆ d moments, it must be identical to ˆ Z , and in particular it must agree with ˆ Z on all higher moments as well.
The fact that k ´ 1 • 2 ˆ d then concludes the proof.
To attain some intuition for the result we imported from Cohen & Yeredor (2011), consider the simple case where ˆ d “ 1 .
The transition matrix ˆ A is then a scalar ˆ a P R , the random variable ˆ Z is deterministically equal to ˆ a , and the teacher’s impulse response is given by the moments
E r
ˆ Z j s “ ˆ a j , j “ 0 , 1 , . .
Since by assumption k ° 2 ˆ d , the fact that L p ⇥ q “ 0 means the random variable corresponding to the student, Z , agrees with the ﬁrst two moments of ˆ Z .
That is, Z satisﬁes E r Z s “ ˆ a and E r Z 2 s “ ˆ a 2 .
This implies that Var r Z s “ E r Z 2 s ´
E r Z s 2 “ 0 , and therefore Z is deterministically equal to ˆ a
, i.e. it is identical to ˆ Z .
The two random variables thus agree on all of their moments, meaning the impulse responses of the student and teacher are the same.
Proof of Theorem 5.
By Lemma 6 (as well as continuity considerations) ⇥ ˚ is balanced.
Lemmathen implies that it extrapolates.
4. A PPROXIMATE C ONVERGENCE L EADS TO A PPROXIMATE E XTRAPOLATION
Theorem 5 in Section 4.1 proves extrapolation in the case where GF converges to a zero loss solution.
Theorem 8 below extends this result by establishing that, as long as the teacher is stable , approximate convergence leads to approximate extrapolation.
Theorem 8.
Assume the conditions of Theorem 5, and that the teacher parameters ˆ ⇥ are stable, i.e. the singular values of ˆ A are in r´ 1 , 1 s .
Assume also that ˆ ⇥ are non-degenerate, in the sense that the input-output mapping they realize is not identically zero.
Then, for any ✏ ° 0 and q P N , there exists  p ✏ , q q ° 0 which admits the following: whenever GF leads the student parameters ⇥ to within (Frobenius) distance  p ✏ , q q from a balanced point ⇥ ˚ satisfying L p ⇥ ˚ q “ 0 , the student ✏ -extrapolates with horizon q .
Proof sketch (for complete proof see Appendix D.3).
In order to show ✏ -extrapolation with horizon q we need to show that @
✏ ° 0 , q P N , D  p ✏ , q q ° 0 such that for any ⇥ 1 satisfying || ⇥ 1 ´ ⇥
˚ || †  , we have for n “ 0 , . . .
, q that
[Equation]
[Equation]
First, by Theorem 5 ⇥ ˚ is an extrapolating solution.
Following the interpretation of discrete random variables of Lemma 7 proof, we can bound the Wasserstein-1 distance between the extrapolating distribution deﬁned by ⇥ ˚ and the approximating distribution deﬁned by ⇥ 1 .
From Proposition 2 in
Wu & Yang (2020), we have that if | C 1 p
A 1 q n B 1 ´ C ˚ p
A ˚ q n B ˚ | §  for n P t 0 , . .
, 2 ˆ d u , then
[Equation]
The second part of the proof ties the error at step p with the W p distance (Biswas & Mackey, 2021):
[Equation]
The ﬁnal step of the proof is based on Panaretos & Zemel (2019) and bounds W p in terms of W 1 , i.e.,
[Equation]
where S is the difference between the maximal and minimal eigenvalues in the teacher and student support (due to our assumption on convergence, S is bounded).
Combining the three steps above, we have that if | C 1 p
A 1 q n B 1 ´ C ˚ p
A ˚ q n B ˚ | §  for n P t 0 , . .
, 2 ˆ d u , then @ p ° n
[Equation]
By setting  †
[Equation]
4ˆ d the RHS of Equation 4.5 satisﬁes
[Equation]
4. B ALANCEDNESS C APTURES N EAR -Z ERO I NITIALIZATION
Theorems 5 and 8 assume that GF emanates from a balanced initialization, i.e. from a point ⇥ “ p A , B , C q satisfying B “ C J .
It was shown in Cohen-Karlik et al.
that theoretical predictions derived assuming balanced initialization faithfully match experiments conducted with near-zero initialization (an initialization commonly used in practice).
Proposition 9 below theoretically supports this ﬁnding, establishing that with high probability, random near-zero initialization leads GF to arrive at an approximately balanced point, i.e. a point ⇥ “ p A , B , C q for which the difference between B and C J is negligible compared to their size.
Proposition 9.
Suppose that: (i) d ° 4 ; (ii) the teacher parameters ˆ ⇥ are balanced and are non-degenerate, in the sense that the input-output mapping they realize is not identically zero; and (iii) the student parameters are learned by applying GF to the loss L p¨q .
Let ˜ ⇥ be a random point in parameter space, with entries drawn independently from the standard normal distribution.
For ✏ ° 0 , consider the case where GF emanates from the initialization ✏ ˜ ⇥ , and denote the resulting curve by ⇥ ✏ p ⌧ q “ p A ✏ p ⌧ q , B ✏ p ⌧ q , C ✏ p ⌧ qq , with ⌧ • 0 .
Then, w.p. at least 0 .
75 , for every {+✏ ° 
+} there exists ⌧ ✏ • 0 such that:
[Equation]
Proof sketch (for complete proof see Appendix D.4).
The idea behind the proof is as follows.
Assume ✏ is sufﬁciently small.
Then, when the entries of ⇥ “ p A , B , C q are on the order of ✏ , we have BB B L p ⇥ q « ´ 2 ˆ C ˆ B ¨ C J and BB C L p ⇥ q « ´ 2 ˆ C ˆ B ¨ B J .
This implies that during the ﬁrst part of the curve ⇥ ✏ p ⌧ q it holds that dd ⌧ p B ✏ p ⌧ q ´ C J ✏ p ⌧ qq « ´ 2 ˆ C ˆ B ¨ p B ✏ p ⌧ q ´ C J ✏ p ⌧ qq and similarly dd ⌧ p B ✏ p ⌧ q ` C J ✏ p ⌧ qq « 2 ˆ C ˆ B ¨ p B ✏ p ⌧ q ` C J ✏
p ⌧ qq .
Since ˆ C ˆ B ° 0 (follows from the teacher parameters being balanced and non-degenerate), the entries of B ✏ p ⌧ q ´ C J ✏ p ⌧ q shrink exponentially fast while those of B ✏ p ⌧ q ` C J ✏ p ⌧ q grow at the same rate.
This exponential shrinkage/growth leads } B ✏ p ⌧ q ´ C J ✏ p ⌧ q}
L } B ✏ p ⌧ q ` C J ✏ p ⌧ q} to become extremely small, more so the smaller ✏ is.
E XPERIMENTS
In this section we present experiments corroborating our theoretical analysis (Section 4).
The latter establishes that, under certain conditions, a linear RNN with state space dimension d extrapolates when learning from a teacher network with state space dimension ˆ d via training sequences of length k , irrespective of how large d is compared to ˆ d and k .
A key condition underlying the result is that k is larger than 2 ˆ d .
Section 5.1 below considers the theoretically analyzed setting, and empirically evaluates extrapolation as k varies.
Its results demonstrate a phase transition, in the sense that extrapolation takes place when k ° 2 ˆ d , in compliance with theory, but fails when k falls below 2 ˆ d , in which case the theory indeed does not guarantee extrapolation.
Section 5.2 displays the same phenomenon with linear RNNs that do not adhere to some of the assumptions made by the theory (in particular the assumption of symmetric transition matrices, and those concerning balancedness).
Finally, Section 5.3 considers non-linear RNNs (speciﬁcally, Gated Recurrent Unit networks Chung et al.
(2014)), and shows that they too exhibit a phase transition in extrapolation as the training sequence length varies.
For brevity, we defer some of the details of our implementation, as well as additional experiments, to Appendix A.
5. T HEORETICALLY A NALYZED S ETTING
Our ﬁrst experiment considers the setting described in Section 3 and theoretically analyzed in Section 4.
As representative values for the state space dimensions of the teacher and (overparameterized) student, we choose ˆ d “ 5 and d “ 40 respectively (higher state space dimensions for the student, namely d “ 100 and d “ 200 , yielded qualitatively identical results).
For a given training sequence length k , the student is learned via GD applied directly to the population loss deﬁned in Equation 3.
(applying GD to the empirical loss deﬁned in Equation 3.3, with N “ 10 , 000 training examples, led to similar results).
Figure 1(a) reports the extrapolation error (quantiﬁed by the ` 8 distance between the impulse response of the learned student and that of the teacher) as a function of k .
As can be seen, extrapolation exhibits a phase transition that accords with our theory: when k ° 2 ˆ d extrapolation error is low, whereas when k falls below 2 ˆ d extrapolation error is high.
5. O THER S ETTINGS W ITH L INEAR R ECURRENT N EURAL N ETWORKS
To assess the generality of our ﬁndings, we experiment with linear RNNs in settings that do not adhere to some of the assumptions made by our theory.
Speciﬁcally, we evaluate settings in which: (i) the teacher is unbalanced, meaning ˆ B ‰ ˆ C J , and its transition matrix ˆ A is non-symmetric; (ii) the student’s transition matrix A is not restricted to be symmetric; (iii) learning is implemented by optimizing the empirical loss deﬁned in Equation 3.3 (rather than the population loss deﬁned in Equation 3.4); and (iv) optimization is based on Adam Kingma & Ba (2014) (rather than GD), emanating from standard near-zero initialization which is generally unbalanced (namely, B ‰ C J ).
Figure 1(b) reports the results of an experiment where the state space dimensions of the teacher and (overparameterized) student are ˆ d “ 10 and d “ 50 respectively (higher state space dimensions for the student, namely d “ 100 and d “ 200 , yielded qualitatively identical results), and where the teacher implements a delay line of ˆ d time steps (for details see Appendix B.2.2).
Similar results obtained with randomly generated teachers are reported in Appendix A. As can be seen, despite the fact that our theory does not apply to the evaluated settings, its conclusions still hold — extrapolation error is low when the training sequence length k is greater than 2 ˆ d , and high when k falls below 2 ˆ d .
[Figure]
accords introduced
5. N ON -L INEAR R ECURRENT N EURAL N ETWORKS
As a ﬁnal experiment, we explore implicit extrapolation with non-linear RNNs, namely GRU networks.
Speciﬁcally, we evaluate the extent to which a student GRU with state space dimension d g “extrapolates when learning from a teacher GRU with state space dimension ˆ d g “ 10 (higher state space dimensions for the student, namely d g “ 200 and d g “ 500 , yielded qualitatively identical results).
The student is learned by optimizing an empirical loss comprising training sequences of length k g , where k g is predetermined.
Optimization is based on Adam emanating from standard nearzero initialization.
Figure 2(a) reports the extrapolation error (quantiﬁed by the ` 8 distance between the response of the learned student and that of the teacher, averaged across randomly generated input [Figure] threshold which extrapolation error is high in sequences) for different choices of k g .
As can be seen, similarly to the case with linear RNNs (see
Sections 5.1 and 5.2), there exists a critical threshold for the training sequence length k g , above which extrapolation error is low and below which extrapolation error is high (note that this critical threshold is around four times the teacher’s state space dimension, whereas with linear RNNs the critical threshold was around two times the teacher’s state space dimension; theoretically explaining this difference is an interesting direction for future work).
Figure 2(b) displays the average output response over different inputs of the teacher alongside those of two students — one trained with sequences of length k g “ 30 , and the other with sequences of length k g “ 60 . [-
-]2 As expected, the impulse response of each student tracks that of the teacher for the ﬁrst k g time steps (where k g is student-dependent).
However, while the student for which k g “ 30 fails to track the teacher beyond k g time steps, the student for which k g “ 60 succeeds, thereby exemplifying implicit extrapolation.
C ONCLUSION
This paper studies the question of extrapolation in RNNs, and more speciﬁcally, of whether a student
RNN trained on data generated by a teacher RNN can capture the behavior of the teacher over sequences longer than those seen in training.
We focus on overparameterized students that can perfectly ﬁt training sequences while producing a wide range of behaviors over longer sequences.
Such a student will fail to extrapolate, unless the teacher possesses a certain structure, and the learning algorithm is biased towards solutions adhering to that structure.
We show — theoretically for linear
RNNs and empirically for both linear and non-linear RNNs — that such implicit extrapolation takes place when the teacher has a low dimensional state space and the learning algorithm is GD.
Existing studies of implicit extrapolation in (linear) RNNs (Emami et al., 2021; Cohen-Karlik et al., 2022) suggest that GD is biased towards solutions with short-term memory.
While low dimensional state space and short-term memory may coincide in some cases, in general they do not, and a solution with low dimensional state space may entail long-term memory.
Our theory and experiments show that in settings where low dimensional state space and short-term memory contradict each other, the implicit extrapolation chooses the former over the latter.
An important direction for future work is extending our theory to non-linear RNNs.
We believe it is possible, in the same way that theories for linear (feed-forward) NNs were extended to account for non-linear NNs (see, e.g., Razin et al.
; Lyu & Li (2019)).
An additional direction to explore is the applicability of our results to the recently introduced S4 model (Gu et al., 2022).
A F URTHER E XPERIMENTS
In this section we provide additional experiments that are not included in the main manuscript due to space constraints.
[Figure]
A. B ALANCED T EACHER
In Section 5.1, we have experimented with our proposed theoretical setup.
In this section we provide additional ﬁgures and experiments.
A.1. U NBALANCED S TUDENT
In this experiment we use the same balanced teacher with ˆ d “ 5 as done in Section 5.1.
Instead of the diagonal student with balanced initialization, we use a general (non-diagonal) student with weights sampled from a Gaussian with scale 10 ´ 5 and d “ 40 .
Results are depicted in Figure 3(a).
A similar phase transition phenomenon to the one in Figure 1 is found also here.
A.1. E FFECT OF THE I NITIALIZATION S CALE
Proposition 9 provides theoretical support for the fact that under near-zero initialization, the learned
RNN tends to balancedness, which according to theorems 5 and 8 guarantees extrapolation.
Below we empirically explore the impact of varying the initialization scale.
We use the same setting as in
Section A.1.1, and repeat the experiment with different initialization scales for the students’ weights.
As can be seen in Figure 4, the extrapolation deteriorates for larger initialization scale, in the sense that it requires longer training sequences for getting good extrapolation error.
This suggests that the condition of small initialization required by our theory is not an artifact of our proof technique, but rather a necessary condition for extrapolation to occur.
A. U NBALANCED T EACHER
In Section 5.2, we have tested the extrapolation with respect to a speciﬁc unbalanced teacher and have observed a similar phase transition as predicted by the theory of Section 4 and empirical evaluation of
Section 5.1.
Here we show that the phase transition is not limited to the speciﬁc teacher discussed by testing with respect to a randomly generated unbalanced (non-diagonal) teacher (see Section B.2.2).
The teacher is set to ˆ d “ 5 and student to d “ 50 .
Results are presented in Figure 3 (b).
Here too we can observe the phase transition phenomena.
A. I MPULSE R ESPONSE F IGURES
In Section 5 we have presented the extrapolation performance in different settings.
In order to better convey the meaning of extrapolating vs non-extrapolating solutions we present here ﬁgures of the impulse response of different models.
We start with the impulse response corresponding to the experiment described in Section 5.1.
The ﬁgure depicts the balanced teacher with ˆ d “ 5 and two selected students (with d “ 40 ), one trained with k “ 10 and the other with k “ 20 .
Section B be seen, track the teacher up to k used in training, for k “ 10 there is extrapolation for larger k whereas k “ 20 tracks the teacher well beyond the sequence length used in training.
We can see that the student trained with k “ 10 tracks the teacher several steps beyond the 10 th time step and then decays to zero.
For k “ 20 we can see near perfect extrapolation for the horizon evaluated.
Next we turn to Section 5.2 and depict the average impulse responses of the “delay teacher” and the students trained with respect to the mentioned teacher.
Since the teacher here has ˆ d “ 10 , a model trained with k “ 8 is trained with respect to the zero impulse response (see Section B.2.2 for details on delay teacher), and as expected results with the ‘zero’ solution.
we can see that for k “ 18 the student diverges from the teacher shortly after the 18 th time step.
For k “ 20 we can see near perfect extrapolation up to the horizon considered.
B I MPLEMENTATION D ETAILS
All the experiments are implemented using PyTorch.
B. O PTIMIZATION
In Section 5.1 we optimize the population loss, which entails minimizing Equation 3.4 with respect to the parameters of the learned model.
We use 15K optimization steps with Adam optimizer and a learning rate of 10 ´ 3 .
In this experiment, the results were not sensitive to the initialization scale of the (balanced) student.
In Section 5.2 and Section 5.3 in the experiments that involve minimizing the empirical loss, we use 50K optimization steps with early stopping (most experiments required less than 10K steps).
The batch size is set to 100 , data is sampled from a Gaussian with zero mean and scale of 1 .
Experiments were not sensitive to most hyper-parameters other than learning rate and initialization scale.
The examination of the effect of initialization scale presented in Section A.1.
is done with learning rate scheduler torch.optim.lr_scheduler.
MultiStepLR using milestones at r 5000 , 10000 , 15000 , 30000 s and a decaying factor of  “ 0 . 1 .
B. T EACHER G ENERATION
One of the main challenges in empirically evaluating extrapolation is that randomly sampling weights from a Gaussian distribution may result with an RNN of lower effective rank (i.e. the resulting RNN may be accurately approximated with another RNN with a smaller hidden dimension).
We will now describe the teacher generation scheme for the different experiments.
B.2. B ALANCED T EACHER G ENERATION
A balanced teacher consists of d entries corresponding to the diagonal teacher and d entries representing ˆ B “ ˆ C J .
In order to avoid cases of rapid decay in the impulse response on the one hand, and exponential growth on the other, we set the eigenvalues to distribute uniformly between 0 . 6 and 1 . 05 .
The values of ˆ B and ˆ C are randomly sampled from a Gaussian around 0.5 and scale 1 and then normalized such that ˆ C ˆ B “ 1 .
B.2. U NBALANCED T EACHER G ENERATION
In this experiment, the teacher has a general (non-symmetrid) matrix ˆ A and ˆ B ‰ ˆ C J .
We set the weights as described next.
Delay Teacher A ‘delay’ teacher has an impulse response of 1 at time step i “ ˆ d ´ 1 , that is, the teacher has an impulse response of p 0 , . . .
q .
In order to generate the mentioned impulse response we set the weights as follows,
[Equation] [Equation]
˚˚˝
[Equation]
˛ ‹‹‚ , and C J “ ¨ ˚˚˝ ...
˛ ‹‹‚ .
(B.1)
Note that B , C above are set to extract the last entry of the ﬁrst row of A i and A is a Nilpotent shift matrix.
It is straightforward to verify that CA i B “ 1 for i “ ˆ d ´ 1 and 0 otherwise.
Random Unbalanced Teacher The second unbalanced teacher is randomly generated.
In order to avoid the caveats mentioned in Section A.1, we randomly sample the diagonal (from a Gaussian with zero mean and scale 0 . 1 ) and super diagonal (from a Gaussian with mean 0 .
7 and scale 0 . 1 ) of A .
We set B, C as in equation B.1.
The structure of A ensures similar properties to that of the delayed teacher, speciﬁcally, that the ﬁrst entries of the impulse response is zero and the teacher is ‘revealed’ only after ˆ d time steps.
B.2. N ON -L INEAR T EACHER G ENERATION
As opposed to the linear teacher discussed in previous sections, when the teacher is a Gated Recurrent
Units (GRU), it is unclear how to generate a non-trivial teacher.
When randomly generating a teacher
GRU the result is either a trivial model that quickly decays to zero or a teacher with an exploding impulse response (depending on the scale of the initialization).
In order to produce a teacher with interesting extrapolation behaviour, we initialize a model with an initialization scale of 10 ´ 6 and train for 1000 step the model to mimic an arbitrarily chosen impulse response.
The result of the mentioned procedure is a teacher GRU with non-trivial behaviour.
Figure 2(b) shows that we get with this non-trivial teacher the phase transition phenomena as described in Section 5.3.
B. E XTRAPOLATION E RROR
The concept of extrapolation is very intuitive, and yet it does not admit any standard error measure.
A proper extrapolation error measure should: (a) capture ﬁne differences between two models with good extrapolation behaviour; and on the other hand, (b) be insensitive to the scale in which two non-extrapolating model explode.
A natural approach which we take here is to report the ` 8 norm difference on the tail of the impulse response.
A model is considered non-extrapolating if the extrapolation error is worse than the extrapolation error of a trivial solution which has an impulse response of zeros.
C A CCUMULATING L OSS
In the main paper the analysis is performed for the loss function deﬁned in Section 3, which corresponds to a regression problem over sequences.
Another important and common loss function is an accumulating loss deﬁned over the full output sequence.
Speciﬁcally, the empirical loss of
Equation 3.3 is replaced with,
[Equation]
In this section we discuss the adaptations required to accommodate our theory with the loss deﬁned in Equation C.1.
C. P OPULATION L OSS
A similar derivation of the population loss described in Appendix D.1.1 can be applied to Equation C.1.
The difference is that an additional summation is introduced and is preserved throughout the analysis to result with,
[Equation]
The loss above can be viewed as a different weighting of the original population loss, i.e. Equation C. can be written as
[Equation]
It is clear that the minimizers of Equation C.2 are the same minimizers of Equation 3.4 (i.e. CA i B “ w i for i “ 0 , . . .
, k ´ 1 ).
Thus Lemma 7 holds with no additional modiﬁcations.
C. A PPROXIMATE E XTRAPOLATION
For Theorem 8 the analysis in the proof makes use of the fact that the difference in the moments deﬁned by the student and teacher is bounded by ?
.
The same is true for the case of the weighted loss, speciﬁcally, if the loss §  , then for all i “ 0 , . . .
, k ´ 1 , a p k
´ i qp
CA i B
´
w i q § ?
.
Since k ´ i • 1 for i “ 0 , . . .
k ´ 1 we have p
CA i B ´
w i q § a p
k
´
i qp CA i B
´
w i q § ?
and the remainder of the proof is the same.
C. I MPLICIT B IAS FOR B ALANCEDNESS
The proof of the implicit bias for balancedness involve the gradients of the population loss deﬁned in Equation 3.4.
For the weighted population loss the gradients differ, but the symmetries are all preserved (the gradient computation boils down to adding an external summation to the terms computed in Section D.1.4.
The same steps described in Section D.4 apply for the weighted loss.
D D EFERRED P ROOFS
Here we provide complete proofs for the results in the paper.
D. A UXILARY P ROOFS
In this section we provide missing proofs from the main paper and additional lemmas to be used in the main proofs.
D.1. P OPULATION L OSS
Lemma D.1 ( Proof of Equation 3.4 ) .
Assume x „ D such that E x „ D r x s “ 0 , E x „ D r xx J s “ I k P
R k,k , where I k is the identity matrix.
y is given by y “ { RNN p x q where { RNN p¨q denotes the output of a teacher RNN, ˆ ⇥ “ p ˆ A , ˆ B , ˆ C q .
Denote w i “ ˆ C ˆ A
i ˆ B , the loss for the student RNN satisﬁes:
[Equation]
Proof of Lemma D.1 .
The population loss for training with sequences of length k is
[Equation]
Reversing the order of summation, expanding the terms,
[Equation]
where the transition from the second to third rows is by our assumption that E x „ D r x i
x
j s “ 1 i “ j .
Therefore we have,
[Equation]
concluding the proof.
D.1. P ERFECT G ENERALIZATION AND F AILED E XTRAPOLATION
Proposition D.2 (Proposition 3 in main paper) .
Assume d ° k , and let ✏ • 0 and q P t k ` 1 , k ` 2 , . . .
u .
Then, for any teacher parameters ˆ ⇥ , there exist student parameters ⇥ with which the population loss in Equation 3.4 equals zero [Equation] and yet the student does not ✏ -extrapolate with horizon q .
Proof.
Consider a student, ⇥ , such that A is symmetric (and therefore has an orthogonal eigendecomposition).
Denote A “ U ⇤ U J .
The impulse response at time step i can be expressed as
CA i B “ CU ⇤ i
U J B .
The latter can be written compactly in matrix form as V g where V is the
Vandermonde matrix with diag p ⇤ q as its values,
[Equation]
and g is deﬁned as g ” p CU q J ˝
U J B .
3 A known result on square Vandermonde matrices is that they are invertible if and only if  i ‰
j , @ i ‰ j .
Given a ﬁxed set of distinct values p  1 , . . .
,  d q and an arbitrary impulse response r P R d , in order for the student to generate the impulse response r (i.e. V g “ r ), one can set the coefﬁcient vector, g “ V ´ 1 r and end up with a symmetric student with r as its impulse response of length d .
Consider a teacher RNN, ˆ ⇥ “ p A , B , C q , we can set and the ﬁrst k entries of r to r i “ ˆ C ˆ A
i ´ 1 ˆ B , @
i “ t 1 , . .
. , k u .
We are therefore left with d ´ k degrees of freedom which yields many different students that correspond to the ﬁrst k entries of the teacher while ﬁtting arbitrary values beyond the k .
D.1. E QUIVALENCE B ETWEEN B ALANCED RNN S WITH S YMMETRIC AND D IAGONAL T RANSITION M ATRICES
Lemma D.3.
A balanced RNN, ⇥ “ p A , B , C q , with a symmetric transition matrix (i.e. B “ C
J and A “ A J ) has an equivalent (i.e. generating the same impulse response) RNN, ⇥ 1 “ p A 1 , B 1 , C 1 q , which is balanced and its transition matrix is diagonal.
Lemma D.3 allows alternating between systems with symmetric and diagonal matrices.
This is useful to simplify the analysis in Section 4.
Proof of Lemma D.3 .
Any symmetric matrix admits an orthogonal eigendecomposition with real (non-imaginary) eigenvalues.
Denote A “ U ⇤ U J .
We can deﬁne
[Equation]
and
[Equation]
The i th index of the impulse response is given by
[Equation]
concluding that ⇥ and ⇥ 1 have the same impulse response of any length.
D.1. G RADIENT D ERIVATION
For completeness and Section D.4, we compute the gradients for the general setting.
Lemma D.4.
Given the population loss
[Equation]
Denote r ` i “ CA i B ´ w i , the derivatives of the loss with respect to B , and C satisfy:
[Equation]
Proof of Lemma D.4 .
Here, we will compute the gradient of the population loss.
Note that for j • 0 , the derivative of CA j B w.r.t to B is given by
[Equation]
Similarly, the derivative of CA j B w.r.t to C is given by
[Equation]
Using these derivatives, we can calculate the derivative of the population loss, (assigning w i “ ˆ B ˆ A i ˆ C ), [Equation] [Equation]
D.
[Equation]
Denoting r ` i “ CA i B ´ w i , and noting that w i is constant (depends on ˆ ⇥ ), we have for X P t B , C u :
[Equation]
Plugging in Equation D.11 and Equation D.12, we have:
[Equation]
1. L EMMA 6 (C ONSERVATION OF B ALANCEDNESS )
Lemma D.5.
[Lemma 6 in main paper] When optimizing equation 3.4 with GF emenating from a balanced initialization ⇥ p 0 q , the parameters ⇥ p ⌧ q are balanced for all ⌧ P R ` .
We prove the above result by ﬁrst showing it for gradient descent and then translating the result to gradient ﬂow.
The gradient descent result is stated below, and generalizes a result that was shown in
Cohen-Karlik et al.
for the memoryless case.
Lemma D.6.
When optimizing equation 3.4 with GD with balanced initial conditions, then @
t P N , ⇥ has a balanced weight conﬁguration, i.e. B t “ C J t .
Proof of Lemma D.6 .
We prove by induction.
By our assumption, the condition holds for t “ 0 .
Assume B t “ C J t , our goal is to show the conditions hold for p B t ` 1 , C t ` 1 q .
In order to show that
B t ` 1 “ C J t ` 1 , we only need to show that B L B B t “ ´
B L B C t ¯ J .
Writing the gradients (Lemma D.4), we have
[Equation]
(D.17) where the inequality follows from the induction assumption and the symmetric structure of A t .
To conclude, the gradients at time t are the same and B t “ C J t by the induction assumption, arriving at
[Equation]
The proof of Lemma 6 follows from Lemma D.6 and the fact that for sufﬁciently small step size GD approximates GF with arbitrary precision (see Theorem 3 in Elkabetz & Cohen, 2021).
D
. C ONSERVATION OF DIFFERENCE OF NORMS
Appendix D.1.5 shows that if weights are initialized to be balanced, this property is conserved throughout optimization.
Here we show under standard initialization schemes, the difference between the norms of B and C is also conserved.
Lemma D.7.
The difference between the norms of B , C is conserved throughout GF, i.e.,
[Equation]
Proof of Lemma D.7 .
We wish to prove that the difference between the norms is conserved over time.
Consider the following expression:
[Equation]
With this notation, we just need to prove that 9 ↵ “ 0 .
The derivative of B , C with respect to time is given by,
[Equation]
[Equation]
Using the interchangeability of derivative and transpose, we have:
[Equation]
Plugging equation D.21 and equation D.22, we get
[Equation]
establishing that ddt
[Equation]
Lemma D.8.
[Lemma 7 in main paper] Suppose that d ° k ° 2 ˆ d , the teacher is balanced, and that the student parameters ⇥ are balanced and satisfy L p ⇥ q “ 0 .
Then ⇥ extrapolates.
Proof of Lemma D.8 .
By Lemma D.3, a balanced RNN with symmetric transition matrix has an equivalent (generating the same impulse response) balanced RNN with a diagonal transition matrix.
We will continue under the assumption of diagonal transition matrices.
Without loss of generality we assume ˆ C ˆ B “ 1 .
Otherwise, the problem can be rescaled by ˆ C ˆ B , which is equivalent to rescaling the initial conditions, and providing no additional information.
From the balanced assumption, we have ˆ C J
“ ˆ B .
Denote ˆ p “ ˆ C J d ˆ B “ ˆ B d ˆ B , and we get ˆ p i [-ě 
-]• 0 and ∞
i ˆ p
i “ 1 , and therefore ˆ p may be interpreted as a distribution over a random variable with ˆ d possible values.
We shall assume that these values are ˆ A 1 , . . .
, ˆ A ˆ d , and denote the corresponding random variable by Z .
Furthermore, we can also interpret elements of the impulse response of ˆ ⇥ as moments of this distribution.
Let us write the n th element of the impulse response as:
[Equation]
where E p r
Z s is the expected value of a random variable Z under the distribution p .
In the same way, we can deﬁne for the learned model ⇥ , a distribution p i “ C i B i , and write the learned impulse response as:
[Equation]
This view provides us with a moment matching interpretation of the learning problem.
Namely, the fact that ⇥ matches the ﬁrst k elements of the teacher impulse response, is the same as saying they agree on the ﬁrst k ´ 1 moments E p r Z j s for j P t 1 , . .
k ´ 1 u .
The question of extrapolation is whether equality in the ﬁrst k ´ 1 moments implies an equality in all other moments.
In (Cohen & Yeredor, 2011, Theorem 1) and in (Wu & Yang, 2020, Lemma 4) it is shown that the ﬁrst 2 ˆ d moments of a discrete random variable taking at most ˆ d different values uniquely deﬁne this random variable.
Therefore, any other discrete random variable identifying with the teacher on 2 ˆ d moments must be the same random variable and therefore identiﬁes on higher moments as well.
Since we assumed k ° 2 ˆ d , this result immediately implies that equality in the ﬁrst k ´ 1 moments implies equality in all other moments. ∞ i C i B i =
∞ i ˆ C i ˆ B i =1. For the case ˆ C ˆ B “ 0 , from our assumption that the teacher is balanced, we have that the condition is met only if ˆ C i “ ˆ B i “ 0 for i “ 1 , . . .
, ˆ d .
Such a teacher has an impulse response of zeros, for k • 1 , a student minimizing the loss must also satisfy CB “ 0 and therefore has the zeros as its impulse response thus extrapolating with respect to the said teacher.
D. T HEOREM 8 (A PPROXIMATE E XTRAPOLATION )
This section is devoted to the proof of Theorem 8 which ties the approximation error of optimization to that of extrapolation.
Lemma D.9 below shows that if the parameters are within close proximity to an extrapolating solution, the generalization error (i.e. loss values) are small.
Lemma D.9.
Assume } X ´ X ˚ } 8 †  for X P t A , B , C u , and ⇥ ˚ “ p A ˚ , B ˚ , C ˚ q is an extrapolating solution (with L p ⇥ ˚ q “ 0 q .
Then, there exists a constant depending on k and ⇥ ˚ , c p k, ⇥ ˚ q , such that L p ⇥ q § c  2 .
Proof of Lemma D.9 .
First we note that ⇥ ˚ attains zero loss and is an extrapolating with respect to the teacher ˆ ⇥ , speciﬁcally,
[Equation]
From the above, the loss can be written with respect to ⇥ ˚ instead of ˆ ⇥ , i.e.
[Equation]
Denote by l P R k the residuals vector with respect to ⇥ ˚ , i.e. l i ” [Equation]
i B ´
[Equation]
´ C ˚ p
A ˚ q i B ˚ .
With these notations, the loss amounts to [Equation]
From equivalence of norms, we have that for v P R n , } v } 8 § } v } 2 § ? n } v } 8 .
The latter means that bounding } v } 8 will provide the desired bound on the loss with a factor of ?
k .
Recall that ⇥ ˚ is ﬁxed, speciﬁcally, } X ˚ } is bounded by some constant m X ˚ for X P t A , B , C u .
Denote m ” max [Equation] m A ˚ , m B ˚ , m C ˚ , 1 u , the i th entry of l can be bounded by
[Equation]
[Equation]
[Equation]
›› (D
[Equation]
From Cauchy-Schwarz inequality, the RHS satisﬁes  } A i ´ 1 B } §  } A } i ´ 1 } B } §  p m `  q i .
The last transition uses,
[Equation]
(D
[Equation]
33) for X P t A , B u .
In a similar manner, the LHS can be bounded by
[Equation]
D [Equation] 35) again, we have used the triangle inequality and Cauchy-Schwarz inequality.
In order to deal with } A i ´
´
p
A ˚ q i ´ 1 } , we make use of the equality x n ´ y n “ p x ´ y qp x {+n ´ 1 ` x n ´ 2 y ` x n ´ 
+}
y 2 ` ¨ ¨ ¨ ` xy n ´ 2 y n ´ 1 .
Assigning A and A ˚ , we have
[Equation]
where the last inequality uses,
[Equation]
Plugging all terms back together,
[Equation]
The above is maximized for i “ k , to result with } l } 8 “ O p  m k q .
Plugging back to the bound on the loss, we have
[Equation]
When k is small which is the case of interest, L p A , B , C q “ O p  2 q .
Theorem D.10.
[Theorem 8 in main paper] Assume the conditions of Theorem [Equation] and that the teacher parameters ˆ ⇥ are stable, i.e. the singular values of ˆ A are in r´ 1 , 1 s .
Assume also that ˆ ⇥ are non-degenerate, in the sense that the input-output mapping they realize is not identically zero.
Then, for any ✏ ° 0 and q P N , there exists  p ✏ , q q ° 0 which admits the following: whenever GF leads the student parameters ⇥ to within (Frobenius) distance  p ✏ , q q from a balanced point ⇥ ˚ satisfying
L p ⇥ ˚ q “ 0 , the student ✏ -extrapolates with horizon q .
Proof of Theorem D.10 .
Given a desired approximation error ✏ ° 0 , and horizon q P t k ` 1 , k ` 2 , . . .
u , our goal is to ﬁnd  p ✏ , q q such that if } X ´ X ˚ } F †  for X P t A , B , C u , then | CA j B ´ ˆ
C ˆ A j ˆ B | † ✏ .
First we note that by Theorem 5, ⇥ ˚ is an extrapolating solution with respect to ˆ ⇥ (i.e. ⇥ ˚ and ˆ ⇥ generate identical impulse responses), therefore we can measure extrapolation with respect to ⇥ ˚ .
In a similar manner to Lemma D.9, we can deﬁne the residual vector l such that its i th entry is deﬁned as
[Equation]
Note that ⇥ is initialized balanced and by Lemma 6, this property is preserved throughout GF.
By
Lemma D.3 both ⇥ and ⇥ ˚ have equivalent balanced systems ¯ ⇥ and ¯ ⇥ ˚ for which the transition matrices are diagonal.
For the sake of readability, we overload notations and carry on under the assumption that A and A ˚ are diagonal with entries t
A 1 , . . .
, A d u and t
A ˚ 1 , . . .
, A ˚ d u on their respective diagonals.
Another observation is that a balanced RNN must satisfy that its ﬁrst element in the impulse response moment is non-negative ( CB “ ∞ c i b i “ ∞ b 2 i • 0 ).
Without loss of generality, we can assume the teacher RNN ˆ ⇥ satisﬁes ˆ C ˆ B “ 1 .
Suppose this is not the case, i.e. ˆ C ˆ B “ {+⌘ ° 
+} (recall we assume the teacher is non-degenerate).
One can re-scale ˆ C and ˆ B by a factor of ⌘ ´ 0 .
5 and result with ˆ C ˆ A
i ´ 1 ˆ B ⌘ on the i th entry of its impulse response.
The re-scaled system satisﬁes ˆ C ˆ B ⌘ “ 1 and provides no extra information.
With the claims above, we may interpret ˆ ⇥ as a discrete random variable, i.e. ˆ p “ p ˆ C 1 ˆ B 1 , . .
, ˆ C ˆ d ˆ B ˆ d q deﬁnes a probability vector, 6 the random variable is assigned the value
ˆ A i with probability
ˆ p i .
Denote the random variable associated with ˆ ⇥ by ˆ z. Recall that ⇥ ˚ is balanced and attains zero loss, hence it also admits the random variable interpretation (denoted as z ˚ ).
Under the probabilistic perspective, the impulse response of the systems correspond to the moment of the random variables, i.e.
[Equation]
We can write the loss in terms of the deﬁned residual vector, L p ⇥ q “ } l } 2 .
From Lemma D.9, we have L p ⇥ q “ O p  2 q .
Speciﬁcally, } l } § c  2 for some constant c P R .
Recall that for v P R n , } v } 8 § } v } 2 , applying to our case, we have } l } 8 § } l } § ? c  .
Speciﬁcally, | l 0 | § } l } 8 § ?
c  and since we assume C ˚ B ˚ “ 1 , we have
[Equation]
which leads to 1 ´ ?
c  § CB [-ď 
-]§ 1 ` ? c  .
Denote CB “ 1 `  where  P r´?
c  , ?
c  s .
A re-scaled student introduces a factor  into the approximation error, that is
[Equation]
Recall that ⇥ ˚ is ﬁxed, speciﬁcally, } X ˚ } is bounded by some constant m X ˚ for X P t A , B , C u .
Denote m ” max t m A ˚ , m B ˚ , m C ˚ , 1 u .
Combining the latter with Cauchy-Schwarz inequality,
[Equation]
where the last inequality uses 0 .
5 † 1 `  for ?
c  † 0 .
Per our claim above, we can assume with out loss of generality that CB “ 1 .
A balanced student can thus be interpreted as a discrete random variable with probability vector p taking on values t A 1 , . . .
, A d u with probability p i .
Recall that under the assumption of a stable teacher, we have that for all i “ 1 , . . .
, ˆ d, | ˆ A i | § 1 .
Denote the joint support of A and ˆ A by S ” t A i u di “î t ˆ A j u ˆ dj “ 1 , and deﬁne s “ max t S u ´ min t S u .
By our assumption on convergence and stability of the teacher, we have s [Equation]
From Proposition 2 in Wu & Yang (2020), we have that if | CA i B
´ ˆ C ˆ A
i ˆ B | §  for i P t 1 , . . .
, 2 ˆ d u , then
[Equation]
By (Panaretos & Zemel, 2019, Section 2.3), we can tie the Wasserstein distance to higher order q th
Wasserstien distance.
For the s deﬁned above, we have that for any q P N :
[Equation]
Denote by m q p p q the q th moment of p , from (Biswas & Mackey, 2021, Equation 1), we have:
[Equation]
Recall that the difference in the i th moments of the student and teacher correspond to the difference in the impulse responses.
So we have that if } X ´ X ˚ } F †  for X P t A , B , C u then @
q P N ,
[Equation]
D. P ROPOSITION 9 (I MPLICIT B IAS FOR B ALANCEDNESS )
The proof of Proposition 9 consists of several steps.
First, we bound with high probability the norms of B and C at initialization (Lemma D.11).
We then derive bounds on the differential equations of ddt p B p t q ` C J p t qq and ddt p B p t q ´ C J p t qq (Lemma D.15).
We show that when the initialization scale tends to zero, the ratio between the differential equations tends to zero.
(Lemma D.14).
Before we turn to prove Proposition 9, we ﬁrst need to bound the initial values for a vector v P R n initialized with N p 0 , ✏ 2 n q .
Lemma D.11.
Assume a vector v P R n with N p 0 , ✏ 2 n q per coordinate.
Then:
[Equation]
Proof of Lemma D.11 .
The proof of D.11 uses known results on the Chi-square distribution Laurent & Massart (2000), applied to our speciﬁc setting to achieve the desired bounds.
We will begin by changing variables, ˜ v i “ v i ¨ ?
n ✏ .
The entries ˜ v i , are standard Gaussian variables.
The squared norm of ˜ v distributes according to the  -squared distribution.
By (Laurent & Massart, 2000, Lemma 1), in our case (plugging in x “ 9 n { 64 ), the following inequalities hold: [Equation]
[Equation]
We simplify these inequalities.
From the ﬁrst bound:
[Equation]
Changing variables back to v :
[Equation]
Similarly, for the second bound:
[Equation]
Taking the complementary probability, we have the desired result of
[Equation]
Proposition D.12.
[Proposition 9 in main paper] Suppose that: (i) d ° 4 ; (ii) the teacher parameters ˆ ⇥ are balanced and are non-degenerate, in the sense that the input-output mapping they realize is not identically zero; and (iii) the student parameters are learned by applying GF to the loss L p¨q .
Let ˜ ⇥ be a random point in parameter space, with entries drawn independently from the standard normal distribution.
For ✏ ° 0 , consider the case where GF emanates from the initialization ✏ ˜ ⇥ , and denote the resulting curve by ⇥ ✏ p ⌧ q “ p A ✏ p ⌧ q , B ✏ p ⌧ q , C ✏ p ⌧ qq , with ⌧ • 0 .
Then, w.p. at least 0 .
75 , for every {+✏ ° 
+} there exists ⌧ ✏ • 0 such that:
[Equation]
(D.61)
The consequence of Proposition 9 is that as ✏ converges to zero, B and C converge towards each other.
For convenience, we refer to the mentioned initialization scheme (where every coordinate in a vector is initialized as N
[Equation]
) as ✏ -normal initialization .
In order to prove the proposition we deﬁne a few relevant terms,
[Equation]
We will in fact prove the stronger, following lemma, for any matrix A , not necessarily symmetric.
Lemma D.13.
Assume w 0 ° 0 , and A , B , C are ✏ -normally initialized.
Then D t such that
[Equation]
The proof of Lemma D.13 follows three steps.
(I)
We show that there is a time in the optimization in which the norms of all parameters are bounded (Lemma D.14); (II) we derive upper (and lower) bounds for the differential equations describing the evolvement of Y, W and A .
Our approximations are limited to the initial phase of training.
Concretely, we show that for 0 § t § 12 w 0 ln
[Equation]
, all norms are bounded.
Thus, it is possible to obtain meaningful bounds on the ODEs of Y and W while
A remains in the magnitude of initialization.
(III) Using the relevant bounds, we show that
Lemma D.14.
Assume ✏ -initialization and bounded teacher.
Then, for all 0 § t § 12 w 0 ln ` 1 ✏
[Equation]
.˘
[Equation]
there exist M 1 , M 2 such that:
[Equation]
and
[Equation]
To prove this, we note that at initialization, A , B and C satisfy these bounds.
From continuity, there exists a maximal time for which they are satisﬁed.
We bound the rate of their growth, and thus show that for all t as described, we are within this region.
Lemma D.15.
Assume w 0 ° 0 and assume A , B , C are ✏ -normally initialized, we have the following bounds for all 0 § t § 12 w 0 ln
[Equation]
and
[Equation]
Lemma D.15 shows that the growth rate of W p t q J W p t q and the decay rate of Y p t q J Y p t q both depend on the sign of w 0 .
In our analysis we assume w° 0 for simplicity of presentation, the same analysis applies for w 0 † 0 with opposite roles for Y and W .
The proof of Lemma D.15 follows from writing the leading terms of the ODE and bounding the remaining terms by their upper bounds in the time considered.
Using these lemmas, we proceed to prove Lemma D.13.
Proof of Lemma D.13 .
Consider the dynamics at time 0 § t § 12 w 0 ln ` 1 ✏ 0 .˘ .
By Lemma D.15, we have
[Equation]
[Equation]
and
[Equation]
We can calculate the limit
[Equation]
From Lemma D.14, } A p t q} F § M 2 ✏ , so we can calculate the limit,
[Equation]
which concludes the proof.
Proof of Lemma D.14 .
At [Equation] “ 0 , by Lemma [Equation]
[Equation]
where we bound M [Equation]
If there exists t s.t. }
A p t q} F ° M 2 ✏ (otherwise, we get trivially the desired bound for } A p t q} ), from continuity and the intermediate value theorem, there exists a maximal t A s.t.
for§ t † t 1 , } A p t q} F † ✏ , } A p t 1 q} F “ M 2 ✏ .
By similar argument for B, C, we can deﬁne t B , t C .
By taking the minimal value between them, we get T ˚ s.t.[-@ 
-]
[Equation]
We start by proving the lemma under the assumption that 0 § t † T ˚ .
Later, we will show that
T ˚ ° T .
Let us recall the derivative of B :
[Equation]
Bounding it by the Cauchy-Schwartz inequality, we have
[Equation]
Denote M “ max i p| w i |q
` M 21 ✏ 1 . 5 , then we have
[Equation]
Recall the conservation law from Lemma D.7 for the norms of B and C :
[Equation]
[Equation]
By Lemma D.11, we have
[Equation]
Therefore, we get
[Equation]
[Equation]
Note also that by assuming M 2 ✏ † 1 , we have
[Equation]
Plugging the various values into equation D.73, we have:
[Equation]
(D.81)
Denoting R “ } B } 2 “ B J B , then 9 R “ 9 B J B ` B J 9 B “ 2 B J 9 B (D.82)
Taking absolute value and then plugging equation D.80 results with | [Equation] [Equation] p} B
[Equation]
4 ✏ q .
(D.83)
Using the deﬁnition of R , we get that
[Equation]
a
[Equation]
(D.84)
Notice that for all t s.t. } B } “ ?
R † 4 ✏ † M 1 ✏ 0 .
75 , we are done.
For any t 2 s.t.
a
R p t 2 q ° 4 ✏ , from continuity and the intermediate value theorem, there exist a maximal t 1 s.t.[-
-]
[Equation] t
[Equation]
t
[Equation]
, a
R p t
[Equation]
q “ 4 ✏ .
For all t† t † t 2 , a
[Equation]
✏ . Thus,
[Equation]
D.85)
Integrating by t and substituting, we get that for all t 1 † t int † t 2 :
[Equation]
Denoting M 21 “ 16 ¨ exp p 4 kM q , we have:
[Equation]
To summarize, we have shown that for all [-
 ă-]0 † t † T , there exists M 1 s.t } B p t q} † M 1 ✏ 0 .
Using an identical proof with the roles of C J , B reversed, we have | C p t q| † M 1 ✏ 0 .
As we have deﬁned t B , t C as the maximal points s.t.
this property occurs for T † t B , t C .
Next, we will bound } A } F .
Recall that } A } 2 F “ Tr p
A J A q , and
[Equation]
[Equation]
where we have used the linearity of trace and its invariance to transpose.
Recall that
[Equation]
D
[Equation]
[Equation]
Multiplying it from the left by A J and then taking trace provides us with
[Equation]
[Equation]
Taking a transpose and then using the cyclic property of trace and, we have for the terms in the sum:
[Equation]
Thus
[Equation]
we end up with
[Equation]
We use it to bound Tr p
A J 9 A q :
[Equation]
Using the Cauchy-Schwartz inequality and then plugging M ° | r ` i | and the bounds found for [Equation]
75 leads to
[Equation]
Putting all together provides us with
[Equation]
Noting that } A } F † 1 and denoting ˜ M 2 “
[Equation]
2 2 M ¨ M 21 leads to
[Equation]
Therefore, we get that
[Equation]
Notice that
[Equation]
[Equation]
Therefore, for any 0 § t 3 § T :
[Equation]
D
[Equation] [Equation]
D
[Equation]
Recall that by Lemma D.
[Equation]
we have for any x ° 0 , ln p x q † x that } A p 0 q} F † 4 d ✏ .
(D.111)
Using the fact that
T “ 2 w 0 ln ˆ ✏ 0 . 5 ˙ † 2 w 0 ✏ ´ 0 .(D.112)
we get that ñ } A p t 3 q} F † ˜ M2 w 0 ✏ ` 4 d ✏ .
(D.113)
Taking M 2 “ ˜ M 2 2 w [-
-]0 ` 4 ✏ , we have for all 0 § t § T that ñ } A p t q} F † M 2 ✏ (D.114)
Once again, recall that we deﬁned t A as the maximal t with this property, thus T † t A .
Since we already have also that T † t B , t C , we get T † min p T A , t B , t C q “ T ˚ .
D.4. B OUNDING W J W AND Y J Y
Proof of Lemma D.15 .
Denote
[Equation]
(D
[Equation]
Recall that
[Equation]
(D [Equation]
We can write the change in Y ,
[Equation]
˘
Denote p
A i q S “ A
i `p
A J q iand p
A i q ¯ S
“ A i ´p
A J q ithe symmetric and anti-symmetric parts of A i .
We can now write
[Equation]
Note also that for i “ 0 , we have A
“ I and its anti-symmetric part is the zero matrix, writing i “separately and assigning equation D.118 into equation D.117,
[Equation]
Let us look at ddt p Y J Y q [Equation] Y [Equation]
J 9 Y . Multiplying equation D.119 from the left with Y J evaluates to
[Equation]
We now turn to bound the terms in the sum.
[Equation]
We can bound each term using Cauchy–Schwarz.
We ﬁrst need to bound } Y } and } W } , which are trivially bounded by
[Equation]
As for the symmetric and anti-symmetric parts of A ,
[Equation]
(D [Equation] 123) where the last inequality follows again from Cauchy–Schwarz (the same considerations apply for p
A i q ¯ S ).
From Cauchy–Schwarz we can bound |
Y J p A
i q ¯ S W | § } Y J }}p A
i q ¯ S } F } W } § } Y }} A } iF } W } , we denote M“ max p M 1 , M 2 q and derive:
[Equation]
(D.124) which is maximized when i “ 1 .
We again bound: M “ max i p| w i |q
` M 21 ✏ 1 . 5 ° | r ` i | .
We can bound the terms in equation D.121 by
[Equation]
D.
[Equation]
Plugging back into equation D.121:
[Equation]
We can also bound r ` 0 “ p CB ´ w 0 q [Equation] w [Equation]
CB |
[Equation]
´
[Equation]
`
[Equation]
C }}
[Equation]
} § ´
[Equation]
` M 21 ✏ 1 . 5 .
Note also that we multiply by Y J Y
so we can bound
[Equation]
putting back together, we get
[Equation]
In particular, there exists M 4 such that
[Equation]
Recall that we were interested in bounding ddt p Y
[Equation]
Y
[Equation]
Denoting z p t q ” Y p t q J Y p t q and x p t q ” W p t q J W p t q , and using Lemma D.16, we have the desired bounds
[Equation]
In particular, we can write
[Equation]
and x
[Equation]
where c i ’s are positive constants.
Note that the derivation of x p t q is exactly the same as z p t q with opposite signs and bounding from below instead.
D.4. I NTEGRAL BOUND OF } Y } 2 AND } W }
Lemma D.16.
Assume
z
† ´ 2 w
{+z ` M 4 ✏ 2 . 5 † 0 , 9 x ° 2 w 
+}
z
´ M 4 ✏ 2 . 5 ° 0 , where w 0 ° 0 .
Then, under the assumptions of Lemma D.11:
[Equation]
Proof of Lemma D.16 .
Assume
z
† ´ 2 w
z ` M 4 ✏ 2 . 5 , where w 0 ° 0 , 2 w 0 z ° M 4 ✏ 2 . 5 .
Similarly, assume 9 x ° 2 w 0 x ´ M 4 ✏ 2 . 5 .
Then:
[Equation]
Integrating both sides by dt, and using integration by substitution, we get:
[Equation]
We note:
[Equation]
Combining equations, we have:
[Equation]
We note, z p 0 q “ Y T p
q Y p 0 q .
By linearity of sum of variances, Y p 0 q ’s entries are distributed according to N p?
2 ¨ ✏ q , by Lemma D.11:
[Equation]
From Cauchy-Schwartz, z p 0 q † 3 ✏ 2 with high probability.
W is distributed as Y , therefore, similarly, x p 0 q ° 12 ✏ 2 .
Assuming M 4 ✏ 2 .
5 † w [-
 ϵ-]0 ✏ 2 2 , we have:
[Equation]
Concluding the proof.