Signal Processing Meets SGD: From Momentum to Filter

In the SGDF algorithm, $s_{t}$ serves as a key indicator, calculated as the exponential moving average of the squared difference between the current gradient $g_{t}$ and its momentum $m_{t}$, acting as a marker for gradient variation with weight-adjusted by $\beta_{2}$. Zhuang et al. [66] first proposed the calculation of $s_{t}$, which is utilized for estimating the fluctuation variance of the stochastic gradient. We derived a correction factor $(1-\beta_{1})(1-\beta_{1}^{2t})/(1+\beta_{1})$ under the assumption that $m_{t}$ and $g_{t}$ are independently and identically distributed (i.i.d.), to accurately estimate the variance of $m_{t}$ using $s_{t}$. Fig. 11 compares performances with and without the correction factor, showing superior results with correction. For the derivation of the correction factor, please refer to Appendix A.2.

At each time step $t$, $g_{t}$ represents the stochastic gradient for our objective function, while $m_{t}$ approximates the historical trend of the gradient through an exponential moving average. The difference $g_{t}-m_{t}$ highlights the gradient’s deviation from its historical pattern, reflecting the inherent noise or uncertainty in the instantaneous gradient estimate, which can be expressed as $p(g_{t}|\mathcal{D})\sim\mathcal{N}(g_{t};m_{t},\sigma_{t}^{2})$ [37].

SGDF utilizes the gain $K_{t}$, where the components of each dimension of the estimated gain range between 0 and 1, to balance the current observed gradient $g_{t}$ and the past corrected gradient $\widehat{m}_{t}$, thus optimizing the gradient estimate. This balance plays a crucial role in noisy or complex optimization scenarios, helping to mitigate noise and achieve stable gradient direction, faster convergence, and enhanced performance. The computation of $K_{t}$, based on $s_{t}$ and $g_{t}-m_{t}$, aims to minimize the expected variance of the corrected gradient $\widehat{g}_{t}$ for optimal linear estimation in noisy conditions. For the method derivation, please refer to Appendix A.1.

SGDF optimizes gradient estimate through the Wiener Filter, enhancing the stability and effectiveness of the optimization process and facilitating faster convergence. By fusing two Gaussian distributions, SGDF significantly reduces the variance of gradient estimates, thereby benefiting in solving complex stochastic optimization problems. In this section, we will delve into how SGDF achieves the reduction of gradient estimate variance.

The properties of SGDF ensure that the estimated gradient is a linear combination of the current noisy gradient observation $g_{t}$ and the first-order moment estimate $\widehat{m}_{t}$. These two components are assumed to have Gaussian distributions, where $g_{i}\sim\mathcal{N}\left(\mu,\sigma^{2}\right)$. Hence, their fusion by the filter naturally ensures that the fused estimate $\widehat{g}_{t}$ is also Gaussian.

Consider two Gaussian distributions for the momentum term $\widehat{m}_{t}$ and the current gradient $g_{t}$:

• •

  The momentum term $\widehat{m}_{t}$ is normally distributed with mean $\mu_{m}$ and variance $\sigma_{m}^{2}$, denoted as $\widehat{m}_{t}\sim\mathcal{N}(\mu_{m},\sigma_{m}^{2})$.

• •

  The current gradient $g_{t}$ is normally distributed with mean $\mu_{g}$ and variance $\sigma_{g}^{2}$, denoted as $g_{t}\sim\mathcal{N}(\mu_{g},\sigma_{g}^{2})$.

The product of their probability density functions is given by:

Through coefficient matching in the exponential terms, we obtain the new mean and variance:

The new mean $\mu^{\prime}$ is a weighted average of the two means, $\mu_{m}$ and $\mu_{g}$, with weights inversely proportional to their variances. This places $\mu^{\prime}$ between $\mu_{m}$ and $\mu_{g}$, closer to the mean with the smaller variance. The new standard deviation $\sigma^{\prime}$ is smaller than either of the original standard deviations $\sigma_{m}$ and $\sigma_{g}$, reflecting the reduced uncertainty in the estimate due to the combination of information from both sources. This is a direct consequence of the Wiener Filter’s optimality in the mean-square error sense. The proof is provided in Appendix A.3.

In Section 3.2, we explained the principle of noise reduction by SGDF. In this section, we illustrate the advantages of SGDF over other variance reduction techniques, ensuring better generalization.

In prior variance reduction techniques [13, 27, 48], the variance is gradually reduced at a rate of $\xi^{t-1},\xi\in(0,1)$. This approach rapidly decreases the initial variance, facilitating a linear convergence rate, as noted in [5]. However, this approach poses a potential problem in that the variance may approach zero, thus depriving the optimization process of necessary stochastic exploration. The Wiener Filter satisfies statistical estimation theory, wherein the Cramér-Rao lower bound (CRLB) [43] provides an important lower bound of the variance in the estimator. Thus, we construct a dynamical system modeled by the Fokker-Planck equation to demonstrate the potential benefits of having a lower bound on the variance for optimization.

The existence of a variance lower bound critically enhances the algorithm’s exploration capabilities, especially in regions of the loss landscape where gradients are minimal. By preventing the probability density function from becoming unbounded, it ensures continuous exploration and increases the probability of converging to flat minima associated with better generalization properties [57]. The proof of Theorem 3.1 is provided in Appendix A.4.

Finally, we provide the convergence property of SGDF as shown in Theorem 3.2 and Theorem 3.3. The assumptions are common and standard when analyzing the convergence of convex and non-convex functions via SGD-based methods [9, 32, 44]. Proofs for convergence in convex and non-convex cases are provided in Appendix B and Appendix C, respectively. In the convergence analysis, the assumptions are relaxed and the upper bound is reduced due to the estimation gain introduced by SGDF, promoting faster convergence.

For the convex case, Theorem 3.2 implies that the regret of SGDF is upper bounded by $O(\sqrt{T})$. In the Adam-type optimizers, it’s crucial for the convex analysis to decay $\beta_{1,t}$ towards zero [32, 66]. We have relaxed the analysis assumption by introducing a time-varying gain $K_{t}$, which can adapt with variance. Moreover, $K_{t}$ converges with variance at the end of training to improve convergence [50].

Theorem 3.3 implies the convergence rate for SGDF in the non-convex case is $O(\log T/\sqrt{T})$, which is similar to Adam-type optimizers [9, 44]. Note that these bounds are derived under the worst possible case, so they may be loose. The components of each dimension of the estimated gain $K_{t}$ converge with variance, causing the upper bound to decrease rapidly. This is why SGDF converges faster than SGD.

In this study, we focus on the following tasks: Image Classification. We employed the VGG [49], ResNet [23], and DenseNet [26] models for image classification tasks on the CIFAR dataset [33]. The major difference between these three network architectures is the residual connectivity, which we will discuss in section 4.4. We evaluated and compared the performance of SGDF with other optimizers such as SGD, Adam, RAdam [37], AdamW [38], MSVAG [2], Adabound [40], Sophia [36], and Lion [10], all of which were implemented based on the official PyTorch. Additionally, we further tested the performance of SGDF on the ImageNet dataset [14] using the ResNet model. Object Detection. Object detection was performed on the PASCAL VOC dataset [18] using Faster-RCNN [45] integrated with FPN. For hyper-parameter tuning related to image classification and object detection, refer to [66]. Image Generation. Wasserstein-GAN (WGAN) [1] on the CIFAR-10 dataset.

Hyperparameter tuning. Following Zhuang et al. [66], we delved deep into the optimal hyperparameter settings for our experiments. In the image classification task, we employed these settings:

• •

  SGDF: We adhered to Adam’s original parameter values: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-8}$.

• •

  SGD: We set the momentum to 0.9, the default for networks like ResNet and DenseNet. The learning rate was searched in the set {10.0, 1.0, 0.1, 0.01, 0.001}.

• •

  Adam, RAdam, MSVAG, AdaBound: Traversing the hyperparameter landscape, we scoured $\beta_{1}$ values in $\{0.5,0.6,0.7,0.8,0.9\}$, probed $\alpha$ as in SGD, while tethering other parameters to their literary defaults.

• •

  AdamW, SophiaG, Lion: Mirroring Adam’s parameter search schema, we fixed weight decay at $5\times 10^{-4}$; yet for AdamW, whose optimal decay often exceeds norms [38], we ranged weight decay over $\left\{10^{-4},5\times 10^{-4},10^{-3},10^{-2},10^{-1}\right\}$.

• •

  SophiaG, Lion: We searched for the learning rate among $\{10^{-3},10^{-4},10^{-5}\}$ and adjusted Lion’s learning rate [36]. Following  [36, 10], we set $\beta_{1}$=0.965, 0.9 and $\beta_{2}$=0.99 as the default parameters.

CIFAR-10/100 Experiments. We initially trained on the CIFAR-10 and CIFAR-100 datasets using the VGG, ResNet, and DenseNet models and assessed the performance of the SGDF optimizer. In these experiments, we employed basic data augmentation techniques such as random horizontal flip and random cropping (with a 4-pixel padding). To facilitate result reproduction, we provide the parameter table for this subpart in Tab. 5. The results represent the mean and standard deviation of 3 runs, visualized as curve graphs in Fig. 2.

As Fig. 2 shows, that it is evident that the SGDF optimizer exhibited convergence speeds comparable to adaptive optimization algorithms. Additionally, SGDF’s final test set accuracy was either better than or equal to that achieved by SGD.

ImageNet Experiments. We utilized the ResNet18 model to train on the ImageNet dataset. Due to extensive computational demands, we did not conduct an exhaustive hyperparameter search for other optimizers and instead opted to use the best-reported parameters from [8, 37]. We applied basic data augmentation strategies such as random cropping and random horizontal flipping. The results are presented in Tab. 1. To facilitate result reproduction, we provide the parameter table for this subpart in Tab. 6. Detailed training and test curves are depicted in Fig. 8. Additionally, to mitigate the effect of learning rate scheduling, we employed cosine learning rate scheduling as suggested by [10, 65] and trained ResNet18, 34, and 50 models. The results are summarized in Tab. 2. Experiments on the ImageNet dataset demonstrate that SGDF has improved convergence speed and achieves similar accuracy to SGD on the test set.

Object Detection. We conducted object detection experiments on the PASCAL VOC dataset [18]. The model used in these experiments was pre-trained on the COCO dataset [35], obtained from the official website. We trained this model on the VOC2007 and VOC2012 trainval dataset (17K) and evaluated it on the VOC2007 test dataset (5K). The utilized model was Faster-RCNN [45] with FPN, and the backbone was ResNet50 [23]. Results are summarized in Tab. 3. To facilitate result reproduction, we provide the parameter table for this subpart in Tab. 6. As expected, SGDF outperforms other methods. These results also illustrate the efficiency of our method in object detection tasks.

Image Generation. The stability of optimizers is crucial, especially when training Generative Adversarial Networks (GANs). If the generator and discriminator have mismatched complexities, it can lead to imbalance during GAN training, causing the GAN to fail to converge. This is known as model collapse. For instance, Vanilla SGD frequently causes model collapse, making adaptive optimizers like Adam and RMSProp the preferred choice. Therefore, GAN training provides a good benchmark for assessing optimizer stability. For reproducibility details, please refer to the parameter table in Table 7.

We evaluated the Wasserstein-GAN with gradient penalty (WGAN-GP) [47]. Using well-known optimizers [4, 61], the model was trained for 100 epochs. We then calculated the Frechet Inception Distance (FID) [24] which is a metric that measures the similarity between the real image and the generated image distribution and is used to assess the quality of the generated model (lower FID indicates superior performance). Five random runs were conducted, and the outcomes are presented in Figure 3. Results for SGD and MSVAG were extracted from Zhuang et al. [66].

Experimental results demonstrate that SGDF significantly enhances WGAN-GP model training, achieving a FID score higher than vanilla SGD and outperforming most adaptive optimization methods. The integration of a Wiener filter in SGDF facilitates smooth gradient updates, mitigating training oscillations and effectively addressing the issue of pattern collapse.

The success of optimization algorithms in deep learning not only depends on their ability to minimize training loss but also critically hinges on the nature of the solutions they converge to. Adam-type algorithms tend to converge to sharp minima, characterized by large eigenvalues of the Hessian matrix [6].

We computed the Hessian spectrum of ResNet-18 trained on the CIFAR-100 dataset for 200 epochs using four optimization methods: SGD, SGDM, Adam, and SGDF. These experiments ensure that all methods achieve similar results on the training set. We employed power iteration [58] to compute the top eigenvalues of Hessian and Hutchinson’s method [59] to compute the Hessian trace. Histograms illustrating the distribution of the top 50 Hessian eigenvalues for each optimization method are presented in Fig. 4.

We visualized the loss landscapes of models trained with SGD, SGDM, SGDF, and Adam using the ResNet-18 model on CIFAR-100, following the method in [34]. All models are trained with the same hyperparameters for 200 epochs, as detailed in Section 4.1. As shown in Fig. 5, SGDF finds flatter minima. Notably, the visualization reveals that Adam is more prone to converge to sharper minima.

We’ve conducted comparative experiments on the CIFAR-100 dataset, evaluating both the vanilla Adam algorithm and Wiener Adam, which substitutes the first-moment gradient estimates in the Adam optimizer with Wiener filter estimates. The results are presented in Tab. 4, and the detailed test curves are depicted in Fig. 10. This suggests that our first-moment filter estimation method has the potential to be applied to other optimization methods.

The performance improvement observed when using filter estimate can be attributed to several factors related to the nature of the Wiener filter and the architecture of the neural networks used.

Noise Reduction. The Wiener Filter is fundamentally designed to minimize the mean square error in the presence of noise. In the context of gradient estimation, this means it can effectively reduce the noise in the gradient updates. The vanilla Adam optimizer calculates the first moment (the mean) of the gradients, which may still include some noise components. By employing the Wiener filter, the algorithm may benefit from cleaner and more reliable gradient signals, thus improving the optimizer’s ability to converge to better minima.

Network Architecture. The impact of the Wiener Filter on gradient estimation varies with the network architecture’s unique handling of features and gradients. For VGG, which lacks residual connections, the Wiener filter significantly enhances performance by providing more accurate gradient estimations, reducing noise-induced errors, and leading to better overall accuracy. In contrast, the benefits are less pronounced in ResNet and DenseNet, which incorporate residual and dense connections, respectively. These architectures naturally promote more stable gradient flows through direct paths or dense layer interconnectivity, diminishing the additional advantage offered by the Wiener filter. Intuitively, we can find the loss landscape with residual connected networks in [34]. Thus, while the Wiener filter improves performance in simpler architectures like VGG, its impact is mitigated in more complex architectures where inherent features already assist in gradient stability.