Neural Network Regression with Beta, Dirichlet, and Dirichlet-Multinomial Outputs

1 INTRODUCTION . Artificial neural networks are typically trained by maximizing the conditional likelihood of output targets given input features . Each target is modeled as a sample from a distribution p ( y|x ) parameterized by the output activity of the neural network , where the choice of parametric distribution is implied by the choice of objective function . Thus , the support of the probability distribution should match the target space , but in practice , this is often not the case . Today , the vast majority of neural network output layers implicitly model the targets as samples from one of four distributions : a binomial , a categorical , a Gaussian , or a Laplacian distribution — respectively corresponding to the binomial cross-entropy loss , multi-class cross-entropy loss , mean squared error , and mean absolute error . These distributions are commonly used even when the target space does not match the support , because the gradient calculations for these distributions are simple ( and easy to compute ) when paired with the appropriate output layer activation functions . These distributions dominate to such a degree that few alternatives are even available in most common deep learning software packages such as Keras ( Chollet et al. , 2015 ) and PyTorch ( Paszke et al. , 2017 ) . Alternatives do exist — using neural networks to parameterize more complex distributions is not new . The standard regression approach can be generalized to a heteroskedastic Gaussian output layer ( Nix & Weigend , 1994 ; Williams , 1996 ) , where the neural network predicts both a mean and a variance for each target . Multi-model distributions can be modeled with a mixture density ( Bishop , 1994 ) . And more recently , the Gamma output layer was proposed to model targets in R > 0 ( Ng et al. , 2017 ) . In principle , any parametric distribution with well-defined gradients could serve as a probabilistic prediction at the output of a neural network model . The approach proposed here is simpler than the one taken by Conditional Variational Autoencoders ( CVAEs ) ( Kingma & Welling , 2013 ; Sohn et al. , 2015 ) . While CVAEs can , in theory , model arbitrary high-dimensional conditional distributions , computing the exact conditional likelihood of a target requires marginalizing over intermediate representations , making exact gradient calculations intractable . Thus , training a CVAE requires approximating the gradients through sampling . In this work we show that restricting the output to a particular class of distributions , namely the Dirichlet or Dirichlet-multinomial compound distributions , enables a calculation of the exact likelihood of the targets and the exact gradients . Interpreting the output of a neural network classifier as a probability distribution has obvious benefits . One can derive different point estimates , define confidence intervals , or integrate over possible outcomes — a necessity for managing risk in decision making . Potentially , it could also lead to better learning — matching the output support to the target space essentially constrains the learning problem by incorporating outside knowledge . Allowing the network to output “ uninformative ” distributions — e.g . a uniform distribution over the support — could make training faster by allowing the network to focus on the easiest training examples first — a self-guided form of curriculum learning . In the present work , we derive gradients for the Beta distribution , Dirichlet distribution , and Dirichlet-multinomial compound distribution . We then propose activation functions that stabilize numerical optimization with stochastic gradient descent . Finally , we demonstrate through experiments that this approach can be used to model three common types of targets : ( 1 ) targets over the multivariate simplex , ( 2 ) real-valued scalar targets with lower and upper bounds , and ( 3 ) nonnegative integer-valued counts ( samples from the Dirichlet-multinomial compound distribution ) . The experiments demonstrate that our approach provides interpretable predictions with learned uncertainty , without decreasing the performance of the point estimates . 2 DIRICHLET OUTPUT LAYERS . 2.1 TARGETS ON THE MULTIDIMENSIONAL SIMPLEX . Consider a supervised learning scenario , in which the goal is to model the relationship between a set of input , target pairs ( x ( t ) , y ( t ) ) , where y ( t ) lies on the d-simplex ∆d , i.e . y ( t ) ∈ Rd with y ( t ) i > 0 ∀ i ∈ { 1 , . . . , d } and ∑d i=1 y ( t ) i = 1 . We can construct a neural network that takes in x ( t ) and outputs a length-d vector α ( t ) = < α ( t ) 1 , . . . , α ( t ) d > , with α ( t ) i ∈ ( 0 , ∞ ) , that parameterizes a Dirichlet distribution over ∆d with probability density function pα ( z ) = ∏d i=1 z αi−1 i B ( α ) ( 1 ) B ( α ) = ∏d i=1 Γ ( αi ) Γ ( ∑d i=1 αi ) ( 2 ) where Γ is the Gamma function , which generalizes the factorial function to real values . Thus , given a set of neural network weights and some input x ( t ) , we have a conditional density function over the domain of the target y ( t ) . The network can be trained to minimize the negative log-likelihood ( NLL ) of the training set , −∑t ln pα ( t ) ( y ( t ) ) , using gradient descent . Dropping the example index t for clarity , we write the log-likelihood for a single example as ln ( pα ( y ) ) = ln ( d∏ i=1 yαi−1i ) − ln ( ∏d i=1 Γ ( αi ) Γ ( ∑d i=1 αi ) ) = d∑ i=1 ln ( yαi−1i ) + ln ( Γ ( d∑ i=1 αi ) ) − d∑ i=1 ln ( Γ ( αi ) ) = d∑ i=1 ( αi − 1 ) ln ( yi ) + ln ( Γ ( d∑ i=1 αi ) ) − d∑ i=1 ln ( Γ ( αi ) ) . ( 3 ) The gradient w.r.t . the network output αi is then ∂ ∂αi ln ( pα ( y ) ) = ln ( yi ) + ∂ ∂αi ln ( Γ ( d∑ j=1 αj ) ) − ∂ ∂αi ln ( Γ ( αi ) ) = ln ( yi ) + ∂ ( ∑d j=1 αj ) ∂αi ψ ( d∑ j=1 αj ) − ψ ( αi ) = ln ( yi ) + ψ ( d∑ j=1 αj ) − ψ ( αi ) ( 4 ) where ψ is the digamma function ψ ( x ) = ∂∂x ln ( Γ ( x ) ) , and the gradient for a particular target is ∂ ∂αi ln ( pα ( z = y ) ) = ln ( yi ) + ψ ( d∑ j=1 αj ) − ψ ( αi ) ( 5 ) Accurate numerical approximations of the digamma function are readily available , and obtaining point estimates from the network output is as simple as computing the mean of the Dirichlet distribution , E [ z ] = α∑d i=1 αi , or the mode α− 1∑d i=1 αi−d when αi > 1 ∀i ∈ { 1 , . . . , d } . The Dirichlet distribution can also be parameterized by a vector µ ∈ ∆d , together with a scalar γ > 0 , where pµ , γ ( z ) = ∏d i=1 z αi−1 B ( α ) ( 6 ) αi = µi × γ ( 7 ) In this alternative parameterization , µ predicts the expectation of the Dirichlet , and the summation to one could be enforced using a softmax activation . This is conceptually similar to the heteroskedastic Gaussian model where the neural network computes the mean and standard deviation of a Gaussian distribution . 2.2 TARGETS SAMPLED FROM A DIRICHLET-MULTINOMIAL COMPOUND DISTRIBUTION . The Dirichlet output layer can also be used to model vectors of non-negative integer targets y ∈ Nd0 as samples from the Dirichlet-multinomial compound distribution . This allows us to treat each target as a collection of related trials , conditioned on the same input . The distribution is parameterized by the Dirichlet parameters α and the number of multinomial trials n. The number of trials n can be fixed for each training example , fixed for the entire data set , or treated as a random variable sampled from a conditional probability distribution parameterized by an additional neural network output . Here we assume that n = ∑ i yi is given for each training example , so that the probability of target y is given by the following : pα , n ( y ) = ∫ ∆d Mulz , n ( y ) Dirα ( z ) dz = ∫ ∆d ( 1 B ( y ) d∏ i=1 zyii ) ( 1 B ( α ) d∏ i=1 zαi−1i ) dz = 1 B ( α ) B ( y ) ∫ ∆d d∏ i=1 zyi+αi−1i dz = B ( α+ y ) B ( α ) B ( y ) ( 8 ) where we use the definition of the multivariate beta function to integrate over the simplex ∆d and marginalize out z . This leads to the log-likelihood ln ( pα , n ( y ) ) = lnB ( α+ y ) − lnB ( α ) − lnB ( y ) = ln ( ∏d i=1 Γ ( αi + yi ) Γ ( ∑d i=1 αi + yi ) ) − ln ( ∏d i=1 Γ ( αi ) Γ ( ∑d i=1 αi ) ) − lnB ( y ) = d∑ i=1 ln ( Γ ( αi + yi ) ) − ln ( Γ ( d∑ i=1 ( αi + yi ) ) ) + ln ( Γ ( d∑ i=1 αi ) ) − d∑ i=1 ln ( Γ ( αi ) ) − lnB ( y ) ( 9 ) and the gradient of the log-likelihood w.r.t . network output αi , ∂ ∂αi ln ( pα , n ( y ) ) = ψ ( αi + yi ) − ψ ( d∑ j=1 ( αj + yj ) ) + ψ ( d∑ j=1 αj ) − ψ ( αi ) . ( 10 ) 2.3 UNIVARIATE TARGETS WITH LOWER AND UPPER BOUNDS . When the targets are real-valued scalars with lower and upper bounds , we can shift and rescale the values to be in the range [ 0 , 1 ] and model the target as a sample from a Beta distribution . The Beta distribution is the d = 2 case of the Dirichlet , and can be used to predict univariate targets y ∈ [ 0 , 1 ] . It can be parameterized by a neural network that outputs two values α , β > 0 , with pα , β ( z ) = z α−1 ( 1− z ) β−1 Γ ( αβ ) Γ ( α ) Γ ( β ) ( 11 ) with log-likelihood ln ( pα , β ( z ) ) = ( α− 1 ) ln ( z ) + ( β − 1 ) ln ( 1− z ) ( 12 ) + ln ( Γ ( αβ ) ) − ln ( Γ ( α ) ) − ln ( Γ ( β ) ) and gradients ∂ ∂α ln ( pα , β ( z ) ) = ln ( z ) + ψ ( αβ ) − ψ ( α ) ( 13 ) ∂ ∂β ln ( pα , β ( z ) ) = ln ( 1− z ) + ψ ( αβ ) − ψ ( β ) . ( 14 ) Alternatively , we can parameterize the Beta using two scalar network outputs : µ ∈ [ 0 , 1 ] and γ > 0. pµ , γ ( z ) = y α−1 ( 1− z ) β−1 Γ ( αβ ) Γ ( α ) Γ ( β ) ( 15 ) α = µ× γ ( 16 ) β = γ − µ× γ ( 17 ) In this parameterization , µ = αα+β = E [ z ] is the expectation of the distribution , and γ = α + β controls the width of the density function .

Summary: The authors use neural networks to parameterize conditional probability distributions. This is well-known and has been applied in the literature since extensions to generalized linear models beyond their canonical link function in the 70s. Their transformation from real-valued network output to, say, strictly positive concentration parameters in a Dirichlet are worth studying; but they don't analyze this in any detail.