Original: P REDICTIVE U NCERTAINTY THROUGH Q UANTIZATION
Revised: P REDICTIVE U NCERTAINTY THROUGH Q UANTIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: High-risk domains require reliable conﬁdence estimates from predictive models.
Revised: High-risk domains require reliable conﬁdence estimates from predictive models.

Original: Deep latent variable models provide these, but suffer from the rigid variational distributions used for tractable inference, which err on the side of overconﬁdence.
Revised: Deep latent variable models provide these, but suffer from the rigid variational distributions used for tractable inference, which err on the side of overconﬁdence.

Original: We propose Stochastic Quantized Activation Distributions (SQUAD), which imposes a ﬂexible yet tractable distribution over discretized latent variables.
Revised: We propose Stochastic Quantized Activation Distributions (SQUAD), which imposes a ﬂexible yet tractable distribution over discretized latent variables.

Original: The proposed method is scalable, self-normalizing and sample efﬁcient.
Revised: The proposed method is scalable, self-normalizing and sample efﬁcient.

Original: We demonstrate that the model fully utilizes the ﬂexible distribution, learns interesting non-linearities, and provides predictive uncertainty of competitive quality
Revised: We demonstrate that the model fully utilizes the ﬂexible distribution, learns interesting non-linearities, and provides predictive uncertainty of competitive quality

Original: [Equation]
Revised: [Equation]

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: +EWjpBsZMttz6wVX5vLRUAGKkIJObzN3wkkmuxQzomE9pSHLTt1YF5gG5q6kx7YtnaaOdXZdd85azWcZsP5vFc/2pvPtVXwGrwBW8AB78ER+AhOQBdgEIPv4Cf4VTuooVpU+3orXV6a57wCJaupP+1S9io=</latexit>
Revised: +EWjpBsZMttz6wVX5vLRUAGKkIJObzN3wkkmuxQzomE9pSHLTt1YF5gG5q6kx7YtnaaOdXZdd85azWcZsP5vFc/2pvPtVXwGrwBW8AB78ER+AhOQBdgEIPv4Cf4VTuooVpU+3orXV6a57wCJaupP+1S9io=</latexit>

Original: In high-risk domains, prediction errors come at high costs.
Revised: In high-risk domains, prediction errors come at high costs.

Original: Luckily such domains often provide a failsafe in the form of rejecting a prediction: self-driving cars perform an emergency stop, doctorscan run another diagnostic test, and industrial processes can be temporarily halted.
Revised: Luckily such domains often provide a fail-safe: self-driving cars perform an emergency stop, doctors run another diagnostic test, and industrial processes are temporarily halted.

Original: For deep learning models, this can be achieved via apredetermined threshold on predicted conﬁdence scores.
Revised: For deep learning models, this can be achieved by rejecting datapoints with a conﬁdence score below a predetermined threshold.

Original: As such, a low error rate can be guaranteed at the cost of rejecting some predictions.
Revised: This way, a low error rate can be guaranteed at the cost of rejecting some predictions.

Original: However, estimating  conﬁdence scores from neural networks that create well-ordered rankings of correct and incorrect predictions remains an active area of research.
Revised: However, estimating high quality conﬁdence scores from neural networks, which create well-ordered rankings of correct and incorrect predictions, remains an active area of research.

Original: Deep Latent Variable Models (DLVMs, ﬁg. 1) approach this by postulating latent variables z for which the uncertainty in p ( z | x ) inﬂuences the conﬁdence in the target prediction.
Revised: Deep Latent Variable Models (DLVMs, ﬁg. 1) approach this by postulating latent variables z for which the uncertainty in p ( z | x ) inﬂuences the conﬁdence in the target prediction.

Original: Recently, efﬁcient inference algorithms have been proposed in the form of variational inference, where an inference neural network is optimized to predict parameters of a variational distribution that approximates an otherwise intractable distribution (Kingma & Welling (2013); Rezende et al.
Revised: Recently, efﬁcient inference algorithms have been proposed in the form of variational inference, where an inference neural network is optimized to predict parameters of a variational distribution that approximates an otherwise intractable distribution (Kingma & Welling (2013); Rezende et al.

Original: ; Alemi et al.
Revised: ; Alemi et al.

Original: ; Achille & Soatto (2016)).
Revised: ; Achille & Soatto (2016)).

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: /2fPiKe4NmO0oVrkeFjgeJwVlIHCaWxCf4tZnjq8ShCzNf2wEu1nXX2ENWDNC3TbpF2q7RTpJ0qvbyYUzfQl1U4KMA7Xc8K8KwKZwU4q0IVFqh5wwpVJVcFxVUVnhfgeRVOC3BahUfuArr6yMD6BuwgRjB9LSHyfZJd8PJfVWRyVSni+/MiItJ+tcOEpsO2ozujv0f+sZvaFAcKDmHDgtfXsNGGtiDjUEGnYu4Xvr1frrxte0SYyeAPzVW1s68aLmbYwaYXItHMXrccXS/eNptxHw9liGJ8cOPf8QWa7hDGsIBmlwdtM3VgXmAL6oaVvjVtzTSzqrPrbnDablqtpvVpr3G4N59rK+AleAU2gQXegEPwARyDHvDAJfgBfoJftfc1WpO15EZ6b2nueQFKq/btD8Xo+cg=</latexit>
Revised: /2fPiKe4NmO0oVrkeFjgeJwVlIHCaWxCf4tZnjq8ShCzNf2wEu1nXX2ENWDNC3TbpF2q7RTpJ0qvbyYUzfQl1U4KMA7Xc8K8KwKZwU4q0IVFqh5wwpVJVcFxVUVnhfgeRVOC3BahUfuArr6yMD6BuwgRjB9LSHyfZJd8PJfVWRyVSni+/MiItJ+tcOEpsO2ozujv0f+sZvaFAcKDmHDgtfXsNGGtiDjUEGnYu4Xvr1frrxte0SYyeAPzVW1s68aLmbYwaYXItHMXrccXS/eNptxHw9liGJ8cOPf8QWa7hDGsIBmlwdtM3VgXmAL6oaVvjVtzTSzqrPrbnDablqtpvVpr3G4N59rK+AleAU2gQXegEPwARyDHvDAJfgBfoJftfc1WpO15EZ6b2nueQFKq/btD8Xo+cg=</latexit>

Original: Variational inference relies on a tractable class of distributions that can be optimized to closely resemble the true distribution (ﬁg. 2), and it’s hypothesized that more ﬂexible classes lead to more faithful approximations and thus better performance (Jordan et al.
Revised: Variational inference relies on a tractable class of distributions that can be optimized to closely resemble the true distribution (ﬁg. 2), and it’s hypothesized that more ﬂexible classes lead to more faithful approximations and thus better performance (Jordan et al.

Original: To explore this hypothesis we propose a novel tractable class of highly ﬂexible variational distributions.
Revised: To explore this hypothesis, we propose a novel tractable class of highly ﬂexible variational distributions.

Original: Considering that neural networks with low-precision activations exhibit good performance (Holi & Hwang (1993);
Revised: Considering that neural networks with low-precision activations exhibit good performance (Holi & Hwang (1993);

Original: Hubara et al.
Revised: Hubara et al.

Original: , we make the modeling assumption that latent variables can be expressed under a strong quantization scheme, without loss of predictive ﬁdelity.
Revised: , we make the modeling assumption that latent variables can be expressed under a strong quantization scheme, without loss of predictive ﬁdelity.

Original: If this assumption holds, it becomes tractable to model a scalar latent variable with a ﬂexible multinomial distribution over the values (ﬁg. 3).
Revised: If this assumption holds, it becomes tractable to model a scalar latent variable with a ﬂexible multinomial distribution over the quantization bins (ﬁg. 3).

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: +9H/6iGIp0hMli6xFRAwlikezgjqcMBZb599iTqZYRBHigXF7ODFu2hkjZnpJUqTtPG2XaStPW2V6cT6nfmguyrCXg7e6nubgaRnOcnBWhnqUo3ZFNCpLLnOKyzI8y8GzMpzm4LQMD/wF9M2BhdUN2EKcEvZSQRQENH3Yxb+q6fiyVCQI5kVkZIJyhzFL+k3PtAZ/r/xzO3EZCTXsw5oDr65grQldSYcjDb1Scje3926x8isXU2knQ9C3T9VNd9VfzK69TTxCsp4utzxTzb82l4uA9NUIxWTvOn87kGi6TTknEtpT7jXt1IFZgS1oak7y1ra108wpz67bzkmz7jTqztHr2v6H+VxbBc/BC7AJHPAG7INP4BB0AAYx+A5+gl+Vd5WgMq5E19LlpXnOM1CwyvkfQc33DQ==</latexit>
Revised: +9H/6iGIp0hMli6xFRAwlikezgjqcMBZb599iTqZYRBHigXF7ODFu2hkjZnpJUqTtPG2XaStPW2V6cT6nfmguyrCXg7e6nubgaRnOcnBWhnqUo3ZFNCpLLnOKyzI8y8GzMpzm4LQMD/wF9M2BhdUN2EKcEvZSQRQENH3Yxb+q6fiyVCQI5kVkZIJyhzFL+k3PtAZ/r/xzO3EZCTXsw5oDr65grQldSYcjDb1Scje3926x8isXU2knQ9C3T9VNd9VfzK69TTxCsp4utzxTzb82l4uA9NUIxWTvOn87kGi6TTknEtpT7jXt1IFZgS1oak7y1ra108wpz67bzkmz7jTqztHr2v6H+VxbBc/BC7AJHPAG7INP4BB0AAYx+A5+gl+Vd5WgMq5E19LlpXnOM1CwyvkfQc33DQ==</latexit>

Original: By repositioning the variational distribution from a potentially limited description of moments – as found in exponential families – to a direct expression of probabilities per value, a variety of beneﬁts arise.
Revised: By re-positioning the variational distribution from a potentially limited description of moments, as found in commonly applied conjugate distributions, to a direct expression of probabilities per value, a variety of beneﬁts arise.

Original: As the output domain is constrained, the method becomes self-normalizing, relieving the model from poorly scalable batch normalization techniques (Ioffe & Szegedy (2015)).
Revised: As the output domain is constrained, the method becomes self-normalizing, relieving the model from hard-to-parallelize batch normalization techniques (Ioffe & Szegedy (2015)).

Original: More interesting priors can be exploredoutside the space of exponential families and the model is able to learn non-linearities per neuron.
Revised: More interesting priors can be explored and the model is able to learn unique activation functions per neuron.

Original: More concretely, the contributions of this work are as follows:
Revised: More concretely, the contributions of this work are as follows:

Original: • We propose a novel variational inference method by leveraging multinomial distributions on quantized latent variables.
Revised: • We propose a novel variational inference method by leveraging multinomial distributions on quantized latent variables.

Original: • We show that the emerging predicted distributions are multimodal, motivating the need for ﬂexible distributions.
Revised: • We show that the emerging predicted distributions are multi-modal, motivating the need for ﬂexible distributions in variational inference.

Original: • We demonstrate that the proposed method applied to the information bottleneck objective computes competitive uncertainty over the predictions and that this manifests in better accuracy under strong risk guarantees.
Revised: • We demonstrate that the proposed method applied to the information bottleneck objective computes competitive uncertainty over the predictions and that this manifests in better performance under strong risk guarantees.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In this work, we explore deep neural networks for regression and classiﬁcation.
Revised: In this work, we explore deep neural networks for regression and classiﬁcation.

Original: We have data-points ( x , y) in a dataset D
Revised: We have data-points consisting of inputs x and targets y in a dataset D

Original: = { ( x i , y i ) | i ∈ [1 , ..., N ] } and postulate latent variables
Revised: = { ( x i , y i ) | i ∈ [1 , ..., N ] } and postulate latent variables

Original: z that represent the data.
Revised: z that represent the data.

Original: We focus on the Information Bottleneck (IB) perspective: ﬁrst proposed by Tishby et al.
Revised: We focus on the Information Bottleneck (IB) perspective: ﬁrst proposed by Tishby et al.

Original: (2000), the information bottleneck objective I ( y , z ; θ ) − βI ( x , z ; θ ) is optimized to maximize the mutual information between z and y , whilst minimizing the mutual information between z and x .
Revised: (2000), the information bottleneck objective I ( y , z ; θ ) − βI ( x , z ; θ ) is optimized to maximize the mutual information between z and y , whilst minimizing the mutual information between z and x .

Original: The objective can be efﬁciently optimized using a variational inference scheme as shown concurrently by both Alemi et al.
Revised: The objective can be efﬁciently optimized using a variational inference scheme as shown concurrently by both Alemi et al.

Original: and Achille & Soatto (2016).
Revised: and Achille & Soatto (2016).

Original: Under the Markov assumption P ( z , x , y ) =
Revised: Under the Markov assumption P ( z , x , y ) =

Original: p ( z | x ) p ( y | x ) p ( x ) , they derive the following lower bound:
Revised: p ( z | x ) p ( y | x ) p ( x ) , they derive the following lower bound:

Original: [Equation]
Revised: [Equation]

Original: where E p ( zn | x n ) is commonly estimated using a single Monte Carlo sample and r ( z ) is a variational distribution to the marginal distribution of z .
Revised: where E p ( z | x n ) is commonly estimated using a single Monte Carlo sample and r ( z ) is a variational approximation to the marginal distribution of z .

Original: In practice r ( z ) is ﬁxed to a simple distribution such as a spherical Gaussian.
Revised: In practice r ( z ) is ﬁxed to a simple distribution such as a spherical Gaussian.

Original: Alemi et al.
Revised: Alemi et al.

Original: and Achille & Soatto (2016) continue to show that the Variational Auto Encoder (VAE) Evidence Lower Bound (ELBO) proposed in Kingma &Ba (2014);
Revised: and Achille & Soatto (2016) continue to show that the Variational Auto Encoder (VAE) Evidence Lower Bound (ELBO) proposed in Kingma &

Original:  Rezende et al.
Revised: Ba (2014); Rezende et al.

Original: (2014) is a special case of the IB bound when
Revised: (2014) is a special case of the IB bound when

Original: y
Revised: y

Original: = x and β = 1 :
Revised: = x and β = 1 :

Original: [Equation]
Revised: [Equation]

Original: where i represents the identity of data-point x i .
Revised: where i represents the identity of data-point x i .

Original: Interestingly, the VAE perspective considers the bound to optimize a variational distribution q ( z | x ) , whilst the IB perspective prescribes that q ( z | x ) in the ELBO is not a variational posterior but the true encoder p ( z | x ) , and instead p ( y | z ) and p ( z ) are the distributions approximated with variational counterparts. 
Revised: Interestingly, the VAE perspective considers the bound to optimize a variational distribution q ( z | x ) , whilst the IB perspective prescribes that q ( z | x ) in the ELBO is not a variational posterior but the true encoder p ( z | x ) , and instead p ( y | z ) and p ( z ) are the distributions approximated with variational counterparts. However, deriving the ELBO on a conditional distribution p ( y | x ) leads to a KL divergence between the approximate posterior q ( z | x ) and the real posterior p ( z | x ) , which is not available at test time, rather than the prior p ( z ) .

Revised: Hence, the information bottleneck perspective provides a more fundamental motivation for the bound in equation 1.

Original: From yet another perspective, equation 1 can be interpreted as a domain-translating beta-VAE (Higgins et al.
Revised: Alternatively, equation 1 can be interpreted as a domain-translating beta-VAE (Higgins et al.

Original: (2016)), where an input image is encoded into a latent space and decoded into the target domain.
Revised: (2016)), where an input image is encoded into a latent space and decoded into the target domain.

Original: The Lagrange multiplier β then controls the trade-off between rate and distortion, as argued by Alemi et al.
Revised: The Lagrange multiplier β then controls the trade-off between rate and distortion, as argued by Alemi et al.

Original: In this work, we follow the IB interpretation of the bound in equation 1 and leave the evaluation of our proposed variational inference scheme in other models such as the VAE for further work.
Revised: However, this does not naturally arise from deriving the ELBO using Jensen’s inequality on the conditional likelihood p ( y | x ) .

Original:  M ETHOD
Revised:  M ETHOD

Revised: In this work, we follow the IB interpretation of the bound in equation 1 and leave the evaluation of our proposed variational inference scheme in other models such as the VAE for further work.

Original: At the heart of our proposal lies the assumption that neuron networks can be effective under strong activation quantization schemes.
Revised: At the heart of our proposal lies the assumption that neuron networks can be effective under strong activation quantization schemes.

Original: For clarity, we start with presenting the derivation for a single layer information bottleneck.
Revised: We start with presenting the derivation if the model in the context of a single latent-layer information bottleneck, following the single data-point loss in equation 1 and dropping the subscript n for clarity, with ﬁgure 4 for visual reference:

Original: We model the encoder using a mean-ﬁeld approximation p ( z | x ) = (cid:81) k p
Revised: [Equation] To impose a ﬂexible, multi-modal distribution over z , we ﬁrst make a mean-ﬁeld assumption p ( z | x ) = (cid:81) k p

Original: ( z k | x ) .
Revised: ( z k | x ) .

Original: We individually quantize the domain of each of the K scalar latent variables z k into a small set v = v 1 , . .
Revised: We then quantize the domain of each of the K scalar latent variables z k [Table] [Table]

Original: . , v C consisting of a small number of C different real values (e.g. C = 11 ).
Revised: [Figure] such that only a small set of potential values remain: z k ∈ v = { v 1 , . . . , v C } ∀ k with e.g. C = 11 , see ﬁg. 3.

Original: We usea neural network f θ ( · ) with a softmax output over the  quantization bins to model the distributions over the  values for z : [Equation]
Revised: We use the same quantization scheme, deﬁned by v , for all the latents. The speciﬁc scheme is considered part of the prior, and is discussed in ﬁgure 5. Although the number of values of z is relatively small, the space of values over all the latent variables grows exponentially.

Original:  To optimize the parameters θ with Stochastic Gradient Descent (SGD), we need to derive a fully differentiable loss for our model. Providing ﬁgure 4 for visual reference, we derive the bound as
Revised: We thus take to a monte-carlo estimation scheme. To optimize the parameters θ with Stochastic Gradient Descent (SGD), we need to derive a fully differentiable sampling scheme that allows us to sample values of z . .

Original: [Table] [Table] [Figure] follows, starting with the single-data-point loss in equation 1:
Revised: To formulate this, we re-parametrize the expectation over z in equation 3 using a set of variables s k ∈ { 1 , . . .

Original: [Equation] We re-parametrize z as a set of bin-indexing variables s k ∈ { 1 , . . . , C }  :
Revised: , C } which index the value vector v , allowing us to use a softmax function to represent the distribution over z :

Original: [Equation]
Revised: [Equation]

Original: We model q θ ( y | z )  with a softmax layer f θ ( z ) = softmax( Wz + b )  ( abusing notation to indicate element-wisevector indexing with v [ s ] ):
Revised: These indexing values s are then used in conjunction with values v as in input for q θ ( y | z ) , which is modelled with a small network f θ ( · ) : ( abusing notation to indicate element-wise indexing with v [ s ] ):

Original: [Equation]
Revised: [Equation]

Original: To enable sampling from the discrete variables s , we use the Gumbel-Max trick (Gumbel (1954)), denoted gumb() , re-parameterizing the expectation E s ∼ p with uniform noise (cid:15) ∼ U (0 , 1) : 
Revised: To enable sampling from the discrete variables s , we use the Gumbel-Max trick (Gumbel (1954)), denoted gumb() , re-parameterizing the expectation E s ∼ p with uniform noise (cid:15) ∼ U (0 , 1) : [Equation]

Original: [Equation]
Revised: [Equation]

Original: As the argmax is not differentiable, we approximate this expectation using the Gumbel-Softmax trick (Maddison et al.
Revised: As the argmax is not differentiable, we approximate this expectation using the Gumbel-Softmax trick (Maddison et al.

Original: ; Jang et al.
Revised: ; Jang et al.

Original: (2016)), which generates samples that smoothly deform into one-hot samples as the softmax temperature τ approaches 0 :
Revised: (2016)), which generates samples that smoothly deform into one-hot samples as the softmax temperature τ approaches 0 .

Revised: Using the inner product (denoted · ) of the approximate one-hot samples and v , we create samples from z :

Original: [Equation]
Revised: [Equation]

Original: In practice, we anneal τ from 1 . 0 to 0 .
Revised: In practice, we anneal τ from 1 . 0 to 0 .

Original: 5 during the training process, as proposed by Yang et al.
Revised: 5 during the training process, as proposed by Yang et al.

Original: (2017) to reduce gradient variance initially, at the risk of introducing bias.
Revised: (2017) to reduce gradient variance initially, at the risk of introducing bias.

Original: For the remainder of this work, we will refer to the latent variables as z in lieu of s , for clarity.

Original: To conclude our derivation, we use a ﬁxed SQUAD distribution to model the  marginalp ( z ) : = r ( z ) as shown in ﬁgure 5.
Revised: To conclude our derivation, we use a ﬁxed SQUAD distribution to model the variational marginal r ( z ) as shown in ﬁgure 5.

Original: We can then derive the KL term analytically following the deﬁnition for discrete distributions.
Revised: We can then derive the KL term analytically following the deﬁnition for discrete distributions.

Original: Using the fact that the KL divergence is additive for independent variables:
Revised: Using the fact that the KL divergence is additive for independent variables, we get our ﬁnal loss:

Original: [Equation]
Revised: [Equation]

Original: The derived lower bound is fully differentiable, allowing it to be optimized with common SGDbased methods.
Revised: For the remainder of this work, we will refer to the latent variables as z in lieu of s , for clarity.

Original: At test time, we can approximate the predictive function p ( y ∗ | x ∗ ) for a new data-point x ∗ by taking
Revised: At test time, we can approximate the predictive function p ( y ∗ | x ∗ ) for a new data-point x ∗ by taking

Original: T samples from the latent variables z i.e. ˆ z t ∼ p ( z | x ∗ ) , and averaging the predictionss for y ∗ :
Revised: T samples from the latent variables z i.e. ˆ z t ∼ p ( z | x ∗ ) , and averaging the predictions for y ∗ :

Original: [Equation]
Revised: [Equation]

Revised: A downside of mean-ﬁeld variational approximations is that latents are assumed to be uncorrelated.

Original: We improve the ﬂexibility of the proposed model by creating a hierarchical set of latent variables.
Revised: We can bring some correlation to the latents, and extend the ﬂexibility of the proposed model, by creating a hierarchical set of latent variables.

Original: The joint  distribution of L layers of latents is then:
Revised: We maintain a mean-ﬁeld assumption for the prior, and the joint posterior distribution of L layers of latents factorizes as follows:

Original: [Equation]
Revised: [Equation]

Original: and is straightforwardly implemented with a simple ancestral sampling scheme.

Original: Finally, we have q θ ( y | z 
Revised: With q θ ( y | z 1 , . .

Original: [Equation]
Revised: . , z L ) =

Original: [Equation] [Equation]
Revised: q θ ( y | z L ) .

Original: = [Equation]
Revised: This is straightforwardly implemented with a simple ancestral sampling scheme.

Original: Interestingly, the strong quantization proposed in our method can itself be considered an additional information bottleneck, as it exactly upper-bounds the number of bits per latent variable.
Revised: Interestingly, the strong quantization proposed in our method can itself be considered an additional information bottleneck, as it exactly upper-bounds the number of bits per latent variable.

Original: Such bottlenecks are theorized to have a beneﬁcial effect on generalization (Tishby et al.
Revised: Such bottlenecks are theorized to have a beneﬁcial effect on generalization (Tishby et al.

Original: ; Achille & Soatto (2016); Alemi et al.
Revised: ; Achille & Soatto (2016); Alemi et al.

Original: and we can directly control this bottleneck by varying the number of quantization bins.
Revised: and we can directly control this bottleneck by varying the number of quantization bins.

Original: The computational complexity of the method, as well as the number of model parameters θ , scale linearly in C , i.e. O ( C ) (with C the number of quantization bins).
Revised: The computational complexity of the method, as well as the number of model parameters θ , scale linearly in C , i.e. O ( C ) (with C the number of quantization bins).

Original: It is thus suitable for large-scale inference.
Revised: It is thus suitable for large-scale inference.

Original: We would like to stress that the proposed method differs from work that leverages the
Revised: We would like to stress that the proposed method differs from work that leverages the

Original: Gumbel-Softmax trick to model categorical latent variables: our proposal models continuous scalar latent variables by quantizing their domain and modeling belief with a multinomial distribution.
Revised: Gumbel-Softmax trick to model categorical latent variables: our proposal models continuous scalar latent variables by quantizing their domain and modeling belief with a multinomial distribution.

Original: Categorical latent variable models would incur a much larger polynomial complexity penalty of
Revised: Categorical latent variable models would incur a much larger polynomial complexity penalty of

Original: O ( C 2 ) .
Revised: O ( C 2 ) .

Original: The method is easily integrated into existing deep neural network architectures, and SQUAD layer implementations are provided at github.com/anonymized .
Revised: The method is easily integrated into existing deep neural network architectures, and SQUAD layer implementations are provided at github.com/anonymized .

Original: Matrix-factorization variant  To improving the tractability of using a large number of quantization bins, we propose a variant of SQUAD that uses a matrix factorization scheme to improve the parameter efﬁciency.
Revised: Matrix-factorization variant In the natural image domain, we anticipate a need for a high amount of information per variable, for which a small number of bins does not sufﬁce. To improving the tractability of using a large number of quantization bins, we propose a variant of SQUAD that uses a matrix factorization scheme to improve the parameter efﬁciency.

Original: Formally, equation equation 3 becomes:
Revised: Formally, equation 4 becomes:

Original: [Equation]
Revised: [Equation]

Original: with full layer weights W (cid:48)(cid:48) and W (cid:48) respectively of shape ( K, B, C ) and ( | X | , K, B ) , where K denotes the number of neurons, B the number of factorization, C number of quantization bins and | X | the input dimensionality.
Revised: with full layer weights W (cid:48)(cid:48) and W (cid:48) respectively of shape ( K, B, C ) and ( | X | , K, B ) , where K denotes the number of neurons, B the number of factorization inputs, C number of quantization bins and | X | the input dimensionality.

Original: The right side of ﬁgure 4 provides visual clariﬁcation of this variant.

Original: To improve the parameter efﬁciency, we can learn W (cid:48) per layer as well, resulting in shape ( | X | , 1 , B ) , which is found to be beneﬁcial for large C in the hyper-parameter search presented in section 5.
Revised: To improve the parameter efﬁciency, we can learn W (cid:48) per layer as well, resulting in shape ( | X | , 1 , B ) , which is found to be beneﬁcial for large C by the hyperparameter search presented in section 5.

Original: We depict this alternative model on the right side of ﬁgure 4 and will refer to it as SQUAD-factorized.
Revised: We depict this alternative model on the right side of ﬁgureand will refer to it as SQUAD-factorized.

Original: We leave further extensions such as Network-in-Network(Lin et al.
Revised: We leave further extensions such as Network-in-Network (Lin et al.

Original: for future work.
Revised: for future work.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Outside the realm of DLVMs, other methods have been explored for predictive uncertainty.
Revised: Outside the realm of DLVMs, other methods have been explored for predictive uncertainty.

Original: Lakshminarayanan et al.
Revised: Lakshminarayanan et al.

Original: (2017) propose deep ensembles: straightforward averaging of predictions from a small set of separately adversarially trained DNNs.{+r(z) r(z) 0 .4 .
 [Figure] (z [Figure]+}
Revised: (2017) propose deep ensembles: straightforward averaging of predictions from a {+r(z) r(z) 0 .4 .
 [Figure] (z [Figure]+}

Revised: z

Original:  Although highly scalable, this method requires retraining a model up to 10 times, which can be inhibitively expensive for large datasets.
Revised: [Figure] [Figure] small set of separately adversarially trained DNNs. Although highly scalable, this method requires retraining a model up to 10 times, which can be inhibitively expensive for large datasets.

Original: Gal & Ghahramani (2015b) propose the use of dropout (Srivastava et al.
Revised: Gal & Ghahramani (2015b) propose the use of dropout (Srivastava et al.

Original: at test time  and present a Bayesian neural network interpretation of this method.
Revised: at test time (Monte Carlo Dropout, or MC-Dropout) and present a Bayesian neural network interpretation of this method.

Original: A follow-up work by Gal et al.
Revised: A follow-up work by Gal et al.

Original: (2017) explores the use of Gumbel-Softmax to smoothly deform the dropout noise to allow optimization of the dropout rate during training.
Revised: (2017) explores the use of Gumbel-Softmax to smoothly deform the dropout noise to allow optimization of the dropout rate during training.

Original: A downside of MC-dropout is the limited ﬂexibility of the ﬁxed bimodal delta-peak distribution imposed on the weights, which requires a large number of samples for good estimates of uncertainty.
Revised: A downside of MC-Dropout (MCD) is the limited ﬂexibility of the ﬁxed bi-modal delta-peak distribution imposed on the weights, which requires a large number of samples for good estimates of uncertainty.

Original: van den Oord et al.
Revised: van den Oord et al.

Original: (2017) propose the use of vector quantization in variational inference, quantizing a multi-dimensional embedding, rather than individual latent variables, and explore this in the context of auto-encoders.[-r(z) r(z) 0 .4 .
 [Figure]-]
Revised: (2017) propose the use of vector quantization in variational inference, quantizing a multi-dimensional embedding, rather than individual latent variables, and explore this in the context of auto-encoders.[-r(z) r(z) 0 .4 .
 [Figure]-]

Original: In the space of learning non-linearities, Su et al.
Revised: In the space of learning non-linearity’s, Su et al.

Original: (2017) explore a ﬂexible non-linearity that can assume the form of most canonical activations.
Revised: (2017) explore a ﬂexible non-linearity that can assume the form of most canonical activations.

Original: More ﬂexible distributions have been explored for distributional reinforcement learning by Dabney et al.
Revised: More ﬂexible distributions have been explored for distributional reinforcement learning by Dabney et al.

Original: (2017) using quantile regression, of which can be seen as a special case of SQUAD where the bin values are learned but have ﬁxed uniform probability.
Revised: (2017) using quantile regression, of which can be seen as a special case of SQUAD where the bin values are learned but have ﬁxed uniform probability.

Original: Categorical distributions on scalar variables have been used to model more ﬂexible
Revised: Categorical distributions on scalar variables have been used to model more ﬂexible

Original: Bayesian neural network posteriors as by Shayer et al.
Revised: Bayesian neural network posteriors as by Shayer et al.

Original: The use of a mixture of diracs distribution to approximate a variety of distributions was proposed by Schrempf et al.
Revised: The use of a mixture of diracs distribution to approximate a variety of distributions was proposed by Schrempf et al.

Original:  R ESULTS
Revised:  R ESULTS

Original: Quantifying the quality of uncertainty estimates of models remains an open problem.
Revised: Quantifying the quality of uncertainty estimates of models remains an open problem.

Original: Various methods have been explored in previous works, such as relative entropy Louizos & Welling (2017); Gal & Ghahramani (2015a), probability calibration, and proper scoring rules Lakshminarayanan et al.
Revised: Various methods have been explored in previous works, such as relative entropy Louizos & Welling (2017); Gal & Ghahramani (2015a), probability calibration, and proper scoring rules Lakshminarayanan et al.

Original: Although interesting in their own right, these metrics do not directly measure a good ranking of predictions, nor indicate applicability in high-risk domains.
Revised: Although interesting in their own right, these metrics do not directly measure a good ranking of predictions, nor indicate applicability in high-risk domains.

Original: Proper scoring rules are the exception, but a model with good ranking ability does not necessarily exhibit good performance on proper scoring rules: any score that provides relative ordering sufﬁces and does not have to reﬂect true calibrated probabilities.
Revised: Proper scoring rules are the exception, but a model with good ranking ability does not necessarily exhibit good performance on proper scoring rules: any score that provides relative ordering sufﬁces and does not have to reﬂect true calibrated probabilities.

Original: In fact, well-ranked conﬁdence scores can be re-calibrated (Niculescu-Mizil & Caruana (2005)) after training to improve performance on proper scoring rules and calibration metrics.
Revised: In fact, well-ranked conﬁdence scores can be re-calibrated (Niculescu-Mizil & Caruana (2005)) after training to improve performance on proper scoring rules and calibration metrics.

Original: In order to evaluate the applicability of the model in high-risk ﬁelds such as medicine, we want to quantify how models perform under a desired risk requirement.
Revised: In order to evaluate the applicability of the model in high-risk ﬁelds such as medicine, we want to quantify how models perform under a desired risk guarantee.

Original: We propose to use the selection
Revised: We propose to use the Selection with

Original: [Table] with guaranteed risk (SGR) method 1 introduced by Geifman & El-Yaniv (2017) to measure this.
Revised: Guaranteed Risk (SGR) method 1 introduced by Geifman & El-Yaniv (2017) to measure this.

Original: In summary, the SGR method deﬁnes a selective classiﬁer using a trained neural network and can guarantee auser-selected desired risk with high probability (e.g. 99%), by selectively rejecting data points with a predicted conﬁdence score below an optimal threshold.
Revised: In summary, the SGR method uses a conﬁdence score, e.g. the estimated probability of the prediction, and rejects the prediction if this score is below a certain threshold. This threshold is optimized on a hold-out set to the tightest possible value that still guarantees a desired minimial error rate with high probability (e.g. 99%). SGR is formally proven, and empirically shown, to ensure this minimal error rate under i.i.d. assumptions.

Original: To limit the inﬂuence of hyper-parameters on the comparison, we use theautomated optimization method TPE (Bergstra et al.
Revised: We propose to use SGR as evaluation metric for conﬁdence estimation methods. Speciﬁcally, on a hold-out set we measure the percentage of predictions that are accepted by SGR under a speciﬁc risk guarantee, i.e. the coverage . To illustrate, one can imagine that a trivial solution to guarantee no errors is to reject all data-points, which has 0% coverage. By lowering the conﬁdence threshold, coverage goes up, but the accuracy of the method goes down.

Original: for both baselines and our models.
Revised: SGR ﬁnds the optimal threshold that ensures the accuracy does not drop below a predetermined threshold, and by measuring the coverage under this threshold we can effectively compute the real world value of a conﬁdence estimation methods. To limit the inﬂuence of hyper-parameters on the comparison, we use a extensive hyper-parameter scheme on all variants and baselines. Speciﬁcally we use the automated optimization method TPE (Bergstra et al. over a maximum of 1000 evaluations per model per experiment.

Original: The hyper-parameters are optimized for coverage at 2% risk ( δ = 0 .
Revised: The hyperparameters are optimized for coverage at 2% risk ( δ = 0 .

Original: 01 ) on fashionMNIST, and sub-sequentially evaluated on notMNIST.
Revised: 01 ) on fashionMNIST (Xiao et al. , and we evaluate the optimal hyperparameters on another dataset, notMNIST, to measure how sensitive the hyperparameters are to a change in the data-distribution.

Original: Larger models are evaluated onSVHN.
Revised: Larger models are evaluated on

Revised: Street View Housing Numbers (SVHN) (Netzer et al. (2011)) using the same optimization scheme.

Revised: For the details on the evaluated search space and chosen hyper-parameters we refer the reader to the appendix.

Original: We compare 2 our model against plain MLPs, MCDropout using Maxout activations 3 (Goodfellow et al.
Revised: We compare 2 our model against plain Multi-Layer Perceptrons (MLPs) with ReLU activations, MCDropout using Maxout activations 3 (Goodfellow et al.

Original: ; Chang & Chen (2015)) and an information bottleneck model using mean-ﬁeld Gaussian distributions.
Revised: ; Chang & Chen (2015)) and an information bottleneck model using mean-ﬁeld Gaussian distributions.

Original: We evaluate the complementary deep ensembles technique (Lakshminarayanan et al.
Revised: We evaluate the complementary deep ensembles technique (Lakshminarayanan et al.

Original: (2017)) for all methods.
Revised: for all methods.

Original: 5. M AIN RESULTS
Revised: 5. M AIN RESULTS

Original: We start our analysis by comparing the predictive uncertainty of 2-layer models with 32 latent variables per layer.
Revised: We start our analysis by comparing the predictive uncertainty of 2-layer models with 32 latent variables per layer.

Original: In ﬁgure 6 we visualize the risk/coverage trade-off achieved using the predicted uncertainty as a selective threshold, and present coverage results in table 1.
Revised: In ﬁgure 6 we visualize the risk/coverage trade-off achieved using the predicted uncertainty as a selective threshold, and present coverage results in table 1.

Original: Overall, we ﬁnd that
Revised: Overall, we ﬁnd that

Original: SQUAD performs signiﬁcantly better than plain MLPs and deep Gaussian IB models, and we tentatively attribute this to the increased ﬂexibility of the multinomial distribution.
Revised: SQUAD performs signiﬁcantly better than plain MLPs and deep Gaussian IB models, and we tentatively attribute this to the increased ﬂexibility of the multinomial distribution.

Original: Compared to a Maxout
Revised: Compared to a Maxout

Original: MCdropout model with a similar number of weights, SQUAD appears to have a slight —though not signiﬁcant— advantage, despite the strong quantization scheme, especially at low risk. Deep ensembles improve results forall methods, which ﬁts the hypothesis that ensembles integrate over a form of weight uncertainty.
Revised: (2017) in that we use Softmax Response (SR) – the probability taken from the softmax output for the most likely class – as the conﬁdence score for all methods.

Original:  El-Yaniv (2017) instead proposed to use the variance of the probabilities for MCdropout, but our experiments showed that SR paints MCdropout in a more favorable light.
Revised: Geifman & El-Yaniv (2017) instead proposed to use the variance of the probabilities for MCdropout, but our experiments showed that SR paints MCDropout in a more favorable light.

Original: All models are optimized with ADAM (Kingma & Ba (2014)), weight initialization as proposed by He et al.
Revised: All models are optimized with ADAM (Kingma & Ba (2014)), weight initialization as proposed by He et al.

Original: (2015), a weight decay of 10 − 5 and adaptive learning rate decay scheme — 10x reduction after 10 epochs of no validation accuracy improvement— and use early stopping after 20 epochs of no improvement.
Revised: (2015), a weight decay of 10 − 5 and adaptive learning rate decay scheme — 10x reduction after 10 epochs of no validation accuracy improvement— and use early stopping after 20 epochs of no improvement.

Original: We found this baseline to perform stronger in comparison to conventional ReLU MCdropout models, under equal number of latent variables. 
Revised: We found this baseline to perform stronger in comparison to conventional ReLU MCdropout models, under equal number of latent variables. [Table] [Table]

Original: [Table]
Revised: [Table]

Revised: MCdropout model with a similar number of weights, SQUAD appears to have a slight —though not signiﬁcant— advantage, despite the strong quantization scheme, especially at low risk guarantees.

Revised: Deep ensembles improve results for all methods, which ﬁts the hypothesis that ensembles integrate over a form of weight uncertainty.

Revised: When comparing the optimal hyperparameters found for fashionMNIST on the similar notMNIST dataset

Revised: [Table]

Revised: we ﬁnd that SQUAD shows strong performance, as shown in table 2.

Revised: This provides some evidence that despite the increase in number hyperparameters in SQUAD

Original: [Table]
Revised: [Table]

Revised: the optimal settings are more robust to a change in the data distribution

Revised: [Table]

Revised: we ﬁnd that the ﬂexible variational distribution is used to its full advantage.

Revised: A) visualizes a subset of interesting stereotypical distributions we hope to ﬁnd in the model.

Revised: (B) summarizes distributions predicted by the model similar to stereotypes, discovered by looking at predicted distributions with low KL.

Revised: Figure

Revised: C) shows how often distributions similar to stereotypes arise, as measured by the KL distance (

Revised: KL is closer to stereotypes)

Original: 5. N ATURAL I MAGES
Revised: 5. N ATURAL I MAGES

Original: To explore larger models trained on natural image datasets, we lightly tune hyper-parameters on 256-latent 2-layer models over 100 TPE evaluations.
Revised: To explore larger models trained on natural image datasets, we re-tune hyper-parameters on 256latent 2-layer models over 100 TPE evaluations.

Original: As SVHN contains natural images in color, we anticipate a need for a higher amount of information per variable.
Revised: As SVHN contains natural images in color, we anticipate a need for a higher amount of information per variable.

Original: We thus explore the effect of the matrix-factorized variant.
Revised: We thus explore the effect of the matrix-factorized variant.

Original: As shown in table 3, SQUAD-factorized outperforms the non-factorized variant.
Revised: As shown in table 3, SQUAD-factorized outperforms the non-factorized variant.

Original: Considering the computational cost at the optimum of a 4-neuron factorization ( B = 4 ) with C =37 quantization bins, the model clocks 3.4 million weights.
Revised: Considering the computational cost at the optimum of a 4-neuron factorization ( B = 4 ) with C =37 quantization bins, the model clocks 3.4 million weights.

Original: In comparison, the optimum for the presented MCdropout results has C =11, using 9.0 million weights.
Revised: In comparison, the optimum for the presented MCdropout results has C =11, using 9.0 million weights.

Original: On an NVIDIA Titan XP, the dropout baseline takes 13s per epoch on average, while SQUAD-factorized spans just 9s.
Revised: On an NVIDIA Titan Xp, the dropout baseline takes 13s per epoch on average, while SQUAD-factorized spans just 9s.

Original: To evaluate the sample efﬁciency of the methods, we compare results at T = 4 samples. 
Revised: To evaluate the sample efﬁciency of the methods, we compare results at T = 4 samples. We ﬁnd that

Original: We ﬁnd that SQUAD’s results suffer less from undersampling than MCdropout.
Revised: SQUAD’s results suffer less from under-sampling the predictive distribution than MCdropout.

Original: We tentatively attribute the sample efﬁciency to the ﬂexible approximating posterior on the activations, which is in stark contrast to the rigid approximating distribution that MCdropout imposes on the weights.
Revised: We tentatively attribute the sample efﬁciency to the ﬂexible approximating posterior on the activations, which is in stark contrast to the rigid approximating distribution that MCdropout imposes on the weights.

Original: In conclusion, SQUAD comes out favorably in a resource-constrained environment.
Revised: In conclusion, SQUAD comes out favorably in a resource-constrained environment.

Original: 5. A NALYSIS OF LATENT VARIABLE DISTRIBUTIONS
Revised: 5. A NALYSIS OF LATENT VARIABLE DISTRIBUTIONS

Original: In order to evaluate if the proposed variational distribution does not simply collapse into single mode predictions, we want to ﬁnd out what type of distributions the model predicts over the latent variables.
Revised: In order to evaluate if the proposed variational distribution does not simply collapse into single mode predictions, we want to ﬁnd out what type of distributions the model predicts over the latent variables.

Original: We visualize the forms of predicted distributions in ﬁgure 7a.
Revised: We visualize the forms of predicted distributions in ﬁgure 7a.

Original: While we capture only a small subset of potential multi-modal behavior, this demonstrates that the model indeed utilizes the distribution to its full potential.
Revised: Although this showcases only a small subset of potential multi-modal behavior that emerges, this demonstrates that the model indeed utilizes the distribution to its full potential.

Original: To provide an intuition on how these predicted distributions emerge, we presenta ﬁgure 9 in the appendix.
Revised: To provide an intuition on how these predicted distributions emerge, we present ﬁgure 9 in the appendix.

Original: In ﬁgure 8 we visualize one of the non-linearities that the method learns for a 1-dimensional input SQUAD-factorized model.
Revised: In ﬁgure 8 we visualize one of the activation functions that the method learns for a 1-dimensional input SQUAD-factorized model.

Original: The learned non-linearities resemble “peaked” sigmoid activations, which can be interpreted as a combination of an RBF kernel and sigmoid.
Revised: The learned activation functions resemble “peaked” sigmoid activations which can be interpreted as a combination of an RBF kernel and sigmoid.

Original: This provides food for thought on how non-linearities for conventional neural networks can be designed, and the effect of using such a nonlinearity will be studied in further work.
Revised: This provides food for thought on how non-linearity’s for conventional neural networks can be designed and the effect of using such a non-linearity can be studied in further work.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: As an interesting side effect the method breaks the gradient ﬂow normally present in deep neural networks. As the input correlates stronger with the learned weights, the output adjusts accordingly. With SQUAD, stronger correlation with learned weights perturbates the belief of values, rather than the values itself. As a beneﬁt, this can potentially decrease the – arguably underdeﬁned (Santurkar et al.
Revised: In this work we have proposed a new ﬂexible class of variational distributions.

Original: (2018)) – problem of internal covariate shift, as the output domain is constrained.

Original: Furthermore, this allows the model to learn the non-linearity. As a downside, this breaks certain assumptions on which best-practice hyperparameter settings have been determined, such as weight initialization and activation functions (He et al. and would require further study. Conclusion In this work, wehave proposed a new ﬂexible class of variational distributions and achieve strong performance when applied  to a deep variational information bottleneck model.
Revised: To measure the effectiveness for real world classiﬁcation we applied the class to a deep variational information bottleneck model.

Original: By placing a quantization-based distribution on the activations, we can computeprincipled uncertainty estimates.
Revised: By placing a quantization-based distribution on the activations we can compute uncertainty estimates over the outputs.

Original: We proposed an evaluation scheme motivated by the need in real-world domains to guarantee a minimal risk.
Revised: We proposed an evaluation scheme motivated by the need in real-world domains to guarantee a minimal risk.

Original: The results presented indicate that SQUAD provides an improvement over plain neural networks, Gaussian IB and MCDropout in high-risk domains and that the ﬂexible distribution is used to its full advantage.
Revised: The results presented indicate that SQUAD provides an improvement over plain neural networks and Gaussian information bottleneck models. In comparison to a MCDropout model, which approximates a Bayesian neural network, we get competitive performance.

Revised: Moreover, qualitatively we ﬁnd that the ﬂexible distribution is used to its full advantage is sample efﬁcient.

Revised: The method learns interesting non-linearity’s, is tractable and scale-able, and as the output domain is constrained, no batch normalization techniques are required.

Original: A PPENDIX 7. E FFECT OF HYPER - PARAMETERS ON COVERAGE :
Revised: Various directions for future work arise. The improvement of ensemble methods over individual models indicates that there remains room for improvement for capturing the full uncertainty of the output, and thus a fully Bayesian approach to SQUAD which would include weight uncertainty, shows promise. The ﬂexible class allows us to deﬁne a wide variety of interesting priors, which provides opportunity to study interesting priors that are hard to deﬁne as a continuous density. Likewise, more effective initialization of parameters for the proposed method requires further attention.

Revised: Orthogonally, the proposed class can be applied to other variational objectives as well, such as the variational auto-encoder.

Revised: Finally, the discretized nature of the variables allows for the analytical computation of other divergences such as mutual information and the Jensen-Shannon divergence, the effectiveness of which remains to be studied.

Revised:  A PPENDIX 7. E FFECT OF HYPER - PARAMETERS ON COVERAGE :

Original: The optimal conﬁguration of hyper-parameters and bin priors have been determined using 700 evaluations selected using TPE.
Revised: The optimal conﬁguration of hyper-parameters and bin priors have been determined using 700 evaluations selected using TPE.

Original: The space of parameters explored is as follows, presented in the hyperopt
Revised: The space of parameters explored is as follows, presented in the hyperopt

Original: API for transparency: # Shared
Revised: API for transparency: # Shared

Original: C: quniform(2, 10, 1) * 2 + 1, dropout rate: uniform(0.01, .95), lr: loguniform(log(0.0001), log(0.01)), batch_size: qloguniform(log(32), log(512), 1) # SQUAD & Gaussian kl_multiplier: loguniform(log(1e-6), log(0.01)), init_scale: loguniform(log(1e-3), log(20)), # SQUAD use_bin_probs: choice([’uni’, ’gaus’]), use_bins: choice([’equal_prob_gaus’, ’linearly_spaced’]), learn_bin_values: choice([ ’per_neuron’, ’per_layer’, ’fixed’]),
Revised: C: quniform(2, 10, 1) * 2 + 1, dropout rate: uniform(0.01, .95), lr: loguniform(log(0.0001), log(0.01)), batch_size: qloguniform(log(32), log(512), 1) # SQUAD & Gaussian kl_multiplier: loguniform(log(1e-6), log(0.01)), init_scale: loguniform(log(1e-3), log(20)), # SQUAD use_bin_probs: choice([’uni’, ’gaus’]), use_bins: choice([’equal_prob_gaus’, ’linearly_spaced’]), learn_bin_values: choice([ ’per_neuron’, ’per_layer’, ’fixed’]),

Original: In ﬁgure 10 we visualize the pairwise effect of these hyper-parameters on the coverage.
Revised: In ﬁgure 10 we visualize the pairwise effect of these hyper-parameters on the coverage.

Original: The optimal conﬁguration found in for the main SQUAD model are: batch size: 244, KL multiplier: 0.0027, learn bin values: per layer, p ( z ) : uniform, v : linearly spread over (-3.5,3.5), lr: 0.0008, C : 15, initialization scale: 3.214.
Revised: The optimal conﬁguration found in for the main SQUAD model are: batch size: 244, KL multiplier: 0.0027, learn bin values: per layer, p ( z ) : uniform, v : linearly spread over (-3.5,3.5), lr: 0.0008, C : 15, initialization scale: 3.214.