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 paper presents a sequence prediction method that incorporates a particular inductive bias -- infrequency of latent state updates -- as measured by an L_0 term on latent state differences. This term is approximated and made more easily optimizable using a method similar to [16], which examined L_0 regularization in non-sequential domains. The paper conducts experiments on several low-dimensional POMDP tasks and then draws the following conclusions: the method generalizes better to out-of-distribution inputs in several partially observable domains than other commonly used RNNs, exhibits long term memory in a domain, and exhibits explainability (in its latent states), sometimes at the cost of reduced prediction performance.