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: This paper proposes a sparse normalization function called evidential softmax (ev-softmax). Specifically, it extends the post-hoc sparsification approach of Itkina et al, to propose a normalization function that can be applied during both training and testing. Notably, the authors claim that ev-softmax preserves the multimodality of the distribution, while providing sparse output distributions. The authors argue that sparsity is likely to help with interpretability and also improve task performance, while multimodality is useful when the output space is multimodal. They also show experiments showing that previous sparse normalization functions collapse multimodality of distributions.   The formulation of the ev-softmax function is accompanied with relevant proofs showing that a continuous family of approximations exist for ev-softmax, which make it possible to train with NLL and KL divergence based losses. Experiments on image generation using a generative model (VAE) and machine translation (using an autoregressive model) show some improvements over previous sparse normalization techniques on distributional accuracy and classification performance.