Efficient Orthogonal Parametrisation of Recurrent Neural Networks Using Householder Reflections

Introduction

Recurrent Neural Networks (RNNs) have been successfully used in many applications involving time series. This is because RNNs are well suited for sequential data as they process inputs one element at a time and store relevant information in their hidden state. In practice, however, training simple RNNs (sRNN) can be challenging due to the problem of exploding and vanishing gradients (Hochreiter et al., 2001). It has been shown that exploding gradients can occur when the transition matrix of an RNN has a spectral norm larger than one (Glorot & Bengio, 2010). This results in an error surface, associated with some objective function, having very steep walls (Pascanu et al., 2013). On the other hand, when the spectral norm of the transition matrix is less than one, the information at one time step tend to vanish quickly after a few time steps. This makes it challenging to learn long-term dependencies in sequential data.

Different methods have been suggested to solve either the vanishing or exploding gradient problem. The LSTM has been specifically designed to help with the vanishing gradient (Hochreiter & Schmidhuber, 1997). This is achieved by using gate vectors which allow a linear flow of information through the hidden state. However, the LSTM does not directly address the exploding gradient problem. One approach to solving this issue is to clip the gradients (Mikolov, 2012) when their norm exceeds some threshold value. However, this adds an extra hyperparameter to the model. Furthermore, if exploding gradients can occur within some parameter search space, the associated error surface will still have steep walls. This can make training challenging even with gradient clipping.

Another way to approach this problem is to improve the shape of the error surface directly by making it smoother, which can be achieved by constraining the spectral norm of the transition matrix to be less than or equal to one. However, a value of exactly one is best for the vanishing gradient problem. A good choice of the activation function between hidden states is also crucial in this case. These ideas have been investigated in recent works. In particular, the unitary RNN (Arjovsky et al., 2016) uses a special parametrisation to constrain the transition matrix to be unitary, and hence, of norm one. This parametrisation and other similar ones (Hyland & Rätsch, 2017; Wisdom et al., 2016) have some advantages and drawbacks which we will discuss in more details in the next section.

The main contributions of this work are as follows:

We first show that constraining the search space of the transition matrix of an RNN to the set of unitary matrices $\mathbf{U}(n)$ is equivalent to limiting the search space to a subset of $\mathbf{O}(2n)$ ($\mathbf{O}(2n)$ is the set of $2n\times 2n$ orthogonal matrices) of a new RNN with twice the hidden size. This suggests that it may not be necessary to work with complex matrices.

We present a simple way to parametrise orthogonal transition matrices of RNNs using Householder matrices, and we derive the expressions of the back-propagated gradients with respect to the new parametrisation. This new parametrisation can also be used in other deep architectures.

We develop an algorithm to compute the back-propagated gradients efficiently. Using this algorithm, we show that the worst case time complexity of one gradient step is of the same order as that of the sRNN.

Related Work

Throughout this work we will refer to elements of the following sRNN architecture.

where $W$, $V$ and $Y$ are the hidden-to-hidden, input-to-hidden, and hidden-to-output weight matrices. $h^{(t-1)}$ and $h^{(t)}$ are the hidden vectors at time steps $t-1$ and $t$ respectively. Finally, $\phi$ is a non-linear activation function. We have omitted the bias terms for simplicity.

Recent research explored how the initialisation of the transition matrix $W$ influences training and the ability to learn long-term dependencies. In particular, initialisation with the identity or an orthogonal matrix can greatly improve performance (Le et al., 2015). In addition to these initialisation methods, one study also considered removing the non-linearity between the hidden-to-hidden connections (Henaff et al., 2016), i.e. the term $Wh^{(t-1)}$ in Equation (1) is outside the activation function $\phi$. This method showed good results when compared to the LSTM on pathological problems exhibiting long-term dependencies.

Another interesting parametrisation (Hyland & Rätsch, 2017) has been suggested which takes advantage of the algebraic properties of the unitary group $\mathbf{U}(n)$. The idea is to use the corresponding matrix Lie algebra $\mathbf{u}(n)$ of skew hermitian matrices. In particular, the transition matrix can be written as $W=\exp\left[\sum_{i=1}^{n^{2}}\lambda_{i}T_{i}\right]$, where $\exp$ is the exponential matrix map and $\{T_{i}\}^{n^{2}}_{i=1}$ are predefined $n\times n$ matrices forming a bases of the Lie algebra $\mathbf{u}(n)$. The learning parameters are the weights $\{\lambda_{i}\}$. The fact that the matrix Lie algebra $\mathbf{u}(n)$ is closed and connected ensures that the exponential mapping from $\mathbf{u}(n)$ to $\mathbf{U}(n)$ is surjective. Therefore, with this parametrisation the search space of the transition matrix spans the whole unitary group. This is one advantage over the original unitary parametrisation (Arjovsky et al., 2016). However, the cost of computing the matrix exponential to get $W$ is $\mathcal{O}(n^{3})$, where $n$ is the dimension of the hidden state. .

Another method (Wisdom et al., 2016) performs optimisation directly of the Stiefel manifold using the Cayley transformation. The corresponding model was called full-capacity unitary RNN. Using this approach, the transition matrix can span the full set of unitary matrices. However, this method involves a matrix inverse as well as matrix-matrix products which have time complexity $\mathcal{O}(n^{3})$. This can be problematic for large neural networks when using stochastic gradient descent with a small mini-batch size.

A more recent study (Vorontsov et al., 2017) investigated the effect of soft versus hard orthogonal constraints on the performance of RNNs. The soft constraint was applied by specifying an allowable range for the maximum singular value of the transition matrix. To this end, the transition matrix was factorised as $W=USV^{\prime}$, where $U$ and $V$ are orthogonal matrices and $S$ is a diagonal matrix containing the singular values of $W$. A soft orthogonal constraint consists of specifying small allowable intervals around 1 for the diagonal elements of $S$. Similarly to (Wisdom et al., 2016), the matrices $U$ and $V$ were updated at each training iteration using the Cayley transformation, which involves a matrix inverse, to ensure that they remain orthogonal.

Complex unitary versus orthogonal

Assuming that the activation function $\phi$ applies to the real and imaginary parts separately, it is easy to show that the update equation of the complex hidden state $h^{(t)}$ of the unitary RNN has the following real space representation

Even when the activation function $\phi$ does not apply to the real and imaginary parts separately, it is still possible to find an equivalent representation in the real space. Consider the activation function proposed by (Arjovsky et al., 2016)

Now we will show that the matrix $\hat{W}$ is orthogonal. By definition of a unitary matrix, we have $WW^{*}=I$ where the ∗ represents the conjugate transpose. This implies that $AA^{\prime}+BB^{\prime}=I$ and $BA^{\prime}-AB^{\prime}=0$. And since we have

it follows that $\hat{W}\hat{W}^{\prime}=I$. Also note that $\hat{W}$ has a special structure - it is a block-matrix.

Parametrisation of the transition matrix

Note that $\mathcal{H}_{1}(u)$ is not necessarily a Householder reflection. However, when $u\in\{1,-1\}$, $\mathcal{H}_{1}(u)$ is orthogonal.

We propose to parametrise the transition matrix $W$ of an RNN using the mappings $\{\mathcal{M}_{k}\}$. When using $m$ reflection vectors $\{\mathbf{u}_{i}\}$, the parametrisation can be expressed as

For the particular case where $m=n$ in the above parametrisation, we have the following result.

Note that Theorem 1 would not be valid if $\mathcal{H}_{1}(\cdot)$ was a standard Householder reflection. In fact, in the two-dimensional case, for instance, the matrix $\left(\begin{smallmatrix}1&0\\ 0&-1\end{smallmatrix}\right)$ cannot be expressed as the product of exactly two standard Householder matrices.

The parametrisation in (8) has the following advantages:

The parametrisation is smoothexcept on a subset of zero Lebesgue measure., which is convenient for training with gradient descent. It is also flexible - a good trade-off between expressiveness and speed can be found by tuning the number of reflection vectors.

The time and space complexities involved in one gradient calculation are, in the worst case, the same as that of the sRNN with the same number of hidden units. This is discussed in the following subsections.

When $m<n$, the matrix $W$ is always orthogonal, as long as the reflection vectors are nonzero. For $m=n$, the only additional requirement for $W$ to be orthogonal is that $u_{1}\in\{-1,1\}$.

When $m=n$, the transition matrix can span the whole set of $n\times n$ orthogonal matrices. In this case, the total number of parameters needed for $W$ is $n(n+1)/2$. This results in only $n$ redundant parameters since the orthogonal set $\mathbf{O}(n)$ is a $n(n-1)/2$ manifold.

Let $\mathcal{L}$ be a scalar loss function and $C^{(t)}\coloneqq Wh^{(t-1)}$, where $W$ is constructed using the $\{\mathbf{u}_{i}\}$ vectors following Equation (8). In order to back-propagate the gradients through time, we need to compute the following partial derivatives

at each time step $t$. Note that in Equation (12) $h^{(t-1)}$ is taken as a constant with respect to $U$. Furthermore, we have $\frac{\partial\mathcal{L}}{\partial U}=\sum_{t=1}^{T}\frac{\partial\mathcal{L}}{\partial U^{(t)}}$, where $T$ is the length of the input sequence. The gradient flow through the RNN at time step $t$ is shown in Figure 1.

Before describing the algorithm to compute the back-propagated gradients $\frac{\partial\mathcal{L}}{\partial U^{(t)}}$ and $\frac{\partial\mathcal{L}}{\partial h^{(t-1)}}$, we first derive their expressions as a function of $U$, $h^{(t-1)}$ and $\frac{\partial\mathcal{L}}{\partial C^{(t)}}$ using the compact WY representation (Joffrain et al., 2006) of the product of Householder reflections.

where $\mbox{striu}(U^{\prime}U)$, and $\mbox{diag}(U^{\prime}U)$ represent the strictly upper part and the diagonal of the matrix $U^{\prime}U$, respectively.

Equation (15) is the compact WY representation of the product of Householder reflections. For the particular case where $m=n$, the RHS of Equation (15) should be replaced by $\left(I-UT^{-1}U^{\prime}\right)\mathcal{H}_{1}(u_{1})$, where $\mathcal{H}_{1}$ is defined in (7) and $U=(\mathbf{u}_{n}|\dots|\mathbf{u}_{2})$.

The following theorem gives the expressions of the gradients $\frac{\partial\mathcal{L}}{\partial U^{(t)}}$ and $\frac{\partial\mathcal{L}}{\partial h^{(t-1)}}$ when $m\leq n-1$ and $h=h^{(t-1)}$.

The proof of Equations (16) and (17) is provided in Appendix A. Based on Theorem 2, Algorithm 1 performs the one-step forward-propagation (FP) and back-propagation (BP) required to compute $C^{(t)}$, $\frac{\partial\mathcal{L}}{\partial U^{(t)}}$, and $\frac{\partial\mathcal{L}}{\partial h^{(t-1)}}$. See Appendix B for more detail about how this algorithm is derived using Theorem 2.

In the next section we analyse the time and space complexities of this algorithm.

2 Time and Space complexity

At each time step $t$, the flop count required by Algorithm 1 is $(13n+2)m$; $6nm$ for the one-step FP and $(7n+2)m$ for the one-step BP. Note that the vector $N$ only needs to be calculated at one time step. This reduces the flop count at the remaining time steps to $(11n+3)m$. The fact that the matrix $U$ has all zeros in its upper triangular part can be used to further reduce the total flop count to $(11n-3m+5)m$; $(4n-m+2)m$ for the one-step FP and $(7n-2m+3)m$ for the one-step BP. See Appendix C for more details.

Note that if the values of the matrices $H$, defined in Algorithm 1, are first stored during a “global” FP (i.e. through all time steps), then used in the BP steps, the time complexityWe considered only the time complexity due to computations through the hidden-to-hidden connections of the network. for a global FP and BP using one input sequence of length $T$ are, respectively, $\approx 3n^{2}T$ and $\approx 5n^{2}T$, when $m\approx n$ and $n\gg 1$. In contrast with the sRNN case with $n$ hidden units, the global FP and BP have time complexities $\approx 2n^{2}T$ and $\approx 3n^{2}T$. Hence, when $m\approx n$, the FP and BP steps using our parametrisation require only about twice more flops than the sRNN case with the same number of hidden units.

Note, however, that storing the values of the matrices $H$ at all time steps requires the storage of $mnT$ values for one sequence of length $T$, compared with $nT$ when only the hidden states $\{h^{(t)}\}_{t=1}$ are stored. When $m\gg 1$ this may not be practical. One solution to this problem is to generate the matrices $H$ locally at each BP step using $U$ and $h^{(t-1)}$. This results in a global BP complexity of $(11n-3m+5)mT$. Table 2 summarises the flop counts for the FP and BP steps. Note that these flop counts are for the case when $m\leq n-1$. When $m=n$, the complexity added due to the multiplication by $\mathcal{H}_{1}(u_{1})$ is negligible.

3 Extension to the Unitary case

Although we decided to focus on the set of real-valued orthogonal matrices, for the reasons given in Section 3, our parametrisation can readily be modified to apply to the general unitary case.

Let $\hat{\mathcal{M}}_{1}$ be the mapping defined as

The image of $\hat{\mathcal{M}}_{1}$ spans the full set of unitary matrices $\mathbf{U}(n)$ and any point on its image is a unitary matrix.

Experiments

All RNN models were implemented using the python library theano (Theano Development Team, 2016). For efficiency, we implemented the one-step FP and BP algorithms described in Algorithm 1 using C codeOur implementation can be found at https://github.com/zmhammedi/Orthogonal_RNN.. We tested the new parametrisation on five different datasets all having long-term dependencies. We call our parametrised network oRNN (for orthogonal RNN). We set its activation function to the ky_ReLU defined as $\phi(x)=\max(\frac{x}{10},x)$. To ensure that the transition matrix of the oRNN is always orthogonal, we set the scalar $u_{1}$ to -1 if $u_{1}\leq 0$ and 1 otherwise after each gradient update. Note that the parameter matrix $U$ in Equation (11) has all zeros in its upper triangular part. Therefore, after calculating the gradient of a loss with respect to $U$ (i.e. $\frac{\partial\mathcal{L}}{\partial U}$), the values in the upper triangular part are set to zero.

For all experiments, we used the adam method for stochastic gradient descent (Kingma & Ba, 2014). We initialised all the parameters using uniform distributions similar to (Arjovsky et al., 2016). The biases of all models were set to zero, except for the forget bias of the LSTM, which we set to 5 to facilitate the learning of long-term dependencies (Koutník et al., 2014).

In this experiment, we followed a similar setting to (Koutník et al., 2014) where we trained RNNs to encode song excerpts. We used the track Manyrista from album Musica Deposita by Cuprum. We extracted five consecutive excerpts around the beginning of the song, each having $800$ data points and corresponding to 18ms with a 44.1Hz sampling frequency. We trained an sRNN, LSTM, and oRNN for 5000 epochs on each of the pieces with five random seeds. For each run, the lowest Normalised Mean Squared Error (NMSE) during the 5000 epochs was recorded. For each model, we tested three different hidden sizes. The total number of parameters $N_{p}$ corresponding to these hidden sizes was approximately equal to $250$, $500$, and $1000$. For the oRNN, we set the number of reflection vectors to the hidden size for each case, so that the transition matrix is allowed to span the full set of orthogonal matrices. The results are shown in Figures 2 and 3. All the learning rates were set to $10^{-3}$. The orthogonal parametrisation outperformed the sRNN and performed on average better than the LSTM.

2 Addition Task

In this experiment, we followed a similar setting to (Arjovsky et al., 2016), where the goal of the RNN is to output the sum of two elements in the first dimension of a two-dimensional sequence. The location of the two elements to be summed are specified by the entries in the second dimension of the input sequence. In particular, the first dimension of every input sequence consists of random numbers between 0 and 1. The second dimension has all zeros except for two elements equal to 1. The first unit entry is located in the first half of the sequence, and the second one in the second half. We tested two different sequence lengths $T=400,800$. All models were trained to minimise the Mean Squared Error (MSE). The baseline MSE for this task is 0.167; for a model that always outputs one.

We trained an oRNN with $n=128$ hidden units and $m=16$ reflections. We trained an LSTM and sRNN with hidden sizes 28 and 54, respectively, corresponding to a total number of parameters $\approx 3600$ (i.e. same as the oRNN model). We chose a batch size of 50, and after each iteration, a new set of sequences was generated randomly. The learning rate for the oRNN was set to 0.01. Figure 4 displays the results for both lags.

The oRNN was able to beat the baseline MSE in less than 5000 iterations for both lags and for two different random initialisation seeds. This is in line with the results of the unitary RNN (Arjovsky et al., 2016).

3 Pixel MNIST

In this experiment, we used the MNIST image dataset. We split the dataset into training (55000 instances), validation (5000 instances), and test sets (10000 instances). We trained oRNNs with $n\in\{128,256\}$ and $m\in\{16,32,64\}$, where $n$ and $m$ are the number of hidden units and reflections vectors respectively, to minimise the cross-entropy error function. We experimented with (mini-batch size, learning rate) $\in\{(1,10^{-4}),(50,10^{-3})\}$.

Table 3 compares the test performance of our best model against results available in the literature for unitary/orthogonal RNNs. Despite having fewer total number of parameters, our model performed better than three out the four models selected for comparison (all having $\geq 16K$ parameters). Figure 5 shows the validation accuracy as a function of the number of epochs of our oRNN model in Table 3. Figure 6 shows the effect of varying the number of reflection vectors $m$ on the performance.

4 Penn Tree Bank

In this experiment, we tested the oRNN on the task of character level prediction using the Penn Tree Bank Corpus. The data was split into training (5017K characters), validation (393K characters), and test sets (442K characters). The total number of unique characters in the corpus was 49. The vocabulary size was 10K and any other words were replaced by the special token k>. The number of characters per instance (i.e. char/line) in the training data ranged between 2 and 518 with an average of 118 char/line. We trained an oRNN and LSTM with hidden units 512 and 183 respectively, corresponding to a total of $\approx 180$K parameters, for 20 epochs. We set the number of reflections to 510 for the oRNN. The learning rate was set to 0.0001 for both models with a mini-batch size of 1.

Similarly to (Pascanu et al., 2013) we considered two tasks: one where the model predicts one character ahead and the other where it predicts a character five steps ahead. It was suggested that solving the later task would require the learning of longer term correlations in the data rather than the shorter ones. Table 4 summarises the test results. The oRNN and LSTM performed similarly to each other on the one-step head prediction task. Whereas on the five-step ahead prediction task, the LSTM was better. The performance of both models on this task was close to the state of the art result for RNNs $3.74$ bpc (Pascanu et al., 2013).

Nevertheless, our oRNN still outperformed the results of (Vorontsov et al., 2017) which used both soft and hard orthogonality constraints on the transition matrix. Their RNN was trained on 99% of the data (sentences with $\leq$ 300 characters) and had the same number of hidden units as the oRNN in our experiment. The lowest test cost achieved was 2.20(bpc) for the one-step-ahead prediction task.

5 Copying task

We tested our model on the copy task described in details in (Gers et al., 2001; Arjovsky et al., 2016). Using an oRNN with the ky_ReLU we were not able to reproduce the same performance as the uRNN (Arjovsky et al., 2016; Wisdom et al., 2016). However, we were able to achieve a comparable performance when using the U activation function (Chernodub & Nowicki, 2016), which is a norm-preserving activation function. In order to explore whether the poor performance of the oRNN was only due to the activation function, we tested the same activation as the uRNN (i.e. the real representation of ReLU defined in Equation (4)) on the oRNN. This did not improve the performance compared to the ky_ReLU case suggesting that the block structure of the uRNN transition matrix, when expressed in the real space (see Section 3), may confer special benefits in some cases.

Discussion

In this work, we presented a new parametrisation of the transition matrix of a recurrent neural network using Householder reflections. This method allows an easy and computationally efficient way to enforce an orthogonal constraint on the transition matrix which then ensures that exploding gradients do not occur during training. Our method could also be applied to other deep neural architectures to enforce orthogonality between hidden layers. Note that a “soft” orthogonal constraint could also be applied using our parametrisation by, for example, allowing $u_{1}$ to vary continuously between -1 and 1.

It is important to note that our method is particularly advantageous for stochastic gradient descent when the mini-batch size is close to 1. In fact, if $B$ is the mini-batch size and $T$ is the average length of the input sequences, then a network with $n$ hidden units trained using other methods (Vorontsov et al., 2017; Wisdom et al., 2016; Hyland & Rätsch, 2017) that enforce orthogonality (see Section 2), would have time complexity $\mathcal{O}(BTn^{2}+n^{3})$. Clearly when $BT\gg n$ this becomes $\mathcal{O}(BTn^{2})$, which is the same time complexity as that of the sRNN and oRNN (with $m=n$). In contrast with the case of fully connected deep forward networks with no weight sharing between layers (# layer $=L$), the time complexity using our method is $\mathcal{O}(BLnm)$ whereas other methods discussed in this work (see Section 2) would have time complexity $\mathcal{O}(BLn^{2}+Ln^{3})$. The latter methods are less efficient in this case since $B\gg n$ is less likely to be the case compared with $BT\gg n$ when using SGD.

From a performance point of view, further experiments should be performed to better understand the difference between the unitary versus orthogonal constraint.

Acknowledgment

The authors would like to acknowledge Department of State Growth Tasmania for partially funding this work through SenseT. We would also like to thank Christfried Webers for his valuable feedback.

References

Appendix A Proofs

(giles2008extended) Let $A$, $B$, and $C$ be real or complex matrices, such that $C=f(A,B)$ where $f$ is some differentiable mapping. Let $\mathcal{L}$ be some scalar quantity which depends on $C$. Then we have the following identity

where $dA$, $dB$, and $dC$ represent infinitesimal perturbations and

where $B=\mbox{striu}(J_{m})+\frac{1}{2}I_{m}$ and $J_{m}$ is the $m\times m$ matrix of all ones.

Calculating the infinitesimal perturbations of $C$ gives

By substituting this back into the expression of $dC$, multiplying the left and right-hand sides by $\overline{C}^{\prime}$, and applying the trace we get

Now using the identity $\mbox{Tr}(AB)=\mbox{Tr}(BA)$, where the second dimension of A agrees with the first dimension of B, we can rearrange the expression of $\mbox{Tr}(\overline{C}^{\prime}dC)$ as follows

To simplify the expression, we will use the short notations

$\mbox{Tr}(\overline{C}^{\prime}dC)$ becomes

Now using the two following identities of the trace

we can rewrite $\mbox{Tr}(\overline{C}^{\prime}dC)$ as follows

By rearranging and taking the transpose of the third and fourth term of the right-hand side we obtain

Taking this fact into account, a similar argument to that used in the proof of Theorem 1 can be used here. ∎

Appendix B Algorithm Explanation

Let $U\coloneqq(v_{i,j})_{\begin{subarray}{c}1\leq i\leq n\end{subarray}}^{1\leq j\leq m}$. Then the element of the matrix $T\coloneqq\mbox{striu}(U^{\prime}U)+\frac{1}{2}\mbox{diag}(U^{\prime}U)$ can be expressed as

where $\delta_{i,j}$ is the Kronecker delta and $\llbracket\cdot\rrbracket$ is the Iversion bracket (i.e. $\llbracket p\rrbracket$ = 1 if p is true and $\llbracket p\rrbracket$ = 0 otherwise).

Finally, note that from Equation (16), we have for $1\leq i\leq n$ and $1\leq k\leq m$

Therefore, when $C=C^{(t)}$ and $h=h^{(t-1)}$ we have

Appendix C Time complexity

Table 5 shows the flop count for different operations in the local backward and forward propagation steps in Algorithm 1.

Appendix D Matlab implementation of Algorithm 1