Sparsely Changing Latent States for Prediction and Planning in Partially Observable Domains

1 Introduction . When does the meeting start ? Where are my car keys ? Is the stove turned off ? Humans memorize lots of information over extended periods of time . In contrast , classical planning methods assume that the state of the environment is fully observable at every time step [ 1 ] . This assumption does not hold for realistic applications , where generative processes are only indirectly observable or entities are occluded . Planning in such Partially Observable Markov Decision Processes ( POMDP ) is a challenging problem , because suitably-structured memory is required for decision making . Recurrent neural networks ( RNNs ) are often used to deal with partial observability [ 2–4 ] . They encode past observations by maintaining latent states , which are iteratively updated . However , continuously updating the latent state causes past information to quickly “ wash out ” . Long-Short Term Memory networks ( LSTM , [ 5 ] ) and Gated Recurrent Units ( GRU , [ 6 ] ) deal with this problem by using internal gates . However , they can not leave their latent states completely unchanged , because 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . small amounts of information continuously leak through the sigmoidal gating functions . Additionally , inputs typically need to pass through the latent state to affect the output , making it hard to disentangle observable from unobservable information within their latent states . Our hypothesis is that many generative latent factors in the physical world are constant over extended periods of time . Thus , there might not be the need to update memory at every time step . For example , consider dropping an object : If the drop-off point as well as some latent generative factors , such as gravity and aerodynamic object properties , are known , iteratively predicting the fall can be reasonably accomplished by a non-recurrent process . Similarly , when an agent picks up a key , it is sufficient to memorize that the key is inside their pocket . However , latent factors typically do change significantly and systematically at particular points in time . For example , the aerodynamic properties of an object change drastically when the falling object shatters on the floor , and the location of the key changes systematically when the agent removes it from their pocket . These observations are related to assumptions used in causality research . A common assumption is that the generative process of a system is composed of autonomous mechanisms that describe causal relationships between the system ’ s variables [ 7–9 ] . When considering Markov Decision Processes , it has been proposed that these mechanisms tend to interact sparsely in time and locally in space [ 10 , 11 ] . Causal models aim at creating dependencies between variables only when there exists a causal relationship between them , in order to improve generalization [ 8 ] . Updating the latent state of a model in every time step , on the other hand , induces the prior assumption that the generative latent state typically depends on all previous inputs . Thus , by suitably segmenting the dependencies of the latent variables over time , one can expect improved generalization across spurious temporal dependencies . Very similar propositions have been made for human cognition . Humans tend to perceive their stream of sensory information in terms of events [ 12–16 ] . Event Segmentation Theory ( EST ) [ 16 ] postulates a set of active event models , which encode event-respective aspects over extended periods of time and switch individually at event transitions . To learn about the transitions and consolidate associated latent event encodings , measurements of surprise and other significant changes in predictive model activities , as well as latent state stability assumptions , have been proposed as suitable inductive event segmentation biases [ 16–22 ] . Explicit relations to causality have been put forward in [ 23 ] . In accordance to EST and our sparsely changing latent factor assumption , we introduce Gated L0 Regularized Dynamics ( GateL0RD ) . GateL0RD applies L0-regularized gates , inducing an inductive learning bias to encode piecewise constant latent state dynamics . GateL0RD thus becomes able to memorize task-relevant information over long periods of time . The main contributions of this work can be summarized as follows . ( i ) We introduce a stochastic , rectified gating function for controlling latent state updates , which we regularize towards sparse updates using the L0 norm . ( ii ) We demonstrate that our network performs as good or better than state-of-the-art RNNs for prediction or control in various partially-observable problems with piecewise constant dynamics . ( iii ) We also show that the inductive bias leads to better generalization under distributional shifts . ( iv ) Lastly , we show that the latent states can be easily interpreted by humans . 2 Background . Let fθ : X ×H → Y ×H be a recurrent neural network ( RNN ) with learnable parameters θ mapping inputs1 xt ∈ X and ht−1 ∈ H the latent ( hidden ) state to the output ŷt ∈ Y and updated latent states ht . The training dataset D consists of sequences of input-output pairs d = [ ( x1 , y1 ) , . . . , ( xT , yT ) ] of length T . In this paper , we consider the prediction and control of systems that can be described by a partially observable Markov decision process ( POMDP ) with state space S , action space A , observations space O , and deterministic hidden transitions S ×A → S.2 1Notation : bold lowercase letters denote vectors ( e.g. , x ) . Vector dimensions are denoted by superscript ( e.g . x = [ x1 , x2 , . . . , xn ] ∈ Rn ) . Time or other additional information is denoted by subscript ( e.g. , xt ) . 2We treat the prediction of time series without any actions as a special case of the POMDP with A = ∅ . 3 L0-regularization of latent state changes We want the RNN fθ to learn to solve a task , while maintaining piecewise constant latent states over time . The network creates a dynamics of latent states ht when applied to a sequence : ( ŷt , ht ) = fθ ( xt , ht−1 ) starting from some h0 . The most suitable measure to determine how much a time-series is piecewise constant is the L0 norm applied to temporal changes . With the change in latent state as ∆ht = ht−1 − ht , we define the L0-loss as LL0 ( ∆h ) = ‖∆h‖0 = ∑ j=1 I ( ∆hj 6= 0 ) , ( 1 ) which penalizes the number of non-zero entries of the vector of latent state changes ∆h . The regularization loss from Eq . 1 can be combine in the usual way with the task objective to yield the overall learning objective L of the network : L ( D , θ ) = Ed∼D [ ∑ t Ltask ( ŷt , yt ) + λLL0 ( ∆ht ) ] ( 2 ) with ( ŷt , ht ) = fθ ( xt , ht−1 ) . The task-dependent loss Ltask ( · , · ) can be , for instance , the meansquared error for regression or cross-entropy loss for classification . The hyperparameter λ controls the trade-off between the task-based loss and the desired latent state regularization . Unfortunately , we can not directly minimize this loss using gradient-based techniques , such as stochastic gradient descent ( SGD ) , due to the non-differentiability of the L0-term . Louizos et al . [ 24 ] proposed a way to learn L0 regularization of the learnable parameters of a neural network with SGD . They achieve this by using a set of stochastic gates controlling the parameters ’ usage . Each learnable parameter θj that is subject to the L0 loss is substituted by a gated version θ′j = Θ ( sj ) θj where Θ ( · ) is the Heaviside step function ( Θ ( s ) = 0 if s ≤ 0 and 1 otherwise ) and s is determined by a distribution q ( s|ν ) with learned parameters ν . Thus , θ′j is only non-zero if sj > 0 . This allows to rewrite the L0 loss ( Eq . 1 ) for θ′ as : LL0 ( θ′ , ν ) = ‖θ′‖0 = ∑ j Θ ( sj ) with s ∼ q ( s ; ν ) , ( 3 ) where parameters ν influence sparsity and are affected by the loss . To tackle the problem of non-differentiable binary gates , we can use a smooth approximation as a surrogate [ 24–26 ] . Alternatively , we can substitute its gradients during the backward pass , for example using the straight-through estimator [ 27 ] , which treats the step function as a linear function during the backward pass , or approximate its gradients as in the REINFORCE algorithm [ 28 ] . To transfer this approach to regularize the latent state dynamics in an RNN , we require an internal gating function Λ ( · ) ∈ [ 0 , 1 ] , which controls whether the latent state is updated or not . For instance : ht = ht−1 + Λ ( s ) ∆h̃t−1 with ∆h̃t−1 = h̃t − ht−1 ( 4 ) where h̃ is the proposed new latent state and s is a stochastic variable depending on the current input and previous latent state and the parameters , i.e . st ∼ q ( st ; xt , ht−1 , ν ) . For brevity , we merge the parameters ν into the overall parameter set , i.e . ν ⊂ θ . For computing Eq . 2 we need to binarize the gate by applying the step function Θ ( Λ ( s ) ) . Thus we can rewrite Eq . 2 as L ( D , θ ) = Ed∼D [ ∑ t Ltask ( ŷt , yt ) + λ ∑ t Θ ( Λ ( st ) ) ] . ( 5 ) LSTMs and GRUs use deterministic sigmoidal gates for Λ in Eq . 4 to determine how to update their latent state . However , it is not straight forward to apply this approach to them ( detailed in Suppl . A ) . Thus , we instead introduce a novel RNN , that merges components from GRUs and LSTMs , to implement the proposed L0 regularization of latent state changes while still allowing the network to make powerful computations . We name our network Gated L0 Regularized Dynamics ( GateL0RD ) . 4 GateL0RD . The core of GateL0RD implements the general mapping ( ŷt , ht ) = fθ ( xt , ht−1 ) using three functions , or subnetworks : ( 1 ) a recommendation network r , which proposes a new candidate latent state , ( 2 ) a gating network g , which determines how the latent state is updated , and ( 3 ) an output function , which computes the output based on the updated latent state and the input . The network is systematically illustrated in Fig . 1a . The overall processing is described by the following equations : st ∼ N ( g ( xt , ht−1 ) , Σ ) ( sample gate input ) ( 6 ) Λ ( s ) : = max ( 0 , tanh ( s ) ) ( new gating function ) ( 7 ) ht = ht−1 + Λ ( st ) ( r ( xt , ht−1 ) − ht−1 ) ( update or keep latent state ) ( 8 ) ŷt = p ( xt , ht ) o ( xt , ht ) , ( compute output ) ( 9 ) where denotes element-wise multiplication ( Hadamard product ) . We start with the control of the latent state in Eq . 8 . Following Eq . 4 , a new latent value is proposed by the recommendation function r ( xt , ht−1 ) and the update is “ gated ” by Λ ( s ) . Importantly , if Λ ( s ) = 0 no change to the latent state occurs . Note that the update in Eq . 8 is in principle equivalent to the latent state update in GRUs [ 6 ] , for which it is typically written as ht = Λ ( s ) r ( xt , ht−1 ) + ( 1− Λ ( s ) ) ht−1 with Λ ( s ) a deterministic sigmoidal gate . Because we aim for piecewise constant latent states , the gating function Λ defined in Eq . 7 needs to be able to output exactly zero . A potential choice would be the Heaviside function , i.e . either copy the new latent state or keep the old one . This , however , does not allow any multiplicative computation . So a natural choice is to combine the standard sigmoid gate of RNNs with the stepfunction : Λ ( s ) = max ( 0 , tanh ( s ) ) which we call ReTanh ( rectified tanh ) 3 . Figure 1b shows the activation function Λ depending on its input . The gate is closed ( Λ ( si ) = 0 ) for all inputs si ≤ 0 . A closed gate results in a latent state that remains constant in dimension i , i.e. , hit = h i t−1 . On the other hand , for si > 0 the latent state is interpolated between the proposed new value and the old one . The next puzzle piece is the input to the gate . Motivated from the L0 regularization in Eq . 1 we use a stochastic input . However , in our RNN setting , it should depend on the current situation . Thus , we use a Gaussian distribution for q with the mean determined by the gating network g ( xt , ht−1 ) as defined in Eq . 6 . We chose a fixed diagonal covariance matrix Σ , which we set to Σi , i = 0.1 . To train our network using backpropagation , we implement the sampling using the reparametrization trick [ 29 ] . We introduce a noise variable and compute the gate activation as st = g ( xt , ht−1 ) + with ∼ N ( 0 , Σ ) . ( 10 ) During testing we set = 0 to achieve maximally accurate predictions . 3Note that tanh ( s ) = 2 · sigmoid ( 2s ) − 1 . Finally the output ŷ is computed from the inputs and the new latent state ht in Eq . 9 . Inspired by LSTMs [ 5 ] , the output is determined by a multiplication of a normal branch ( p ( xt , ht ) ) and a sigmoidal gating branch ( o ( xt , ht ) ) . We thus enable both additive as well as multiplicative effects of xt and ht on the output , enhancing the expressive power of the piecewise constant latent states . In our implementation , all subnetworks are MLPs . r , p use a tanh output activation ; o uses a sigmoid ; g has a linear output . p , o are one-layer networks . By default , r , g are also one-layer networks . However , when comparing against deep ( stacked ) RNNs , we increase the number of layers of r and g to up to three ( cf . Suppl . B ) . We use the loss defined in Eq . 5 . GateL0RD is fully differentiable except for the Heaviside step function Θ in Eq . 5 . A simple approach to deal with discrete variables is to approximate the gradients by a differentiable estimator [ 25–27 ] . We employ the straight-through estimator [ 27 ] , which substitutes the gradients of the step function Θ by the derivative of the linear function ( see Fig . 1c ) . We use GateL0RD as a memory module of a more general architecture illustrated in Fig . 1d . The network input is preprocessed by a feed-forward network fpre ( xt ) . Similarly , its output is postprocessed by an MLP fpost ( ŷt ) ( i.e . a readout layer ) before computing the loss . The latent state h0 of GateL0RD could be initialized by 0 . However , improvements can be achieved if the latent state is instead initialized by a context network finit , a shallow MLP that sets h0 based on the first input [ 30 , 31 ] . In the Supplementary Material we ablate various components of GateL0RD , such as the gate activation function Λ ( Suppl . C.1 ) , the gate stochasticity ( Suppl . C.2 ) , the context network finit ( Suppl . C.3 ) , the multiplicative output branch o ( Suppl . C.4 ) , and compare against L1/L2-variants ( Suppl . C.5 ) .

Summary: This work proposes an RNN with a gating mechanism that encourages its latent state to have sparse updates by incorporating stochastic gating units trained using an approximation of L0 regularization of the state updates. The authors propose to use this RNN in partially observable problems in which the latent dynamics are known to change slowly and long-term memory is important. The authors show in a number of partially observed tasks that their proposed RNN and training objective outpeforms well known RNN architectures when predicting the environment dynamics.