Lipschitzness Effect of a Loss Function on Generalization Performance of
  Deep Neural Networks Trained by Adam and AdamW Optimizers
The generalization performance of deep neural networks with regard to the
optimization algorithm is one of the major concerns in machine learning. This
performance can be affected by various factors. In this paper, we theoretically
prove that the Lipschitz constant of a loss function is an important factor to
diminish the generalization error of the output model obtained by Adam or
AdamW. The results can be used as a guideline for choosing the loss function
when the optimization algorithm is Adam or AdamW. In addition, to evaluate the
theoretical bound in a practical setting, we choose the human age estimation
problem in computer vision. For assessing the generalization better, the
training and test datasets are drawn from different distributions. Our
experimental evaluation shows that the loss function with a lower Lipschitz
constant and maximum value improves the generalization of the model trained by
Adam or AdamW.
Mohammad Lashkari, Amin Gheibi
Lipschitzness Effect of a Loss Function on Generalization Performance of Deep Neural Networks Trained by Adam and AdamW Optimizers

###### Abstract

The generalization performance of deep neural networks with regard to the optimization algorithm is one of the major concerns in machine learning. This performance can be affected by various factors. In this paper, we theoretically prove that the Lipschitz constant of a loss function is an important factor to diminish the generalization error of the output model obtained by Adam or AdamW. The results can be used as a guideline for choosing the loss function when the optimization algorithm is Adam or AdamW.

In addition, to evaluate the theoretical bound in a practical setting, we choose the human age estimation problem in computer vision. For assessing the generalization better, the training and test datasets are drawn from different distributions. Our experimental evaluation shows that the loss function with lower Lipschitz constant and maximum value improves the generalization of the model trained by Adam or AdamW.

+
Footnote †: journal: Amirkabir University of Technology (Tehran Polytechnic)

+
Footnote †: journal: Amirkabir University of Technology (Tehran Polytechnic)

+
Footnote †: journal: Amirkabir University of Technology (Tehran Polytechnic)

## 1 Introduction

The adaptive moment estimation (Adam) algorithm is one of the most widely used optimizers for training deep learning models. Adam is an efficient algorithm for stochastic optimization, based on adaptive estimates of first-order and second-order moments of gradient [1]. The method is computationally efficient and has little memory usage. Adam is much more stable than stochastic gradient descent (SGD) and the experiments of work [1] show that it is faster than previous stabilized versions of SGD, such as SGDNesterov [2], RMSProp [3] and AdaGrad [4] to minimize the loss function in the training phase. It is recently used in several machine learning problems and performs well. Thus, any improvement in the generalization performance of a model trained by Adam is essential.

One of the main concerns in machine learning is the generalization performance of deep neural networks (DNNs). A generalization measurement criterion is the generalization error which is defined as the difference between the true risk and the empirical risk of the output model [5]. One established way to address the generalization error of machine learning models in order to derive an upper bound for it, is the notion of uniform stability [5; 6; 7]. Roughly speaking, the uniform stability measures the difference in the error of the output model caused by a slight change in the training set. The pioneering work of [6], shows that if a deterministic learning algorithm is more stable, then the generalization error of the ultimate model achieves a tighter upper bound. In the following work of [7], Hardt _et al._ extend the notion of uniform stability to randomized learning algorithms to drive an upper bound for the expected generalization error of a DNN trained by SGD. They prove that SGD is more stable, provided that the number of iterations is sufficiently small. In the recent work of [5], Ali Akbari _et al._ derive a high probability generalization error bound instead of an expected generalization error bound. They demonstrate that if SGD is more uniformly stable, then the generalization error bound is tighter. They also proved the direct relationship between the uniform stability of SGD and loss function properties i.e. its Lipschitzness, resulting in the generalization error connection to the Lipschitz constant of a loss function.

In our work, Adam is central instead of SGD. We distinguish the relationship between the uniform stability of Adam and Lipschitzness of a loss function. Through this way, we connect the generalization error of a DNN trained by Adam to the properties of a loss function including its Lipschitzness. Subsequently, we assess the generalization performance of a DNN trained by AdamW optimizer which decouples weight decay from estimates of moments to make the regularization technique more effective [8]. We connect the uniform stability of AdamW and the
generalization error of a DNN trained by it to Lipschitzness of a loss function. In Experiments, we evaluate our theoretical results in the human age estimation problem.

Human age estimation is one of the most significant topics in a wide variety of applications such as age-specific advertising, customer profiling, or recommending new things. However, we are facing many challenges to solve this problem. Face makeup, insufficient light, skin color, and unique features of each person are the factors that can affect the accuracy of the model. Based on these reasons, collecting more data cannot necessarily reduce the generalization error of the final model. The practical results show that choosing a stable loss function can improve the accuracy of the model trained by Adam or AdamW.

## 2 Related Work

There is a variety of approaches to derive upper bounds for the generalization error including algorithmic stability [5; 6; 7; 9; 10; 11], robustness [12; 13], and PAC-Bayesian Theory [14; 15]. Each of these approaches theoretically analyzes some effective factors and gives the researchers some information which can enhance the generalization performance of deep learning models.

The notion of uniform stability was firstly introduced in [6] for deterministic algorithms. It was extended to randomized algorithms in the work of [7] to derive an expected upper bound for the generalization error which is directly related to the number of training epochs of SGD.

Recently, based on the uniform stability definition of SGD, the generalization error of a DNN trained by it, with high probability, is upper-bounded by a vanishing function which is directly related to the Lipschitz constant and the maximum value of a loss function [5]. In our work, we analyze the uniform stability for Adam and AdamW and its relationship with the Lipschitz constant of a loss function. We show that the loss function proposed in [5] stabilizes the training process and reduces the generalization error when the optimization algorithm is Adam or AdamW.

## 3 Preliminaries

Let \(X\) and \(Y\subseteq\mathbb{R}^{\mathbb{M}}\) be the input and output spaces of a problem respectively, and \(F\) be the set of all mappings from \(X\) to \(Y\). A learning problem is to find \(f^{\theta}:X\to Y\), parameterized by \(\theta\in H\) where \(H\) is the set of all possible values for the neural network parameters. Assume \(\ell:Y\times Y\rightarrow\mathbb{R}^{+}\) denotes the loss function of the problem. The goal of a learning algorithm is to minimize the true risk \(R_{true}(f^{\theta})\coloneqq\mathbb{E}_{\mathrm{(x,y)}}\left[\ell(f^{\theta }(\mathrm{x}),y)\right]\) where \(\mathrm{(x,y)}\in X\times Y\):

\[f^{\theta}_{true}=\operatorname*{argmin}_{f^{\theta}\in F}R_{true}(f^{\theta}). \tag{1}\]

Since the distribution of \(X\times Y\) is unknown; \(f^{\theta}_{true}\) cannot be found in the equation (1). Hence, we have to estimate the true risk. Let \(S\in(X\times Y)^{N}\) be the training set. The true risk is estimated by the empirical risk \(R_{emp}(f^{\theta})\coloneqq\frac{1}{N}\sum_{i=1}^{N}\ell(f^{\theta}(\mathrm{ x_{i}}),\mathrm{y_{i}})\) in which, \(N=|S|\) and \(\mathrm{(x_{i},y_{i})}\in S\). In current deep learning algorithms, training the model means minimizing \(R_{emp}(f^{\theta})\). In the rest of this paper, in the theorems and proofs, the loss function is denoted by \(\ell(\mathrm{y},\mathrm{y})\) where \(\mathrm{\hat{y}}\) is the prediction vector and \(\mathrm{y}\) is the target vector.

**Definition 3.1** (Partition): _Suppose that \(S\) is a training set of size \(N\). Let \(1<k<N\) be a number that \(N\) is divisible by \(k\) (if it is not possible, we repeat a sample enough to make divisibility possible). A partition of \(S\), which we denote by \(B_{S}=\{B_{1},B_{2},\ldots,B_{k}\}\), is a set of \(k\) subsets of \(S\) such that every sample is in exactly one set and the size of each subset is \(\frac{N}{k}\)._

We use the definition 3.1 to formalize the training process of deep learning models mathematically. Assume \(S\) is the training set and \(B_{S}=\{B_{1},B_{2},\ldots,B_{k}\}\) is a partition of it. Each element of \(B_{S}\) represents a mini-batch of \(S\). Without loss of generality we suppose that in each iteration of the optimization algorithm, a mini-batch \(B_{i}\in B_{S}\) is randomly selected to the parameters be updated. This is done by the algorithm using a random sequence \(R=(r_{1},r_{2},\ldots,r_{T})\) of indices of elements in \(B_{S}\), where \(T\) is the number of iterations. We use \(f^{\theta}_{B_{S},R}\) to denote the output model of the optimization algorithm, applied to a partition \(B_{S}\) and a random sequence \(R\).

**Definition 3.2** (Generalization Error): _Given a partition \(B_{S}\) of a training set \(S\) and a sequence \(R\) of random indices of \(B_{S}\) elements, the generalization error of \(f^{\theta}_{B_{S},R}\) trained by an arbitrary optimization algorithm, is defined as \(E(f^{\theta}_{B_{S},R})=R_{true}(f^{\theta}_{B_{S},R})-R_{emp}(f^{\theta}_{B_ {S},R})\)._
**Definition 3.3** (Lipschitzness): _Let \(Y\subseteq\mathbb{R}^{\mathrm{M}}\) be the output space of a problem. A loss function \(\ell(\hat{\mathrm{y}},\mathrm{y})\) is \(\gamma\)-Lipschitz with regard to its first argument, if \(\forall\,\mathrm{y}_{1},\mathrm{y}_{2}\in Y\), we have:_

\[\left|\ell(\mathrm{y}_{1},\mathrm{y})-\ell(\mathrm{y}_{2},\mathrm{y})\right| \leq\gamma\left\|\mathrm{y}_{1}-\mathrm{y}_{2}\right\|,\]

_where \(\left\|.\right\|\) is the \(L_{2}\) norm._

As mentioned before, uniform stability of the optimization algorithm is effective to the generalization performance of the ultimate model \(f^{\theta}_{B_{\mathrm{S}},R}\)[5]. We follow the uniform stability definition of work [7] to link Lipschitzness of the loss function to the generalization error of \(f^{\theta}_{B_{\mathrm{S}},R}\). For simplicity, moving forward, we denote \(f^{\theta}_{B_{\mathrm{S}},R}\) by \(f_{B_{\mathrm{S}},R}\) and \(E(f^{\theta}_{B_{\mathrm{S}},R})\) by \(E(f_{B_{\mathrm{S}},R})\).

Along with the notion of uniform stability which we define in Section 5, another concept called bounded difference condition (BDC) affects the generalization error [5]:

**Definition 3.4** (Bdc): _Consider two numbers \(k,T\in\mathbb{N}\). If \(G:\{1,2,\ldots,k\}^{T}\rightarrow\mathbb{R}^{+}\), is a measurable function and for \(R,R^{\prime}\in Dom(G)\) which are different only in two elements, constant \(\rho\) exists such that_

\[\sup_{R,R^{\prime}}\left|G(R^{\prime})-G(R)\right|\leq\rho,\]

_then, \(G(.)\) holds bounded difference condition (BDC) with the constant \(\rho\). We use the \(\rho\)-BDC expression to denote that a function holds this condition with the constant \(\rho\)._

In Definition 3.4, we assumed the slight change in the input to be the difference in two elements, which we will see its reason in the proof of the theorems. Intuitively, if a function satisfies the above condition, its value does not differ much due to a slight change in the input. Such functions are dense around their expectation with respect to the input random sequence \(R\)[16].

## 4 Formulation of Age Estimation Problem

Our problem in the experimental part is human age estimation. Let \((\mathrm{x},y)\) be a training sample where \(\mathrm{x}\) is the input image of a person's face and \(y\in\mathbb{N}\) is the corresponding age label. Due to the correlation of the neighboring ages, classification methods based on single-label learning [17] are not efficient because these methods ignore this correlation. Also, regression-based models are not stable to solve this problem [5].

According to the aforementioned reasons, another method based on label distribution learning (LDL) framework which was firstly introduced in the work of [18], is used for this problem [5]. In this method \(y\) is replaced by \(\mathrm{y}=[y_{1},y_{2},\ldots,y_{\mathrm{M}}]\in\mathbb{R}^{\mathrm{M}}\) where \(y_{i}\) is the probability of facial image \(\mathrm{x}\) belongs to class \(i\). As usual, \(\mathrm{y}\) is assumed to be a normal distribution, centering at \(y\) and standard deviation \(\sigma\) which controls the spread of the distribution [18]. Therefore, the output space, \(Y\) is a subset of \(\mathbb{R}^{\mathrm{M}}\) and our objective is to find \(f^{\theta}\) which maps \(\mathrm{x}\) to \(\mathrm{y}\in Y\).

### Loss Functions for Age Estimation Problem

Let \((\mathrm{x},\mathrm{y})\in S\) be a training instance where \(\mathrm{x}\) represents the facial image and \(\mathrm{y}\in\mathbb{R}^{\mathrm{M}}\) is the corresponding label distribution. Consider \(\hat{\mathrm{y}}=f^{\theta}(\mathrm{x})\), representing the estimated label distribution by \(f^{\theta}\). To obtain \(f^{\theta}\), a convex loss function named Kullbeck-Leibler (KL) divergence has been widely utilized. The KL loss function is defined as below:

\[\ell_{KL}(\hat{\mathrm{y}},\mathrm{y})=\sum_{m=1}^{\mathrm{M}}y_{m}\log(\frac{ y_{m}}{\hat{y}_{m}}).\]

As an alternative to KL, another convex loss function called Generalized Jeffries-Matusita (GJM) distance has been proposed in [5] under the LDL framework, defined as

\[\ell_{GJM}(\hat{\mathrm{y}},\mathrm{y})=\sum_{m=1}^{\mathrm{M}}y_{m}\left|1- \left(\frac{\hat{y}_{m}}{y_{m}}\right)^{\alpha}\right|^{\frac{1}{2}},\]

where \(\alpha\in(0,1]\). According to the experiments of [5], the best value of \(\alpha\) for good generalization is \(0.5\). It has been proved that if \(\alpha=0.5\), then the Lipschitz constant and the maximum value of GJM are less than the Lipschitz constant and the maximum value of KL respectively 1[5].
## 5 Uniform Stability and Generalization Error Analysis

The notion of uniform stability was firstly introduced in [6] for deterministic learning algorithms. They demonstrate that smaller stability measure of the learning algorithm, the tighter generalization error is. However, their stability measure is limited to deterministic algorithms and is not appropriate for randomized learning algorithms such as Adam. Therefore, we follow [5, 7] to define the uniform stability measure for randomized optimization algorithms generally:

**Definition 5.1** (Uniform Stability): _Let \(S\) and \(S^{\prime}\) denote two training sets drawn from a distribution \(\mathbb{P}\). Suppose that \(B_{S}\) and \(B_{S^{\prime}}\) of equal size k, are two partitions of \(S\) and \(S^{\prime}\) respectively, which are different in only one element (mini-batch). Consider a random sequence \(R\) of \(\{1,2,\ldots k\}\) to select a mini-batch at each iteration of an optimization algorithm, \(A_{opt}\). If \(f_{B_{S},R}\) and \(f_{B_{S^{\prime}},R}\) are output models obtained by \(A_{opt}\) with the same initialization, then \(A_{opt}\) is \(\beta\)-uniformly stable with regard to a loss function \(\ell\), if_

\[\forall S,S^{\prime}\ \ \sup_{(\mathrm{x},\mathrm{y})}\mathbb{E}_{R}\left[| \ell(f_{B_{S^{\prime}},R}(\mathrm{x}),\mathrm{y})-\ell(f_{B_{S},R}(\mathrm{x} ),\mathrm{y})|\right]\leq\beta.\]

To evaluate the uniform stability of Adam and AdamW in order to prove its link to loss function properties, a lemma named **Growth recursion** which has been stated in [7] for SGD is central to our analysis. In the following, we state this lemma for an arbitrary iterative optimization algorithm, but before stating the lemma, we need some definitions. As we know, gradient-based optimization algorithms are iterative, and in each iteration, the network parameters are updated. Consider \(H\) as the set of all possible values for the network parameters. Let \(A_{opt}\) be an arbitrary iterative optimization algorithm that runs \(T\) iterations. In the \(t\)-th iteration, the update that is computed in the last command of the loop for the network parameters, is a function \(A^{t}:H\to H\) mapping \(\theta_{t-1}\) to \(\theta_{t}\) for each \(1\leq t\leq T\). We call \(A^{t}\) the **update rule** of \(A_{opt}\). Let's define two characteristics of an update rule: The update rule, \(A^{t}(.)\) is \(\sigma\)**-bounded** if

\[\sup_{\theta\in H}\left\|\theta-A^{t}(\theta)\right\|\leq\sigma, \tag{2}\]

and it is \(\tau\)**-expensive** if

\[\sup_{\theta,\,\theta^{\prime}\in H}\frac{\left\|A^{t}(\theta)-A^{t}(\theta^{ \prime})\right\|}{\left\|\theta-\theta^{\prime}\right\|}\leq\tau, \tag{3}\]

where \(\left\|.\right\|\) is the \(L_{2}\) norm.

**Lemma 5.2** (Growth recursion): _[_7_]_ _Given two training set \(S\) and \(S^{\prime}\), suppose that \(\theta_{0},\theta_{1},\ldots,\theta_{T}\) and \(\theta^{\prime}_{0},\theta^{\prime}_{1},\ldots\theta^{\prime}_{T}\) are two updates of network parameters with update rules \(A^{t}_{S}\) and \(A^{t}_{S^{\prime}}\), running on \(S\) and \(S^{\prime}\) respectively such that for each \(1\leq t\leq T\), \(\theta_{t}=A^{t}_{S}(\theta_{t-1})\) and \(\theta^{\prime}_{t}=A^{t}_{S^{\prime}}(\theta^{\prime}_{t-1})\). If \(A^{t}_{S}\) and \(A^{t}_{S^{\prime}}\) are both \(\tau\)-expensive and \(\sigma\)-bounded, then for \(\Delta_{t}=\left\|\theta_{t}-\theta^{\prime}_{t}\right\|\), we have:_

* _If_ \(A^{t}_{S}=A^{t}_{S^{\prime}}\) _then_ \(\Delta_{t}\leq\tau\Delta_{t-1}\)_._
* _If_ \(A^{t}_{S}\neq A^{t}_{S^{\prime}}\) _then_ \(\Delta_{t}\leq\Delta_{t-1}+2\sigma\)__2_._ Footnote 2: In the work of [7], this inequality has been written as \(\Delta_{t}\leq\min(1,\tau)\Delta_{t-1}+2\sigma\) which is less than \(\Delta_{t}+2\sigma\) that we just need in the proofs of the theorems.

We state the proof of Lemma 5.2 in Appendix A. In Subsection 5.1, we discuss the uniform stability of Adam to upper-bound the generalization error of a DNN trained by it. Subsequently, In Subsection 5.2, we state different theorems for the uniform stability of AdamW and the generalization error because AdamW exploits decoupled weight decay, and its update parameters statement is different from Adam.

### Adam Optimizer

Let \(\ell(f^{\theta};B)\) represents the computation of a loss function on an arbitrary mini-batch, \(B=\{(\mathrm{x}_{i},\mathrm{y}_{i})\}_{i=1}^{b}\), which we use at each iteration to update parameters in order to minimize \(R_{emp}(f^{\theta})\):

\[\ell(f^{\theta};B)=\frac{1}{b}\sum_{i=1}^{b}\ell(f^{\theta}(\mathrm{x}_{i}), \mathrm{y}_{i}),\]

in which \(\theta\) is the parameters and \(b\) is the batch size. Let \(g(\theta)=\nabla_{\theta}\ell(f^{\theta};B)\) where \(\nabla_{\theta}\) is the gradient. For \(t\geq 1\) suppose that \(m_{t}\), \(v_{t}\) are estimates of the first and second moments respectively:

\[m_{t} =\beta_{1}\cdot m_{t-1}+(1-\beta_{1})\cdot g(\theta_{t-1});\ m_{0 }=0, \tag{4}\] \[v_{t} =\beta_{2}\cdot v_{t-1}+(1-\beta_{2})\cdot g^{2}(\theta_{t-1});\ v _{0}=0, \tag{5}\]
where \(\beta_{1},\beta_{2}\in(0,1)\) are exponential decay rates and the multiply operation is element-wise. Let \(\widehat{m}_{t}=m_{t}/(1-\beta_{1}^{t})\) and \(\widehat{v}_{t}=v_{t}/(1-\beta_{2}^{t})\) be the bias-corrected estimates; Adam computes the parameters update using \(\widehat{m}_{t}\) adapted by \(\widehat{v}_{t}\):

\[\theta_{t}=\theta_{t-1}-\eta\cdot\frac{\widehat{m}_{t}}{(\sqrt{\widehat{v}_{t}} +\epsilon)},\]

where \(\eta\) is the learning rate and \(\epsilon=10^{-8}\). Based on what we discussed so far, to evaluate the uniform stability of Adam, we need to formulate its update rule. Given \(\beta_{1},\beta_{2}\in(0,1)\) for each \(1\leq t\leq T\) let

\[\hat{M}(m_{t-1},\theta) =\frac{\beta_{1}\cdot m_{t-1}+(1-\beta_{1})\cdot g(\theta)}{1- \beta_{1}^{t}}, \tag{6}\] \[\hat{V}(v_{t-1},\theta) =\frac{\beta_{2}\cdot v_{t-1}+(1-\beta_{2})\cdot g^{2}(\theta)}{ 1-\beta_{2}^{t}}, \tag{7}\]

where \(m_{t-1}\) and \(v_{t-1}\) are the biased estimates for the first and second moments of the gradient at the previous step respectively as we explained in the equations (4) and (5). Adam's update rule is obtained as follows:

\[A^{t}(\theta)=\theta-\eta\cdot\left(\frac{\hat{M}(m_{t-1},\theta)}{\sqrt{\hat {V}(v_{t-1},\theta)}+\epsilon}\right), \tag{8}\]

where \(\eta\) is the learning rate and the division operation is element-wise. We use the following lemma in the proof of Theorem 5.4:

**Lemma 5.3**.: _Let \(m_{t-1}=\beta_{1}\cdot m_{t-2}+(1-\beta_{1})\cdot g(\theta_{t-2})\) such that \(\beta_{1}\in(0,1)\) is constant and \(m_{0}=0\). Let \(\ell(\mathrm{y},\mathrm{y})\) be \(\gamma\)-Lipschitz. Then for all \(t\geq 1\) and \(\theta\in H\), we have \(\left\|\hat{M}(m_{t-1},\theta)\right\|\leq\gamma\)._

The proof of Lemma 5.3 is available in Appendix A. Now we can state the theorems which link the generalization error with the loss function properties. In Theorem 5.4 we assess the stability measures including the uniform stability and in Theorem 5.5, we drive an upper bound for the generalization error of a DNN trained by Adam.

**Theorem 5.4**.: _Assume Adam is executed for \(T\) iterations with a learning rate \(\eta\) and batch size \(b\) to minimize the empirical risk in order to obtain \(f_{B_{g},R}\). Let \(\ell(\mathrm{y},\mathrm{y})\) be convex and \(\gamma\)-Lipschitz. Then, Adam is \(\beta\)-uniformly stable with regard to the loss function \(\ell\), and for each \((\mathrm{x},\mathrm{y})\), \(\ell(f_{B_{g},R}(\mathrm{x}),\mathrm{y})\) holds the \(\rho\)-BDC with respect to \(R\). Consequently, we have_

\[\beta\leq\frac{2\eta}{c}\cdot\frac{bT\gamma^{2}}{N},\quad\rho\leq\frac{8\eta}{ c}\cdot\left(\frac{b\gamma}{N}\right)^{2},\]

_in which \(c\in(0,1)\) is a constant number and \(N\) is the size of the training set._

**Proof.** Consider Adam's update rule, \(A^{t}(.)\) in the equation (8). In order to prove that \(A^{t}(.)\) satisfies the conditions of Lemma 5.2, \(\sigma\)-boundedness and \(\tau\)-expensiveness of \(A^{t}(.)\) are needed to be evaluated. From the formula (2), we have:

\[\left\|\theta-A^{t}(\theta)\right\|=\left\|\eta\cdot\left(\frac{\hat{M}(m_{t- 1},\theta)}{\sqrt{\hat{V}(v_{t-1},\theta)}+\epsilon}\right)\right\|\]

where \(m_{t-1}\) and \(v_{t-1}\) are the biased estimates for \(\mathbb{E}\left[g\right]\) and \(\mathbb{E}\left[g^{2}\right]\geq 0\) in the \(t\)-th step respectively. Therefore:

\[\left\|\eta\cdot\left(\frac{\hat{M}(m_{t-1},\theta)}{\sqrt{\hat {V}(v_{t-1},\theta)}+\epsilon}\right)\right\| \leq\eta\cdot\left\|\frac{\hat{M}(m_{t-1},\theta)}{\epsilon}\right\| \tag{9}\] \[\leq\frac{\eta\gamma}{\epsilon}. \tag{10}\]

Because \(\epsilon>0\) and \(\hat{V}(v_{t-1},\theta)\geq 0\), we deduced the inequality (9). In the inequality (10), Lemma 5.3 has been applied, which implies that, \(A^{t}(.)\)\(\sigma\)-bounded such that \(\sigma\leq\frac{\eta\gamma}{\epsilon}\). Now, we check the \(\tau\)-expensiveness condition: we know that for all \(\theta\in H\), \(\frac{\hat{M}(m_{t-1},\theta)}{\sqrt{\hat{V}(v_{t-1},\theta)}}\simeq\pm 1\) because \(|\mathbb{E}[g]|/\sqrt{\mathbb{E}[g^{2}]}\leq 1\). On the other hand \(\ell(\mathrm{y},\mathrm{y})\) is convex. Thus, for two updates of network parameters \(\theta_{t-1}\) and \(\theta_{t-1}^{\prime}\) in an arbitrary iteration \(t\) with the same initialization,
by choosing a sufficiently small learning rate, the two vectors \(\frac{\hat{M}(m_{c-1},\theta_{t-1})}{\sqrt{V(v_{t-1},\theta_{t-1})}}\) and \(\frac{\hat{M}(m_{c-1},\theta_{t-1}^{\prime})}{\sqrt{V(v_{t-1},\theta_{t-1}^{ \prime})}}\) are approximately equal. Thus, by substituting \(A^{t}(.)\) in the formula (3), it is concluded that, \(A^{t}(.)\) is \(1\)-expensive.

Let \(B_{S}\) and \(B_{S^{\prime}}\) having equal size \(k\), be two partitions of training sets \(S\) and \(S^{\prime}\) respectively, such that \(B_{S}\) and \(B_{S^{\prime}}\) are different in only one mini-batch. Let \(\theta_{0},\theta_{1},\ldots,\theta_{T}\) and \(\theta_{0}^{\prime},\theta_{1}^{\prime},\ldots\theta_{T}^{\prime}\) be two parameters updates obtained from training the network by Adam with update rules \(A_{S}^{t}\) and \(A_{S^{\prime}}^{t}\) respectively where \(A_{S}^{t}\) runs on \(B_{S}\) and \(A_{S^{\prime}}^{t}\) runs on \(B_{S^{\prime}}\) with the same random sequence \(R\) such that \(\theta_{0}=\theta_{0}^{\prime}\). Let two mini-batches \(B\) and \(B^{\prime}\) have been selected for updating the parameters in the \(t\)-th iteration. If \(B=B^{\prime}\), then \(A_{S^{\prime}}^{t}=A_{S}^{t}\) else \(A_{S^{\prime}}^{t}\neq A_{S}^{t}\). \(B=B^{\prime}\) occurs with probability \(1-\frac{1}{k}\) and the opposite occurs with probability \(\frac{1}{k}\). At the beginning of the proof, we demonstrated that \(A^{t}(.)\) (for an arbitrary training set) is \(\sigma\)-bounded and \(1\)-expensive. Let \(\Delta_{t}=\|\theta_{t}-\theta_{t}^{\prime}\|\), from Lemma 5.2, we have:

\[\Delta_{t} \leq(1-\frac{1}{k})\Delta_{t-1}+\frac{1}{k}\left(\Delta_{t-1}+ \frac{2\eta\gamma}{\epsilon}\right)\] \[=\Delta_{t-1}+\frac{1}{k}\cdot\frac{2\eta\gamma}{\epsilon}.\]

We know \(k=\frac{N}{b}\). Therefore, solving the recursive relation gives

\[\Delta_{T}\leq\Delta_{0}+2T\eta\cdot\frac{\gamma}{k\epsilon}=2\eta\cdot\frac {bT\gamma}{N\epsilon}.\]

Let \(\theta_{T,i}\) are the effective parameters of \(\theta_{T}\) on the \(i\)-th neuron of the last layer with \(M\) neurons. notation \(\langle.,.\rangle\) is inner product and \(\left[f(i)\right]_{i=1}^{M}\) for an arbitrary function \(f\), denotes the vector \(\left[f(1),f(2),\ldots,f(M)\right]\). Now we proceed to prove Adam's uniform stability. According to Definition 5.1, we have:

\[\mathbb{E}_{R}\left(\left|\ell(f_{B_{S^{\prime}},R}(\mathrm{x}), \mathrm{y})-\ell(f_{B_{S},R}(\mathrm{x}),\mathrm{y})\right|\right)\] \[\leq\mathbb{E}_{R}\left(\gamma\left\|f_{B_{S^{\prime}},R}( \mathrm{x})-f_{B_{S},R}(\mathrm{x})\right\|\right)\] \[=\gamma\mathbb{E}_{R}\left(\left\|\left[\left(\theta_{T,i}^{ \prime},\mathrm{x}\right)\right]_{i=1}^{M}-\left[\left(\theta_{T,i},\mathrm{x }\right)\right]_{i=1}^{M}\right\|\right)\] \[\leq\gamma\mathbb{E}_{R}\left(\|\theta_{T}^{\prime}-\theta_{T}\|\right) \tag{11}\] \[=\gamma\mathbb{E}_{R}\left[\Delta_{T}\right]\] \[\leq 2\eta\cdot\frac{bT\gamma^{2}}{N\epsilon}. \tag{12}\]

In the inequality (11), we assumed \(\|\mathrm{x}\|\leq 1\); that is the re-scaling technique that is common in computer vision. In the last inequality, \(\epsilon\) is a constant number between \(0\) and \(1\).

After showing the relation between the uniform stability of Adam and the Lipschitz constant of the loss function, we evaluate the bounded difference condition for the loss function with respect to the random sequence and a fixed training set. Suppose that \(R\) and \(R^{\prime}\) are two random sequences of batch indices to update the parameters in which only the location of two indices has been changed; that is if \(R=(\ldots,i,\ldots,j,\ldots)\) then \(R^{\prime}=(\ldots,j,\ldots,i,\ldots)\). Without loss of generality, assume \(1\leq i\leq\frac{k}{2}\) and \(\frac{k}{2}+1\leq j\leq k\). The probability of selecting two identical batches in the \(t\)-th iteration is \(1-\frac{4}{Tk^{2}}\). Thus, two updates of neural network parameters as \(\theta_{0}^{R},\theta_{1}^{R},\ldots,\theta_{T}^{R}\) and \(\theta_{0}^{R^{\prime}},\theta_{1}^{R^{\prime}},\ldots,\theta_{T}^{R^{\prime}}\) are made with the same initialization, \(\theta_{0}^{R}=\theta_{0}^{R^{\prime}}\). Let \(\Delta_{t}=\left\|\theta_{t}^{R}-\theta_{t}^{R^{\prime}}\right\|\). From Lemma 5.2, we have:

\[\Delta_{T}\leq\frac{8}{Tk^{2}}\cdot\frac{\eta T\gamma}{\epsilon}=\frac{8}{k^{2}} \cdot\frac{\eta\gamma}{\epsilon}.\]

According to Definition 3.4, we have:

\[\left|\ell(f_{B_{S},R^{\prime}}(\mathrm{x}),\mathrm{y})-\ell(f_{B _{S},R}(\mathrm{x}),\mathrm{y})\right|\] \[\leq\gamma\left\|f_{B_{S},R^{\prime}}(\mathrm{x})-f_{B_{S},R}( \mathrm{x})\right\|\] \[=\gamma\left\|\left[\left(\theta_{T,i}^{R^{\prime}},\mathrm{x} \right)\right]_{i=1}^{M}-\left[\left(\theta_{T,i}^{R},\mathrm{x}\right)\right]_{i =1}^{M}\right\|\] \[\leq\gamma\left\|\theta_{T}^{R^{\prime}}-\theta_{T}^{R}\right\| \tag{13}\] \[=\gamma\Delta_{T}\] \[\leq\frac{8}{k^{2}}\cdot\frac{\eta\gamma^{2}}{\epsilon}. \tag{14}\]

The inequality (13) has been obtained similar to (12). Replacing \(k\) by \(\frac{N}{b}\) in the inequality (14) leads to the inequality in the proposition.
**Theorem 5.5**.: _Let \(\ell(\mathrm{y},\mathrm{y})\) with the maximum value of \(L\) be convex and \(\gamma\)-Lipschitz. Assume Adam is run for \(T\) iterations with a learning rate \(\eta\) and batch size \(b\) to obtain \(f_{B_{S},R}\). Then we have the following upper bound for \(E(f_{B_{S},R})\) with probability at least \(1-\delta\):_

\[E(f_{B_{S},R})\leq\frac{2\eta}{c}\left(4\left(\frac{b\gamma}{N}\right)^{2} \sqrt{T\ log(2/\delta)}+\frac{bT\gamma^{2}}{N}\left(1+\sqrt{2N\log(2/\delta)} \right)\right)+L\sqrt{\frac{\log(2\delta)}{2N}}, \tag{15}\]

_in which \(c\in(0,1)\) is a constant number and \(N\) is the size of the training set._

**Proof.** In the work of [5], an upper bound for the generalization error of the output model trained by any optimization algorithm \(A_{opt}\) is established with probability at least \(1-\delta\), under the condition \(A_{opt}\) satisfies uniform stability measure with bound \(\beta\) and for each \(\mathrm{(x,y)}\), \(\ell(f_{B_{S},R}\mathrm{(x,y)}\mathrm{)}\) holds the \(\rho\)-BDC with regard to \(R\)3:

Footnote 3: In the assumptions of the main theorem in the work of [5], it has been stated that the model trained by stochastic gradient descent, but by studying the proof, we realize that their argument can be extended to any iterative algorithm that is \(\beta\)-uniformly stable because, in their proof, the upper bound has been derived independently of the update rule of stochastic gradient descent. The proof is available at [http://proceedings.mlr.press/v139/akbari21a/akbari21a-supp.pdf](http://proceedings.mlr.press/v139/akbari21a/akbari21a-supp.pdf).

\[E(f_{B_{S},R})\leq\rho\sqrt{T\log(2/\delta)}+\beta(1+\sqrt{2N\log(2/\delta)}) +L\sqrt{\frac{\log(2/\delta)}{2N}}. \tag{16}\]

By combining Theorem 5.4 and the inequality (16), we have the following upper bound with probability \(1-\delta\):

\[E(f_{B_{S},R})\leq\frac{2\eta}{c}\left(4\left(\frac{b\gamma}{N}\right)^{2} \sqrt{T\ log(2/\delta)}+\frac{bT\gamma^{2}}{N}\left(1+\sqrt{2N\log(2/\delta)} \right)\right)+L\sqrt{\frac{\log(2\delta)}{2N}}, \tag{17}\]

where \(c\in(0,1)\) is a constant number.

\(\Box\)

Theorem 5.5 shows how the generalization error bound of deep learning models trained by Adam depends on the Lipschitz constant \(\gamma\) and the maximum value \(L\). Furthermore, the inequality (15), implies the sensitivity of the generalization error to the batch size; when the batch size grows, \(E(f_{B_{S},R})\) increases. On the other hand, from the basics of machine learning, we know, if the batch size is too small, the parameters update is very noisy. Thus, an appropriate value should be considered for the batch size according to the training set size.

As we mentioned in Section 4 for the KL and GJM losses we have \(\gamma_{KL}\leq\gamma_{GJM}\) and \(L_{KL}\leq L_{GJM}\)[5]. Hence, following Theorem 5.5 we have the following corollary:

**Corollary 5.6**.: _Let \(f^{KL}_{B_{S},R}\) and \(f^{GJM}_{B_{S},R}\) be the output models trained by Adam optimizer using the KL and GJM loss functions respectively and the partition \(B_{S}\) obtained from the training set \(S\). We have_

\[E(f^{GJM}_{B_{S},R})\leq E(f^{KL}_{B_{S},R}).\]

**Proof.** We know if \(\alpha=0.5\), then the Lipschitz constant and the maximum value of GJM are less than the Lipschitz constant and the maximum value of KL respectively [5]. So, under the same settings for hyper-parameters of Adam and the same initialization, from Theorem 5.5, we have:

\[E(f^{GJM}_{B_{S},R})\leq E(f^{KL}_{B_{S},R}).\]

\(\Box\)

### AdamW Optimizer

The objective of regularization techniques is to control the network parameters' domain in order to prevent the over-fitting issue. \(L_{2}\)-regularization which exploits the \(L_{2}\) norm of the parameters vector, is more practical than \(L_{1}\) because it keeps the loss function differentiable and convex. In continuation, we study \(L_{2}\)-regularization and note its effect on SGD and Adam. The lack of significant effect of this technique on Adam led to AdamW 4[8].
Let \(\ell^{reg}(f^{\theta};B)\) be a regularized loss function computed on a mini-batch, \(B=\{(\mathrm{x}_{i},\mathrm{y}_{i})\}_{i=1}^{b}\):

\[\ell^{reg}(f^{\theta};B)=\frac{1}{b}\left(\sum_{i=1}^{b}\ell(f^{\theta}(\mathrm{ x}_{i}),\mathrm{y}_{i})+\frac{\lambda}{2}\left\|\theta\right\|^{2}\right), \tag{18}\]

where \(\left\|.\right\|\) is the \(L_{2}\) norm, \(\lambda\in\mathbb{R}^{+}\) is the weight decay and \(b\) is the batch size. According to the equation (18), to compute the parameters update in SGD, we have:

\[\theta_{t}=\left(1-\frac{\eta\lambda}{b}\right)\theta_{t-1}-\frac{\eta}{b} \sum_{i=1}^{b}\nabla_{\theta}\ell(f^{\theta}(\mathrm{x}_{i}),\mathrm{y}_{i}).\]

In SGD, minimizing the regularized loss function can improve the output model generalization. However, this technique, cannot be effective in Adam because it uses adaptive gradients to update the parameters [8]. In AdamW, the weight decay hyperparameter was decoupled from optimization steps taking gradient of the loss function. Let \(\widehat{m}_{t}\) and \(\widehat{v}_{t}\) denote the bias-corrected estimates illustrated in Subsection 5.1. The parameters update is computed as follows:

\[\theta_{t}=\theta_{t-1}-\alpha_{t}\left(\eta\cdot\frac{\widehat{m}_{t}}{( \sqrt{\widehat{v}_{t}}+\epsilon)}+\lambda\theta\right) \tag{19}\]

where \(\alpha_{t}\) is the schedule multiplier. The equation (19) exhibits that AdamW updates the parameters in a different way than Adam. Hence, we need to state theorems specific to AdamW for the stability and the generalization error. Consider \(\hat{M}(m_{t-1},\theta)\) and \(\hat{V}(v_{t-1},\theta)\) in the equations (6) and (7). According to update parameters statement of adamW in formula 19, The AdamW's update rule is defined as

\[A_{W}^{t}(\theta)=\theta-\alpha_{t}\left(\eta\cdot\frac{\hat{M}(m_{t-1}, \theta)}{\sqrt{\hat{V}(v_{t-1},\theta)}+\epsilon}+\lambda\theta\right). \tag{20}\]

where \(0\mid\mathrm{a}_{t}\lambda<1\) because otherwise, the update occurs in a wrong direction which means it goes away from the minimum. Consider the set of all possible values for the network parameters, \(H\subset\mathbb{R}^{K}\). Without loss of generality we can assume \(H\) is bounded 5. Let \(\left\|\theta\right\|_{\mathrm{sup}}=\sup_{\theta\in H}\left\|\theta\right\|\):

Footnote 5: We know the number of iterations, T is finite. Therefore the set of visible values for parameters in the training stage is finite. So we can assume the set of all possible values is an infinite bounded superset of visible values.

**Theorem 5.7**.: _Assume AdamW is executed for \(T\) iterations with a learning rate \(\eta\), batch size \(b\), weight decay \(\lambda\), and schedule multiplier \(\alpha_{t}\) to minimize the empirical risk in order to obtain \(f_{B_{S},R}\). Let \(\ell(\hat{\mathrm{y}},\mathrm{y})\) be convex and \(\gamma\)-Lipschitz. Then, Adam is \(\beta\)-uniformly stable with regard to the loss function \(\ell\), and for each \((\mathrm{x},\mathrm{y})\), \(\ell(f_{B_{S},R}(\mathrm{x}),\mathrm{y})\) holds the \(\rho\)-BDC with respect to \(R\). Consequently, we have_

\[\beta\leq\frac{2bT}{N}\sum_{t=1}^{T}\alpha_{t}\left(\frac{\eta\gamma^{2}}{c}+ \gamma\lambda\left\|\theta\right\|_{\mathrm{sup}}\right),\quad\rho\leq\frac{8 b^{2}}{N^{2}}\sum_{t=1}^{T}\alpha_{t}\left(\frac{\eta\gamma^{2}}{c}+\gamma \lambda\left\|\theta\right\|_{\mathrm{sup}}\right),\]

_in which \(c\in(0,1)\) is a constant number and \(N\) is the size of the training set._

**Proof.** First, we check the \(\sigma\)-boundedness of \(A_{W}^{t}(\theta)\):

\[\left\|\theta-A_{W}^{t}(\theta)\right\| =\left\|\alpha_{t}\left(\eta\cdot\frac{\hat{M}(m_{t-1},\theta)}{ \sqrt{\hat{V}(v_{t},\theta)}+\epsilon}+\lambda\theta\right)\right\|\] \[\leq\left\|\alpha_{t}\cdot\frac{\hat{M}(m_{t-1},\theta)}{\epsilon} \right\|+\alpha_{t}\lambda\left\|\theta\right\|\] \[=\alpha_{t}\left(\eta\cdot\frac{\left\|\hat{M}(m_{t-1},\theta) \right\|}{\epsilon}+\lambda\left\|\theta\right\|\right)\] \[\leq\alpha_{t}\left(\frac{\eta\gamma}{\epsilon}+\lambda\left\| \theta\right\|_{\mathrm{sup}}\right). \tag{21}\]
By applying Lemma 5.3, we concluded the inequality (21), which shows that \(A_{W}^{t}(\theta)\) is \(\sigma\)-bounded. Now we evaluate the \(\tau\)-expensiveness of AdamW. According to the formula (3), we have

\[\frac{\left\|A_{W}^{t}(\theta)-A_{W}^{t}(\theta^{\prime})\right\|}{ \left\|\theta-\theta^{\prime}\right\|}\] \[=\frac{\left\|-\alpha_{t}\left(\eta\cdot\frac{\bar{M}(m_{t-1}, \theta)}{\sqrt{\bar{V}(v_{t-1},\theta)+\epsilon}}+\lambda\theta\right)+\alpha_ {t}\left(\eta\cdot\frac{\bar{M}(m_{t-1},\theta^{\prime})}{\sqrt{\bar{V}(v_{t-1 },\theta^{\prime})+\epsilon}}+\lambda\theta^{\prime}\right)+\theta-\theta^{ \prime}\right\|}{\left\|\theta-\theta^{\prime}\right\|}. \tag{22}\]

As said in the proof of Theorem 5.4, for every \(\theta\in H\), we have \(\frac{\bar{M}(m_{t-1},\theta)}{\sqrt{\bar{V}(v_{t-1},\theta)}}\simeq\pm 1\) because \(|\mathbb{E}[g]|/\sqrt{\mathbb{E}[g^{2}]}\leq 1\). Therefore, the equation (22) is written as follows:

\[\frac{\left\|-\alpha_{t}\lambda\theta+\alpha_{t}\lambda\theta^{ \prime}+\theta-\theta^{\prime}\right\|}{\left\|\theta-\theta^{\prime}\right\|} =\frac{\left\|\alpha_{t}\lambda(\theta^{\prime}-\theta)+\theta- \theta^{\prime}\right\|}{\left\|\theta-\theta^{\prime}\right\|}\] \[=\frac{\left|1-\alpha_{t}\lambda\right\|\left\|\theta-\theta^{ \prime}\right\|}{\left\|\theta-\theta^{\prime}\right\|}\] \[=\left|1-\alpha_{t}\lambda\right|<1. \tag{23}\]

AdamW update rule in the equation (20) implies that \(0<\alpha_{t}\lambda<1\) which its consequent is the inequality (23). with an analogous demonstration to what we did in the proof of Theorem 5.4, i.e. considering update sequences and using Lemma 5.2 in order to evaluate the uniform stability and bounded difference condition according to their definitions, we conclude the following inequalities:

\[\beta \leq\frac{2bT}{N}\sum_{t=1}^{T}\alpha_{t}\left(\frac{\eta\gamma^{ 2}}{\epsilon}+\gamma\lambda\left\|\theta\right\|_{\sup}\right),\] \[\rho \leq\frac{8b^{2}}{N^{2}}\sum_{t=1}^{T}\alpha_{t}\left(\frac{\eta \gamma^{2}}{\epsilon}+\gamma\lambda\left\|\theta\right\|_{\sup}\right).\]

**Theorem 5.8**.: _Let \(\ell(\hat{\mathrm{y}},\mathrm{y})\) with the maximum value of \(L\) be convex and \(\gamma\)-Lipschitz. Assume AdamW is run for \(T\) iterations with a learning rate \(\eta\), batch size \(b\), weight decay \(\lambda\), and schedule multiplier \(\alpha_{t}\) to obtain \(f_{B_{S},R}\). Then we have the following upper bound for \(E(f_{B_{S},R})\) with probability at least \(1-\delta\):_

\[E(f_{B_{S},R})\leq\frac{2b}{N}\sum_{t=1}^{T}\alpha_{t}\left(\frac{\eta\gamma^{ 2}}{c}+\gamma\lambda\left\|\theta\right\|_{\sup}\right)\left(\frac{4b}{N} \sqrt{T\log(2/\delta)}+T\sqrt{2N\log(2/\delta)}\right)+L\sqrt{\frac{\log(2/ \delta)}{2N}}. \tag{24}\]

_in which \(c\in(0,1)\) is a constant number and \(N\) is the size of the training set._

**Proof.** By combining the equation (16) and Theorem 5.7 we conclude the proposition.

\(\Box\)

The inequality (24) implies that the generalization error growth of a DNN trained by AdamW, is directly related to the Lipschitz constant and the maximum value of a loss function. Following Theorem 5.8 we have the following corollary for the KL and GJM loss functions:

**Corollary 5.9**.: _Let \(f_{B_{S},R}^{KL}\) and \(f_{B_{S},R}^{GJM}\) be the output models trained by AdamW optimizer using the KL and GJM loss functions respectively using the partition \(B_{S}\) obtained from the training set \(S\). We have_

\[E(f_{B_{S},R}^{GJM})\leq E(f_{B_{S},R}^{KL}).\]

**Proof.** The proposition is concluded by Theorem 5.8 and an analogous argument to Corollary 5.6.
## 6 Experimental Evaluation

### Datasets

We use \(4\) datasets, including UTKFace [19], AgeDB [20], MegaAge-Asian [21], and FG-NET [22] to evaluate age estimation performance. UTKFace dataset contains \(23,708\) facial images, providing enough samples of all ages, ranging from \(0\) to \(116\) years-old. AgeDB contains \(16,488\) in-the-wild images in the age range from \(0\) to \(100\) years-old. MegaAge-Asian has been already split into MegaAge-Train and Mega-Age-Test datasets, containing \(40,000\) and \(3,945\) images respectively, belonging to Asian people with the age label in the range from \(1\) to \(69\) years-old. FG-NET dataset contains \(1,002\) facial images in the age range of \(0\) to \(69\) years. This dataset covers variations in pose, age expression, resolution, and lighting conditions. By collecting the samples from UTKFace, MegaAge-Train, and AgeDB datasets whose ages are in the range from \(0\) to \(100\) years-old, we create a new dataset called UAM, which includes \(80,174\) images. We use UTKFace and UAM as the training sets. FG-NET, MegaAge-Test, and \(10\%\) randomly selected from AgeDB called AgeDB-Test, are left as the test sets.

### Settings

All images are pre-processed by the following procedures: face detection and alignment are done by prepared modules in OpenCV package. All images are reshaped to the size of \(256\times 256\) and standard data augmentation techniques, including random cropping and horizontal flipping, are carried out during the training phase. We use two neural network architectures VGG16 [23] and ResNet50 [24], pre-trained on ImageNet [25] and VGGFace2 [26] datasets respectively, to estimate human age. VGGFace2 dataset was created with the aim of estimating human pose and age. With the same seed, the last layer of these models is replaced with a M-neurons dense layer with random weights. The last layer of VGG16 is trained on UTKFace in \(5\) epochs and the last layer of ResNet50 is trained on UAM in \(15\) epochs. M is set to \(116\) in VGG16 and \(101\) in ResNet50 model. We train the models via Adam and AdamW with learning rate \(2\times 10^{-5}\) for KL and \(10^{-4}\) for GJM 6. The batch size and AdamW's weight decay are set to \(64\) and \(0.9\) respectively. We set \(\beta_{1}=0.9\) and \(\beta_{2}=0.999\) for both Adam and AdamW as the authors of [1] and [8] suggested.

Footnote 6: In our experiments, when we set the learning rate to \(2\times 10^{-5}\) for the GJM loss, the ultimate model at the last epoch remained under-fit.

### Evaluation Metrics and Results

As the first observation, we measure the generalization error estimate in the training steps of ResNet50 trained by Adam and AdamW which is defined as

\[\hat{E}(f_{B_{S},R})=|R_{train}(f_{B_{S},R})-R_{val}(f_{B_{S},R})|,\]

where \(f_{B_{S},R}\) is the output model, \(R_{train}(f_{B_{S},R})\), \(R_{val}(f_{B_{S},R})\) are the average of loss values on the training and validation sets respectively. The results of this experiment are shown in Figure 1 and Figure 2. In the first epochs, the models are still under-fit and the loss is far from its minimum; therefore, \(\hat{E}(f_{B_{S},R})\) does not give us critical information about the generalization error, but in the rest of epochs, when the experimental loss of the models approaches its minimum, \(\hat{E}(f_{B_{S},R})\) can represent the generalization error. As can be seen in Figure 0(a) and Figure 0(a), after epoch \(5\) or \(6\) the generalization error estimate of the models trained by Adam and AdamW using the GJM loss function is lower than the models trained using the KL loss.
In addition, we measure the generalization performance in terms of Mean Absolute Error (MAE) and Cumulative Score (CS). Consider the training set \(S\), and the test set \(S_{test}\in(X\times Y)^{D}\). Let \((\mathrm{x}_{k},y_{k})\in S_{test}\) represents a test example where \(y_{k}\in\mathbb{R}\) is the label of \(k\)-th example of the test set. Since we use label distribution learning, for each \((\mathrm{x},\mathrm{y})\in S\), \(\mathrm{y}\in\mathbb{R}^{\mathrm{M}}\) is the probability distribution corresponding to \(\mathrm{x}\). Therefore, in the evaluation phase, the output of the model per the test example \(\mathrm{x}_{k}\) is the predicted probability distribution \(\hat{\mathrm{y}}_{k}=[\hat{y}_{k,1},\hat{y}_{k,2},\ldots,\hat{y}_{k,\mathrm{M}}]\). MAE is defined as \(\frac{1}{D}\sum_{k=1}^{D}|\hat{l}_{k}-l_{k}|\) where \(\hat{l}_{k}\) is the index of the largest element of \(\hat{\mathrm{y}}_{k}\) and \(l_{k}\) is the true label. CS is defined as \(\frac{D_{I}}{D}\times 100\%\) where \(D_{I}\) is the number of test samples such that \(|\hat{l}_{k}-l_{k}|<I\). Commonly, the value of \(I\) is set to 5 [5][27].

The results are reported in Tables 1-3. The ResNet50 models are more accurate than the VGG16 models because the version of VGG16 is pre-trained on ImageNet dataset which is not suitable for age estimation. Tables 1- 3 show that when we train a DNN by Adam or AdamW, the GJM loss performs better than the KL loss.