Full-Capacity Unitary Recurrent Neural Networks

Introduction

Deep feed-forward and recurrent neural networks have been shown to be remarkably effective in a wide variety of problems. A primary difficulty in training using gradient-based methods has been the so-called vanishing or exploding gradient problem, in which the instability of the gradients over multiple layers can impede learning . This problem is particularly keen for recurrent networks, since the repeated use of the recurrent weight matrix can magnify any instability.

This problem has been addressed in the past by various means, including gradient clipping , using orthogonal matrices for initialization of the recurrence matrix , or by using pioneering architectures such as long short-term memory (LSTM) recurrent networks or gated recurrent units . Recently, several innovative architectures have been introduced to improve information flow in a network: residual networks, which directly pass information from previous layers up in a feed-forward network , and attention networks, which allow a recurrent network to access past activations . The idea of using a unitary recurrent weight matrix was introduced so that the gradients are inherently stable and do not vanish or explode . The resulting unitary recurrent neural network (uRNN) is complex-valued and uses a complex form of the rectified linear activation function. However, this idea was investigated using, as we show, a potentially restricted form of unitary matrices.

The two main components of our contribution can be summarized as follows:

1) We provide a theoretical argument to determine the smallest dimension $N$ for which any parameterization of the unitary recurrence matrix does not cover the entire set of all unitary matrices. The argument relies on counting real-valued parameters and using Sard’s theorem to show that the smooth map from these parameters to the unitary manifold is not onto. Thus, we can show that a previously proposed parameterization cannot represent all unitary matrices larger than $7\times 7$. Thus, such a parameterization results in what we refer to as a restricted-capacity unitary recurrence matrix.

2) To overcome the limitations of restricted-capacity parameterizations, we propose a new method for stochastic gradient descent for training the unitary recurrence matrix, which constrains the gradient to lie on the differentiable manifold of unitary matrices. This approach allows us to directly optimize a complete, or full-capacity, unitary matrix. Neither restricted-capacity nor full-capacity unitary matrix optimization require gradient clipping. Furthermore, full-capacity optimization still achieves good results without adaptation of the learning rate during training.

To test the limitations of a restricted-capacity representation and to confirm that our full-capacity uRNN does have practical implications, we test restricted-capacity and full-capacity uRNNs on both synthetic and natural data tasks. These tasks include synthetic system identification, long-term memorization, frame-to-frame prediction of speech spectra, and pixel-by-pixel classification of handwritten digits. Our proposed full-capacity uRNNs generally achieve equivalent or superior performance on synthetic and natural data compared to both LSTMs and the original restricted-capacity uRNNs .

In the next section, we give an overview of unitary recurrent neural networks. Section 3 presents our first contribution: the theoretical argument to determine if any unitary parameterization has restricted-capacity. Section 4 describes our second contribution, where we show how to optimize a full-capacity unitary matrix. We confirm our results with simulated and natural data in Section 5 and present our conclusions in Section 6.

Unitary recurrent neural networks

The uRNN proposed by Arjovsky et al. consists of the following nonlinear dynamical system that has real- or complex-valued inputs ${\bf x}_{t}$ of dimension $M$, complex-valued hidden states ${\bf h}_{t}$ of dimension $N$, and real- or complex-valued outputs ${\bf y}_{t}$ of dimension $L$:

Arjovsky et al. propose the following parameterization of the unitary matrix ${\bf W}$:

where ${\bf D}$ are diagonal unitary matrices, ${\bf R}$ are Householder reflection matrices , $\mathcal{F}$ is a discrete Fourier transform (DFT) matrix, and ${\bf P}$ is a permutation matrix. The resulting matrix ${\bf W}_{u}$ is unitary because all its component matrices are unitary. This decomposition is efficient because diagonal, reflection, and permutation matrices are $\mathcal{O}(N)$ to compute, and DFTs can be computed efficiently in $\mathcal{O}(N\log N)$ time using the fast Fourier transform (FFT). The parameter vector $\theta_{u}$ consists of $7N$ real-valued parameters: $N$ parameters for each of the $3$ diagonal matrices where ${D}_{i,i}=e^{j\theta_{i}}$ and $2N$ parameters for each of the $2$ Householder reflection matrices, which are real and imaginary values of the complex reflection vectors ${\bf u}_{i}$: ${\bf R}_{i}={\bf I}-2\frac{{\bf u}_{i}{\bf u}_{i}^{H}}{\langle{\bf u}_{i},{\bf u}_{i}\rangle}$.

Estimating the representation capacity of structured unitary matrices

In this section, we state and prove a theorem that can be used to determine when any particular unitary parameterization does not have capacity to represent all unitary matrices. As an application of this theorem, we show that the parameterization (3) does not have the capacity to cover all $N\times N$ unitary matrices for $N>7$. First, we establish an upper bound on the number of real-valued parameters required to represent any $N\times N$ unitary matrix. Then, we state and prove our theorem.

The set of all unitary matrices is a manifold of dimension $N^{2}$.

If a family of $N\times N$ unitary matrices is parameterized by $P$ real-valued parameters for $P<N^{2}$, then it cannot contain all $N\times N$ unitary matrices.

Proof: We consider a family of unitary matrices that is parameterized by $P$ real-valued parameters through a smooth map $g:\mathcal{P}(P)\rightarrow\mathcal{U}(N^{2})$ from the space of parameters $\mathcal{P}(P)$ to the space of all unitary matrices $\mathcal{U}(N^{2})$. The space $\mathcal{P}(P)$ of parameters is considered as a $P$-dimensional manifold, while the space $\mathcal{U}(N^{2})$ of all unitary matrices is an $N^{2}$-dimensional manifold according to lemma 3.1. Then, if $P<N^{2}$, Sard’s theorem implies that the image $g(\mathcal{P})$ of $g$ is of measure zero in $\mathcal{U}(N^{2})$, and in particular $g$ is not onto. Since $g$ is not onto, there must exist a unitary matrix ${\bf W}\in\mathcal{U}(N^{2})$ for which there is no corresponding input ${\bf P}\in\mathcal{P}(P)$ such that ${\bf W}=g({\bf P})$. Thus, if $P$ is such that $P<N^{2}$, the manifold $\mathcal{P}(P)$ cannot represent all unitary matrices in $\mathcal{U}(N^{2})$.

We now apply Theorem 3.2 to the parameterization (3). Note that the parameterization (3) has $P=7N$ real-valued parameters. If we solve for $N$ in $7N<N^{2}$, we get $N>7$. Thus, the parameterization (3) cannot represent all unitary matrices for dimension $N>7$.

Optimizing full-capacity unitary matrices on the Stiefel manifold

Under this canonical inner product on the tangent space, the gradient in the Stiefel manifold of the loss function $f$ with respect to the matrix ${\bf W}$ is ${\bf A}{\bf W}$, where ${\bf A}={\bf G}^{H}{\bf W}-{\bf W}^{H}{\bf G}$ is a skew-Hermitian matrix and ${\bf G}$ with ${G}_{i,j}=\frac{\delta f}{\delta W_{i,j}}$ is the usual gradient of the loss function $f$ with respect to the matrix ${\bf W}$ . Using these facts, Tagare suggests a descent curve along the Stiefel manifold at training iteration $k$ given by the matrix product of the Cayley transformation of ${\bf A}^{(k)}$ with the current solution ${\bf W}^{(k)}$:

where $\lambda$ is a learning rate and ${\bf A}^{(k)}={{\bf G}^{(k)}}^{H}{\bf W}^{(k)}-{{\bf W}^{(k)}}^{H}{\bf G}^{(k)}$. Gradient descent proceeds by performing updates ${\bf W}^{(k+1)}={\bf Y}^{(k)}(\lambda).$ Tagare suggests an Armijo-Wolfe search along the curve to adapt $\lambda$, but such a procedure would be expensive for neural network optimization since it requires multiple evaluations of the forward model and gradients. We found that simply using a fixed learning rate $\lambda$ often works well. Also, RMSprop-style scaling of the gradient ${\bf G}^{(k)}$ by a running average of the previous gradients’ norms before applying the multiplicative step (6) can improve convergence. The only additional substantial computation required beyond the forward and backward passes of the network is the $N\times N$ matrix inverse in (6).

Experiments

All models are implemented in Theano , based on the implementation of restricted-capacity uRNNs by , available from https://github.com/amarshah/complex_RNN. All code to replicate our results is available from https://github.com/stwisdom/urnn. All models use RMSprop for optimization, except that full-capacity uRNNs optimize their recurrence matrices with a fixed learning rate using the update step (6) and optional RMSprop-style gradient normalization.

First, we compare the performance of full-capacity uRNNs to restricted-capacity uRNNs and LSTMs on two tasks with synthetic data. The first task is synthetic system identification, where a uRNN must learn the dynamics of a target uRNN given only samples of the target uRNN’s inputs and outputs. The second task is the copy memory problem, in which the network must recall a sequence of data after a long period of time.

For the task of system identification, we consider the problem of learning the dynamics of a nonlinear dynamical system that has the form (1), given a dataset of inputs and outputs of the system. We will draw a true system ${\bf W}_{sys}$ randomly from either a constrained set $\mathcal{W}_{u}$ of restricted-capacity unitary matrices using the parameterization ${\bf W}_{u}(\theta_{u})$ in (3) or from a wider set $\mathcal{W}_{g}$ of restricted-capacity unitary matrices that are guaranteed to lie outside $\mathcal{W}_{u}$. We sample from $\mathcal{W}_{g}$ by taking a matrix product of two unitary matrices drawn from $\mathcal{W}_{u}$.

We use a sequence length of $T=150$, and we set the input dimension $M$ and output dimension $L$ both equal to the hidden state dimension $N$. The input-to-hidden transformation ${\bf V}$ and output-to-hidden transformation ${\bf U}$ are both set to identity, the output bias ${\bf c}$ is set to ${\bf 0}$, the initial state is set to ${\bf 0}$, and the hidden bias ${\bf b}$ is drawn from a uniform distribution in the range $[-0.11,-0.09]$. The hidden bias has a mean of $-0.1$ to ensure stability of the system outputs. Inputs are generated by sampling $T$-length i.i.d. sequences of zero-mean, diagonal and unit covariance circular complex-valued Gaussians of dimension $N$. The outputs are created by running the system (1) forward on the inputs.

We compare a restricted-capacity uRNN using the parameterization from (3) and a full-capacity uRNN using Stiefel manifold optimization with no gradient normalization as described in Section 4. We choose hidden state dimensions $N$ to test critical points predicted by our arguments in Section 3 of ${\bf W}_{u}(\theta_{u})$ in (3): $N\in\{4,6,7,8,16\}$. These dimensions are chosen to test below, at, and above the critical dimension of $7$.

For all experiments, the number of training, validation, and test sequences are $20000$, $1000$, and $1000$, respectively. Mean-squared error (MSE) is used as the loss function. The learning rate is $0.001$ with a batch size of $50$ for all experiments. Both models use the same matrix drawn from $\mathcal{W}_{u}$ as initialization. To isolate the effect of unitary recurrence matrix capacity, we only optimize ${\bf W}$, setting all other parameters to true oracle values. For each method, we report the best test loss over $100$ epochs and over $6$ random initializations for the optimization.

The results are shown in Table 1. “${\bf W}_{sys}$ init.” refers to the initialization of the true system unitary matrix ${\bf W}_{sys}$, which is sampled from either the restricted-capacity set $\mathcal{W}_{u}$ or the wider set $\mathcal{W}_{g}$.

Notice that for $N<7$, the restricted-capacity uRNN achieves comparable or better performance than the full-capacity uRNN. At $N=7$, the restricted-capacity and full-capacity uRNNs achieve relatively comparable performance, with the full-capacity uRNN achieving slightly lower error. For $N>7$, the full-capacity uRNN always achieves better performance versus the restricted-capacity uRNN. This result confirms our theoretical arguments that the restricted-capacity parameterization in (3) lacks the capacity to model all matrices in the unitary group for $N>7$ and indicates the advantage of using a full-capacity unitary recurrence matrix.

1.2 Copy memory problem

The experimental setup follows the copy memory problem from , which itself was based on the experiment from . We consider alternative hidden state dimensions and extend the sequence lengths to $T=1000$ and $T=2000$, which are longer than the maximum length of $T=750$ considered in previous literature.

In this task, the data is a vector of length $T+20$ and consists of elements from 10 categories. The vector begins with a sequence of 10 symbols sampled uniformly from categories 1 to 8. The next $T-1$ elements of the vector are the ninth ’blank’ category, followed by an element from the tenth category, the ‘delimiter’. The remaining ten elements are ‘blank’. The task is to output $T+10$ blank characters followed by the sequence from the beginning of the vector. We use average cross entropy as the training loss function. The baseline solution outputs the blank category for $T+10$ time steps and then guesses a random symbol uniformly from the first eight categories. This baseline has an expected average cross entropy of $\frac{10\log\left(8\right)}{T+20}$.

2 Speech data

We now apply restricted-capacity and full-capacity uRNNs to real-world speech data and compare their performance to LSTMs. The main task we consider is predicting the log-magnitude of future frames of a short-time Fourier transform (STFT). The STFT is a commonly used feature domain for speech enhancement, and is defined as the Fourier transform of short windowed frames of the time series. In the STFT domain, a real-valued audio signal is represented as a complex-valued $F\times T$ matrix composed of $T$ frames that are each composed of $F=N_{win}/2+1$ frequency bins, where $N_{win}$ is the duration of the time-domain frame. Most speech processing algorithms use the log-magnitude of the complex STFT values and reconstruct the processed audio signal using the phase of the original observations.

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$.We use the TIMIT dataset . According to common practice , we use a training set with 3690 utterances from 462 speakers, a validation set of 400 utterances, an evaluation set of 192 utterances. Training, validation, and evaluation sets have distinct speakers. Results are reported on the evaluation set using the network parameters that perform best on the validation set in terms of the loss function over three training trials. All TIMIT audio is resampled to $8$kHz. The STFT uses a Hann analysis window of $256$ samples ($32$ milliseconds) and a window hop of $128$ samples ($16$ milliseconds).

The LSTM requires gradient clipping during optimization, while the restricted-capacity and full-capacity uRNNs do not. The hidden state dimensions $N$ of the LSTM are chosen to match the number of parameters of the full-capacity uRNN. For the restricted-capacity uRNN, we run models that match either $N$ or number of parameters. For the LSTM and restricted-capacity uRNNs, we use RMSprop with a learning rate of $0.001$, momentum $0.9$, and averaging parameter $0.1$. For the full-capacity uRNN, we also use RMSprop to optimize all network parameters, except for the recurrence matrix, for which we use stochastic gradient descent along the Stiefel manifold using the update (6) with a fixed learning rate of $0.001$ and no gradient normalization.

Results are shown in Table 5.2, and Figure 2 shows example predictions of the three types of networks. Results in Table 5.2 are given in terms of the mean-squared error (MSE) loss function and several metrics computed on the time-domain signals, which are reconstructed from the predicted log-magnitude and the original phase of the STFT. These time-domain metrics are segmental signal-to-noise ratio (SegSNR), short-time objective intelligibility (STOI), and perceptual evaluation of speech quality (PESQ). SegSNR, computed using , uses a voice activity detector to avoid measuring SNR in silent frames. STOI is designed to correlate well with human intelligibility of speech, and takes on values between 0 and 1, with a higher score indicating higher intelligibility . PESQ is the ITU-T standard for telephone voice quality testing , and is a popular perceptual quality metric for speech enhancement . PESQ ranges from 1 (bad quality) to 4.5 (no distortion).

Note that full-capacity uRNNs generally perform better than restricted-capacity uRNNs with the same number of parameters, and both types of uRNN significantly outperform LSTMs.

As another challenging long-term memory task with natural data, we test the performance of LSTMs and uRNNs on pixel-by-pixel MNIST and permuted pixel-by-pixel MNIST, first proposed by and used by to test restricted-capacity uRNNs. For permuted pixel-by-pixel MNIST, the pixels are shuffled, thereby creating some non-local dependencies between pixels in an image. Since the MNIST images are $28\times 28$ pixels, resulting pixel-by-pixel sequences are $T=784$ elements long. We use 5000 of the 60000 training examples as a validation set to perform early stopping with a patience of 5. The loss function is cross-entropy. Weights with the best validation loss are used to process the evaluation set. The full-capacity uRNN uses RMSprop-style gradient normalization.

Learning curves are shown in Figure 3, and a summary of classification accuracies is shown in Table 3. For the unpermuted task, the LSTM with $N=256$ achieves the best evaluation accuracy of $98.2\%$. For the permuted task, the full-capacity uRNN with $N=512$ achieves the best evaluation accuracy of $94.1\%$, which is state-of-the-art on this task. Both uRNNs outperform LSTMs on the permuted case, achieving their best performance after fewer traing epochs and using an equal or lesser number of trainable parameters. This performance difference suggests that LSTMs are only able to model local dependencies, while uRNNs have superior long-term memory capabilities. Despite not representing all unitary matrices, the restricted-capacity uRNN with $N=512$ still achieves impressive test accuracy of $93.3\%$ with only $1/16$ of the trainable parameters, outperforming the full-capacity uRNN with $N=116$ that matches number of parameters. This result suggests that further exploration into the potential trade-off between hidden state dimension $N$ and capacity of unitary parameterizations is necessary.

Unitary recurrent matrices prove to be an effective means of addressing the vanishing and exploding gradient problems. We provided a theoretical argument to quantify the capacity of constrained unitary matrices. We also described a method for directly optimizing a full-capacity unitary matrix by constraining the gradient to lie in the differentiable manifold of unitary matrices. The effect of restricting the capacity of the unitary weight matrix was tested on system identification and memory tasks, in which full-capacity unitary recurrent neural networks (uRNNs) outperformed restricted-capacity uRNNs from as well as LSTMs. Full-capacity uRNNs also outperformed restricted-capacity uRNNs on log-magnitude STFT prediction of natural speech signals and classification of permuted pixel-by-pixel images of handwritten digits, and both types of uRNN significantly outperformed LSTMs. In future work, we plan to explore more general forms of restricted-capacity unitary matrices, including constructions based on products of elementary unitary matrices such as Householder operators or Givens operators.

Acknowledgments: We thank an anonymous reviewer for suggesting improvements to our proof in Section 3 and Vamsi Potluru for helpful discussions. Scott Wisdom and Thomas Powers were funded by U.S. ONR contract number N00014-12-G-0078, delivery orders 13 and 24. Les Atlas was funded by U.S. ARO grant W911NF-15-1-0450.