Title: GATO: GATES ARE NOT THE ONLY OPTION

Abstract: Recurrent Neural Networks (RNNs) facilitate prediction and generation of structured temporal data such as text and sound. However, training RNNs is hard. Vanishing gradients cause difficulties for learning long-range dependencies. Hidden states can explode for long sequences and send unbounded gradients to model parameters, even when hidden-to-hidden Jacobians are bounded. Models like the LSTM and GRU use gates to bound their hidden state, but most choices of gating functions lead to saturating gradients that contribute to, instead of alleviate, vanishing gradients. Moreover, performance of these models is not robust across random initializations. In this work, we specify desiderata for sequence models. We develop one model that satisfies them and that is capable of learning long-term dependencies, called GATO. GATO is constructed so that part of its hidden state does not have vanishing gradients, regardless of sequence length, and so that its state does not explode. We study GATO on copying and arithmetic tasks with long dependencies and on modeling intensive care unit and language data. Training GATO is more stable across random seeds and learning rates than GRUs and LSTMs.

1. INTRODUCTION:

RNNs allow us to model structured data like text, image, and sound from domains such as art, healthcare, and environmental sciences. These models and their deep variants form the basis for tasks such as generation and prediction (Pascanu et al., 2013) . Modeling joint distributions of sequences facilitates the generation of images (Oord et al., 2016) , videos (Denton & Fergus, 2018) , drawings (Ha & Eck, 2017) , and jazz music (Huang et al., 2019) . Modeling conditional distributions enables applications like predicting patient outcomes in medical centers with electronic health records (Che et al., 2018) . In all of these cases, the data generation process may involve dependencies that range over hundreds of time steps (e.g. number of video frames or health over a lifetime) where early events may determine the distribution of later events (Kamiya et al., 2015; Denton & Fergus, 2018) .

The basic RNN (Elman, 1990) suffers from vanishing and exploding gradient problems (Bengio et al., 1994) that cause difficulties for learning long-range dependencies. The problems stem from extreme eigenvalues of parameter matrices in the RNN. Such models can also suffer from unbounded growth or shrinkage of the hidden state over long sequence lengths. This can happen even when the hidden-to-hidden Jacobian does not vanish or explode and contributes additional difficulties for learning recurrent and decoder parameters. Careful parameterization and initialization of RNNs is crucial for their success in learning long-range dependencies using gradient-based optimization.

The majority of proposed solutions make use of RNNs with gated units, notably LSTMs (Hochreiter & Schmidhuber, 1997) , GRUs (Cho et al., 2014) and Recurrent Highway Networks (RHNs) (Zilly et al., 2017) . These models and deep variants have been the basis for many state-of-the-art empirical results. Gated architectures address gradient issues by bounding the hidden state and designing the recurrence to copy the state forward in time. The GRU and RHN update their state with a weighted average of bounded updates and the current state, where the weighted average in GRU is a convex combination. The LSTM uses an auxiliary cell state with linear recursive updates to copy the hidden state forward. However, gated models are difficult to optimize in practice, in part due to the saturating gradients of gating non-linearities.

An alternative to gated models is to restrict hidden transformations to be unitary (Arjovsky et al., 2016; Wisdom et al., 2016; Jing et al., 2017; Mhammedi et al., 2017) . Unitary matrices can be safely raised to powers without vanishing gradients since they have modulus 1 eigenvalues. Besides unitary transformations, Zhang et al. (2018) controls the singular values of the transformation using SVD to avoid vanishing and exploding gradient problems. But, these approaches often impose computational overhead to ensure the constraint and require non-trivial parameterizations. These approaches only address hidden-to-hidden gradients. In this work we discuss other properties, such as exploding hidden states in forward propagation when the input is constant for long periods.

We propose a sequence model called GATO. The key insight in GATO is that only part of the hidden state must have non-vanishing gradients for states in the far past to influence the current state. GATO splits its hidden state into two channels. At each time step, each channel is updated by a function of the input and current hidden state. To ensure that one of the channels will not have vanishing gradients, GATO limits the functions for both channels to only depend on one of the channels. This results in an identity-block structure in the hidden-hidden Jacobian. GATO then uses non-saturating functions, such as the cosine, to bound the hidden state upon input to the decoder to stabilize gradients for model parameters. We use a simple version of GATO where the dimensions within the channels do not interact and show that it can model arbitrary conditional distributions through memorization. Partitioned hidden states have been mentioned in (Oliva et al., 2017) and (Zhang et al., 2018) . Oliva et al. (2017) can guarantee the identity Jacobian when α = 1, but in this case the state stays constant.

We study GATO copying and adding on long sequences. GATO solves the copying task while the others do not (Figure 1 ). On adding, the GRU experiences difficulty as sequence length increases while GATO does not (Figure 2 ). We also test GATO on health and language data. Using a critical care database collected in Boston (Johnson et al., 2016) , GATO predicts patient length-of-stay in the intensive care unit more accurately than GRUs and LSTMs across seeds and learning rates. GATO achieves lower held-out perplexity on language modeling using the Penn Treebank dataset. We explore this in a setting with no methods such as dropout or gradient clipping. GATO is stable to train across random initializations compared to GRUs, LSTMs, and other recently proposed RNNs such as the SRU (Oliva et al., 2017) , EURNN (Jing et al., 2017) , and RHN (Zilly et al., 2017) .

2. BACKGROUND ON SEQUENCE MODELS:

Sequence modeling targets learning the sequence of conditional distributions that constitute the joint. Let x = (x 1 , x 2 , . . . , x T ) be a sequence of T observations with joint distribution Q(x). We consider the standard probabilistic model p θ (x) with parameter θ, where each sequence element depends on all previous elements. Let L t at timestep t be the negative conditional log probability of the next sequence element according to the model:

p θ (x) = p θ (x 1 ) T -1 t=1 p θ (x t+1 |x ≤t ); L t (θ) = -log p θ (x t+1 |x ≤t ); L(θ) = t L t (θ).

To minimize L(θ) on samples from Q, we need expressive p θ to capture complicated conditionals.

2.1. SEQUENCE MODELING WITH NEURAL NETWORKS:

Sequence models often represent conditionals using a hidden state that evolves as a sequence of data is observed. Using a function D to map from states to distribution parameters, we define the conditional probability p(x t+1 |x ≤t ) with parameters D(h t ). To represent the true conditionals of Q, h t must summarize the information in x 1 . . . x t useful for predicting x >t . Typically, D is chosen to be a feed-forward neural network. Let U be hidden-to-hidden parameters, W be data-to-hidden parameters, and σ be an element-wise non-linearity. RNNs proceed by parameterizing the dynamics of the hidden state via the following discrete-time recurrence:

EQUATION

Though RNNs can represent a broad class of sequences (Schäfer & Zimmermann, 2006; Siegelmann & Sontag, 1995) , they face problems in modeling long sequences where an early part of the sequence influences a later part. On one hand, modeling such data requires the use of powerful functions, but on the other, gradient-based optimization of such functions is unstable and requires manual tuning.

2.2. RNNS HAVE PROBLEMS:

Vanishing and Exploding Gradients. Consider two states h t and h T with times t T . Exploding gradients occur when h T changes dramatically for small changes in h t , while vanishing gradients occur when the influence of h t on h T decreases exponentially with gap T -t.

Consider the RNN as in Equation ( 1) and recall that σ is an element-wise non-linearity. Let L T be the loss at time T . Let D t = diag(σ (Uh t-1 + Wx t )) be the Jacobian for σ. Suppose that λ is a lower-bound on the absolute-value of the Jacobian D's entries. Then if the smallest singular value of U is larger than 1/λ, ∂h T /∂h t blows up exponentially as time gap T -t increases. Suppose instead that λ is an upper-bound on the absolute-value of D's entries. Then if the largest singular value of U is smaller than 1/λ, ∂h T /∂h t decays exponentially with time gap T -t. The Jacobian ∂h T /∂h t involved in L T /∂U is

∂L T ∂U = T t=1 ∂L T ∂h T ∂h T ∂h t ∂h t ∂U ,

and can stop losses L T from backpropagating to U. These issues cause short-term dependencies to be favored over long-term ones (Bengio et al., 1994; Chung et al., 2014) .

Exploding Forward Computation. The hidden state can also grow without bound or vanish for long sequences. Consider the following recurrence where σ denotes sigmoid:

z (1) = σ(W (1) x t+1 + U (1) h t ); z (2) = σ(W (2) x t+1 + U (2) h t ); h = tanh(W (3) x t+1 + U (3) h t ); h t+1 = z (1) h t + z (2) h.

The state h can grow without bound. Two examples of systems with good hidden-hidden gradients but that have unbounded forward propagation are RHN (Zilly et al., 2017) and (Efficient) Unitary RNN (Arjovsky et al., 2016; Jing et al., 2017) . They are described in Appendix G and Appendix H. These models NAN during forward propagation in our experiments.

2.3. LONG SHORT-TERM MEMORY AND GATED RECURRENT UNITS:

To improve gradients, Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) and Gated Recurrent Units (GRU) (Cho et al., 2014) augment RNNs with additional computation that helps bound the state as well as reduce gradient vanishing. The LSTM retains a cell state that copies the hidden state forward in time (Appendix E). The GRU uses a convex combination of the current state and a bounded proposal. Its state is bounded when initialized correctly (Appendix F). This variety of neural sequence models has produced state-of-the-art performance in machine translation (Sutskever et al., 2014; Luong et al., 2015) and image captioning (Xu et al., 2015) . However, these models are difficult to optimize. Practitioners must experiment with optimization techinques including gradient clipping, randomly dropping-out hidden state updates (Krueger et al., 2016) , and even training discriminators to keep generated and data sequences close (Lamb et al., 2016) . Empirical investigation into the behavior of the LSTM gates has also shown that forget and input gate values tend toward 0.5, contradicting explicit interpretations as gates (Li et al., 2018b) . Recent analysis has shown that LSTMs still have gradient issues that cause them to favor learning short-term dependencies in data (Kanuparthi et al., 2019) . See Appendix J for a brief discussion.

3. DESIDERATA FOR SEQUENCE MODELS:

We now develop criteria for building sequence models and consider several important gradients.

Hidden-to-Hidden gradients. For sequence models to learn long-term dependencies, part of the hidden state at time t needs to have non-decaying influence on states at later times T . Therefore only some, not all, of the hidden units need to avoid vanishing gradients. Our first criterion is:

1. Part (a block) of the Jacobian ∂h T /∂h t does not depend on time gap T -t.

Decoder Gradients. Let γ be a non-linear function and consider a loss at time t defined through a linear multi-output decoder D(h t ) = Wγ(h t ). Then the gradient of the loss with respect to W is:

∂L t /∂W = (∂L t /∂D)(∂D/∂W) = (∂L t /∂D)γ(h t )

To ensure that ||∂L t /∂W|| does not become too large, one property that γ needs to have is: 2. γ is bounded to ensure that decoder gradients ∂D/∂W do not grow with time.

Loss-to-Hidden Gradients. ∂L t /∂h t for one time step t also depends on the choice of γ through its gradient γ . It is defined through the decoder by

∂L t /∂h t = (∂L t /∂D)(∂D/∂h t ) = W (∂L t /∂D) γ (h t )

This means that:

3. γ (h t ) should not converge to 0 as ||h t || grows (if h t does grow).

4. The magnitude ||γ || evaluated at h t is bounded even when ||h t || increases.

These criteria require bounded, differentiable functions whose gradients do not go to 0 as their input increases. There is a relationship between the particular form of h's update and the choice of γ. For example, if the update contracts to a constant c from both the positive and negative direction, then γ should be well-behaved at c (e.g. non-zero gradients). These criteria rule out traditional choices of non-linear functions that saturate, such as sigmoid and tanh, and functions that grow such as SoftPlus and ReLU. Periodic functions such as cos and sin meet these criteria.

Forward Computation. In tasks such as language modeling, the hidden state can be processed for thousands of steps through the corpus. 

4. GATO:

We propose a new sequence model that meets the desiderata called GATO. GATO is capable of learning long-range dependencies and is stable to train, exhibiting little variance across learning rates and initializations. GATO requires fewer parameters to perform well on tasks such as copying, arithimetic, classification, and language modeling. We describe the model and analyze its gradients.

4.1. GATO MODEL:

Building on the ideas in Section 3, GATO partitions the RNN hidden state into two parts, one of which is guaranteed to have good long-term gradients by construction. The updates of the hidden state impose a structure on the hidden-to-hidden Jacobian wherein a sub-matrix is the identity I. This means that parts of ∂h T /∂h t do not dependent on T -t and therefore do not vanish with time.

Partitioned Hidden State. The hidden state h t has two channels r t and s t . At each time step, both channels are updated by a non-linear function of the previous state. To ensure that one part of the state, s t , has time-independent hidden-to-hidden gradients, the updates do not depend on s. Let f φ and g ψ be non-linear functions with parameters φ and ψ. GATO's hidden state evolves as:

r t+1 = f φ (r t , x t+1 ); s t+1 = s t + g ψ (r t , x t+1 ); h t+1 = [r t+1 , s t+1 ].

This is like an RNN where the hidden state update is only a function of part of the state, and where the updates of the other part are defined by residual connections (He et al., 2016; Wang & Tian, 2016) .

In our experiments, r and s are the same size. In Appendix A we experiment with s smaller than r.

4.2. GATO HAS GOOD GRADIENTS:

Hidden-to-Hidden Gradients. As shown in Section 2.2, RNNs suffer from vanishing gradients. We show that the GATO hidden state recurrence does not have vanishing gradients for part of the hidden state. First consider ∂s T /∂s t and recall the evolution of r and s:

r T = f φ (f φ (...f φ (r t , x t+1 ))); s T = s t + T -1 k=t g ψ (r k , x k+1

). The right hand side for r T and the second term of the right hand side for s T do not depend on s t , so ∂s T /∂s t is the identity matrix I. Next, ∂r T /∂s t = 0 because s is not included among the arguments to the update f φ (r t , x t ). The whole Jacobian matrix of hidden states h T with respect to h t is:

∂h T ∂h t = ∂s T ∂st ∂s T ∂rt ∂r T ∂st ∂r T ∂rt =   I T -1 k=t ∂g ψ (r k ,x k+1 ) ∂r k ∂r k ∂rt 0 T -1 =t ∂f φ (r ,x +1 ) ∂r 

 .

(2)

Since ∂s T /∂s t = I does not depend on time difference T -t, part of the hidden state, s T , retains its influence from s t (criterion 1). Though ∂s T /∂r t and ∂r T /∂r t are complicated, ∂s T /∂s t is guaranteed not to vanish so gradients can be propagated back through time. In Appendix A, we explore the effect of removing this feature and find that model performance suffers.

Decoder and Loss-to-Hidden Gradients. We have just shown that the hidden state channel s t 's gradient does not vanish. As discussed in Section 3, there are additional optimization issues that can be introduced by the hidden state. We focus on s t . In the case of loss defined through decoder D(s t ) = Wγ(s t ), γ must be chosen to ensure that ||∂L t /∂W|| and ||∂L T /∂s t || do not shrink or grow, regardless of ||s t ||. In the experiments, GATO ensures this by using γ(s t ) = cos(s t ). We experiment with sin in Appendix B.2.

4.3. NON-INTERACTING HIDDEN UNITS:

Consider a version of GATO where interactions are limited to pairs of units. Let j index pairs of hidden units and let the hidden state size be 2J. Then the updates for each pair (r (j) , s (j) ) are:

r (j) t+1 = f (j) φ (r (j) t , x t+1 ); s (j) t+1 = s (j) t + g (j)

ψ (r

EQUATION

Non-interacting units in RNN hidden states have been explored previously (Li et al., 2018a) . We show that this version of GATO can represent arbitrary conditionals by memorization. Consider a sequence model of scalar data with conditionals defined via hidden states p(x t+1 |x ≤t ) = p(x t+1 ; D(s t , r t )).

For this model to represent arbitrary conditionals, the hidden state needs to store x ≤t . One way for GATO in Equation ( 3) to memorize its inputs x would be for each r (j) to learn how to count and each function g

(j)

φ to store its one-dimensional input x in s (j) when the count r (j) equals j. A powerful decoder can then read the memorized sequence from the hidden state and compute a conditional distribution. This construction requires the size J to be at least the length of the sequence T .

4.4. TWO EXAMPLES OF NON-INTERACTING GATO:

Let x t ∈ R D denote the input at time t. Let h t = [r t , s t ] denote the hidden state with r t , s t ∈ R J . Let A, B, C ∈ R J×D and a, b, c ∈ R J be parameters. Let σ be a nonlinear function whose maximum absolute value is 1. Let λ ∈ R.

One layer GATO.

s t+1 = s t + SoftPlus(Ax t+1 + a r t ); r t+1 = λ • σ(Bx t+1 + b r t ) r t + tanh(Cx t+1 + c r t ).

(4)

Two layer GATO. Let F(r, x) denote a neural network with inputs r, x such that each output dimension F j is a one hidden layer neural network with ReLU activation and hidden size k that only depends on r j and x. Then two-layer GATO is defined as:

EQUATION

We call the λ • σ(.) term the regularizer. In both models, setting λ to be a number less than one helps meet criterion 5 in Section 3. To see this, assume λ = 0.7 and that σ takes its maximum value 1:

r t+1 = 0.7 • r t + tanh(Bx t+1 + b r t ).

Now assume that the tanh takes on either of its extremes, for example 1. Then the update is r = 0.7 • r + 1 and any dimension of r that is greater than 10/3 will be shrunk towards 10/3. We set σ to the sigmoid in the main experiments (Section 5) and to tanh in Appendix B.1.

5. EXPERIMENTS:

We study GATO on two standard sequence tasks, copying and performing arithmetic on long sequences (Hochreiter & Schmidhuber, 1997; Arjovsky et al., 2016) . We compare against several kinds of RNNs, ranging from well-established to recently proposed alternatives. In these tasks, the models train directly on new samples from the data distribution at each batch. There is no notion of overfitting on these experiments. These experiments show where common sequence models fail.

We then use GATO for hospital length of stay prediction using electronic health records and on language modeling. For these tasks, we report held-out metrics (accuracy and perplexity), though we do not focus on regularizing sequence models in this work. We explore this in a setting with no gradient clipping (which could stabilize training) or dropout (which could reduce overfitting). We run experiments across three random seeds and choices of learning rates.

5.1. MODEL AND BASELINE SPECIFICATIONS:

For baselines, we use PyTorch modules nn.GRU and nn.LSTM (Paszke et al., 2017) . Since GATO has fewer parameters than these models for a given hidden size, we consider two alternative baselines. We use diagonal weight matrices for LSTM and GRU (GATO's weight matrices are also diagonal).

We call these models GRU-DIAG and LSTM-DIAG. When we set them to have the same number of parameters as GATO rather than the same hidden size, we call them GRU-DIAG (same param) and LSTM-DIAG (same param). We also compare against SRU (Appendix I), RHN (Appendix G), and EURNN (Appendix H). For GATO, we choose σ in the r update regularization term to be the sigmoid and choose λ = 0.7. In our experiments, we find λ = 0.5 performed similar to 0.7.

We give exact numbers of RNN parameters for each experiment. Embedding and decoder parameters are constant across models and excluded in counts. More details including initialization are described in Appendix C. All models are trained across three random seeds for all tasks. The error bands in all plots represent the minimum and maximum across seeds with a bold line representing the mean.

5.2. EXPERIMENT 1: LONG COPY TASK:

Task Description. We model sequences of the form ABA, where A is a sequence of length 20 with tokens uniformly drawn from {1, . . . , 10} and B is 100 constant BLANK tokens. The model must predict the second A as a copy of the first. The hidden state must memorize the first 20 tokens and keep them in memory while processing the 100 BLANK's. This is a harder variant of the copy task in (Arjovsky et al., 2016) . We do not use a special token to mark the beginning of the second A region.

Experiment Setup. All models use an input embedding size of 4, recurrent hidden size of 1024, and a fully-connected decoder with one hidden layer of size 256. We use two layer GATO. GATO has 138, 752 parameters,while the GRU and LSTM have 3, 118, 080 and 4, 222, 976, respectively. We draw new sequences from the data distribution for each step of training. We evaluate on a held-out set of 1000 sequences. We let each model see 1 million training points.

Results. We show the results in Figure 1 . For the LSTM and GRU, the probability of the tokens to be copied is near chance (0.10) across the three seeds. GATO learns the task with at least .50 probability averaged across the held-out set for all tokens to be copied. GRU-DIAG beats other baselines when excluding EURNN. EURNN NAN'ed on one seed (not shown) and was unstable for the other two. Figure 1 : GATO learns to copy long sequences with an order of magnitude fewer parameters than the GRU and LSTM. We plot average probability of copied tokens on a held-out set. Error bands correspond to random seeds. The GRU and LSTM output probabilities near chance for the copied tokens. GRU-DIAG with same number of parameters as GATO beats other baselines when excluding EURNN. EURNN NANed on one seed (not shown) and was unstable for the other two.

5.3. EXPERIMENT 2: ADDING PROBLEM:

Task Description. This experiment follows the adding task described in (Hochreiter & Schmidhuber, 1997) . In this task, the input consists of two sequences of length T . Each element of the first sequence is sampled from the uniform distribution U[0, 1]. The second sequence plays the role of an indicator. It is all 0's except for two 1's. The location of the first 1 is sampled uniformly from the first half of the sequence, while the location of the second 1 is sampled uniformly from the second half of the sequence. The task is to predict the sum of the two numbers in the first sequence whose locations correspond to the locations of the 1's in the second sequence. The mean squared error of predicting the expectation, 1, is 0.167. This forms the baseline.

Experiment Setup. We compare all models with the baseline and we use two-layer GATO. Loss is measured by mean squared error. We do not use embeddings for this task, so the input size at each timestep is 2 (uniform variate and indicator). All models use recurrent hidden size 512 and a decoder with one hidden layer of size 256. For recurrent parameters, GATO has 43,264 parameters while the GRU and LSTM have 792,576 and 1,056,768 respectively. To make a prediction, we apply the decoder to the RNN hidden state after processing the sequence. We sample new data points at each step of training. We evaluate on a held-out set of 1000 examples. We run experiment with sequence lengths T = 100, 200, 400, 750.

Results. Our results are shown in Figure 2 . After training with 200,000 examples, the mean squared error of the LSTM and both LSTM-DIAGs was still around the baseline. GRU and the GRU-DIAGs learned for T=100 and 200. On T=400, GRU had a NAN error on one seed (denoted by infinite loss). The GRU did not learn the task on T=750, but GRU-DIAG did. Neither SRU nor RHN could learn the task, with missingness in the plot resulting from infinite and NAN losses. GATO (magenta line) learned this task at all lengths and was stable across random seeds, also learning with fewer examples.

Figure 2 : GATO is more robust across different seeds and sequence lengths. We show the mean squared errors of the adding problem for sequence lengths T = 100, 200, 400, 750 over number of training examples. LSTM failed to learn the task. For length T = 400, the GRU numerically failed on one seed (NAN error). We denote this by infinite error. GRU-DIAG could learn the task with more samples than GATO. SRU, EURNN, and RHN could not learn this task (NAN and/or worse than chance). NANs and errors greater than 0.5 (significantly worse than chance) are not shown (e.g. EURNN, T=750).

5.4. EXPERIMENT 3: LENGTH OF STAY IN THE INTENSIVE CARE UNIT (ICU):

Task Description and Experimental Setup. This experiment uses MIMIC-III, a publicly-available Boston-based critical care database. The task is to predict if a patient stays more than 1 week in the ICU unit using the first 48 hours of their vitals. There are 53,211 patients in total, 27,088 of whom stay more than one week. We use the vitals extracted for Interpolation-Prediction Networks (Shukla & Marlin, 2019) . See Appendix C for more details. We split the data into a training set of size 43,000, validation set of size 5,106, and test set of 5,105. We use cross entropy loss. To make a prediction, we apply a decoder to the RNN hidden state after processing the sequence of vital signs. We use a recurrent hidden size 512 and decoder hidden size 1024 for all models. We use one-layer GATO.

Results. We report average test accuracy across seeds in Table 1 . GATO has higher accuracy than GRU, LSTM, and their diagonal variants on this task across two learning rates. RHN performs in between the GRU and LSTM variants. SRU performs best on this task, but has more parameters for the same hidden size. The ranges in performance for GATO and SRU overlap. Mikolov et al. (2010) . We choose hidden size 1300, embedding size 650 and a linear decoder for all the models for this task. We use one-layer GATO in this experiment. Table 2 shows model perplexity. GATO achieves similar perplexity to GRU. RHN gets infinite loss because of its unbounded hidden state and the long sequence lengths on this task. Excluding the infinite loss, SRU and EURNN have the highest test perplexity. By applying principles from GATO, we bound RHN's hidden state and call it BRHN. BRHN performs similarly to GATO and GRU.

6. CONCLUSION AND FUTURE WORK:

In this work, we study criteria for building sequence models that can capture long-term dependencies.

We introduce one such model, GATO. GATO uses partitioned hidden states. One part of the hidden state is modified, but never used to update future states. This ensures that part of the hidden state does not have vanishing gradients. We explore a variant of GATO with non-interacting hidden units and show it is sufficient in our experiments. We also introduce criteria for avoiding unbounded forward propagation in sequence models. Empirically, such explosion occurs in our baseline models.

GATO performs as well as GRUs and LSTMs in four typical sequence modeling problems: copying, arithmetic tasks with long-term dependencies, classification, and language modeling. We made similar comparisons against GRU and LSTM by matching the number of parameters instead of the hidden size. All alternative models (SRU,RHN,EURNN), demonstrate significant instabilities. In RHN, we identify that its hidden state is unbounded and use principles from this work to improve its performance. Future directions include understanding when forward propagation is stable and how the principles used to develop GATO can be used to design feed-forward neural networks. Since several years of research have been devoted to investigating regularization of GRUs and LSTMs, an important direction is to see whether the same techniques apply to alternate models such as GATO, SRU, EURNN, and others. Studying combinations of GATO and other (deep) variants of RNNs such as Phase LSTM (Neil et al., 2016) , Dilated RNN (Chang et al., 2019) , Skip RNN (Campos et al., 2018) , Fast-slow RNN (Mujika et al., 2017) and Deep RNN (Pascanu et al., 2013) is another interesting direction.

A GATO CANNOT LEARN ADDING WITHOUT LONG-TERM GRADIENT:

Recall that GATO partitions its hidden state as h t = [s t , r t ] and that s t and r t are updated by:

r t+1 = f φ (x t+1 , r t ); s t+1 = s t + g ψ (x t+1 , r t ).

We experiment with two ablations of GATO for the adding task and show that the long-term gradients through ∂s T /∂s t are necessary for its performance.

1. Set the update for s to 0. This makes GATO close to a standard RNN for half of its state.

2. Set the update for s to g ψ . This is a standard RNN with half of its state but where the decoder can use g ψ to compute an additional function of the state r.

Both of these ablations lose GATO's identity block matrix in the Jacobian ∂h T /∂h t . In these experiments, GATO does not learn the adding task for sequence length 750 above chance for either ablation, across three seeds. We perform a third ablation by keeping s, but reduce its size: 3. Set s to be 1/4 of the hidden state instead of 1/2.

We find that 1/4 is less stable for the adding task (Figure 3 ) and performs worse for Penn TreeBank (Table 3 ). This suggests that having a larger fraction of the state devoted to long-term gradient propagation is important. Though s is guaranteed to have good gradients, care must still be taken when choosing the update for r as in any recurrent model. Recall the update for r:

EQUATION

We experiment by setting the σ in the regularizing term to be tanh instead of sigmoid. Performance is similar to sigmoid on Copy (Figure 5 ), Add (Figure 6 ), MIMIC (Table 4 ), and Penn TreeBank (Table 4 )

B.2 DECODER NON-LINEARITY FOR S:

The criteria in Section 3 suggest bounding the hidden state with periodic functions before decoding.

Our experiments used cos. We believe that these criteria are necessary, but not necessarily sufficient. We also try sin and observe no change in performance on Copy (Figure 7 ), Add (Figure 8 ), MIMIC (Table 5 ), and Penn TreeBank (Table 5 ). It is unclear which other functions meet these criteria. A next step would explore finite sums of cos and sin. 

C EXPERIMENTAL DETAILS:

C.1 MIMIC-III Features. To predict hospital length of stay on the MIMIC-III dataset, we use the same features extracted for Interpolation-Prediction Networks (Shukla & Marlin, 2019) . The features include blood oxygen level, heart rate, respiratory rate, systolic blood pressure, diastolic blood pressure, temperature, oxygen saturation, and others. Not all features are recorded at each observation time. We use 25 features including 12 vitals, 12 additional missingness indicators for those features, and a measurement time.

Continuous Embeddings. All models apply a linear transformation followed by an element-wise tanh function to the continuous input before it is passed to the RNN. The linear transformation is learned alongside other model parameters.

C.2 INITIALIZATION:

Recall that experiment 1 is copying, experiment 2 is adding, experiment 3 is hospital stay prediction, and experiment 4 is language modeling on Penn Treebank. We use different initializations for GATO in experiments 1 and 2 than we do for experiments 3 and 4. In experiments 1 and 2, we use initialization U[-0.1, 0.1] for all RNN parameters in GATO. In experiments 3 and 4, we use initialization U[-0.1, 0.1] for all RNN parameters besides the weight matrices in the linear layers used to project x. For those, we use the default PyTorch initializations. For embeddings and decoders for all models, as well as all GRU and LSTM parameters, we use default PyTorch initializations. We use initialization U[-0.1, 0.1] for parameters in RHN, PyTorch initializations for the linear layers in SRU and the linear layer for inputs in EURNN. For the representation of unitary matrices parameters in EURNN, we use standard normal initializations.

C.3 LEARNING RATE AND OPTIMIZER:

We use Adam optimizer for all experiments. We use learning rate 4e-3 for experiments 1 and 2, and 1e-4 for experiment 4. We check both learning rates 4e-3 and 1e-4 for experiment 3. In experiment 2, we start with learning rate 0.004 and check the training loss after seeing every 10,000 examples. If the training loss is larger than for the previous 10,000 examples, we divide the learning rate by 2.

C.4 VALIDATION AND TEST SCORES:

On MIMIC-III and Penn Treebank, we report test set accuracy corresponding to each model's predictions when they achieved their best validation accuracy. Perplexity on the language modeling task is defined as in Merity et al. (2018) . 

C.7 LANGUAGE MODEL TRAINING:

We follow the training strategy in (Merity et al., 2018) . With probability 0.95, we take a batch with sentence length sampled from N (70, 5 2 ) (then rounded to an integer) and with probability 0.05 we take a batch with sentence length sampled from N (35, 5 2 ) (then rounded to an integer).

C.8 OTHER TRICKS:

We do not use dropout, gradient clipping and other tricks in all the experiments.

D VANISHING AND EXPLODING GRADIENTS:

Consider the RNN as in Equation ( 1) and recall that σ is an element-wise non-linearity:

h t+1 = σ(Uh t + Wx t+1 ).

Let L T be the loss at time T . Let D t = diag(σ (Uh t-1 + Wx t )) be the Jacobian for σ.

Then the gradient of the loss L T with respect to U is:

∂L T ∂U = T t=1 ∂L T ∂h T ∂h T ∂h t ∂h t ∂U = T t=1 T k=t U D k ∂L T ∂h T h t-1 .

Suppose that λ is a lower-bound on the absolute-value of the Jacobian D k 's entries. Then if the smallest singular value of U is larger than 1/λ, ∂h T /∂h t blows up exponentially with time gap T -t. This is called the exploding gradient problem. Suppose instead that λ is an upper-bound on the absolute-value of D k 's entries. Then if the largest singular value of U is smaller than 1/λ, ∂h T /∂h t decays exponentially with time gap T -t. This is called the vanishing gradient problem.

The term ∂h T /∂h t in the gradient calculation for ∂L T /∂U is therefore problematic for models like the RNN. Such issues can stop losses L T at later time steps T from backpropagating to U to adjust the computation of earlier states h t for t T .

E LSTM:

The LSTM (Hochreiter & Schmidhuber, 1997 ) maintains an auxilary cell state that helps copy the hidden state forward in time as a sequence is processed. c t is called the cell state, and h t is called the hidden state. Letting f t , i t , g t and o t be non-linear functions of h t-1 and x t and letting denote element-wise product, the LSTM hidden state computation for c and h is defined by

EQUATION

In the update of the cell state, the forget gate f t determines the influence of previous cell states on the next. The input gate i t determines the influence of new information on the cell. f t , i t , o t and g t are each defined via usual RNN recurrences:

g t = tanh(W gh h t-1 + W gx x t + b g ), f t = σ(W f h h t-1 + W f x x t + b f ), i t = σ(W ih h t-1 + W ix x t + b i ), o t = σ(W oh h t-1 + W ox x t + b o ).

F GATED RECURRENT UNITS:

The GRU (Cho et al., 2014) is similar to the LSTM but with fewer gates and parameters. First, two RNN-like functions r and z of inputs x t and h t-1 are computed. They are called the reset and update gates, respectively. Then, a third RNN-like function with inputs x t , h t-1 , and r computes a proposed state h. Finally, the new state is a convex combination of the previous and proposed states. Let σ be the sigmoid and φ be any non-linearity like the tanh. Then the GRU computation is: r = σ (W r x t + U r h t-1 ) z = σ (W z x t + U z h t-1 ) h = φ ((Wx t + U (r h t-1 )))

h t = z h t-1 + (1 -z) h

The expression for h t is a convex combination of h t-1 and proposed update h. Therefore, if φ is chosen to be bounded by 1 and if h starts with all coordinates in (0, 1), the state will stay bounded.

G RECURRENT HIGHWAY NETWORK (RHN):

One-layer RHNs (Zilly et al., 2017) are another alternative gated recurrent networks. Define h t , x t to be hidden states and inputs at time t. The update in the RHNs computation is:

r = σ (W r x t + U r h t-1 ) z = σ (W z x t + U z h t-1 ) h = tanh (Wx t + Uh t-1 ) h t = r h t-1 + z h

The RHN does not bound its hidden state. Based on the criterion 5 in Section 3, we bound the hidden state using a similar structure as the GRU. We call the new structure Bounded Recurrent Highway Network (BRHN). The BRHN updates as follows: r = σ (W r x t + U r h t-1 ) h = tanh (Wx t + Uh t-1 ) h t = r h t-1 + (1 -r) h

BRHN performs similarly to GATO and GRU on Penn TreeBank, while RHN numerically fails.

H EFFICIENT UNITARY RECURRENT NEURAL NETWORK (EURNN):

Efficient Unitary Recurrent Neural Network (EURNN) (Jing et al., 2017) is a fast and tunable version of the Unitary Recurrent Network (Arjovsky et al., 2016) . The forward propagation in EURNN follows the RNN framework:

h t = f (Wx t + Uh t-1 )

Define D to be diagonal matrix and F 1 , . . . , F L to be a sequence of rotation matrices. EURNN uses a parameterization to make U unitary:

U = DF 1 . . . F L

If L equals to the number of rows of U, it can capture the full space of the unitary matrices. Following Jing et al. (2017) , we choose L = 2 to have faster computation and use the same nonlinearity

f (z i ) = z i |z i | * ReLU(|z i | + b i ),

where b i is a trainable parameter. Though EURNN has bounded hidden-to-hidden gradients, we know from Section 3 that these are not the only gradients that matter. The forward propagation of the hidden state can blow up, which may still cause gradient problems. It is a good research question whether it is possible to have both unitary hidden-to-hidden transformations and a bounded hidden state.

:



Figures and Tables:

Figure3: Both sizes of s work well for the add task. 50% is more stable, suggesting long-term gradient propagation is important for this task. This is consistent with the findings in Appendix A where 0% s failed on this task.

Figure 4: Both sizes of s work well for the copy task.

Figure 5: Tanh performs similarly to sigmoid for regularizing non-linearity on copy.

Figure 6: Tanh performs similarly to sigmoid for regularizing non-linearity on add.

Figure 7: sin for decoder non-linearity performs similarly to sos on copy task.

Figure 8: sin for decoder non-linearity performs similarly to cos on adding task.

BATCH SIZEWe use batch size32, 64, 256, and  20 in experiment 1, 2, 3, and 4 respectively.C.6 WEIGHT DECAY We use 1.2e-6 weight decay (the default number inMerity et al. (2018)) in experiments 3 and 4. Since there is no generalization problem in experiments 1 and 2 (because data samples are drawn from the true generating distribution), we do not use weight decay in those experiments.

Let r t denote the part of h t involved in the recurrence to compute h t+1 . Then if ||r t || is large at time t during training, ||∂L t+1 /∂U|| will be large and can cause additional issues in learning recurrent parameters U. 5. ||r t || in general should not become large as t increases.

Table content: <table/>

GATO performs well on test accuracy with both learning rates 4e-3 and 1e-4. We show the accuracy of length of stay prediction averaged across seeds. SRU performs best on this task. The ranges in performance for GATO and SRU overlap. *EURNN had an NAN error on all seeds shortly following the reported accuracies.Description and Results. A language model predicts the probability distribution of the next word in a sentence based on the previous words. We test GATO on the Penn Treebank dataset preprocessed by

Table content: <table><tr><td>Model</td><td colspan="3"># Parameters Test Accuracy, LR 4e-3 Test Accuracy, LR 1e-4</td></tr><tr><td>GATO</td><td>20,736</td><td>0.693</td><td>0.646</td></tr><tr><td>GRU</td><td>827,904</td><td>0.519</td><td>0.569</td></tr><tr><td>LSTM</td><td>1,103,872</td><td>0.650</td><td>0.569</td></tr><tr><td>GRU-DIAG</td><td>41,472</td><td>0.505</td><td>0.506</td></tr><tr><td>LSTM-DIAG</td><td>55,296</td><td>0.633</td><td>0.506</td></tr><tr><td>GRU-DIAG (same params)</td><td>21,504</td><td>0.570</td><td>0.505</td></tr><tr><td>LSTM-DIAG (same params)</td><td>22,400</td><td>0.569</td><td>0.506</td></tr><tr><td>SRU</td><td>315,282</td><td>0.700</td><td>0.683</td></tr><tr><td>EURNN</td><td>14,848</td><td>0.658*</td><td>0.610*</td></tr><tr><td>RHN</td><td>827,904</td><td>0.619</td><td>0.570</td></tr><tr><td colspan="3">5.5 EXPERIMENT 4: LANGUAGE MODELING ON PENN TREEBANK</td><td/></tr></table>

GATO outperforms GRU and LSTM. We calculate the perplexity averaged over three seeds on Penn Treebank language modeling in the unregularized setting. GATO achieves similar perplexity to GRU. RHN gets infinite loss. Excluding the RHN, SRU and EURNN have the highest test perplexity. *By applying principles from GATO, we bound RHN's hidden state to make it perform similarly to GATO and GRU (called BRHN).

Table content: <table><tr><td>Model</td><td colspan="2"># Parameters Test Perplexity Score</td></tr><tr><td>GATO</td><td>1,273,350</td><td>112.85</td></tr><tr><td>GRU</td><td>7,612,800</td><td>115.78</td></tr><tr><td>LSTM</td><td>10,150,400</td><td>122.86</td></tr><tr><td>GRU-DIAG</td><td>2,546,700</td><td>124.44</td></tr><tr><td>LSTM-DIAG</td><td>3,395,600</td><td>119.09</td></tr><tr><td>GRU-DIAG (same params)</td><td>1,273,350</td><td>124.65</td></tr><tr><td>LSTM-DIAG (same params)</td><td>1,306,000</td><td>122.01</td></tr><tr><td>SRU</td><td>2,198,560</td><td>149.97</td></tr><tr><td>EURNN</td><td>850,200</td><td>768.96</td></tr><tr><td>RHN</td><td>7,612,800</td><td>∞</td></tr><tr><td>BRHN*</td><td>5,075,200</td><td>116.11</td></tr></table>

Size of s can matter. For MIMIC, using 1/4 of the hidden state for s works slightly better than using 1/2 as in the main experiments. For Penn Treebank, 1/2 was better than 1/4, suggesting long-term gradient propagation is important for this task.

Table content: <table><tr><td>Model</td><td colspan="3">MIMIC Test Acc., LR 4e-3 MIMIC Test Acc., LR 1e-4 Penn Test PPL</td></tr><tr><td>1/2 state (original)</td><td>0.693</td><td>0.646</td><td>112.85</td></tr><tr><td>1/4 state</td><td>0.696</td><td>0.651</td><td>122.293</td></tr><tr><td colspan="3">B ABLATION STUDY: GATO'S NON-LINEARITIES</td><td/></tr></table>

Tanh performs similarly to sigmoid for regularizing non-linearity on MIMIC and Penn TreeBank.

Table content: <table><tr><td>Model</td><td colspan="3">MIMIC Test Acc., LR 4e-3 MIMIC Test Acc., LR 1e-4 Penn Test PPL</td></tr><tr><td>sigmoid (original)</td><td>0.693</td><td>0.646</td><td>112.85</td></tr><tr><td>tanh</td><td>0.693</td><td>0.646</td><td>118.955</td></tr></table>

sin for decoder non-linearity performs similarly to cos on MIMIC and Penn TreeBank

Table content: <table><tr><td>Model</td><td colspan="3">MIMIC Test Acc., LR 4e-3 MIMIC Test Acc., LR 1e-4 Penn Test PPL</td></tr><tr><td>cos (original)</td><td>0.693</td><td>0.646</td><td>112.85</td></tr><tr><td>sin</td><td>0.692</td><td>0.648</td><td>112.89</td></tr></table>

Peer review (iclr2020):
The paper proposes a new RNN architecture designed to overcome vanishing/exploding gradient problems and to improve long-term memory for sequence modelling. The main ideas are (i) to split the hidden state into two parts, one of which does not influence the recurrence relation, and can therefore not blow up or contract by self-feedback; and (ii) to use periodic functions, in particular the cosine, as non-linearity in the decoder, so that the output is bounded but does not saturate. The paper puts forward a fairly systematic analysis of the gradients in RNNs. The analysis appears correct, and is in fact quite similar to considerations in earlier RNN work (which is correctly cited), and forms the basis for the proposed GATO unit. There are two loose ends in this part: 1) the cosine non-linearity results from a purely negative selection - the function should be bounded, but not saturating. The paper does not even ask the question which periodic function might be a good choice. 2) While the method is presented as a grand theory, with the only constraint that a part of the hidden state does not influence the recurrence function; the actual implementation and experiments are limited to the narrow special case of *non-interacting* GATO, where the *passive* variables make up exactly half of the hidden vector, and the update of each individual hidden variable is influenced only by a single variable from the previous state. So there are in fact no empirical results, not even on toy data, for the general case that the paper claims to introduce. In the experiments, there are two artificial problems (copying, adding) for sequences of symbols. These are illustrative and sensible to verify and analyse the behaviour of GATO in a controlled setting, but rather far from most real sequence modelling tasks. In a third experiment the task is to classify whether or not patients will stay in intensive care for >1 week, based on a (seemingly private) database of time series of vital parameters. Unfortunately, the results for that experiment do not go beyond the usual *ours is better*. While the numbers clearly support GATO, it would have been nice to look a bit closer and pinpoint what makes the difference. Also, the comparison is a bit loose. It would have been better to add additional baselines where also LSTM and GRU are restricted to *non-interacting*, element-wise recurrence (as far as technically feasible). As it stands, it is unfair to claim *we can do it with much fewer parameters* - perhaps LSTM / GRU could, too. In fact, it could even be that the task is just simple, so that more restricted model with fewer parameters generally perform better - I do not claim this is the case, but the experiments do not rule it out and, hence, do not confirm that the clever GATO recurrence makes the difference. A small gap is also that even the *real* experiment might not be completely realistic. Nowadays it is a popular strategy to use deep LSTMS / GRUs, i.e., stack multiple levels of recurrence, possibly with temporal sub-sampling, to better capture long-term relations. While this is potentially even more brittle, because of the additional gradient flow across layers, it does seem to work. But deep RNNs are not tested (in fact, not even mentioned) in the paper. Overall, in spite of a few loose ends, I find both the design and the practical performance of the proposed GATO unit very interesting, and potentially valuable for the ICLR crowd. This is one of the more convincing RNN papers I have recently read.