Evidential Softmax for Sparse Multimodal Distributions in Deep Generative Models

1 Introduction . Learning deep generative models over discrete probability spaces has enabled state-of-the-art performance across tasks in computer vision [ 1 ] , natural language processing [ 2 ] – [ 4 ] , and robotics [ 5 ] – [ 7 ] . The probability distributions of deep generative models are often obtained from a neural network using the softmax transformation . Since the softmax function has full support , the logarithm of the output probabilities is well-defined , which is important for common probabilistic loss functions such as the negative log likelihood , cross entropy , and Kullback-Leibler ( KL ) divergence . Several applications highlight the importance of achieving sparsity in high dimensional dense distributions to make downstream tasks computationally feasible and perhaps even interpretable [ 8 ] – [ 11 ] . In discrete variational autoencoders ( VAEs ) , for instance , calculating the posterior probabilities exactly requires marginalization over all possible latent variable assignments , which is intractable for large latent spaces [ 9 ] , [ 10 ] . Stochastic approaches to circumvent exact marginalization include the score function estimator [ 12 ] , [ 13 ] , which suffers from high variance , and continuous relaxations of the discrete latent variable , such as Gumbel-Softmax [ 14 ] , [ 15 ] , which introduce bias . Alternatively , high-dimensional latent spaces can be tractably marginalized using normalization functions that produce sparse probability distributions . Several sparse alternatives to softmax have been proposed , including sparsemax [ 8 ] , α-entmax [ 16 ] , and sparsehourglass [ 17 ] . However , such approaches often collapse the multimodality of distributions , resulting in unimodal probability distributions [ 10 ] ( see Fig . 1 ) . When marginalizing over discrete latent spaces , maintaining this multimodality is crucial for applications such as image generation , where the latent space is multimodal in nature [ 1 ] , and machine translation , where multiple words can be valid translations and require 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . context to distinguish between the optimal word choice [ 9 ] . Furthermore , these approaches either require the construction of alternative loss functions to the negative log-likelihood , which is undefined for zero-valued output probabilities [ 8 ] , [ 16 ] – [ 18 ] , or they are applied as a post-hoc transformation at test time [ 10 ] . For the latter case , the latent distributions obtained through these existing sparse normalization functions and their resulting posteriors are not trained directly . The post-hoc sparsification procedure introduced by Itkina et al . [ 10 ] provides a method for obtaining sparse yet multimodal distributions at test time . We generalize this method to a sparse normalization function termed evidential softmax ( ev-softmax ) which can be applied during both training and test time . We define continuous approximations of this function that have full support and are thus compatible with typical loss terms , such as negative log-likelihood and KL divergence . We use this full-support approximation of ev-softmax to directly optimize discrete generative models for objective functions with probabilistic interpretations . This approach is in contrast with post-hoc sparsification methods that optimize a proxy model [ 10 ] and methods that use non-probabilistic objective functions , such as sparsemax [ 8 ] . We evaluate these discrete generative models at test time with the ev-softmax function to obtain sparse yet multimodal latent distributions . Contributions : This paper proposes a strategy for training neural networks with sparse probability distributions that is compatible with negative log-likelihood and KL-divergence computations . We generalize the evidential sparsification procedure developed by Itkina et al . [ 10 ] to the ev-softmax function . We prove properties of ev-softmax and its continuous approximation , which has full support . Our approach outperforms existing dense and sparse normalization functions in distributional accuracy across image generation and machine translation tasks . 2 Background . Notation : We denote scalars , vectors , matrices , and sets as a , a , A , and A , respectively . The vectorized Kronecker delta δi is defined to be 1 at the ith index and 0 at all the other indices . The K-dimensional simplex is written as ∆K : = { ξ ∈ RK+1 : 1T ξ = 1 , ξi ≥ 0 } . The KL divergence of probability distribution p from distribution q is denoted as KL ( p || q ) . 2.1 Discrete Variational Autoencoders . To model a data generating distribution , we consider data x ∈ Rd , categorical input variable y ∈ Y , categorical latent variable z ∈ Z , and parameters θ ∈ Rp . In the VAE model , z is sampled from a prior pθ ( z ) , and we generate x conditioned on z as pθ ( x | z ) . The conditional variational autoencoder ( CVAE ) [ 19 ] extends this framework by allowing for the generative process to be controlled by input y . The latent variable z is thus conditioned on the input y in the prior pθ ( z | y ) . The VAE estimates parameters φ and an auxiliary distribution qφ ( z | x ) to calculate the evidence lower bound ( ELBO ) of the log-likelihood for the observed data [ 20 ] : log pθ ( x ) ≥ Eqφ ( z|x ) log pθ ( x | z ) − KL ( qφ ( z | x ) || pθ ( z ) ) ( 1 ) The CVAE estimates parameters φ of an auxiliary distribution qφ ( z | x , y ) to obtain the ELBO [ 19 ] : log pθ ( x | y ) ≥ Eqφ ( z|x , y ) log pθ ( x | z ) − KL ( qφ ( z | x , y ) || pθ ( z | y ) ) . ( 2 ) The subsequent discussion focuses on VAEs , and analogous expressions for CVAEs may be obtained by replacing pθ ( x ) , qφ ( z | x ) , and pθ ( z ) with pθ ( x | y ) , qφ ( z | x , y ) , and pθ ( z | y ) , respectively . Maximizing the VAE ELBO using gradient-based methods involves computing the term ∇φEqφ ( z|x ) log pθ ( x | z ) in Eq . ( 1 ) , which requires marginalization over the auxiliary distribution qφ ( z | x ) . Many Monte Carlo approaches estimate this quantity without explicit marginalization over all possible latent classes z ∈ Z . The score function estimator ( SFE ) is an unbiased estimator which uses the identity ∇φEqφ ( z|x ) log pθ ( x | z ) = Eqφ ( z|x ) pθ ( x | z ) ∇φ log qφ ( z | x ) to estimate the gradient through samples drawn from qφ ( z | x ) . However , the SFE is prone to high variance , and in practice , sampling-based approaches generally require variance reduction methods such as baselines , control variates , and Rao-Blackwellization techniques [ 21 ] . Alternatively , a low-variance gradient estimate can be obtained by the reparameterization trick , which separates the parameters of the distribution from the source of stochasticity [ 20 ] . The Gumbel-Softmax and Straight-Through reparameterizations rely on a continuous relaxation of the discrete latent space [ 14 ] , [ 15 ] . In contrast to these stochastic approaches , our method , like sparsemax [ 8 ] , α-entmax [ 16 ] , and sparsehourglass [ 17 ] , focuses on obtaining sparse auxiliary distributions qφ ( z | x ) to enable exact marginalization over all latent classes z ∈ Z for which the auxiliary distribution has nonzero probability mass . 2.2 Normalization Functions . A common normalization function that maps vectors v ∈ RK to probability distributions in ∆K−1 is the softmax function : SOFTMAX ( v ) k ∝ evk . ( 3 ) One limitation of this function is the resulting probability distribution always has full support , i.e . SOFTMAX ( v ) k 6= 0 for all v and k. Thus , if the distributions qφ in Eq . ( 1 ) and Eq . ( 2 ) are outputs of the softmax function , then calculating the expectations over qφ exactly requires marginalizing over all one-hot vectors z ∈ RK , which can be computationally intractable when K is large [ 9 ] , [ 10 ] . Recent interest in sparse alternatives has led to the development of many new sparse normalization functions [ 8 ] , [ 10 ] , [ 16 ] – [ 18 ] . The sparsemax function projects inputs onto the simplex to achieve a sparse distribution [ 8 ] . The α-entmax family of functions generalizes the softmax and sparsemax functions through a learnable parameter α which tunes the level of sparsity [ 16 ] . Sparsehourglass is another generalization of sparsemax with a controllable level of sparsity [ 17 ] . Despite the controls , we find that these methods achieve sparsity at the expense of collapsing valid modes . 2.3 Post-Hoc Evidential Sparsification . Of particular interest is a sparse normalization function that maintains multimodality in probability distributions . Balancing sparsity and multimodality has been shown to be beneficial for tasks such as image generation [ 10 ] or machine translation [ 16 ] . Itkina et al . [ 10 ] propose a post-hoc sparsification method for discrete latent spaces in CVAEs which preserves multimodality in distributions . The method calculates evidential weights as affine transformations of an input feature vector and interprets these weights as evidence either for a singleton latent class { zk } or its complement { zk } depending on the sign of the weights . This interpretation , with origins in evidential theory [ 22 ] , [ 23 ] , allows for the distinction between conflicting evidence and lack of evidence , which is normally not possible with ordinary multinomial logistic regression . Conflicting evidence can be interpreted as multimodality whereas lack of evidence can correspond to sparsity [ 10 ] . For a given dataset and model , the resulting post-hoc sparsification is as follows . A CVAE is trained on a dataset using the softmax normalization and the ELBO in Eq . ( 2 ) . At test time , for a given query y from the dataset , the learned weights of the last layer of the encoder β̂ ∈ R ( J+1 ) ×K , and the corresponding output of the last hidden layer φ ( y ) ∈ RJ , the evidential weights are : wk = β̂0k − 1 K K∑ ` =1 β̂0 ` + J∑ j=1 [ φj ( β̂jk − 1 K K∑ ` =1 β̂j ` ) ] ( 4 ) for each k ∈ { 1 , . . . , K } . Evidence in support of a latent class zk then corresponds to w+k = max ( 0 , wk ) and evidence against it as w − k = max ( 0 , −wk ) . The resulting sparse distribution is : p ( zk | y ) ∝ 1 { mk 6= 0 } ewk , ( 5 ) where the evidential belief mass mk is : mk = e −w−k ew+k − 1 +∏ ` 6=k ( 1− e−w − ` ) ( 6 ) for each k ∈ { 1 , . . . , K } . Further information can be found in Appendix A . 3 The Evidential Softmax Transformation . Although the evidential sparsification procedure is successful at providing sparse yet multimodal normalizations of scores , it can only be performed post-hoc , without the possibility of being used at training time . Furthermore , it requires O ( JK2 ) computations to calculate all the weights for each input , with J hidden features in the last layer and K latent classes . We propose a transformation , which we call evidential softmax ( ev-softmax ) , on inputs v ∈ RK . We claim that this transformation is equivalent to that of the post-hoc evidential sparsification procedure with the advantages of having a simple closed form and requiring only O ( K ) computations . Let v = 1K ∑K i=1 vi be the arithmetic mean of the elements of v. We define our transformation as , EVSOFTMAX ( v ) k ∝ 1 { vk ≥ v } evk . ( 7 ) Note that the post-hoc sparsification procedure in Section 2.3 takes input vectors φ ∈ RJ and the weights β̂ ∈ R ( J+1 ) ×K of the neural network ’ s last layer whereas the ev-softmax function takes as input v = β̂Tφ ∈ RK . We formalize the equivalence of ev-softmax with the post-hoc sparsification procedure as follows : Theorem 3.1 . For any input vector φ ∈ RJ , weight matrix β̂ ∈ RJ×K , define the score vector v = β̂Tφ and variables wk , w±k , mk , and p ( zk ) as in Section 2.3 for all k ∈ { 1 , . . . , K } . If the probability distribution of v is non-atomic , then EVSOFTMAX ( v ) k = p ( zk ) for each k. The proof of this theorem uses the following two lemmas : Lemma 3.2 . Let v = 1K ∑K i=1 vi . Then vk − v = wk for all k ∈ { 1 , . . . , K } . Proof . The result follows from algebraic manipulation of Eq . ( 4 ) . Full details are in Appendix B. Lemma 3.3 . For mk as defined in Eq . ( 6 ) , 1 { mk 6= 0 } = 1 { vk > v } for all k ∈ { 1 , . . . , K } . Proof . Observe that the condition 1 { mk = 0 } = 1− 1 { mk 6= 0 } is equivalent to the conjunction of the following two conditions : 1 ) wk ≤ 0 and 2 ) there exists some ` 6= k such that w ` ≥ 0 . From Lemma 3.2 , these statements are equivalent to 1 ) vk ≤ v and 2 ) there exists some ` 6= k such that v ` ≥ v. Note that if vk ≤ v , then either v ` = v or there must exist some ` 6= k such that v ` > v. In both cases , there exists ` 6= k such that v ` ≥ v. Thus , we see that vk ≤ v implies there exists some ` 6= k such that v ` ≥ v. We conclude that 1 { mk = 0 } = 1 { vk ≤ v } and the result follows . We briefly sketch the proof of Theorem 3.1 and refer the reader to Appendix B for the full derivation . Lemma 3.2 implies that ewk ∝ evk . In addition , combining Lemma 3.3 with the non-atomicity of the distribution of v yields that 1 { mk 6= 0 } = 1 { vk ≥ v } ( Lemma B.3 ) . Eq . ( 7 ) follows from these two results .

Summary: The paper introduces a normalizing function called evidential softmax (ev-softmax). Based on principles of evidential theory, ev-softmax is able to assign zero probability to classes that lack evidence of the data. It can be seen as a generalization of the work done by Itkina et al. (2020), the crucial difference being that ev-softmax can be applied at training time. The experiment section shows improved performance compared to other sparse methods, on a variety of tasks.