Likelihood-free Inference with Mixture Density Network

An ANN is a computing system inspired by the biological neural network. The unit of an ANN is called a neuron (see Figure 1), and an ANN consists of many neurons. Each neuron transforms the input from other neurons and gives an output

where $x$ is the input of a neuron, $f(\cdot)$ is a nonlinear function that is usually called the activation function, and $w$ and $b$ are parameters to be learned by the network. The parameters of an ANN model can be optimized by minimizing the loss function $\mathcal{L}$ in the training process. The loss function used in this work will be described in Section 2.3. In our analysis, we consider the randomized leaky rectified linear unit (RReLU; Xu et al. 2015) as the activation function:

where $a$ is a random number sampled from a uniform distribution $U(l,u)$, where $l,u\in[0,1)$. Here we adopt the default settings of $l=1/8$ and $u=1/3$ in PyTorch¹¹1https://pytorch.org/docs/stable/index.html. A discussion about the effect of activation functions on the result will be shown in Section 5.3.

The general structure of an ANN is composed of many layers, and each layer contains many neurons. There are three types of layers in a network: an input layer, one or more hidden layers, and an output layer. Data is fed to the input layer, and then the information is passed through each hidden layer and finally computed from the output layer. Furthermore, the batch normalization technique (Ioffe & Szegedy, 2015) is applied before each nonlinear layer to facilitate optimization and speedup convergence.

In general, there is no theory to determine exactly how many neurons should be used in each hidden layer. Here we adopt the model architecture used in Wang et al. (2020a), i.e. the number of neurons in each hidden layer is decreased in proportion to the layer index. The number of neurons in the $i$th hidden layer is therefore

where $N_{\text{in}}$ is the number of neurons in the input layer and $F$ is a decreasing factor defined by

where $N_{\text{out}}$ is the number of neurons in the output layer and $n$ is the number of hidden layers. In our analysis, we consider a network with three hidden layers. We discuss the effect of the number of hidden layers on the result of estimation in Section 5.2.

A mixture model is a probabilistic model that assumes all data points are generated from a mixture of a finite number of distributions with unknown parameters, where the distribution can be of any type. In Figure 2 we show an example of a mixture model of a parameter. For measurement $\bm{d}$ and parameters $\bm{\theta}$, the probability density of $\bm{\theta}$ with $K$ components has the form

where $\omega_{i}$ is a mixture weight representing the probability that $\bm{\theta}$ belongs to the $i$th component. The mixture weights are nonnegative and their sum is normalized, i.e. $\sum_{i=1}^{K}\omega_{i}=1$.

The normal (or Gaussian) distribution acts as the foundation for many modeling approaches in statistics. Therefore, in the following analysis, we mainly consider the Gaussian distribution as the unit of the mixture model. In Section 6, we also discuss the applications of the beta mixture model in the parameter estimation. Therefore, for the case of only one parameter, Equation (5) becomes

For multiple parameters, Equation (5) becomes

where $N$ is the number of parameters.

An MDN is a combination of an ANN and a mixture model. It learns a mapping between the measurement $\bm{d}$ and parameters of the mixture model. For an MDN with Gaussian components specifically, it learns a mapping between the measurement $\bm{d}$ and the parameters of the Gaussian mixture model ($\omega$, $\bm{\mu}$, and $\bm{\Sigma}$ (or $\sigma$ for the case of only one parameter)). Figure 3 shows the general structure of an MDN with Gaussian components. The left side of the MDN accepts the measurement as input and then the parameters of the mixture model are computed from the right side.

The MDN model shown in the Figure 3 is for the case of single observational data. But in cosmology we usually need to combine multiple data sets to estimate cosmological parameters. Therefore, following Wang et al. (2020a), we design a multibranch MDN for a joint constraint on parameters using multiple data sets, as shown in Figure 4. There are multiple branches on the left side of the model that can accept multiple observational data sets as input. Then, the information will be combined in the middle of the model (the yellow part of Figure 4), and finally the MDN outputs the parameters of the mixture model.

The parameters (e.g. $w$ and $b$ in Equation (1)) of an MDN can be learned by minimizing the loss function (Kruse, 2020):

where the loss is averaged over the minibatch input samples to the network. Specifically, for an MDN with Gaussian components and the case of only one parameter, Equation (8) becomes

Because small values in the exponent will give very small values in the logarithm that can lead to numerical underflow, we therefore, use the Log-Sum-Exp trick and carry out the calculations in the logarithmic domain for numerical stability. Thus, Equation (9) can be rewritten as

where

For the case with multiple parameters, Equation (8) becomes

Considering the numerical stability, this equation can be written as

where

Note that for a valid multivariate Gaussian distribution, $\bm{\Sigma}_{i}$ and $\bm{\Sigma}^{-1}_{i}$ here must be positive-definite matrices. Therefore, the precision matrix $\bm{\Sigma}^{-1}_{i}$ can be characterized by its upper Cholesky factor $\bm{U}_{i}$:

where $\bm{U}_{i}$ is an upper triangular matrix with strictly positive diagonal entries. There are $N(N+1)/2$ nonzero entries for $\bm{U}_{i}$, which is much less than that of $\bm{\Sigma}^{-1}_{i}$ ($N^{2}$ entries). Thus, if the MDN learns $\bm{\Sigma}^{-1}_{i}$ instead of $\bm{U}_{i}$, there will be more neurons in the network, which will increase the training time. Also, the output of the network may also make $\bm{\Sigma}^{-1}_{i}$ a non-positive-definite matrix. Therefore, in order to make the convergence of the network more stable, the MDN is actually learns $\bm{U}_{i}$. Outputs of the network can be either positive or negative numbers. Thus, to ensure the diagonal entries of $\bm{U}_{i}$ are positive, we use a Softplus function to enforce the positiveness of the diagonal entries:

where $\widetilde{\bm{U}}_{i}$ is the output of the network, and

where $\beta$ is a parameter and we set $\beta=1$ throughout the paper. This offers an efficient way to calculate the matrix in Equation (14):

Therefore,

In addition,

Therefore, Equation (14) can be finally written as

For parameter estimation with the MDN, we will first train the MDN with simulated data and then estimate parameters with the observational data. In this section, we describe the processes of training and parameter estimation.

To test the capability of the MDN for estimating parameters from one data set, we constrain $w$ and $\Omega_{\rm m}$ of the $w$CDM model using the Pantheon SN-Ia (Scolnic et al., 2018), which contain 1048 data points for the redshift range [0.01, 2.26]. The distance modulus of the Pantheon SN-Ia can be rewritten as

where $m_{B,{\rm corr}}^{*}=m_{B}^{*}+\alpha\times x_{1}-\beta\times c+\Delta_{B}$ is the corrected apparent magnitude, and $M_{B}$ is the $B$-band absolute magnitude (Scolnic et al., 2018). The variable $\alpha$ is the coefficient of the relation between luminosity and stretch, $\beta$ is the coefficient of the relation between luminosity and color, and $\Delta_{B}$ is a distance correction based on predicted biases from simulations.

The measurement of the Pantheon SN-Ia is the corrected apparent magnitudes. Therefore, the input of the MDN is $m_{B,{\rm corr}}^{*}$ for the observational SN-Ia, and $\mu+M_{B}$ for the simulated data generated by the $w$CDM model. Specifically, we have the distance modulus

and the luminosity distance

where $c$ is the speed of light, and $H_{0}$ is the Hubble constant. The $E(z)$ function describes the evolution of the Hubble parameter:

where $\Omega_{\rm m}$ is the matter density parameter and $w$ is the equation of state of dark energy. Equation (26) assumes the spatial flatness. Then we have

As we already know, the absolute magnitude of SN-Ia $M_{B}$ is strongly degenerated with the Hubble constant $H_{0}$. Therefore, we combine $M_{B}$ and $H_{0}$ into a new parameter and constrain it with the cosmological parameters simultaneously. To do this, we define a dimensionless luminosity distance

Then Equation (27) can be rewritten as

where $\mu_{c}\equiv 5\log_{10}\left(c/H_{0}/{\rm Mpc}\right)+M_{B}+25$ is a new nuisance parameter to be estimated with the cosmological parameters simultaneously.

We first constrain $w$, $\Omega_{\rm m}$, and the nuisance parameter $\mu_{c}$ simultaneously using the Pantheon SN-Ia data without considering systematic uncertainties. Here two methods are considered in our analysis: standard MCMC and the MDN. For the MCMC method, we use the public package emcee (Foreman-Mackey et al., 2013) for parameter estimation, which is a Python module that achieves the MCMC method. During the constraining procedures, an MCMC chain with 100,000 steps is generated after burn-in. The best-fitting values with 1$\sigma$ errors of the parameters are calculated from the MCMC chains with coplot,²²2https://github.com/Guo-Jian-Wang/coplot as shown in Table 1. We also draw the corresponding one-dimensional and two-dimensional contours in Figure 5. For the MDN method, we consider three Gaussian components and using the setting illustrated in Sections 2.1, 2.3, and 2.4. There are 3000 samples in the training set and 500 samples in the validation set. After training, we feed the corrected apparent magnitudes $m_{B,{\rm corr}}^{*}$ to the well-trained MDN model to obtain the parameters of the mixture model, and then get an ANN chain of the cosmological parameters. The results are shown in Table 1 and Figure 5. Obviously, the results of the MDN method are almost the same as those of the MCMC method.

We show the losses of the training and validation sets in Figure 6. The loss of the training set is large at the beginning of the training process, and gradually decreases and stabilizes with the epoch. For the validation set, the trend of loss is similar to that of the training set. After 400 epochs, the loss of the validation set is smaller than that of the training set. Obviously, there is no overfitting for the MDN. The reason why the loss of the validation set is smaller than that of the training set is that multiple sets of noise are added to each sample, which means that multiple sets of noise will be generated and added to a sample to ensure that the MDN knows that measurement may differ due to the presence of noise. Our analysis shows that the adding of multiple sets of noise to measurement is beneficial to the stability of the MDN and will make the final estimated cosmological parameters more accurate, especially for the multibranch MDN.

Furthermore, we consider the systematic uncertainties in our analysis to test the capability of the MDN for dealing with observations with covariance. We use the systematic covariance matrix $\bm{C}_{\rm sys}$ (Scolnic et al., 2018) to generate noise then add it to the training set (see Section 2.4.1). We then constrain the parameters with the MCMC and MDN (with three Gaussian components) methods. The results are shown in Table 2 and Figure 7. The difference between the MDN results and the MCMC results can then be calculated using

where $\sigma=\sqrt{\sigma^{2}_{\rm MCMC}+\sigma^{2}_{\rm MDN}}$, and $P_{\rm MCMC}$ and $P_{\rm MDN}$ are the best-fitting parameters for the MCMC and MDN methods, respectively. Specifically, the differences for the parameters are $0.174\sigma$, $0.015\sigma$, and $0.076\sigma$, respectively, which are slightly larger for the parameter $w$.

We can see that the shapes of both the one-dimensional and two-dimensional contours based on the MDN method are slightly different from those of the MCMC method. The reason is that the joint probability distribution of $w$ and $\Omega_{\rm m}$ deviates from the multivariate Gaussian distribution, which makes it more difficult to model their joint probability distribution. We take two approaches to solve this problem. One is to increase the number of components ($p_{i}$ function in Equation (5)), which diversifies the function forms that the MDN can fit. Specifically, we use nine Gaussian components and reestimate the cosmological parameters with the MDN, then plot the one-dimensional and two-dimensional contours in Figure 8. One can see that the MDN result converges to the MCMC result. Quantitatively, the differences between the MDN results and the MCMC results are $0.051\sigma$, $0.046\sigma$, and $0.025\sigma$. These deviations are smaller than those of using three components, especially for parameter $w$. Another way is to more efficiently sample the parameter space. As we illustrated in Section 2.4.1, samples are generated uniformly in the $\pm 5\sigma_{p}$ range of the posterior distribution, which might not be the most efficient way. We will explore different sampling methods in future works.

As we have shown in Section 3.2 the number of components of the mixture model may affect the constraints of parameters. Here, we further discuss the effect of the number of components on the result. Specifically, with the same procedure, we estimate $w$ and $\Omega_{\rm m}$ of the $w$CDM model with Pantheon SN-Ia data by considering one, three, five, and seven Gaussian components. The two-dimensional contours of the parameters are shown in Figure 10.

For the case of one component, the constraint is

This result converges to the MCMC result (see Table 2), and both the best-fitting values and $1\sigma$ errors are similar to those of the MCMC method. However, the two-dimensional contour (the upper left panel of Figure 10) is quite different from the contour of the MCMC method. The reason, as we illustrated in Section 3.2, is that the joint probability distribution of $w$ and $\Omega_{\rm m}$ deviates significantly from the multivariate Gaussian distribution.

For the cases with three, five, and seven components, the results all converge to the MCMC result, and all the shapes of the two-dimensional contours are similar to the MCMC result. This suggests that for $w$ and $\Omega_{\rm m}$ of the $w$CDM model, three Gaussian components can characterize the joint probability distribution. Therefore, we recommend using multiple components, such as three or more, when estimating parameters of other models. However, it is not recommended to use a very large number of components (e.g. more than 10), especially for multiple-parameters cases, because the MDN will be unstable and thus it will become difficult to constrain these components.

To test the effect of hyperparameters in the following sections, we simulate SN-Ia data based on a future Wide-field Infrared Survey Telescope (WFIRST) experiment (Spergel et al., 2015) in a flat $w$CDM model with the following fiducial values of parameters: $w=-1$, $\Omega_{\rm m}=0.3$, and $H_{0}=70~{}\rm km\ s^{-1}\ Mpc^{-1}$. The total number of SN-Ia samples is 2725 in the range of $0.1<z<1.7$. We then estimate $w$ and $\Omega_{\rm m}$ using the MDN method with three Gaussian components. Finally, we calculate the mean deviation between the MDN results and the fiducial values:

where $N$ is the number of cosmological parameters, $\theta_{i,\rm fid}$ is the fiducial parameter, and $\theta_{i,\rm pred}$ and $\sigma_{i,\rm pred}$ are the predicted best-fitting value and error of the cosmological parameter, respectively.

First, we test the effect of the number of hidden layers of the MDN model. We design five different MDN structures with the number of hidden layers varying from one to five. We consider RReLU here as the activation function. We train the MDN models with 3000 samples, and after 2000 epochs, we obtain the corresponding five ANN chains by feeding the simulated observations. We then calculate the mean deviations between the MDN results and the fiducial values. The mean deviation as a function of the number of hidden layers is shown in the upper left panel of Figure 11. The mean deviations of these five cases are $0.072\sigma$, $0.033\sigma$, $0.032\sigma$, $0.094\sigma$, and $0.082\sigma$, respectively. We can see that there is a large deviation for MDNs with fewer or more hidden layers. An MDN with one hidden layer may not have enough layers to build a mapping between the input data and the output parameters, while an MDN with four or five hidden layers will have more difficulty learning the mixture model. We note that the training time will increase as the number of hidden layers increases. As a result, MDNs with many hidden layers become unpragmatic because of the training time.

The settings of the activation function may also affect the parameter estimation. Here we select several commonly used activation functions and test their effects on the performance of the MDN model: rectified linear unit (ReLU), exponential linear unit (ELU), RReLU, sigmoid function, and hyperbolic tangent (Tanh). The ReLU activation function is defined as (see also Nair & Hinton 2010)

The ELU is (Clevert et al., 2016)

where $\alpha$ is a parameter and we set it to 1. The sigmoid function is defined by (see also Han & Moraga 1995)

and the hyperbolic tangent is (Malfliet, 1992)

The MDN model contains three hidden layers in our work. Using the data simulated in Section 5.2, we train another five MDN models with 3000 samples to estimate $w$ and $\Omega_{\rm m}$. After 2000 epochs, we obtain five ANN chains, and then calculate the corresponding deviations. The mean deviation of the parameters between the MDN results and the fiducial values as a function of the activation function is shown in the upper right panel of Figure 11. One can see that the MDN with the sigmoid function has the maximum deviation of $0.158\sigma$, while the deviations of the other MDNs are all smaller than $0.067\sigma$. Therefore, in our specific application of the MDN, the ReLU, ELU, RReLU, and Tanh activation functions can all be used for the MDN method.

To test the effect of the number of samples in the training set, we train several MDN models with the number of samples varying from 1000 to 10,000. The MDN model used here contains three hidden layers, the RReLU nonlinear function is taken as the activation function, and the MDN is trained with 2000 epochs. The mean deviation of parameters between the MDN results and the fiducial values is shown in the lower left panel of Figure 11. We can see that the mean deviation varies in the interval $[0.010\sigma,0.103\sigma]$ as the number of training samples changes. But most of the deviations are less than $0.055\sigma$. Therefore, in our specific application of the MDN, thousands of samples can make an MDN model well trained. Note that the time of training of an MDN model is related to the number of samples in the training set. Therefore, considering the performance and training time of the MDN, the number of samples should be selected reasonably.

The number of epochs may also affect the performance of the MDN. To test this, we train several MDN models with the epoch varying from 500 to 5000. The MDN models used here contain three hidden layers, the activation function is RReLU, and 3000 samples are used to train the models. With the same procedure, we train and calculate deviations between the MDN results and the fiducial values. The mean deviation as a function of the number of epochs is shown in the lower right panel of Figure 11. We can see that as the number of epochs increases, the mean deviation oscillates in the range of $[0.015\sigma,0.059\sigma]$. Also, in our case, there is little improvement by training for more than 2000 epochs. Since the training time is also related to the number of epochs, the number of epochs should also be selected reasonably.

The analysis above is focused on the Gaussian distribution as the unit of the mixture model. Since the unit of the mixture model can be any type of distribution, we use a beta distribution as the unit of the mixture model in this section. A beta mixture model with $K$ components has the form

where $\omega$ is the mixture weight, $\alpha,\beta>0$ are two shape parameters to be determined, and $\Gamma(\cdot)$ is the gamma function defined as

Here we use the Softplus function (Equation (17)) to guarantee $\alpha$ and $\beta$ are positive. Note that the beta distribution is defined on the interval [0, 1]; thus we normalize the parameters in the training set (i.e. the target) to the range [0, 1] by using the min$-$max normalization technique:

The parameters of the beta mixture model can be estimated by minimizing the loss function

For numerical stability, we also use the Log-Sum-Exp trick; then Equation  (40) can be rewritten as

where

Since the beta distribution is for one random variable, thus to test the capability of the beta mixture model for parameter inference, we use the Pantheon SN-Ia to constrain $\Omega_{\rm m}$ of the $w$CDM model. In our analysis, the Hubble constant is set to $H_{0}=70~{}\rm km\ s^{-1}\ Mpc^{-1}$, and the absolute magnitude is set to $M_{B}=-19.3$. We consider two cases of the equation of state of dark energy: $w=-1$ and $w=-0.5$, respectively. The case $w=-0.5$ considered here is completely set manually, which is to test the case where the posterior distribution of the parameter is a truncated distribution under the physical limit. We constrain $\Omega_{\rm m}$ using the MCMC method and the MDN method. For the MDN method, we consider two mixture models: the Gaussian mixture model and the beta mixture model.

For the case of $w=-1$, we first constrain $\Omega_{\rm m}$ using the MCMC method. After generating an MCMC chain, the best-fitting value with $1\sigma$ error of $\Omega_{\rm m}$ can be calculated (see Table 4). The corresponding marginalized distribution is shown in Figure 12 with a green dashed line. We then use the MDN method with three Gaussian components to constraint $\Omega_{\rm m}$ and find the result (the blue solid line in Figure 12) is almost the same as that of the MCMC method for both the best-fitting value and the $1\sigma$ error. We also constrain $\Omega_{\rm m}$ using the MDN method with one beta component and obtain the same result of the MCMC method.

With the same procedure, we constrain $\Omega_{\rm m}$ for the case of $w=-0.5$. The result of the MCMC method is shown in Table 4 and the marginalized distribution is shown in Figure 13 with a green dashed line. We can see the distribution is a truncated distribution due to the physical limit. Then, we constrain $\Omega_{\rm m}$ with the MDN method with Gaussian components and find both the best-fitting value and the $1\sigma$ error deviate slightly from the MCMC result, and the shape of the distribution (the blue solid line in Figure 13) differs slightly from that of the MCMC method. Finally, the result of $\Omega_{\rm m}$ based on the MDN method with beta components is almost the same as that of the MCMC method, and the shape of the distribution (the red solid line in Figure 13) is consistent with the one based on the MCMC method.

These results show that MDN models with beta components are able to constrain parameters, and even have advantages over MDN models with Gaussian components in the case of truncated distribution. In the case of truncated distribution, the reason for the beta mixture model performing better than the Gaussian mixture model is that the beta distribution has more shape features (see Figure 14 for examples), which enables it to fit better for the posterior distribution of parameters.

We now switch to multiparameter estimation. If parameters are not independent, the beta mixture model cannot be used when doing the MDN parameter estimation. But if the correlations between parameters are negligible, the beta mixture model can be used for multiple-parameter tasks to obtain the one-dimensional distribution of each parameter. In this case, the probability density in Equation (6.1) becomes

where $N$ is the number of parameters. To test this, we constrain $w$, $\Omega_{\rm m}$, and the nuisance parameter $\mu_{c}$ using the Pantheon SN-Ia. The Softplus function is used as the activation function considering the stability of the MDN. The one-dimensional distributions of these parameters are shown in Figure 15. We can see that the results based on the MDN with beta components are almost the same as those of the MCMC method. Our analysis shows the instability of the MDN when using Equation (43) as the probability density. The reason may be that there are correlations between parameters, but they are not considered by the MDN. Besides, there are too many components ($N\times K$ components in total) for the probability density in Equation (43), especially for cases of many cosmological parameters, which will also increase the instability of the MDN. Therefore, the beta mixture model based on Equation (43) should be further investigated.

As a comparison, we also constrain $w$, $\Omega_{\rm m}$, and $\mu_{c}$ without considering correlations between them, by using an MDN with a Gaussian mixture model. In this case, the probability density in Equation (2.2) becomes

The one-dimensional distributions of these three parameters are shown as green dotted lines in Figure 15. Obviously, this result is similar to that based on the beta mixture model. We note that the Gaussian mixture model performs poorly for parameters with truncated distribution (see Figure 13), and correlations between parameters will not be obtained. Therefore, the Gaussian mixture model based on Equation (2.2) is recommended for multiparameter estimates.