Uncertainty Estimation and Calibration with Finite-State Probabilistic RNNs

1 INTRODUCTION . Machine learning models are well-calibrated if the probability associated with the predicted class reflects its correctness likelihood relative to the ground truth . The output probabilities of modern neural networks are often poorly calibrated ( Guo et al. , 2017 ) . For instance , typical neural networks with a softmax activation tend to assign high probabilities to out-of-distribution samples ( Gal & Ghahramani , 2016b ) . Providing uncertainty estimates is important for model interpretability as it allows users to assess the extent to which they can trust a given prediction ( Jiang et al. , 2018 ) . Moreover , well-calibrated output probabilities are crucial in several use cases . For instance , when monitoring medical time-series data ( see Figure 1 ( a ) ) , hospital staff should also be alerted when there is a low-confidence prediction concerning a patient ’ s health status . Bayesian neural networks ( BNNs ) , which place a prior distribution on the model ’ s parameters , are a popular approach to modeling uncertainty . BNNs often require more parameters , approximate inference , and depend crucially on the choice of prior ( Gal , 2016 ; Lakshminarayanan et al. , 2017 ) . Applying dropout both during training and inference can be interpreted as a BNN and provides a more efficient method for uncertainty quantification ( Gal & Ghahramani , 2016b ) . The dropout probability , however , needs to be tuned and , therefore , leads to a trade-off between predictive error and calibration error . Sidestepping the challenges of Bayesian NNs , we propose an orthogonal approach to quantify the uncertainty in recurrent neural networks ( RNNs ) . At each time step , based on the current hidden ( and cell ) state , the model computes a probability distribution over a finite set of states . The next state of the RNN is then drawn from this distribution . We use the Gumbel softmax trick ( Gumbel , 1954 ; Kendall & Gal , 2017 ; Jang et al. , 2017 ) to perform Monte-Carlo gradient estimation . Inspired by the effectiveness of temperature scaling ( Guo et al. , 2017 ) which is usually applied to trained models , we learn the temperature τ of the Gumbel softmax distribution during training to control the concentration of the state transition distribution . Learning τ as a parameter can be seen as entropy regularization ( Szegedy et al. , 2016 ; Pereyra et al. , 2017 ; Jang et al. , 2017 ) . The resulting model , which we name ST-τ , defines for every input sequence a probability distribution over state- * Equal contribution . †Work done at NEC Laboratories Europe . 1https : //github.com/nec-research/st_tau transition paths similar to a probabilistic state machine . To estimate the model ’ s uncertainty for a prediction , ST-τ is run multiple times to compute mean and variance of the prediction probabilities . We explore the behavior of ST-τ in a variety of tasks and settings . First , we show that ST-τ can learn deterministic and probabilistic automata from data . Second , we demonstrate on real-world classification tasks that ST-τ learns well calibrated models . Third , we show that ST-τ is competitive in out-of-distribution detection tasks . Fourth , in a reinforcement learning task , we find that ST-τ is able to trade off exploration and exploitation behavior better than existing methods . Especially the outof-distribution detection and reinforcement learning tasks are not amenable to post-hoc calibration approaches ( Guo et al. , 2017 ) and , therefore , require a method such as ours that is able to calibrate the probabilities during training . 2 UNCERTAINTY IN RECURRENT NEURAL NETWORKS . 2.1 BACKGROUND . An RNN is a function f defined through a neural network with parameters w that is applied over time steps : at time step t , it reuses the hidden state ht−1 of the previous time step and the current input xt to compute a new state ht , f : ( ht−1 , xt ) → ht . Some RNN variants such as LSTMs have memory cells ct and apply the function f : ( ht−1 , ct−1 , xt ) → ht at each step . A vanilla RNN maps two identical input sequences to the same state and it is therefore not possible to measure uncertainty of a prediction by running inference multiple times . Furthermore , it is known that passing ht through a softmax transformation leads to overconfident predictions on out-of-distribution samples and poorly calibrated probabilities ( Guo et al. , 2017 ) . In a Bayesian RNN the weight matrices w are drawn from a distribution and , therefore , the output is an average of an infinite number of models . Unlike vanilla RNNs , Bayesian RNNs are stochastic and it is possible to compute average and variance for a prediction . Using a prior to integrate out the parameters during training also leads to a regularization effect . However , there are two major and often debated challenges of BNNs : the right choice of prior and the efficient approximation of the posterior . With this paper , we side-step these challenges and model the uncertainty of an RNN through probabilistic state transitions between a finite number of k learnable states s1 , ... , sk . Given a state ht , we compute a probability distribution over the learnable states . Hence , for the same state and input , the RNN might move to different states in different runs . Instead of integrating over possible weights , as in the case of BNNs , we sum over all possible state sequences and weigh the classification probabilities by the probabilities of these sequences . Figure 2 illustrates the proposed approach and contrasts it with vanilla and Bayesian RNNs . The proposed method combines two building blocks . The first is state-regularization ( Wang & Niepert , 2019 ) as a way to compute a probability distribution over a finite set of states in an RNN . State-regularization , however , is deterministic and therefore we utilize the second building block , the Gumbel softmax trick ( Gumbel , 1954 ; Maddison et al. , 2017 ; Jang et al. , 2017 ) to sample from a categorical distribution . Combining the two blocks allows us to create a stochastic state RNN which can model uncertainty . Before we formulate our method , we first introduce the necessary two building blocks . Deterministic State-Regularized RNNs . State regularization ( Wang & Niepert , 2019 ) extends RNNs by dividing the computation of the hidden state ht into two components . The first component is an intermediate vector ut ∈ Rd computed in the same manner as the standard recurrent component , ut = f ( ht−1 , ct−1 , xt ) . The second component models probabilistic state transitions between a finite set of k learnable states s1 , . . . , sk , where si ∈ Rd , i ∈ [ 1 , k ] and which can also be written as a matrix St ∈ Rd×k . St is randomly initialized and learnt during backpropagation like any other network weight . At time step t , given an ut , the transition over next possible states is computed by : θt = ϕ ( St , ut ) , where θt = { θt,1 , ... , θt , k } and ϕ is some pre-defined function . In Wang & Niepert ( 2019 ) , ϕ was a matrix-vector product followed by a SOFTMAX function that ensures ∑k i=1 θt , i = 1 . The hidden state ht is then computed by ht = g ( θt ) · S > t , ht ∈ Rd , ( 1 ) where g ( · ) is another function , e.g . to compute the average . Because Equation ( 1 ) is deterministic , it can not capture and estimate epistemic uncertainty . Monte-Carlo Estimator with Gumbel Trick . The Gumbel softmax trick is an instance of a pathwise Monte-Carlo gradient estimator ( Gumbel , 1954 ; Maddison et al. , 2017 ; Jang et al. , 2017 ) . With the Gumbel trick , it is possible to draw samples z from a categorical distribution given by paramaters θ , that is , z = ONE HOT ( arg maxi [ γi + log θi ] ) , i ∈ [ 1 . . . k ] , where k is the number of categories and γi are i.i.d . samples from the GUMBEL ( 0 , 1 ) , that is , γ = − log ( − log ( u ) ) , u ∼ UNIFORM ( 0 , 1 ) . Because the arg max operator breaks end-to-end differentiability , the categorical distribution z can be approximated using the differentiable softmax function ( Jang et al. , 2017 ; Maddison et al. , 2017 ) . This enables us to draw a k-dimensional sample vector α ∈ ∆k−1 , where ∆k−1 is the ( k − 1 ) -dimensional probability simplex . 2.2 STOCHASTIC FINITE-STATE RNNS ( ST-τ ) Our goal is to make state transitions stochastic and uitilize them to measure uncertainty : given an input sequence , the uncertainty is modeled via the probability distribution over possible statetransition paths ( see right half of Figure 2 ) . We can achieve this by setting ϕ to be a matrix-vector product and using θt to sample from a Gumbel softmax distribution with temperature parameter τ . Applying Monte Carlo estimation , at each time step t , we sample a distribution over state transition probabilities αt from the Gumbel softmax distribution with current parameter τ , where each state transition has the probability αt , i = exp ( ( log ( θt , i ) + γi ) /τ ) ∑k j=1 exp ( ( log ( θt , j ) + γj ) /τ ) , i ∈ [ 1 . . . k ] . ( 2 ) The resulting αt = { αt,1 , ... , αt , k } can be seen as a probability distribution that judges how important each learnable state st is . The new hidden state ht can now be formed either as an average , ht = ∑k i=1 αt , isi ( the “ soft ” Gumbel estimator ) , or as a one-hot vector , ht = ONE HOT ( arg maxi [ log ( θt , i ) + γi ] ) . For the latter , gradients can be estimated using the straightthrough estimator . Empirically , we found the average to work better . By sampling from the Gumbel softmax distribution at each time step , the model is stochastic and it is possible to measure variance across predictions and , therefore , to estimate epistemic uncertainty . For more theoretical details we refer the reader to Appendix 2.3 . The parameter τ of the Gumbel softmax distribution is learned during training ( Jang et al. , 2017 ) . This allows us to directly adapt probabilistic RNNs to the inherent uncertainty of the data . Intuitively , the parameter τ influences the concentration of the categorical distribution , that is , the larger τ the more uniform the distribution . Since we influence the state transition uncertainty with the learned temperature τ , we refer to our model as ST-τ . We provide an ablation experiment of learning τ versus keeping it fixed in Appendix E. Figure 3 illustrates the proposed model . Given the previous hidden state of an RNN , first an intermediate hidden state ut is computed using a standard RNN cell . Next , the intermediate representation ut is multiplied with k learnable states arranged as a matrix St , resulting in θt . Based on θt , samples are drawn from a Gumbel softmax distribution with learnable temperature parameter τ . The sampled probability distribution represents the certainty the model has in moving to the other states . Running the model on the same input several times ( drawing Monte-Carlo samples ) allows us to estimate the uncertainty of the ST-τ model .

Summary: This paper proposes a method to quantify the uncertainty for RNN. Different from  the traditional Bayesian RNN, the proposed method is more efficient. At each  time, based on the current hidden state and memory, it generates a probability  distribution over the state transition paths on the transition probability by  using the Gumbel softmax function. The next state is computed based on the weighted average of the sampled states and its uncertainty can be  qualified by the sample variance. The hyper-parameter tau of the Gumbel function  is learnt from data to better capture the inherent uncertainty in the data.