Fast Saturating Gate for Learning Long Time Scales with
Recurrent Neural Networks

In this section, we review gated RNNs and their time scale interpretation (Tallec and Ollivier 2018). We begin with an LSTM (Hochreiter and Schmidhuber 1997), which is one of the most popular RNNs. An LSTM has a memory cell $c_{t}\in\mathbb{R}^{n}$ and hidden state $h_{t}\in\mathbb{R}^{n}$ inside, which are updated depending on the sequential input data $x_{t}$ at each time step $t=1,2,\cdots$ by

where $W_{*},U_{*}$ and $b_{*}$ are weight and bias parameters for each $*\in\{f,i,c,o\}$. The sigmoid function $\sigma$ is defined as

$f_{t},i_{t},o_{t}\in[0,1]^{n}$ are called forget, input, and output gates, respectively. They were initially motivated as a binary mechanism, i.e., switching on and off, allowing information to pass through (Gers, Schmidhuber, and Cummins 2000). The forget gate has been reinterpreted as the representation for time scales of memory cells (Tallec and Ollivier 2018). Following that study, we simplify Eq. (1) by assuming $\tilde{c}_{t}=0$ for an interval $t\in[t_{0},t_{1}]$. Then, we obtain

where $\bar{f}=(\prod_{s=t_{0}+1}^{t_{1}}f_{s})^{\frac{1}{t_{1}-t_{0}}}$ is the (entry-wise) geometric mean of the values of the forget gate. Through Eq. (8), the memory cell $c_{t}$ loses its information on data up to time $t_{0}$ exponentially, and the entry of $\bar{f}$ represents its (averaged) decay rate. This indicates that, in order to capture long-term dependencies of the sequential data, the forget gate is desired to take values near 1 on average. We refer the associated time constant¹¹1An exponential function $F(t)=e^{-\alpha t}$ of time $t$ decreases by a factor of $1/e$ in time $T=1/\alpha$, which is called the time constant. $T=-1/\log{\bar{f}}$ as the time scale of units, which has been empirically shown to illustrate well the temporal behavior of LSTMs (Mahto et al. 2021).

The above argument applies not only to an LSTM, but also to general gated RNNs including a GRU (Cho et al. 2014) with state update of the form

where $h_{t},f_{t},i_{t}$ denotes the state, forget gate, and input gate, respectively, and $\tilde{h}_{t}$ is the activation to represent new information at time $t$. Here again, the forget gate $f_{t}$ takes a role to control the time scale of each unit of the state.

The sigmoid function $\sigma(z)$ in the gating mechanism requires large $z$ to take a value near 1 as the output. On the other hand, the derivative $\sigma^{\prime}(z)$ takes exponentially small values for $z\gg 0$ (Fig. 1). Thus, when a gated model needs to learn large gate values such as $0.99$ with gradient methods, parameters in the gate cannot be effectively updated due to gradient vanishing. This is called saturation of bounded activation functions (Gulcehre et al. 2016). The behavior of gate functions on the saturating regime is important for gated RNNs because forget gate values need to be large to represent long time scales as explained in Section 2.1. That is, gated RNNs must face saturation of the forget gate to learn long time scales. Thus, it is hypothesized that saturation causes difficulty in training gated RNNs for data with extremely long time scales (Chandar et al. 2019; Gu et al. 2020b).

Recall Eq. (8), which describes the time scales of the memory cell $c_{t}$ of an LSTM via exponential decay. Let the memory cell at time $t_{1}$ be $c_{t_{1}}=\lambda c_{t_{0}}$ with $\lambda\in[0,1]$. Requiring long time scales corresponds to getting $\lambda$ close to 1. Therefore, we can consider a long-time-scale learning problem as minimizing a loss function $L$ that measures discrepancy of $c_{t_{1}}$ and $\lambda_{*}c_{t_{0}}$ where $\lambda_{*}\in[0,1]$ is a desired value close to 1. We take $L$ as the absolute loss for example. Then, we obtain

using Eq. (8). Let $z_{t}=W_{f}x_{t}+U_{f}h_{t-1}+b_{f}$, so that $f_{t}=\sigma(z_{t})$. Since we are interested in the averaged value of $f_{t}$, we consider $z_{t}$ to be time-independent, that is, $z_{t}=z$ in the same way as Tallec and Ollivier (2018). The problem is then reduced to a problem to obtain $z$ that minimizes

We consider this as the minimal problem to analyze the learnability of the forget gate for long time scales. Note that since the product $c_{t_{1}}=\bar{f}^{t_{1}-t_{0}}c_{t_{0}}$ is taken element-wise, we can consider this as a one-dimensional problem. Furthermore, the global solution can be explicitly written as $z=\sigma^{-1}(\lambda_{*}^{1/(t_{1}-t_{0})})$ where $\sigma^{-1}$ is an inverse of $\sigma$.

Next, we consider the learning dynamics of the model on the aforementioned problem Eq. (14). RNNs are usually trained with gradient methods. Learning dynamics with gradient methods can be analyzed considering learning rate $\to 0$ limit known as gradient flow (Harold and George 2003). Therefore, we consider the following gradient flow

using the loss function introduced above. Here, $\tau$ denotes a time variable for learning dynamics, which should not be confused with $t$ representing the state transition. Our aim is to investigate the convergence rate of a solution of the differential equation Eq. (15) when $\sigma$ in the forget gate is replaced with another function $\phi$.

To investigate the effect of choice of gate functions on the convergence rate, we first define the candidate set $\mathcal{F}$ of bounded functions for the gate function.

$\mathcal{F}$ is a natural class of gating activation functions including $\sigma$. As we explained in Section 2.2, gated RNNs suffer from gradient vanishing due to saturation when learning long time scales. To clarify the issue, we first show that saturation is inevitable regardless of the choice of $\phi\in\mathcal{F}$.

Nevertheless, choices of $\phi$ significantly affect the efficiency of the training. When the target $\lambda_{*}$ takes an extreme value near boundaries, the efficiency of training should depend on the asymptotic behavior of $\phi(z)$ for $z\gg 0$, that is, the rate at which $\phi(z)$ converges as $z\to\infty$. We call the convergence rate of $\phi(z)$ as $z\to\infty$ as the order of saturation. More precisely, we define the notion as follows²²2Our definition for asymptotic order is slightly different from the usual one which adopts $\limsup_{z\to\infty}\frac{g(z)}{1-\phi(z)}<\infty$, since it is more suitable for analyzing training efficiency.:

Intuitively, the order of saturation of $O(g(z))$ means that the convergence rate of $\phi$ to 1 is bounded by the decay rate of $g$ up to constant multiplication of $z$. For example, the sigmoid function $\sigma$ satisfies $e^{-az}/(1-\sigma(z))\to 0$ as $z\to\infty$ for any $a>1$, thus has the exponential order of saturation $O(e^{-z})$. The convexity condition for a higher order of saturation is rather technical, but automatically satisfied for typical functions, see Appendix C.2. If $\phi$ has a higher order of saturation (or saturates faster) than another function $\tilde{\phi}$, then $\phi(z)$ converges faster than $\tilde{\phi}(z)$ as $z\to\infty$, and $\phi^{\prime}(z)$ becomes smaller than $\tilde{\phi}^{\prime}(z)$. In this sense, training with $\tilde{\phi}$ seems more efficient than $\phi$ in the above problem. However, this is not the case as we discuss in the next section.

To precisely analyze learning behavior, we trace the learning dynamics of the output value $f=\phi(z)$ since our purpose is to obtain the desired output value rather than the input $z$. We transform the learning dynamics (Eq. (15)) into that of $f$ by

To treat Eq. (16) as purely of $f$, we define a function $g_{\phi}(f)$ of $f$ by $g_{\phi}(f):=\phi^{\prime}(\phi^{-1}(f)),$ so that Eq. (16) becomes

Our interest is in the dynamics of $f$ near the boundary, i.e., the limit of $f\to 1$. We have the following result:

Theorem 4.4 indicates that a higher order of saturation accelerates the move of the output $f$ near boundaries in accordance with Eq. (17) since $g_{\phi}(f)$ takes larger values. Thus, contrarily to the intuition in Section 4.2, a higher order of saturation leads to more efficient training for target values near boundaries. We demonstrate this effect using two activation functions, the sigmoid function $\sigma(z)$ and normalized softsign function $\sigma_{\rm ns}(z)=(\mathrm{softsign}(z/2)+1)/2$ where $\mathrm{softsign}(z)=z/(1+|z|)$. $\sigma_{\rm ns}$ is the softsign function modified so that $0\leq\sigma_{\rm ns}(z)\leq 1$ and $\sigma_{\rm ns}^{\prime}(0)=\sigma^{\prime}(0)$. $\sigma$ has a higher order of saturation than $\sigma_{\rm ns}$ since $\sigma$ has the order of saturation of $O(e^{-z})$ and $\sigma_{\rm ns}$ has $O(z^{-1})$ (see Fig. 1). We plot the learning dynamics of gradient flow for the problem in Fig. 2. Since $\sigma$ has a higher order of saturation than $\sigma_{\rm ns}$, the gate value $f$ of $\sigma_{\rm ns}$ converges slower to the boundary. Fig. 2 also shows the dynamics of gradient descent with the learning rate 1. While gradient descent is a discrete approximation of gradient flow, it behaves similar to gradient flow.

Explicit convergence rates. Beyond Theorem 4.4, we can explicitly calculate effective bounds of the convergence rate for the problem when the activation function is the sigmoid function $\sigma(z)$ or normalized softsign function $\sigma_{\rm ns}(z)$.

Proposition 4.5 shows the quantitative effect of difference in the order of saturation on the convergence rates. We fit the bounds to the learning curves with the gradient flow in Fig. 2. The convergence rates of the learning are well approximated by the bounds. These asymptotic analysis highlights that choices of the function $\phi$ significantly affects efficiency of training for long time scales.

We consider modification of the usual sigmoid function to another function $\phi\in\mathcal{F}$ for the forget gate in a gated RNN. Function $\phi$ should satisfy the following conditions.

1. (i)

   $\phi$ has a higher order of saturation than $\sigma$,

2. (ii)

   $\phi(z)\approx\sigma(z)$ for $z\approx 0$,

3. (iii)

   $\phi$ is symmetric in a sense that $\phi(-z)=1-\phi(z)$.

Condition (i) comes from the argument in the previous section that fast saturating functions learn values near boundaries efficiently. Conditions (ii) and (iii) indicate that the function $\phi(z)$ behaves similarly to $\sigma(z)$ around $z=0$. In order to avoid possible harmful effects due to the modification, we do not want to change the behavior of the function away from the saturating regime. Hence, we require these conditions. The requirements are analogous to those by Gu et al. (2020b, Section 3.4) for the gate adjustment. The first condition can be viewed as a theoretical refinement of their heuristic modification.

We explore gate functions satisfying the above conditions. Recall that the sigmoid function $\sigma(z)$ has the exponential order of saturation. From condition (i) in the previous section, we explore functions saturating superexponentially. Since any superexponential order can be written as $O(e^{-\alpha(z)})$ with a function satisfying $\alpha(z)>z$ for large $z$, it is enough to consider a function of the form $\phi(z)=\sigma(\alpha(z))$ for such $\alpha$. The desirable properties in Section 5.1 are rephrased as follows in terms of $\alpha$: (i) $\alpha(z)\gg z$ for $z\gg 0$, (ii) $\alpha^{\prime}(0)=1$, and (iii) $\alpha(-z)$ = $-\alpha(z)$ for $z\in\mathbb{R}$. Such functions can be found as examples in the form $\alpha(z)=z+p(z)$ where $p$ is a polynomial consisting of only odd higher degree terms, such as $\alpha(z)=z+z^{3}$. Since a higher degree term has a larger effect on the order of saturation, it mitigates gradient vanishing of the gate function more in accordance with Theorem 4.4. Thus, we take a limit of the degree to infinity, which leads to a simple function expression

Therefore, we adopt $\alpha(z)=\sinh(z)$ for the alternative function $\phi=\sigma\circ\alpha$ and obtain the fast gate

This simple expression enables us to implement it with only one additional function. Note that the above form is an example of gate functions which satisfy desirable properties in Section 5.1 and there are infinitely many other possible choices. The above particular form is one of the simplest choices, and not a limitation of our method. For discussion of other candidates, see Appendix D.

We analyze the fast gate $\phi(z)$ and compare it with other gate functions. First, the order of saturation of the fast gate is $O(e^{-e^{z}})$ since $e^{-e^{az}}/(1-\phi(z))\to 0$ as $z\to\infty$ for any $a>1$. We briefly describe the method of the refine gate (Gu et al. 2020b), which was proposed to avoid gradient vanishing of the gate function. This method exploits an auxiliary gate $r_{t}\in[0,1]^{n}$ to modify the forget gate value $f_{t}$ to $g_{t}=r_{t}(1-(1-f_{t})^{2})+(1-r_{t})f_{t}^{2}$. When a large output value is desired for the forget gate value, the auxiliary gate is expected to take $r_{t}\approx 1$ to push $f_{t}$ to $g_{t}\approx 1-(1-f_{t})^{2}$. Therefore, from the asymptotic view point, this method modifies the order of saturation of the gate function from $O(e^{-z})$ to $O(e^{-2z})$. Compared with the refine gate, the fast gate has a much higher order of saturation. We also analyze the asymptotic convergence rates of solving the toy problem in Section 4. We summarize the results in Tab. 1. See Appendix C.3 for detailed derivation. Since the fast gate has a doubly exponential order of saturation, the order of convergence rate of learning long time scales is faster than the sigmoid and refine gates which have an exponential order of saturation (see also Fig. 2 for comparison of the convergence rates). In addition, the fast gate does not require additional parameters whereas the refine gate does. Therefore, the fast gate is computationally more efficient than the refine gate.

We evaluate the learnability for long time scales across various methods on two synthetic tasks, adding and copy, following previous studies (Hochreiter and Schmidhuber 1997; Arjovsky, Shah, and Bengio 2016). While these tasks are simple and easy to solve for short sequences, they get extremely difficult for gated RNNs to solve when the sequence length grows to hundreds or thousands.

Setup. We compare the fast gate with the refine gate because Gu et al. (2020b) reported that the refine gate achieved the best performance among other previous gate variants. We also include the normalized softsign function $\sigma_{\rm ns}(z)$ in Section 4 as a referential baseline to test the compatibility of our theory. We use these gate functions in a single-layer LSTM. Since the initialization of the forget gate bias $b_{f}$ is critical to model performance (Tallec and Ollivier 2018), we set it so that $\phi(b_{f})=1/(1+e^{-1})$ is satisfied for each gate function $\phi$ (with the bias for an additional gate function initialized by 0 for the refine gate), which amounts to $b_{f}=1$ in the usual sigmoid case (Gers, Schmidhuber, and Cummins 2000; Greff et al. 2016). We also compare performance of these gate variants to that of chrono-initialization (Tallec and Ollivier 2018), a method to initialize parameters to represent long time scales for the sigmoid gate. In addition to the LSTM, we include JANET (Van Der Westhuizen and Lasenby 2018) and NRU (Chandar et al. 2019) as baselines. JANET is one of the simplest gated RNNs specialized to learn long time scales by omitting gates other than the forget gate and applying chrono-initialization. NRU uses non-saturating activation functions to write or erase to a memory cell. We train and evaluate each model three times by varying the random seed. See Appendix E.3 for detailed setting.

Results. The mean squared error on the adding task of sequence length 5000 and the accuracy on the copy task of sequence length 500 during training are shown in Fig. 3 and 4, respectively. While NRU requires the least number of parameter updates on the adding task, the training diverges on the copy tasks due to gradient explosion (Pascanu, Mikolov, and Bengio 2013). This is because the state in the NRU evolves on an unbounded region; thus, a small parameter update can drastically change the behavior of the model. We could not fix this instability even by reducing the clipping threshold for gradient by a factor of $10$. We hypothesize that the training of the NRU tends to be more unstable on the copy task because this task has higher dimensional nature than the adding task in the sense of input dimension (10 vs 2) and the number of tokens to memorize (10 vs 2). Among the gate functions, the fast gate converges the fastest. This is due to the higher order of saturation: the fast gate has the order of saturation $O(e^{-e^{z}})$ whereas the refine gate has $O(e^{-2z})$, thus learns long time scales more efficiently. The normalized softsign gate completely fails to learn since it has a lower order of saturation $O(z^{-1})$ than the sigmoid function. Thus, the performance of the models is well explained by the difference in the order of saturation in Tab. 1. This indicates that our theoretical analysis matches practical settings despite the fact that it builds on a simplified learning problem. The result also shows that modification of the gate function is more effective for learning long-term dependencies than other methods such as chrono-initialization and JANET.

We further observe the growth of time-scale distribution of the memory cell in the LSTM on the adding task. The time scale of $i$-th unit in the memory cell is measured using the bias term of the forget gate by $-1/\log\phi(b_{f,i})$ (Tallec and Ollivier 2018; Mahto et al. 2021). We show the statistics of time scales over all 128 units at each iteration of the training in Fig. 5. The fast gate represents much longer time scales than other gate functions after training, which validates our method. While chrono-initialization set time scales so that they are uniformly distributed within the range $[1,5000]$, they do not change at all during training due to saturation of the usual sigmoid function $\sigma(z)$. Since the fast gate can learn to adapt to even longer time scales than such initialization, it is effective to approximate arbitrary desired time scales which is usually unknown a priori.

Next, we evaluate the fast gate on the sequential image recognition task, where a pixel value is applied into recurrent models at each time step (Le, Jaitly, and Hinton 2015).

Setup. We use the usual order sequential MNIST (sMNIST) task and permuted order version (psMNIST) to introduce more complex and long-term temporal dependencies. We also use the sequential CIFAR-10 (sCIFAR) task, which involves higher dimensional inputs and longer sequence length (i.e., $3\times 1024$) than MNIST. We train the LSTM with the various gates, JANET, and NRU. See Appendix E.4 for training details. We set the hidden dimension as 512 on all models and the dimension of the memory cell in NRU as 144 following the original paper. We evaluate the averaged test accuracy with the standard deviation over three random seeds. We also report the computational time to train each model to compare the computational efficiency of the LSTM with the fast gate to other models.

Results. The results are shown in Tab. 2. The fast gate performs the best among various gates while the normalized softsign gate poorly performs. This is because the order of saturation of activation functions in the forget gate directly affects the learnability for long time scales. The LSTM with the fast gate also outperforms all the other baselines. Note that as chrono-initialization gives too heavy-tailed time scale distribution (Gu et al. 2020b), the chrono-LSTM and JANET performs worse than simply initializing the gate bias as $b_{f}=1$ (Sigmoid in the table). Since large model size tends to result in high accuracy on these tasks (Voelker, Kajić, and Eliasmith 2019; Erichson et al. 2021), we also provide results of smaller models in Appendix F.3 for comparison. While the NRU achieves the highest accuracies on the psMNIST task in Tab. 5 in the appendix, the NRU performs worse than the LSTMs in Tab. 2. Therefore, performance of the LSTM seems to scale better than the NRU in model size. The refine gate requires more than 1.5 times longer processing time than the fast gate for training since it involves an auxiliary gate computation with additional parameters. The NRU requires longer processing time than the LSTMs due to its complicated state update rule. The learning curve for the psMNIST task in Fig. 6 shows that the fast gate reaches high accuracy in as few epochs as the refine gate. Combining with the results on the processing time per epoch, we conclude that the fast gate learns complex time scales the most efficiently among the gate functions.

In natural language processing, performance of language models can suffer from difficulty in learning long time scales because predicting statistically rare words involves long time scales (Mahto et al. 2021). It is expected that using a forget gate function with a higher order of saturation improves learning to predict such rare words. We validate this effect in the following experiment.

Setup. We train and evaluate three-layer LSTM language models following a previous study (Mahto et al. 2021)³³3https://github.com/HuthLab/multi-timescale-LSTM-LMs. on the Penn Treebank (PTB) dataset, replacing every sigmoid forget gate in the baseline LSTM with the fast gate. We compare the LSTM with the fast gate against the LSTM with the sigmoid and refine gates and also NRU by replacing all LSTM modules with NRU modules under the same experimental setting. To evaluate the model performance on data involving different ranges of time scales, the test dataset is divided into four bins depending on their frequencies in the training dataset: more than 10,000, 1000-10,000, 100-1000, and fewer than 100 occurrences.

Results. The results are shown in Tab. 3. The result for the NRU is not in the table since the training diverges. Both refine and fast gates improve perplexity for less frequently observed words compared to the sigmoid gate. The model with the fast gate also achieves the lowest total perplexity. Since frequently observed words involve short-term dependencies, this result indicates that the fast gate improves model performance by learning a wide range of time scales that appear in practical tasks. Tab. 3 also shows the training time taken for one epoch. We observe that the refine gate has larger computational overhead than the fast gate, although the LSTM with the refine gate has the same number of parameters as other LSTMs by using the gate-tying trick (Appendix E.2). This is due to the two-stage computation for the gate value via the auxiliary gate, which cannot be parallelized, thus leads to slow computation. In summary, the fast gate can improve performance on data involving extremely long time scales without sacrificing the performance on data involving short time scales and with less computational overhead.