Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs

Introduction

Deep Neural Networks (LeCun et al., 2015) have been successful on numerous difficult machine learning tasks, including image recognition(Krizhevsky et al., 2012; Donahue et al., 2015), speech recognition(Hinton et al., 2012) and natural language processing(Collobert et al., 2011; Bahdanau et al., 2014; Sutskever et al., 2014). However, deep neural networks can suffer from vanishing and exploding gradient problems(Hochreiter, 1991; Bengio et al., 1994), which are known to be caused by matrix eigenvalues far from unity being raised to large powers. Because the severity of these problems grows with the the depth of a neural network, they are particularly grave for Recurrent Neural Networks (RNNs), whose recurrence can be equivalent to thousands or millions of equivalent hidden layers.

Several solutions have been proposed to solve these problems for RNNs. Long Short Term Memory (LSTM) networks (Hochreiter & Schmidhuber, 1997), which help RNNs contain information inside hidden layers with gates, remains one of the the most popular RNN implementations. Other recently proposed methods such as GRUs(Cho et al., 2014) and Bidirectional RNNs (Berglund et al., 2015) also perform well in numerous applications. However, none of these approaches has fundamentally solved the vanishing and exploding gradient problems, and gradient clipping is often required to keep gradients in a reasonable range.

A recently proposed solution strategy is using orthogonal hidden weight matrices or their complex generalization (unitary matrices) (Saxe et al., 2013; Le et al., 2015; Arjovsky et al., 2015; Henaff et al., 2016), because all their eigenvalues will then have absolute values of unity, and can safely be raised to large powers. This has been shown to help both when weight matrices are initialized to be unitary (Saxe et al., 2013; Le et al., 2015) and when they are kept unitary during training, either by restricting them to a more tractable matrix subspace (Arjovsky et al., 2015) or by alternating gradient-descent steps with projections onto the unitary subspace (Wisdom et al., 2016).

In this paper, we will first present an Efficient Unitary Neural Network (EUNN) architecture that parametrizes the entire space of unitary matrices in a complete and computationally efficient way, thereby eliminating the need for time-consuming unitary subspace-projections. Our architecture has a wide range of capacity-tunability to represent subspace unitary models by fixing some of our parameters; the above-mentioned unitary subspace models correspond to special cases of our architecture. We also implemented an EUNN with an earlier introduced FFT-like architecture which efficiently approximates the unitary space with minimum number of required parameters(Mathieu & LeCun, 2014b).

We then benchmark EUNN’s performance on both simulated and real tasks: the standard copying task, the pixel-permuted MNIST task, and speech prediction with the TIMIT dataset (Garofolo et al., 1993). We show that our EUNN algorithm with an $\mathcal{O}(N)$ hidden layer size can compute up to the entire $N\times N$ gradient matrix using $\mathcal{O}(1)$ computational steps and memory access per parameter. This is superior to the $\mathcal{O}(N)$ computational complexity of the existing training method for a full-space unitary network (Wisdom et al., 2016) and $\mathcal{O}(\log N)$ more efficient than the subspace Unitary RNN(Arjovsky et al., 2015).

Background

A recurrent neural network takes an input sequence and uses the current hidden state to generate a new hidden state during each step, memorizing past information in the hidden layer. We first review the basic RNN architecture.

Consider an RNN updated at regular time intervals $t=1,2,...$ whose input is the sequence of vectors ${\bf x}^{(t)}$ whose hidden layer ${\bf h}^{(t)}$ is updated according to the following rule:

where $\sigma$ is the nonlinear activation function. The output is generated by

where ${\bf b}$ is the bias vector for the hidden-to-output layer. For $t=0$, the hidden layer $\mathbf{h^{(0)}}$ can be initialized to some special vector or set as a trainable variable. For convenience of notation, we define ${\bf z}^{(t)}={\bf U}{\bf x}^{(t)}+{\bf W}{\bf h}^{(t-1)}$ so that ${\bf h}^{(t)}=\sigma({{\bf z}^{(t)}})$.

2 The Vanishing and Exploding Gradient Problems

When training the neural network to minimize a cost function $C$ that depends on a parameter vector ${\bf a}$, the gradient descent method updates this vector to ${\bf a}-\lambda\frac{\partial C}{\partial{\bf a}}$, where $\lambda$ is a fixed learning rate and $\frac{\partial C}{\partial{\bf a}}\equiv\nabla C$. For an RNN, the vanishing or exploding gradient problem is most significant during back propagation from hidden to hidden layers, so we will only focus on the gradient for hidden layers. Training the input-to-hidden and hidden-to-output matrices is relatively trivial once the hidden-to-hidden matrix has been successfully optimized.

In order to evaluate $\frac{\partial C}{\partial W_{ij}}$, one first computes the derivative $\frac{\partial C}{\partial h^{(t)}}$ using the chain rule:

where ${\bf D}^{(k)}={\rm diag}\{\sigma^{\prime}({\bf U}{\bf x}^{(k)}+{\bf W}{\bf h}^{(k-1)})\}$ is the Jacobian matrix of the pointwise nonlinearity. For large times $T$, the term $\prod{\bf W}$ plays a significant role. As long as the eigenvalues of ${\bf D}^{(k)}$ are of order unity, then if ${\bf W}$ has eigenvalues $\lambda_{i}\gg 1$, they will cause gradient explosion $|\frac{\partial C}{\partial{\bf h}^{(T)}}|\rightarrow\infty$, while if ${\bf W}$ has eigenvalues $\lambda_{i}\ll 1$, they can cause gradient vanishing, $|\frac{\partial C}{\partial{\bf h}^{(T)}}|\rightarrow 0$. Either situation prevents the RNN from working efficiently.

Unitary RNNs

In a breakthrough paper, Arjovsky, Shah & Bengio (Arjovsky et al., 2015) showed that unitary RNNs can overcome the exploding and vanishing gradient problems and perform well on long term memory tasks if the hidden-to-hidden matrix in parametrized in the following unitary form:

Here ${\bf D}_{1,2,3}$ are diagonal matrices with each element $e^{i\omega_{j}},j=1,2,\cdots,n$. ${\bf T}_{1,2}$ are reflection matrices, and ${\bf T}=I-2\frac{\widehat{\bf v}\widehat{\bf v}^{\dagger}}{||\widehat{\bf v}||^{2}}$, where $\widehat{\bf v}$ is a vector with each of its entries as a parameter to be trained. $\bm{\Pi}$ is a fixed permutation matrix. $\mathcal{F}$ and $\mathcal{F}^{-1}$ are Fourier and inverse Fourier transform matrices respectively. Since each factor matrix here is unitary, the product ${\bf W}$ is also a unitary matrix.

This model uses $\mathcal{O}(N)$ parameters, which spans merely a part of the whole $\mathcal{O}(N^{2})$-dimensional space of unitary $N\times N$ matrices to enable computational efficiency. Several subsequent papers have tried to expand the space to $\mathcal{O}(N^{2})$ in order to achieve better performance, as summarized below.

2 Full Space Unitary RNNs

In order to maximize the power of Unitary RNNs, it is preferable to have the option to optimize the weight matrix ${\bf W}$ over the full space of unitary matrices rather than a subspace as above. A straightforward method for implementing this is by simply updating ${\bf W}$ with standard back-propagation and then projecting the resulting matrix (which will typically no longer be unitary) back onto to the space of unitary matrices. Defining ${\bf G}_{ij}\equiv\frac{\partial{C}}{\partial{W_{ij}}}$ as the gradient with respect to ${\bf W}$, this can be implemented by the procedure defined by (Wisdom et al., 2016):

This method shows that full space unitary networks are superior on many RNN tasks (Wisdom et al., 2016). A key limitation is that the back-propation in this method cannot avoid $N$-dimensional matrix multiplication, incurring $\mathcal{O}(N^{3})$ computational cost.

Efficient Unitary Neural Network (EUNN) Architectures

In the following, we first describe a general parametrization method able to represent arbitrary unitary matrices with up to $N^{2}$ degrees of freedom. We then present an efficient algorithm for this parametrization scheme, requiring only $\mathcal{O}(1)$ computational and memory access steps to obtain the gradient for each parameter. Finally, we show that our scheme performs significantly better than the above mentioned methods on a few well-known benchmarks.

Any $N\times N$ unitary matrix ${\bf W}_{N}$ can be represented as a product of rotation matrices $\{{\bf R}_{ij}\}$ and a diagonal matrix ${\bf D}$, such that ${\bf W}_{N}={\bf D}\prod_{i=2}^{N}\prod_{j=1}^{i-1}{\bf R}_{ij}$, where ${\bf R}_{ij}$ is defined as the $N$-dimensional identity matrix with the elements $R_{ii}$, $R_{ij}$, $R_{ji}$ and $R_{jj}$ replaced as follows (Reck et al., 1994; Clements et al., 2016):

where $\theta_{ij}$ and $\phi_{ij}$ are unique parameters corresponding to $\mathbf{R_{ij}}$. Each of these matrices performs a $U(2)$ unitary transformation on a two-dimensional subspace of the N-dimensional Hilbert space, leaving an $(N-2)$-dimensional subspace unchanged. In other words, a series of $U(2)$ rotations can be used to successively make all off-diagonal elements of the given $N\times N$ unitary matrix zero. This generalizes the familiar factorization of a 3D rotation matrix into 2D rotations parametrized by the three Euler angles. To provide intuition for how this works, let us briefly describe a simple way of doing this that is similar to Gaussian elimination by finishing one column at a time. There are infinitely many alternative decomposition schemes as well; Fig. 1 shows two that are particularly convenient to implement in software (and even in neuromorphic hardware (Shen et al., 2016)). The unitary matrix ${\bf W}_{N}$ is multiplied from the right by a succession of unitary matrices ${\bf R}_{Nj}$ for $j=N-1,\cdots,1$. Once all elements of the last row except the one on the diagonal are zero, this row will not be affected by later transformations. Since all transformations are unitary, the last column will then also contain only zeros except on the diagonal:

The effective dimensionality of the the matrix ${\bf W}$ is thus reduced to $N-1$. The same procedure can then be repeated $N-1$ times until the effective dimension of ${\bf W}$ is reduced to 1, leaving us with a diagonal matrix:Note that Gaussian Elimination would make merely the upper triangle of a matrix vanish, requiring a subsequent series of rotations (complete Gauss-Jordan Elimination) to zero the lower triangle. We need no such subsequent series because since ${\bf W}$ is unitary: it is easy to show that if a unitary matrix is triangular, it must be diagonal.

where ${\bf D}$ is a diagonal matrix whose diagonal elements are $e^{iw_{j}}$, from which we can write the direct representation of ${\bf W}_{N}$ as

This parametrization thus involves $N(N-1)/2$ different $\theta_{ij}$-values, $N(N-1)/2$ different $\phi_{ij}$-values and $N$ different $w_{i}$-values, combining to $N^{2}$ parameters in total and spans the entire unitary space. Note we can always fix a portion of our parameters, to span only a subset of unitary space – indeed, our benchmark test below will show that for certain tasks, full unitary space parametrization is not necessary. Our preliminary experimental tests even suggest that a full-capacity unitary RNN is even undesirable for some tasks.

2 Tunable space implementation

The representation in Eq. 12 can be made more compact by reordering and grouping specific rotational matrices, as was shown in the optical community (Reck et al., 1994; Clements et al., 2016) in the context of universal multiport interferometers. For example (Clements et al., 2016), a unitary matrix can be decomposed as

is a block diagonal matrix, with $N$ angle parameters in total, and

with $N-1$ parameters, as is schematically shown in Fig. 1a. By choosing different values for $L$ , ${\bf W}_{N}$ will span a different subspace of the unitary space. Specifically,when $L=N$, ${\bf W}_{N}$ will span the entire unitary space.

Following this physics-inspired scheme, we decompose our unitary hidden-to-hidden layer matrix ${\bf W}$ as

3 FFT-style approximation

Inspired by (Mathieu & LeCun, 2014a), an alternative way to organize the rotation matrices is implementing an FFT-style architecture. Instead of using adjacent rotation matrices, each $\mathbf{F}$ here performs a certain distance pairwise rotations as shown in Fig. 1b:

The rotation matrices in $\mathbf{F}_{i}$ are performed between pairs of coordinates

where $p=\frac{N}{2^{i}}$, $k\in\{0,...,2^{i-1}\}$ and $j\in\{1,...,p\}$. This requires only $\log(N)$ matrices, so there are a total of $N\log(N)/2$ rotational pairs. This is also the minimal number of rotations that can have all input coordinates interacting with each other, providing an approximation of arbitrary unitary matrices.

4 Efficient implementation of rotation matrices

To implement this decomposition efficiently in an RNN, we apply vector element-wise multiplications and permutations: we evaluate the product $\mathbf{F}{\bf x}$ as

where $*$ represents element-wise multiplication, $\mathbf{F}$ refers to general rotational matrices such as $\mathbf{F}_{A/B}$ in Eq. 14 and $\mathbf{F}_{i}$ in Eq. 16. For the case of the tunable-space implementation, if we want to implement $\mathbf{F}_{A}^{(l)}$ in Eq. 14, we define $\mathbf{v}$ and the permutation as follows:

For the FFT-style approach, if we want to implement $\mathbf{F}_{1}$ in Eq 16, we define $\mathbf{v}$ and the permutation as follows:

In general, the pseudocode for implementing operation $\mathbf{F}$ is as follows:

Note that $\mathbf{ind_{1}}$ and $\mathbf{ind_{2}}$ are different for different $\mathbf{F}$.

5 Nonlinearity

We use the same nonlinearity as (Arjovsky et al., 2015):

where the bias vector $\mathbf{b}$ is a shared trainable parameter, and $|z_{i}|$ is the norm of the complex number $z_{i}$.

where $|z_{i}|$ is the absolute value of the real number $z_{i}$.

We empirically find that this nonlinearity function performs the best. We believe that this function possibly also serves as a forgetting filter that removes the noise using the bias threshold.

Experimental tests of our method

In this section, we compare the performance of our Efficient Unitary Recurrent Neural Network (EURNN) with

an LSTM RNN (Hochreiter & Schmidhuber, 1997),

a Partial Space URNN (Arjovsky et al., 2015), and

a Projective full-space URNN (Wisdom et al., 2016).

All models are implemented in both Tensorflow and Theano, available from https://github.com/jingli9111/EUNN-tensorflow and https://github.com/iguanaus/EUNN-theano.

We compare these networks by applying them all to the well defined Copying Memory Task (Hochreiter & Schmidhuber, 1997; Arjovsky et al., 2015; Henaff et al., 2016). The copying task is a synthetic task that is commonly used to test the network’s ability to remember information seen $T$ time steps earlier.

Specifically, the task is defined as follows (Hochreiter & Schmidhuber, 1997; Arjovsky et al., 2015; Henaff et al., 2016). An alphabet consists of symbols $\{a_{i}\}$, the first $n$ of which represent data, and the remaining two representing “blank” and “start recall”, respectively; as illustrated by the following example where $T=20$ and $M=5$:

In the above example, $n=3$ and $\{a_{i}\}=\{A,B,C,-,:\}$. The input consists of $M$ random data symbols ($M=5$ above) followed by $T-1$ blanks, the “start recall” symbol and $M$ more blanks. The desired output consists of $M+T$ blanks followed by the data sequence. The cost function $C$ is defined as the cross entropy of the input and output sequences, which vanishes for perfect performance.

We use $n=8$ and input length $M=10$. The symbol for each input is represented by an $n$-dimensional one-hot vector. We trained all five RNNs for $T=1000$ with the same batch size 128 using RMSProp optimization with a learning rate of 0.001. The decay rate is set to 0.5 for EURNN, and 0.9 for all other models respectively. (Fig. 2). This results show that the EURNN architectures introduced in both Sec.4.2 (EURNN with N=512, selecting L=2) and Sec.4.3 (FFT-style EURNN with N=512) outperform the LSTM model (which suffers from long term memory problems and only performs well on the copy task for small time delays $T$) and all other unitary RNN models, both in-terms of learnability and in-terms of convergence rate. Note that the only other unitary RNN model that is able to beat the baseline for $T=1000$ (Wisdom et al., 2016) is significantly slower than our method.

Moreover, we find that by either choosing smaller $L$ or by using the FFT-style method (so that $\mathbf{W}$ spans a smaller unitary subspace), the EURNN converges toward optimal performance significantly more efficiently (and also faster in wall clock time) than the partial (Arjovsky et al., 2015) and projective (Wisdom et al., 2016) unitary methods. The EURNN also performed more robustly. This means that a full-capacity unitary matrix is not necessary for this particular task.

2 Pixel-Permuted MNIST Task

The MNIST handwriting recognition problem is one of the classic benchmarks for quantifying the learning ability of neural networks. MNIST images are formed by a 28$\times$28 grayscale image with a target label between 0 and 9.

To test different RNN models, we feed all pixels of the MNIST images into the RNN models in 28$\times$28 time steps, where one pixel at a time is fed in as a floating-point number. A fixed random permutation is applied to the order of input pixels. The output is the probability distribution quantifying the digit prediction. We used RMSProp with a learning rate of 0.0001 and a decay rate of 0.9, and set the batch size to 128.

As shown in Fig. 3, EURNN significantly outperforms LSTM with the same number of parameters. It learns faster, in fewer iteration steps, and converges to a higher classification accuracy. In addition, the EURNN reaches a similar accuracy with fewer parameters. In Table. 2, we compare the performance of different RNN models on this task.

3 Speech Prediction on TIMIT dataset

We also apply our EURNN to real-world speech prediction task and compare its performance to LSTM. The main task we consider is predicting the log-magnitude of future frames of a short-time Fourier transform (STFT) (Wisdom et al., 2016; Sejdić et al., 2009). We use the TIMIT dataset (Garofolo et al., 1993) sampled at 8 kHz. The audio .wav file is initially diced into different time frames (all frames have the same duration referring to the Hann analysis window below). The audio amplitude in each frame is then Fourier transformed into the frequency domain. The log-magnitude of the Fourier amplitude is normalized and used as the data for training/testing each model. In our STFT operation we uses a Hann analysis window of 256 samples (32 milliseconds) and a window hop of 128 samples (16 milliseconds). The frame prediction task is as follows: given all the log-magnitudes of STFT frames up to time $t$, predict the log-magnitude of the STFT frame at time $t+1$ that has the minimum mean square error (MSE). We use a training set with 2400 utterances, a validation set of 600 utterances and an evaluation set of 1000 utterances. The training, validation, and evaluation sets have distinct speakers. We trained all RNNs for with the same batch size 32 using RMSProp optimization with a learning rate of 0.001, a momentum of 0.9 and a decay rate of 0.1.

The results are given in Table. 3, in terms of the mean-squared error (MSE) loss function. Figure. 4 shows prediction examples from the three types of networks, illustrating how EURNNs generally perform better than LSTMs. Furthermore, in this particular task, full-capacity EURNNs outperform small capacity EURNNs and FFT-style EURNNs.

Conclusion

We have presented a method for implementing an Efficient Unitary Neural Network (EUNN) whose computational cost is merely $\mathcal{O}(1)$ per parameter, which is $\mathcal{O}(\log N)$ more efficient than the other methods discussed above. It significantly outperforms existing RNN architectures on the standard Copying Task, and the pixel-permuted MNIST Task using a comparable parameter count, hence demonstrating the highest recorded ability to memorize sequential information over long time periods.

It also performs well on real tasks such as speech prediction, outperforming an LSTM on TIMIT data speech prediction.

We want to emphasize the generality and tunability of our method. The ordering of the rotation matrices we presented in Fig. 1 are merely two of many possibilities; we used it simply as a concrete example. Other ordering options that can result in spanning the full unitary matrix space can be used for our algorithm as well, with identical speed and memory performance. This tunability of the span of the unitary space and, correspondingly, the total number of parameters makes it possible to use different capacities for different tasks, thus opening the way to an optimal performance of the EUNN. For example, as we have shown, a small subspace of the full unitary space is preferable for the copying task, whereas the MNIST task and TIMIT task are better performed by EUNN covering a considerably larger unitary space. Finally, we note that our method remains applicable even if the unitary matrix is decomposed into a different product of matrices (Eq. 12).

This powerful and robust unitary RNN architecture also might be promising for natural language processing because of its ability to efficiently handle tasks with long-term correlation and very high dimensionality.

Acknowledgment

We thank Hugo Larochelle and Yoshua Bengio for helpful discussions and comments.

This work was partially supported by the Army Research Office through the Institute for Soldier Nanotechnologies under contract W911NF-13-D0001, the National Science Foundation under Grant No. CCF-1640012 and the Rothberg Family Fund for Cognitive Science.

References