Jumpy Recurrent Neural Networks

1 INTRODUCTION . It is said that change happens slowly and then all at once . Billiards balls move across a table before colliding and changing trajectories ; water molecules cool slowly and then undergo a rapid phase transition into ice ; and economic systems enjoy periods of stability interspersed with abrupt market downturns . That is to say , many time series exhibit periods of relatively homogeneous change divided by important events . Despite this , recurrent neural networks ( RNNs ) , popular for time series modeling , treat time in uniform intervals – potentially wasting prediction resources on long intervals of relatively constant change . One reason for this is that standard RNNs are sequence models without an explicit notion of time . Instead , the amount of time represented by a single RNN update is implicitly set by the training data . For example , a model trained on sequences of daily average temperatures has an implicit time step of a day . For a fixed computational budget , this introduces a trade-off between fidelity and temporal range . A model trained at a resolution of one time step per minute would require over 10K iterations to make a prediction for one week in the future . At the other end of the spectrum , a one-week resolution model could achieve this in a single step but could not provide information about the intervening days . As such , selecting a point on this spectrum is a troublesome design decision . In this work , we present Jumpy RNNs , a simple recurrent architecture that takes update steps at variable , data-dependent time-scales while being able to provide dense predictions at intervening points . The core innovation is to define the hidden state as a continuous , piece-wise linear function of time . Specifically , each Jumpy RNN step predicts not only a hidden state hi , but also a hidden velocity ḣi and a span of time ∆ over which the linear latent dynamics h ( t ) = hi + ḣi ( t− i ) should be applied . Our model then jumps forward in time by ∆ before updating again . Any intermediate time step can be produced by decoding the corresponding hidden state h ( t ) . During training , our model learns to use these functions to span the non-uniform time durations between key events , where key events emerge as time points where linear latent extrapolation is FC Jumpy RNN Prediction GRU FC FC GRU Baseline RNN Prediction 111 … 111 5 2 10 1.5 True Dynamics FC GRU Figure 1 : Predicting the dynamics of two billiards balls ( left ) using a baseline RNN cell ( center ) and a Jumpy RNN cell ( right ) . Whereas the baseline model produces a hidden state ht at each time step , our jumpy model predicts a continuous-time hidden state , over a predicted interval ∆i . This allows it to skip over long spans of predictable motion and focus on key events such as collisions . ineffective . In Figure 1 , for example , we see that our model updates at the collision points between the two balls and the walls . During time spans when the balls are undergoing constant motion , our model does not perform cell updates . In contrast , a standard RNN must tick uniformly through time . We demonstrate our proposed model in several physical dynamics prediction tasks . We show Jumpy RNNs achieve comparable performance to the baseline while being between three and twenty times more efficient to sample . This includes settings with non-linear pixel-based observations . Further , we show that our model outperforms RNNs with any fixed step length , showing the importance of data-dependent step sizes . Finally , we demonstrate that a learned Jumpy RNN dynamics model can be leveraged as an efficient forward predictor in a planning domain . Our key contributions are to : – Identify a trade-off between temporal resolution and the computational expense of RNNs , – Propose Jumpy RNNs , which make jumpy predictions and interpolate between them , – Show empirically that Jumpy RNNs are efficient and effective at jumpy time series prediction . 2 JUMPY RECURRENT NEURAL NETWORKS . Consider a continuous-time function x ( t ) sampled at uniform time steps to form the sequence x0 , x1 , . . . , xT . We study the problem of generative modeling , where given an initial set of observations , the goal is to auto-regressively predict a likely continuation . Standard RNN . The per-step RNN computation during auto-regressive sequence generation is : ht = RNNCell ( ht−1 , φ ( x̂t ) ) , x̂t+1 = f ( ht ) where φ ( · ) and f ( · ) are observation encoders and decoders respectively , and ht is the hidden state . This computation is performed at each time-step even when the dynamics of x ( t ) are simple ( or even constant ) for long stretches of time . One way to linearly reduce the computation required to predict into the future is to increase the time-span ∆ between RNN ticks . Standard RNNs , however , then lose the ability to predict at clock times in between the ∆ time steps . This introduces a trade-off between predictive resolution and the computational cost of predicting far into the future . 2.1 JUMPY RNN ARCHITECTURE . Continuous Hidden Dynamics with Constant Jumps . Our first step toward resolving the trade-off is to upgrade the standard RNN so that it can learn to linearly interpolate a continuous-time hidden state h ( t ) between updates . Let ∆ be the time between RNN ticks such that RNN tick i ( starting at i = 0 ) corresponds to continuous time point τi = i∆ . For update i , the RNN predicts both a hidden state hi and hidden velocity ḣi that describes how the hidden state h ( t ) evolves over the time interval [ τi , τi + ∆ ] . Specifically , the operation of this linear-dynamics RNN with constant jump is given by : hi , ḣi = RNNCell ( [ hi−1 , ḣi−1∆ ] , φ ( x̂ ( τi ) ) ) ( 1 ) h ( t ) = hi + ( t− τi ) ḣi for t ∈ [ τi , τi + ∆ ] ( 2 ) x̂ ( t ) = f ( h ( t ) ) ( 3 ) where [ · , · ] denotes concatenation and x̂ ( t ) is the continuous time output prediction that can be immediately produced on demand for any time in between ticks . Under this model , the hidden state evolves as a piece-wise linear function of time , with transitions between linear functions occurring at RNN updates . This linearity constraint provides a strong and natural inductive bias and allows for fast interpolation . Importantly , the hidden-state linearity does not translate to linearity in output space , as the learned decoder f ( · ) can be an arbitrarily complex non-linear function . While the above model can interpolate between time steps , it is still constrained by a constant jump width ∆ . Manually selecting ∆ is difficult since the ideal value depends on the complexity of x ( t ) and the encoder and decoder capacity . Further , x ( t ) may vary in complexity over time such that a single jump width is insufficient . Jumpy RNNs . The Jumpy RNN upgrades the above model to dynamically adjust ∆ throughout the sequence as shown in Figure 2 . Specifically , we now predict a duration ∆i in addition to linear latent dynamics [ hi , ḣi ] at each RNN step . This duration encodes the time span over which the corresponding linear latent dynamics accurately approximate x ( t ) . RNN ticks now occur at variable points in time so that RNN step i occurs at time point τi = ∑i−1 j=0 ∆j . Now our update is : hi , ḣi = RNNCell ( [ hi−1 , ḣi−1∆i−1 ] , φ ( x̂ ( τi ) ) ) ( 4 ) ∆i = LeakyReLU ( Whi + b ) + 1 ( 5 ) h ( t ) = hi + ( t− τi ) ḣi for t ∈ [ τi , τi + ∆i ] ( 6 ) where again x̂ ( t ) = f ( h ( t ) ) . In practice , this approach can be implemented on top of any existing RNN variant , for example , GRUs . Training the jumpy model , however , introduces new complications since the optimal dynamic steps ∆∗i are unknown and interact with the rest of the hidden state . This is partly addressed by the inductive bias of Eq . 5 , which encourages our model to always jump forward in time by at least one time step . The next section details our approach to effective training . 2.2 TRAINING JUMPY RNNS . Even though our model is defined over continuous time , the training data for time-series models typically consists of observation sequences of the form x0 , x1 , x2 , . . . , xT . For the remainder we will assume that there is one time unit between each sample , i.e . xt = x ( t ) . We now describe how to train Jumpy RNNs in this setting by constructing a loss function specifically suited to our jumpy model . Supervising Predictions . Supervising for prediction accuracy simply requires evaluating the model at each discrete time step and applying the loss : Lx = ∑T t=0 ` ( xt , x̂ ( t ) ) , where ` ( · , · ) is a loss between predicted and actual values . In concept , this loss alone may be sufficient for optimizing accuracy , however , it does not encourage the model to take large jumps ( predict large ∆i ) when possible . In fact , it may encourage small jumps as an easier optimization path . Thus , below we augment the loss to encourage jumpiness . Supervising Dynamic Jumps . We interpret the the jump width ∆i as the duration that the linear dynamics can accurately approximate x ( t ) . To formalize this , we say a training point xt is accurately approximated if ` ( xt , x̂ ( t ) ) < , where is a hyperparameter . Ideally , we would like the Jumpy RNN to produce the maximum jump width that maintains this constraint . More formally , given fixed network weights the optimal jump width ∆∗i at step i is the solution to : max ∆i≥1 ∆i s.t . ` ( xt , x̂ ( t ) ) < ∀ { xt | t ∈ [ τi , τi + ∆i ] } , ( 7 ) which seeks the largest jump that does not violate the loss threshold for any intervening observation . This can be solved using a simple forward line search that starts at τi and moves forward through the samples until our model ’ s prediction loss under the current dynamics exceeds . If this occurs on the first step ( implying no time passes between ticks ) , we set ∆∗i to 1 . In Appendix A.1 we show this can be efficiently vectorized for an entire batch as part of a single pass through the sequence . We compute the optimal ∆∗i on-the-fly during training and jump forward accordingly – effectively ignoring the predicted ∆i as in standard teacher-forcing style training . To update the model to predict the correct ∆i , we augment the prediction loss Lx with a jump size loss L∆ = ∑ i ||∆i −∆∗i ||22 . At inference , the predicted ∆i is used to determine jump length . Interpreting Error Threshold . The error threshold is a key hyperparameter introduced in our approach . Examining Equation 7 , acts as a trade-off between jumpiness and approximation error . In the extreme of setting = 0 , the model predictions will likely never be below , ∆i will remain at 1 , and our model will reduce back to a standard RNN that updates at each time step . At the other extreme , choosing =∞ forces our model to predict the entire time series with a single linear latent transition . In practice , we found that setting to the final training loss of the baseline model yields 3-10x jumpy models with the same test error as the baseline model .

Summary: This paper proposes a recurrent network architecture for future prediction where the hidden states (and the outputs) aren't updated step by step as done traditionally. Instead, the network models the hidden state dynamics as being piecewise linear over varying time spans. It learns to produce the linear dynamics together with each time span, and can "jump" to the next time span according to its own predictions. Hidden states for any time step within a span can be easily obtained by interpolation using the predicted linear dynamics. Experiments on a series of synthetic benchmarks are used to demonstrate that the model can learn to utilize this structure to reduce the amount of computation.