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 work proposes a novel method to estimate uncertainties in recurrent neural networks. The proposed model explicitly computes a probability distribution over a set of discrete hidden states given the current hidden state in an RNN. Leveraging the Gumbel softmax trick, the proposed method performs MC gradient estimation. A temperature parameter is also learned to control the concentration of state transition distribution. To estimate uncertainty of a given input, the proposed model is run multiple times to draw samples for estimating the mean and variance. Experiments are conducted in a variety of sequential prediction problems, including a reinforcement learning task, demonstrating the effectiveness of the proposed uncertainty estimation method.