$\alpha$-Divergence Loss Function for Neural Density Ratio Estimation

First, we review the definition of $f$-divergences.

Many divergences are specific cases obtained by selecting a suitable generator function $f$. For example, $f(u)=u\cdot\log u$ corresponds to KL-divergence.

Then, we derive the variational representations of $f$-divergences using the Legendre transform of the convex conjugate of a twice differentiable convex function $f$, $f^{*}(\psi)=\sup_{u\in\mathbb{R}}\{\psi\cdot u-f(u)\}$ [21]:

where the supremum is taken over all measurable functions $\phi:\Omega\rightarrow\mathbb{R}$ with $E_{Q}[\,|f^{\prime}(\phi)|\,]<\infty$ and $E_{P}[\,|f^{*}(f^{\prime}(\phi))|\,]<\infty$. The maximum value is achieved at $\phi(\mathbf{x})=q(\mathbf{x})/p(\mathbf{x})$.

By replacing $\phi$ with a neural network model $\phi_{\theta}$, the optimal function for Equation (1) is trained through back-propagation using an $f$-divergence loss function, such that

where $\phi_{\theta}$ is a real-valued function, the superscript variable “$(R,S)$” can be omitted when unnecessary and instead represented as $\mathcal{L}_{f}(\cdot)$. As shown in Table 2, we list pairs of convex functions and the corresponding loss functions $\mathcal{L}_{f}(\phi_{\theta})$ in Equation (2) for several $f$-divergences.

When $f$-divergence loss functions $\mathcal{L}_{f}(\phi_{\theta})$, as defined in Equation (2), are not lower-bounded, overfitting can occur during optimization. For example, we observe the case of the Pearson $\chi^{2}$ loss function, $\mathcal{L}_{\mathrm{chi}\text{-}\mathrm{sq}}=-2\cdot\hat{E}_{Q}\big{[}\phi_{\theta}\big{]}+\hat{E}_{P}\big{[}\phi_{\theta}^{2}\big{]}$, as follows. Since the term $-2\cdot\hat{E}_{Q}\big{[}\phi_{\theta}\big{]}$ is not lower-bounded, it can approach negative infinity, causing the entire loss function to diverge to negative infinity as $\phi_{\theta}(\mathbf{x}_{i}^{q})\rightarrow\infty$ for $\mathbf{x}_{i}^{q}\in\hat{\mathbf{X}}_{Q[S]}$. Consequently, $\mathcal{L}_{\mathrm{chi}\text{-}\mathrm{sq}}\rightarrow-\infty$ when $\phi_{\theta}(\mathbf{x}_{i}^{q})\rightarrow\infty$ for some $\mathbf{x}_{i}^{q}\in\hat{\mathbf{X}}_{Q[S]}$. As shown in Table 2, both the KL-divergence and Pearson $\chi^{2}$ loss functions are not lower-bounded, and hence, are prone to overfitting during optimization. This phenomenon is referred to as train-loss hacking by Kato and Teshima [14].

Neural network parameters are updated using the accumulated gradients from each mini-batch. It is desirable for these gradients to be unbiased, i.e., $E\big{[}\nabla_{\theta}\mathcal{L}_{f}(\theta)\big{]}=\nabla_{\theta}E\big{[}\mathcal{L}_{f}(\theta)\big{]}$ holds. However, the equality between $E\big{[}\nabla_{\theta}\mathcal{L}_{f}(\theta)\big{]}$ and $\nabla_{\theta}E\big{[}\mathcal{L}_{f}(\theta)\big{]}$ requires the uniform integrability of $\mathcal{L}_{f}(\theta)$, i.e., $\lim_{K\rightarrow\infty}\sup_{\theta}E\big{[}|\mathcal{L}_{f}(\theta)|\cdot I(\mathcal{L}_{f}(\theta)>K)\big{]}=0$. The uniform integrability condition is typically violated when the loss function exhibits heavy-tailed behavior, which often occurs for the standard $f$-divergence loss functions derived solely from Equation (2). Consequently, the standard loss functions frequently result in biased gradients.

To illustrate this, consider employing the KL-divergence loss function for optimizing a shift parameter of $\phi_{\theta}=|x-\theta|$, where $\theta\in(0,1)$ and $x\in[0,1]$. Intuitively, in the above example, biased gradients occur because the KL-divergence loss gradients contain terms inversely proportional to the estimated density ratios, making their expectation diverge. Specifically, the loss function is obtained as $\mathcal{L}_{KL}(\phi_{\theta})=-\hat{E}_{Q}\big{[}\log\phi_{\theta}\big{]}+\hat{E}_{P}\big{[}\phi_{\theta}\big{]}-1$, and the gradient is expressed as $\nabla_{\theta}\mathcal{L}_{KL}(\phi_{\theta})=-\hat{E}_{Q}\big{[}\nabla_{\theta}(\log\phi_{\theta})\big{]}+\hat{E}_{P}\big{[}\nabla_{\theta}(\phi_{\theta})\big{]}$. Then, we have $\frac{\partial}{\partial\theta}E\big{[}\log\phi_{\theta}(x)\big{]}=\frac{\partial}{\partial\theta}\int_{0}^{1}\log|x-\theta|\,dx=-\log(1-\theta)$ and $E\big{[}\frac{\partial}{\partial\theta}\log\phi_{\theta}(x)\big{]}=\int_{0}^{\theta}\frac{1}{\theta-x}dx+\int_{\theta}^{1}\frac{1}{x-\theta}dx=\infty$. Consequently, we generally observe that $\nabla_{\theta}E\big{[}\mathcal{L}_{KL}(\phi_{\theta})\big{]}\neq E\big{[}\nabla_{\theta}\mathcal{L}_{KL}(\phi_{\theta})\big{]}$.

To mitigate this issue, Belghazi et al. [3] introduced a bias-reduction method for stochastic gradients in KL-divergence loss functions.

The vanishing gradient problem is a well-known issue in optimizing GANs [2]. We suggest that this problem occurs when the following two conditions are met: (i) the loss function results in minimal updates of model parameters, and (ii) updating the model parameters leads to negligible changes in the model’s outputs. Thus, the problem emerges when the following equation holds:

where $\mathbf{0}$ denotes a zero vector of the same dimension as the model gradient.

In Equation (3), condition (i) describes the loss function’s gradient vanishing, while condition (ii) ensures that the vanishing gradient condition persists. Specifically, the following three observations clarify their relationship: First, condition (i) does not necessarily imply condition (ii); Second, condition (i) alone does not guarantee its own persistence; Third, condition (ii) ensures the continued validity of condition (i).

To illustrate the first observation, consider the KL-divergence loss. Because its gradient is obtained as $-\hat{E}_{Q}\big{[}\nabla_{\theta}\phi_{\theta}/\phi_{\theta}\big{]}+\hat{E}_{P}\big{[}\nabla_{\theta}\phi_{\theta}\big{]}$ as shown in Table 2, (i) reduces to $E_{Q}\big{[}\nabla_{\theta}\phi_{\theta}/\phi_{\theta}\big{]}=E_{P}\big{[}\nabla_{\theta}\phi_{\theta}\big{]}$. Then, (i) does not ensure (ii). To understand the second and third observations, note that (ii) is both necessary and sufficient for preventing updates to the model parameters. Thus, if condition (ii) does not hold, the model’s predictions may change, potentially leading to the breakdown of condition (i). On the other hand, if condition (ii) holds, the model updates do not alter the outputs, leaving the gradient condition (i) unchanged, thereby causing the vanishing gradient problem to persist.

Consider the scenario where estimated density ratios become extremely small or large, fulfilling sufficient conditions for Equation (3) to hold. Table 2 presents the gradient formulas for the divergence loss functions (as provided in Table 2) along with their asymptotic behavior of the loss gradients as $\phi_{\theta}\rightarrow 0$ or $\phi_{\theta}\rightarrow\infty$. These results demonstrate that major $f$-divergence loss functions satisfy the conditions for Equation (3), showing that $E\big{[}\nabla_{\theta}\mathcal{L}_{f}(\phi_{\theta})\big{]}\rightarrow c_{1}\cdot E_{Q}\big{[}\nabla_{\theta}\phi_{\theta}\big{]}+c_{2}\cdot E_{P}\big{[}\nabla_{\theta}\phi_{\theta}\big{]}$, where $c_{1}$ and $c_{2}$ are constants, as $\phi_{\theta}\rightarrow 0$ or $\phi_{\theta}\rightarrow\infty$. In summary, all the divergence loss functions in Tables 2 and 2 can experience vanishing gradients when the estimated density ratio approaches extremal estimates.

The sample complexity of the KL-divergence is $O(e^{KL(Q||P)})$, which implies that

where $\widehat{KL^{N}}(Q||P)$ represents an arbitrary KL-divergence estimator for a sample size $N$ using a variational representation of the divergence, and $KL(Q||P)$ represents the true value of KL-divergence [26, 31, 19]. That is, when using KL-divergence loss functions, the sample size of the training data must increase exponentially as the true amount of KL-divergence increases in order to sufficiently train a neural network. To address this issue, existing methods divide the estimation of high divergence values into multiple smaller divergence estimations [28].

Here, we define $\alpha$-divergence (Amari’s $\alpha$-divergence), which is a subgroup of $f$-divergence, as [1]:

where $\alpha\in\mathbb{R}\setminus\{0,1\}$. From Equation (5), Hellinger divergence is obtained when $\alpha=1/2$, and $\chi^{2}$ when $\alpha=-1$.

Then, we obtain the following variational representation of $\alpha$-divergence :

From the right-hand side of Equation (6), we obtain a standard $\alpha$-divergence loss function as

Because $\nabla_{\theta}E_{Q}\big{[}\phi_{\theta}^{\alpha}\big{]}\allowbreak\neq\allowbreak E_{Q}\big{[}\nabla_{\theta}(\phi_{\theta}^{\alpha})\big{]}$ and $\nabla_{\theta}E_{P}\big{[}\phi_{\theta}^{\alpha-1}\big{]}\allowbreak\neq\allowbreak E_{P}\big{[}\nabla_{\theta}(\phi_{\theta}^{\alpha-1})\big{]}$ are generally observed when $\alpha<2$, the standard $\alpha$-divergence loss function with $\alpha<2$ has biased gradients.

To obtain unbiased gradients for any $\alpha$, we rewrite the terms $\phi_{\theta}^{\alpha}$ and $\phi_{\theta}^{\alpha-1}$ of the equation in Gibbs density form. Then, we have another variational representation of $\alpha$-divergence.

Subsequently, we obtain our loss function for DRE, called $\alpha$-Divergence loss function ($\alpha$-Div).

The superscript “$(R,S)$” is dropped when unnecessary and instead expressed as $\mathcal{L}_{\alpha\text{-Div}}(T_{\theta}\,;\,\alpha)$.

We train a neural network with $\alpha$-Div as described in Algorithm 1. In practice, neural networks rarely achieve the global optimum in Equation (8). The following theorem suggests that normalizing the estimated values, $q(\mathbf{x})/p(\mathbf{x})=e^{-T_{\theta}(\mathbf{x})}/\hat{E}_{P}\left[e^{-T_{\theta}}\right]$, improves the optimization of the neural networks.

$\alpha$-Div avoids the train-loss hacking problem when $\alpha$ is within $(0,1)$. Table 3 summarizes the lower-boundedness status of $\alpha$-Div for each case: $\alpha<0$, $0<\alpha<1$, or $\alpha>1$. $\alpha$-Div is lower-bounded when $0<\alpha<1$, whereas it is not lower-bounded when $\alpha>1$ or $\alpha<0$. Thus, selecting $\alpha$ from the interval (0, 1) effectively prevents the train-loss hacking problem.

By rewriting $\phi_{\theta}^{\alpha}$ in the Gibbs density form $e^{\alpha\cdot T_{\theta}}$, we mitigate the heavy-tailed behavior that often breaks uniform integrability in the standard $\alpha$-divergence loss functions. Hence, we present Theorem 5.1, which guarantees the unbiasedness of $\alpha$-Div’s gradients.

In Section 6.2, we empirically confirm that this unbiasedness is crucial for stable and effective optimization, as it prevents gradient estimates from drifting in the presence of heavy-tailed data.

When $\alpha$ is within $(0,1)$, $\alpha$-Div avoids the gradient vanishing issue during training. Below, we describe why gradient vanishing does not occur in this case.

First, we obtain the gradients of the standard $\alpha$-divergence loss in Equation (7) and $\alpha$-Div:

Next, consider the case where the estimated probability ratios, $\phi_{\theta}$ and $e^{T_{\theta}}$, are either nearly zero or very large for some point $\mathbf{x}$. Because equations such that $E_{Q}[e^{T_{\theta}}]\rightarrow 0\Leftrightarrow E_{P}[e^{T_{\theta}}]\rightarrow 0$ and $E_{Q}[e^{T_{\theta}}]\rightarrow\infty\Leftrightarrow E_{P}[e^{T_{\theta}}]\rightarrow\infty$ follow from the assumption that $p(\mathbf{x})>0\Leftrightarrow q(\mathbf{x})>0$ for all $\mathbf{x}\in\Omega$, the behavior of $E\big{[}\nabla_{\theta}\mathcal{L}_{\alpha\text{-Div}}(T_{\theta};\alpha)\big{]}$ under certain regularity conditions for $T_{\theta}$, as $E_{P}[e^{T_{\theta}}]\rightarrow 0$ or $E_{P}[e^{T_{\theta}}]\rightarrow\infty$, is summarized in Table 4.

In all cases except for $\alpha$-Div with $0<\alpha<1$, vanishing of the loss gradients is observed, such that $E\big{[}\nabla_{\theta}\mathcal{L}_{\alpha\text{-standard}}(\phi_{\theta})\big{]}\rightarrow\mathbf{0}$ or $E\big{[}\nabla_{\theta}\mathcal{L}_{\alpha\text{-Div}}(T_{\theta})\big{]}\rightarrow\mathbf{0}$. This implies that, during optimization, neural networks may remain stuck at extreme local minima when their density ratio estimations are either $0$ or $\infty$. However, this issue is avoided when $\alpha$ is within the interval $(0,1)$. Additionally, choosing $\alpha$ within the interval $(0,1)$ mitigates numerical instability arising from large differences in gradient values of the loss function for $\alpha>1$ and $\alpha<0$, which is represented as $E\big{[}\nabla_{\theta}\mathcal{L}_{\alpha\text{-Div}}(T{\theta};\alpha)\big{]}\rightarrow\infty-\infty$ in Table 4.

We present the exact upper bound on the sample size required for minimizing $\alpha$-Div in Theorem 5.2, which corresponds to Equation (4) for KL-divergence loss functions. The sample size requirement for minimizing $\alpha$-Div is upper-bounded depending on the value of $\alpha$. Intuitively, this property arises from the boundedness of Amari’s $\alpha$-divergence: $0\leq D_{\alpha}\leq 1/(\alpha\cdot(1-\alpha))$.

Unfortunately, despite the sample requirement stated in Equation (15), we empirically find that the estimation accuracy for $\alpha$-Div and KL-divergence loss functions is roughly the same as that of KL-divergence loss functions in downstream tasks of DRE, including KL-divergence estimation, as discussed in Section 6.3.

We empirically confirmed the stability of optimization using $\alpha$-Div, as discussed in Section 5.1. This includes addressing the potential divergence of training losses for $\alpha>1$ and $\alpha<0$, and observing the stability of optimization when $\alpha$ is within $(0,1)$. Subsequently, we conducted experiments using synthetic datasets to examine the behavior of training losses during optimization across different values of $\alpha$ at each learning step.

Experimental Setup. First, we generated 100 training datasets from two 5-dimensional normal distributions, $P=\mathcal{N}(\mu_{p},I_{5})$ and $Q=\mathcal{N}(\mu_{q},\Sigma_{q})$, where $\mu_{p}=\mu_{q}=(0,0,\ldots,0)$, and $I_{5}$ denotes the $5$-dimensional identity matrix. The covariance matrix $\Sigma_{q}=(\sigma_{ij})_{i=1}^{5}$ is defined as $\sigma_{ii}=1$, and $\sigma_{ij}=0.8$ for $i\neq j$. Subsequently, we trained neural networks using the synthetic datasets by optimizing $\alpha$-Div for $\alpha=-3.0,\allowbreak-2.0,\allowbreak-1.0,\allowbreak 0.2,\allowbreak 0.5,\allowbreak 0.8,\allowbreak 2.0,\allowbreak 3.0$, and $4.0$, while measuring training losses at each learning step. For each value of $\alpha$, 100 trials were performed. Finally, we reported the median of the training losses at each learning step, along with the ranges between the 45th and 55th percentiles and between the 2.5th and 97.5th percentiles.

Results. Figure 2 presents the training losses of $\alpha$-Div across the learning steps for $\alpha=-2.0,\allowbreak 3.0$, and $0.5$. Results for other values of $\alpha$ are provided in Section D.1 in the Appendix. The figures on the left ($\alpha=-2.0$) and in the center ($\alpha=3.0$) show that the training losses diverged to negative infinity when $\alpha<0$ or $\alpha>1$. In contrast, the figure on the right ($\alpha=0.5$) demonstrates that the training losses successfully converged. These results highlight the stability of $\alpha$-Div’s optimization when $\alpha$ is within the interval $(0,1)$, as discussed in Section 5.1.

Unbiased gradients of loss functions are expected to optimize neural network parameters more effectively than biased gradients, since they update the parameters in ideal directions at each iteration. We empirically compared the efficiency of minimizing training losses between the proposed loss function and the standard $\alpha$-divergence loss function derived from Equation (12), which highlighted the effectiveness of the unbiased gradients of the proposed loss function. Additionally, we observed that the estimated density ratios using the standard $\alpha$-divergence loss function diverged to large positive values, suggesting that the gradients of the standard $\alpha$-divergence loss function vanished. In contrast, $\alpha$-Div exhibited stable estimation. This finding aligns with the discussion in Section 5.3.

Experimental Setup. We first generated 100 training datasets from two normal distributions, $P=\mathcal{N}(\mu_{p},I_{5})$ and $Q=\mathcal{N}(\mu_{q},I_{5})$, where $I_{5}$ denotes the 5-dimensional identity matrix. The means were set as $\mu_{p}=(-5/2,0,0,0,0)$ and $\mu_{q}=(5/2,0,0,0,0)$. We then trained neural networks using three different loss functions: the standard $\alpha$-divergence loss function defined in Equation (12), $\alpha$-Div, and deep direct DRE (D3RE) [14]. Training losses were measured at each learning step. D3RE addresses train-loss hacking issues associated with Bregman divergence loss functions, as described in Section 3.4, by mitigating the lower-unboundedness of loss functions. Specifically, for D3RE, we employed the neural network-based Bregman divergence Least Squares Importance Fitting (nnBD-LSIF) loss function, which ensures unbiased gradients and stable optimization. The hyperparameter for nnBD-LSIF was set to $C=2$. For both the standard $\alpha$-divergence loss and $\alpha$-Div, we used $\alpha=0.5$. Finally, we reported the median training losses at each learning step, along with ranges between the 45th and 55th percentiles and between the 2.5th and 97.5th percentiles.

In Section 3.5, we examined the $\alpha$-divergence loss function, hypothesizing that its boundedness could address sample size issues in high KL-divergence data. Theorem 5.2, based on this boundedness, suggests that $\alpha$-Div can be minimized regardless of the true KL-divergence, indicating its potential for effective DRE with high KL-divergence data. To validate this hypothesis, we assessed DRE and KL-divergence estimation accuracy using both $\alpha$-Div and a KL-divergence loss function.

However, we observed that the RMSE of DRE using $\alpha$-Div increased significantly with higher KL-divergence, similar to the KL-divergence loss function.

Additionally, both methods yielded nearly identical KL-divergence estimations. These findings suggest that the accuracy of DRE and KL-divergence estimation is primarily influenced by the true amount of KL-divergence in the data rather than the $\alpha$-divergence.

Experimental Setup. We generated 100 training and 100 test datasets, each containing 10,000 samples. The datasets were drawn from two normal distributions, $P=\mathcal{N}(\mu_{p},\sigma^{2}\cdot I_{3})$ and $Q=\mathcal{N}(\mu_{q},4^{2}\cdot I_{3})$, where $\mu_{p}=(-3/2,-3/2,-3/2)$ and $\mu_{q}=(3/2,3/2,3/2)$, with $I_{3}$ denoting the 3-dimensional identity matrix. The values of $\sigma$ were set to 1.0, 1.1, 1.2, 1.4, 1.6, 2.0, 2.5, and 3.0. Correspondingly, the ground truth KL-divergence values of the datasets were 31.8, 25.6, 21.0, 14.5, 10.4, 5.8, 3.1, and 1.8 nats¹¹1A ’nat’ is a unit of information measured using the natural logarithm (base $e$), reflecting the increasing $\sigma^{2}$ values. The true density ratios of the test datasets are known for this experimental setup. We trained neural networks on the training datasets by optimizing both $\alpha$-Div with $\alpha=0.5$ and the KL-divergence loss function. After training, we measured the root mean squared error (RMSE) of the estimated density ratios using the test datasets. Additionally, we estimated the KL-divergence of the test datasets based on the estimated density ratios using a plug-in estimator. Finally, we reported the median RMSE of the DRE and the estimated KL-divergence, along with the interquartile range (25th to 75th percentiles), for both the KL-divergence loss function and $\alpha$-Div.

Results. Figure 3 shows the experimental results. The $x$-axis represents the true KL-divergence values of the test datasets, while the y-axes of the graphs display the RMSE (left) and estimated KL-divergence (right) for the test datasets. We empirically observed that the RMSE for DRE using $\alpha$-Div increased significantly as the KL-divergence of the datasets increased. A similar trend was observed for the KL-divergence loss function. Additionally, the KL-divergence estimation results were nearly identical between both methods. These findings indicate that the accuracy of DRE and KL-divergence estimation is primarily determined by the magnitude of KL-divergence in the data and is less influenced by $\alpha$-divergence. Therefore, we conclude that the approach discussed in Section 5.4 offers no advantage over the KL-divergence loss function in terms of the RMSE for DRE with high KL-divergence data. However, we believe that these empirical findings contribute to a deeper understanding of the accuracy of downstream tasks in DRE using $f$-divergence loss functions.