Learning RNNs with Commutative State Transitions

1 INTRODUCTION . Many of the successes of deep learning can be attributed to a choice of architecture that fits the application domain . A key example of this are convolutional neural nets ( Krizhevsky et al. , 2012 ) that employ extensive parameter tying across the image , reflecting the shift invariance of visual labels . It is clear that such inductive bias is necessary for any domain if one needs to learn from a limited number of examples . This has prompted much recent interest in the question of modeling architectures that capture invariances of particular domains . One such example is the setting where the input to the network is a set of objects and the true output does not depend on the order of the objects in the set . In this case , the function that we seek to learn is permutation invariant . It would thus be advantageous to retrict learning to models that are permutation invariant by design . This was the observation in the work on Deep Sets ( Zaheer et al. , 2017 ) , which gave necessary and sufficient conditions for a network to be permutation invariant . The underlying idea in many of the permutation invariant architectures is to apply the same network N1 to all items in the set , then sum the outputs and apply another network N2 to the summation . Clearly this is permutation invariant . Similarly other models using attention ( Vinyals et al. , 2016 ; Lee et al. , 2019 ) are permutation invariant since they only rely on comparisons between elements in the set . Although the above approaches results in invariant models , these are not clearly the smallest models that can compute a given invariant function . For example , consider the problem of computing the permutation invariant function fpx1 , . . . , xnq “ maxi xi . In order to calculate it with a DeepSet approach , one would need to set N1pxq “ xk for large k and then N2pzq “ z1 { k. This would approximate the ` 8 metric and therefore the max function . However , in order to implement these two functions N1 , N2 with ReLU networks , one would need Opk q to achieve accuracy ( e.g. , see Telgarsky , 2017 , Lemma 3.4 ) , and thus calculating this seemingly simple function to good accuracy would take a fairly large network . Clearly , if one uses max as the aggregation function in DeepSets , this could be avoided , but the same would then be true for other aggregation functions . On the other hand , the max function can be easily implemented via a simple RNN . Consider an RNN with the following state update st ` 1 “ max pst´1 , xtq and set s0 “ ´8 . Then clearly we would have sn “ maxi “ 1 xi . Namely , we will implement the permutation invariant max function . Furthermore , the st ` 1 state update can be very easily be implemented using a one hidden layer ReLU network , because maxps , xq “ ReLUps´ xq ` ReLUpxq ´ ReLUp´xq . The above example demonstrates that RNNs can in some cases be a natural computational model for permutation invariant functions . The goal of this paper is to ask which RNNs compute permutation invariant functions , and how they can be learned . Clearly , most RNNs are not permutation invariant , as their state has the capability of tracking temporal patterns , and is thus sensitive to changes in ordering . Here we show that by restricting RNN state updates to commutative-associative functions , RNNs do become invariant , and the resulting set of functions is very expressive . After establishing the commutative-associative constraints , we turn to ask how one can learn RNNs that satisfy the latter constraint . We show that for ReLU activation functions the commutative constraint corresponds to a closed form regularizer on the parameters of the RNN and suggest learning with this regularizer in order to achieve commutativity . We show empirically that learning with commutative regularization leads to architectures that are permutation invariant in practice and that outperform DeepSets on several benchmarks . Taken together , our results highlight the importance of the function class used for learning permutation invariant functions , and the important role of recursive computation for these tasks . 2 PROBLEM FORMULATION . We consider the problem of learning functions that map a sequence of inputs x1 , . . . , xn to an output y . We note we could have also addressed the problem of generating an output sequence y1 , . . . , yn but use a single output for simplicity . We say that such a function fpx1 , . . . , xnq is permutation invariant if for any permutation π of rns and for all x inputs we have : fpx1 , . . . , xnq “ fpxπp1q , . . . , xπpnqq ( 2.1 ) We next ask : how should an RNN be constructed such that it is guaranteed to be permutation invariant . We consider standard RNNs , parameterized as follows . Let st P Rr be the state vector at time t. Then the state update rule is : st ` 1 “ Aσ pWψpxtq ` Θstq ( 2.2 ) where ψpxq P Rd is a “ pre-processing ” network ( of arbitrary architecture ) , σ is a pointwise nonlinearity such as ReLU or tanh , and A , W , Θ are matrices which transform to the appropriate dimensions , specifically , A P Rrˆk , W P Rkˆd and Θ P Rkˆr . We note that we add a matrix A to allow the non-linearity to be taken in Rk such that k ‰ r. Finally , the output is obtained via : y “ gpstq ( 2.3 ) where g is any multilayer neural net ( the specific architecture will not be important for our derivation ) . To simplify the derivation of permutation invariance we make the assumption that the dimension of the RNN state , and the dimension of the transformed input ψpxq are the same , i.e . r “ d.1 Next , we introduce the following notation for the RNN state transition : st ` 1 “ Aσ pWψpxtq ` Θstq “ ψpxtq ˝ st ( 2.4 ) Thus , the ˝ operation takes two vectors in Rd and outputs a vector in Rd . Namely ˝ : Rd , Rd Ñ Rd.2 With this notation it is clear that : st “ p. . . pps0 ˝ ψpx1qq ˝ ψpx2qq . . . ˝ ψpxtqq ( 2.5 ) We use parentheses to emphasize that the order of applying ˝ does matter and these operations are neither commutative nor associative for general RNNs . 1It is always possible to achieve by padding either s or x by zeros as necessary . 2For the reminder of this paper we omit the ψ notation as it does not affect our derivations . 3 PERMUTATION INVARIANT RNNS . In this section we relate properties of the state transition operator ˝ to the permutation invariance properties of the RNN . Without loss of generality we assume that the ψpxq “ x , since this is just a pre-processing step . Definition 3.1 . An RNN is permutation invariant if the function fpx1 , . . . , xnq “ gpstq implemented by the RNN is permutation invariant . We begin with basic definitions . Definition 3.2 . The operator ˝ is commutative if for all x , x1 it holds that x ˝ x1 “ x1 ˝ x Definition 3.3 . The operator ˝ is associative if for all x , x1 , x̂ it holds that px ˝ x1q˝x̂ “ x˝px1 ˝ x̂q Definition 3.4 . If an RNN has a ˝ operator that is both commutative and associative we refer to it as a Commutative-Associative RNNs . Theorem 3.5 . If an RNN is Commutative-Associative then it is permutation invariant . Namely st does not depend on the order of x1 , . . . , xn . Proof . The associativity property of the ˝ operator implies that we can remove the parentheses in Eq . 2.5 and commutativity implies that we can then switch the order of the xi arbitrarily . Thus we have that for all permutations π : st “ s0 ˝ x1 ˝ x2 . . . ˝ xt “ s0 ˝ xπp1q ˝ xπp2q . . . ˝ xπt ( 3.1 ) Therefore the state is invariant to the input order , and so is the output y . A complementary question is that of universality . Namely , are the above permutation invariant RNNs sufficiently expressive to capture any permutation invariant function ? Theorem 3.6 . Let fpx1 , . . . , xn ) be a permutation invariant function . Then it can be implemented with Commutative-Associative RNNs . Proof . Since the function f is permutation invariant , it follows from Theorem 2 in the DeepSets paper ( Zaheer et al. , 2017 ) that there exist two functions φ and ρ such that fpx1 , . . . , xnq “ ρ p ř i φpxiqq . We can now implement this architecture with a Commutative-Associative RNNs as follows : set ψ “ φ , g “ ρ and x˝x1 “ x ` x1 . Clearly this implements the DeepSet function and is CommutativeAssociative because ˝ is just addition . We conclude that Commutative-Associative RNNs implement permutation invariant functions and are also sufficiently expressive to capture all permutation invariant functions . 4 COMMUTATIVE REGULARIZATION . In the previous section we highlighted the importance of having an RNN operator ˝ that is both commutative and associative . This suggests that in order to learn a permutation invariant RNN we would like to learn RNNs under the constraint that ˝ is associative and commutative . In this section and the remainder of the paper , we focus on the latter , namely introducing a constraint ( or equivalently , regularizer ) that ˝ is commutative . Formally , the constraint that ˝ is commutative corresponds to requiring : x ˝ x1 “ x1 ˝ x @ x , x1 P Rd ( 4.1 ) Recall that ˝ depends on the matrices A , W , Θ and thus the above constraint translates to a non-linear constraint on these three matrices . However , it is not clear how to write this constraint in a way that facilitates optimization . This will be our goal in what follows . We begin by replacing requirement 4.1 with a more convenient proxy . Lemma 4.1 . If there exists a distribution , D , over Rd , Rd with non-zero density over its support,3 then Ex , x1 „ D ” } x ˝ x1 ´ x1 ˝ x } 22 ı “ 0 ðñ P ” x ˝ x1 “ x1 ˝ x , @ x , x1 P Rd ı “ 1 ( 4.2 ) Proof . The proof relies on the fact that for a continuous non-negative random variable , X , is holds that ( see Section A ) : E ” X ı “ 0 ðñ PrX “ 0s “ 1 ( 4.3 ) We now use 4.3 in order to complete the proof . Denote by R the random variable Rpx , x1q ” } x ˝ x1 ´ x1 ˝ x } 22 . Obviously , R is non-negative by definition . Note also that for any pair x , x1 we have } x ˝ x1 ´ x1 ˝ x } 22 “ 0 ðñ x ˝ x1 “ x1 ˝ x . It follows that : E ” R ı “ 0 ðñ P ” R “ 0 ı “ 1 ðñ P ” x ˝ x1 “ x1 ˝ x , @ x , x1 P Rd ı “ 1 ( 4.4 ) The above lemma already gives the strong result that with probability one the network is commutative . In fact a stronger result is possible , stating that the network is commutative for all inputs . We only need to add the restriction of the network A , W , Θ is Lipschitz . This will hold for example , if we restrict A , W , Θ to any compact set ( e.g. , see similar argument in Arjovsky et al. , 2017 ) . Lemma 4.2 . Assume the network defined by A , W , Θ has a Lipschitz constant bounded by K , and that Eq . 4.2 holds under the conditions therein . Then @ x , x1 it holds that x ˝ x1 “ x1 ˝ x . Proof . Denote fpx , x1q “ } x ˝ x1 ´ x1 ˝ x } 22 . Then f also has a Lipschitz constant bounded by L that is a function of K. Assume in contradiction that there exists x0 , x10 such that fpx0 , x10q “ for some ą 0 . Let Bpc , rq denote an ` 2 ball centered at c with radius r. Then because f is Lipschitz it follows that for all z P Bprx0 , x10s , 2L q it holds that fpzq ě 2 . Let γ denote the minimal density value at Bprx0 , x10s , 2L q . By assumption on D we have that γ ą 0 . It therefore follows that P ” fpXq ě 0.5 ı ě P ” z P Bprx0 , x10s , 2L ı q ě 0.5γ ą 0 and therefore P ” fpXq “ 0 ı ă 1 and we have a contradiction with Lemma 4.1 . As a conclusion of lemma 4.1 , we can express E ” } x ˝ x1 ´ x1 ˝ x } 22 ı in terms of matrices A , W , Θ for a specific distribution and derive conditions under which it is equal to zero . Definition 4.3 . Consider an RNN parameterized by matrices A , W , Θ as in 2.2 with ReLU activation . ∆pA , W , Θq ” Ex , x1 „ D ” } x ˝ x1 ´ x1 ˝ x } 22 ı ( 4.5 ) Below we show that ∆pA , W , Θq can be expressed as a simple function of A , W , Θ when using the ReLU non-linearity . Let D be a multivariate Gaussian distribution with zero mean and covariance identity . In order to calculate ∆pA , W , Θq we will make use of the following integral ( Cho & Saul , 2009 ) : E ” σpu ¨ xqσpv ¨ xq ı “ 1 π } u } 2 } v } 2 psin pθu , vq ` pπ ´ θu , vq cos pθu , vqq ” gpu , vq ( 4.6 ) It turns out that ∆pA , W , Θq is a simple function of the above gpu , vq function , as stated next . Let wi denote the ith row of W , θi the ith row of Θ and ai the ith column of A . Also , denote the following horizontal stacking of the vectors as ui “ rθi wis , vj “ rwj θjs ( 4.7 ) and the matrices with these as columns by U and V respectively . 3e.g . @ x , x1 P Rd , fDpx , x1q ą 0 Theorem 4.4 . Assume x , x1 are independent Gaussian vectors , each with a unit covariance . Then : ∆pA , W , Θq “ h ÿ i “ 1 h ÿ j “ 1 ai ¨ aj ´ gpui , ujq ´ 2gpui , vjq ` gpvi , vjq ¯ ( 4.8 ) Proof . Define the stacked vector : z “ ˆ x1 x ˙ . The function ∆pA , W , Θq is given by : E „ › › › Aσ ` Wx ` Θx1 ˘ ´Aσ ` Wx1 ` Θx ˘ › › › 2 2 ( 4.9 ) We can write 4.9 using definitions in Eq . 4.7 as : E „ ´ Aσ ` UTz ˘ ´Aσ ` V Tz ˘ ¯2 “ E ” σ ` zTU ˘ ATAσ ` UTz ˘ ı ´ 2E ” σ ` zTU ˘ ATAσ ` V Tz ˘ ı ` E ” σ ` zTV ˘ ATAσ ` V Tz ˘ ı Denoting ũ “ σ pUzq , ṽ “ σ pV zq , we can now write : E ” ũTATAũ´ 2ũTATAṽ ` ṽTATAṽ ı . ( 4.10 ) Letting ATA “ Q , we can now write the above as : E « h ÿ i “ 1 h ÿ j “ 1 ũiQij ũj ´ 2 h ÿ i “ 1 h ÿ j “ 1 ũiQij ṽj ` h ÿ i “ 1 h ÿ j “ 1 ṽiQij ṽj ff ( 4.11 ) and more compactly h ÿ i “ 1 h ÿ j “ 1 E ” Qij pũiũj ´ 2ũiṽj ` ṽiṽjq ı ( 4.12 ) writing back ũi “ σpui ¨ zq and ṽj “ σpvj ¨ zq we have h ÿ i “ 1 h ÿ j “ 1 QijE ” σpui ¨ zqσpuj ¨ zq ´ 2σpui ¨ zqσpvj ¨ zq ` σpvi ¨ zqσpvj ¨ zq ı ( 4.13 ) Using Eq . 4.6 , we have that Eq . 4.13 results in the expression in Eq . 4.8 . Theorem 4.4 provides an interesting characterization of when an RNN is commutative . As mentioned earlier , an RNN is commutative iff ∆pA , W , Θq “ 0 . Since the theorem gives a closed-form expression for ∆ it provides a measure of non-commutativity which we will optimize in what follows . The function in Eq . 4.8 does not provide clear structual constraints on the matrices A , W , Θ but a rather elaborate dependence between their elements . The function ∆ indeed achieves a value of zero for the ReLU implementation of the max function discussed in the intro , as well as ( infinitely ) many other commutative state transition functions .

Summary: The paper starts with presenting an RNN formulation and essentially writing out the sequence of RNN applications. Not surprisingly, if these applications were associative and commutative the RNN would be permutation invariant.  Then a condition for commutativity is formulated in terms of an expectation of a difference.  Based on a prior result, it is shown that the expectation can be computed in closed form.  Although it is not shown if an RNN regularized that way is permutation invariant, since the associativity is not demonstrated, empirically it is shown that it may be already of use.