PoF: Post-Training of Feature Extractor for Improving Generalization

The relationship between the local loss landscape and the generalization ability of a minimum has been discussed extensively in theoretical and empirical literature (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Dziugaite & Roy, 2017; Jiang* et al., 2020; Dinh et al., 2017). The previous section described an intuitive picture explaining why a flatter minimum likely generalizes. Another view is given by the bits-back argument (Hinton & Camp, 1993; Honkela & Valpola, 2004). It states models that are stable against weight perturbations can be described with fewer bits. According to the minimum description length (MDL) (Rissanen, 1978) or similar criteria, models that can be represented with a smaller number of bits are expected to have better generalization abilities. Both the geometrical and description-length points of views suggest flatter minima are preferable for generalization.

With deep neural networks, there are a number of studies that aim to measure the flatness of the loss landscape. For example, Keskar et al. (2017) measured the flatness by the worst loss around the minima. Li et al. (2018) visualizes loss landscapes with findings that sharp minimizers tend to have larger generalization error. Loss landscape visualization was carried out with minimization trajectories (Goodfellow et al., 2015) and with parameter interpolations (Im et al., 2016).

Hessian matrices are sometimes utilized to quantify flatness with eigenvalues (Sagun et al., 2016; Wu et al., 2017; Zhang et al., 2018) and spectral norm (Yao et al., 2018). However, measuring the flatness is still an open problem in general due to high dimensionality and architectural complexity of deep networks.

For a given network architecture, different optimization methods tend to reach solutions that are different in terms of flatness. Stochastic gradient descent (SGD) is known to be biased towards flat minima (Jastrzebski et al., 2018; Maddox et al., 2020). The stochastic noise in the parameter updates leads to parameter convergence in distribution with similar loss from a Langevin-dynamics based modeling (Chaudhari & Soatto, 2018). Stochastic Weight Averaging (SWA) was proposed to find a flatter point near a minimizer reached by SGD, but SWA itself does not work like an optimizer that can find a wider basin. Recently, optimization algorithms that seek flatter minima have attracted attention. The sharpness-aware minimization (SAM) (Foret et al., 2021) based on the PAC-Bayes generalization bound (Langford & Shawe-Taylor, 2003) has been shown to be effective for various image classification tasks. There are some extensions of SAM including adaptive optimization methods (Chen et al., 2020b; Kwon et al., 2021). These sharpness-aware methods apply parameter perturbations within some radius at each iteration, aiming to make the perturbed region flat. PoF also adopts a parameter perturbation, whose range is determined in data-driven manner to reduce a specially designed effective loss as described in the next section.

Curvature estimation of a loss surface by approximated Hessian matrices or Fisher information matrices (Roux et al., 2008; Botev et al., 2017; Grosse & Martens, 2016; Martens & Grosse, 2015; Pauloski et al., 2021) is related to the flatness-based optimization. A second-order term appears in the formulation of SAM, but it is simply dropped to reduce the computational cost from a practical perspective. PoF is designed to implicitly reduce some Hessian components without a need to directly compute Hessian or Fisher information matrices.

Some previous work discussed co-adaptation prevention between layers, in particular, a feature extractor and a classifier. Prevention of co-adaptation among neurons likely brings a positive effect on generalization (Hinton et al., 2012). FOCA (Sato et al., 2019) avoids between-layer co-adaptation by using many random weak classifiers during optimization. Moayed and Mansoori (2020) proposed a method to adaptively assign dropout rate according to the co-adaption pressure. Wei et al. (2020) makes a series of weak classifiers to decouple co-adaptation. Our work can be viewed as a method to weaken co-adaptation between the feature extractor and the classifier.

PoF is a perturbation-based method to seek a flatter minimum, i.e., loss gradients are computed with respect to shifted parameters at each iteration. As illustrated in Fig. 1 (c), we assume that the deterioration in loss caused by the positional discrepancy between the training and test minimizers is most severe in the classifier parameter space. PoF aims to update the feature extractor so that the returned features yield a flat loss landscape to an appropriate extent in the classifier parameter space. Making the classifier parameters or the corresponding decision boundary “loose” is analogous to the concept of the maximum margin method (see Fig. 3).

PoF determines the direction of perturbation in a data-driven manner. Perturbing in a spherically uniform way would be very inefficient, because Hessian spectra are in most cases dominated by a very small number of eigenstates compared to the parameter dimension, and the rest of the Hessian components are negligible. Thus, perturbing along the direction of the eigenvector corresponding to the maximum eigenvalue of the Hessian would be much more efficient. However, it is practically infeasible to compute a Hessian and its eigenvectors even at some interval of iterations due to high computational cost. PoF avoids to directly compute Hessians and adopts a much more computationally efficient approach. Figure 2 (a) shows counts of randomly sampled mini-batches whose gradients maximally correlate with $i$-th eigenvector of the Hessian. This toy experiment shows that mini-batch gradients highly likely correlate the principal eigenvector. Though we are unsure to what extent this tendency holds, we simply assume this tendency generally holds. Based on the observation, we adopt an approach where a perturbation is taken along a mini-batch gradient direction in the classifier parameter space.

PoF determines the range of perturbation in a data-driven manner as well. It adopts a linear search method along the direction of (negative) mini-batch gradients evaluated at $\theta_{0}$ to find the nearest minimum in the 1D subspace, i.e.,

where $\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$ is a special short-hand notation of

Here, $\xi_{\mathcal{B}}^{\star}$ is the Euclidean distance to the minimum of the mini-batch loss in the 1D subspace. The classifier with the perturbed parameters $\theta_{0}-\xi\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$ behaves as a somewhat weak classifier for the entire training dataset when $\xi\simeq\xi_{\mathcal{B}}^{\star}$. The idea here is to make this kind of a somewhat weak classifier stronger for an arbitrary mini-batch $\mathcal{B}$ by optimizing the feature extractor. In this way, the perturbed region is expected to become a flat basin. Next, we investigate an appropriate perturbation range of $\xi$. Figure 2 (b) shows probability distributions of $\xi_{\mathcal{B}}^{\star}$ for different mini-batches after an orthodox SGD training using a toy dataset. It indicates that the peak point of the test distribution is roughly twice as that of the training distribution. Then, a naive strategy would be to enlarge the perturbation range by setting $\xi=\gamma\xi_{\mathcal{B}}^{\star}$ with $\gamma\simeq 2$ to compensate the training-test distribution gap.

One could use an iterative gradient descent method instead of linear search to find a nearest minimum of a mini-batch loss. But, the computational cost would become much higher in such a case. Given that the classifier parameters $\theta_{0}$ already reached a local minimum, it would be a decent assumption that a mini-batch loss landscape is locally convex. Computationally efficient linear search usually suffices in practice.

for some $m\gg 1$. An update of feature-extractor parameters, $\Delta\phi$, that approximately minimizes Eq. (4), is given by

with arbitrarily chosen mini-batches $\mathcal{B}$ and $\tilde{\mathcal{B}}$. Here, we artificially set $\partial\hat{\mathcal{L}^{\prime}_{\mathcal{B}}}/\partial\phi\triangleq 0$ so that the update reduces the empirical loss with respect to the perturbed classifier $\theta_{0}-\gamma\xi_{\mathcal{B}}^{\star}\hat{\mathcal{L}^{\prime}_{\mathcal{B}}}$. Modification of Eq. (5) by adding a momentum term, etc., is possible. This update makes the perturbed classifier stronger, contributing to make the local loss landscape flatter, as illustrated in Fig. 3. When $\gamma=0$, the update rule is equivalent to SGD. As an extension of the method, one could randomize $\gamma$ in a predefined range such as $\gamma\in[0,2]$, or employ some scheduling function to $\gamma$ such as linear growth from 0 to 2. A pseudocode of our algorithm is provided in Algorithm 1. One trick that we employ here is that a mini-batch is re-sampled right before computing the feature-extractor update in Eq. (5). This avoids too much over-fitting to a particular mini-batch at each iteration.

PoF updates only feature extractors, and the classifier parameters $\theta_{0}$ are kept unchanged. From our experience, this strategy works well on real datasets. But, it could be possible that the position of the minimum drifts from $\theta_{0}$ during PoF. To avoid this type of drifting, one may add an SGD update step for the classifier parameters $\theta_{0}$ after the feature-extractor parameter update. From our experience, this strategy does not affect the final performance much on real datasets, so we did not apply this strategy for the experiments reported in the experimental section (Sec. 4).

A similar concept to PoF was proposed by Sato et al. , where they proposed a supervised representation learning method, FOCA, for optimizing a feature extractor with respect to weak-classifier ensemble (Sato et al., 2019, 2021). To obtain a weak classifier, FOCA applies gradient descent iterations using a given mini-batch from random initialization, while PoF effectively finds a nearest local minimum of the mini-batch loss starting from the training-loss minimum. Not only does PoF enable post-training, but PoF can effectively find a flatter minimum, as is explained in the next subsection.

As for word choice, we adopt “post-training” rather than “fine-tuning” throughout this paper. The latter word is commonly used in transfer learning settings, whereas PoF intends to improve in-distribution performance.

Next, we present a mathematical analysis about the relationship of the proposed algorithm and the loss landscape. Let us approximate the mini-batch loss landscape along a direction of mini-batch loss gradients $\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$ as

where $\mathcal{H}_{\mathcal{B}}(\phi,\theta_{0})$ is the second-order derivatives of the mini-batch loss defined in the classifier parameter space,

Here, we model the landscape as a quadratic function of $\xi\in\mathbb{R}$. As is obvious from Eq. (6), $\xi=\xi_{\mathcal{B}}^{\star}$ minimizes the mini-batch loss in the linear subspace whose basis is given by $\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$. It is also assumed that the minimum is zero, which is not a bad approximation, given that $\theta_{0}$ already minimizes the entire training loss and the mini-batch loss is further optimized by linear search. Next, we expand the loss function of the entire training set $\mathcal{D}$ as

where $\mathcal{H}_{\mathcal{D}}(\phi,\theta_{0})$ is the Hessian computed from the entire training set in the classifier parameter space. Since the training loss is assumed to be locally minimized at $\theta=\theta_{0}$, there is no first-order term of $\xi$ in Eq. (8). Setting $\xi=\gamma\xi_{\mathcal{B}}^{\star}$, Eq. (8) becomes

Let us call this approximated loss as the effective loss. The algorithm effectively reduces these quantities for arbitrary choices of $\mathcal{B}$ by searching appropriate feature-extractor parameters $\phi$, on which $\mathcal{L}_{\mathcal{D}}$ and $\mathcal{H}_{\mathcal{D}}$ depend. The effective loss consists of the zero-th and the second order derivatives of the training loss. If $\xi_{\mathcal{B}}^{\star}$ goes small, the zero-th order term becomes a dominant term; oppositely, if $\xi_{\mathcal{B}}^{\star}$ goes large, the second order term becomes a dominant term. The quantity $\frac{\gamma^{2}\xi_{\mathcal{B}}^{\star 2}}{2}\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}^{\top}\mathcal{H}_{\mathcal{D}}\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$ is a rough estimate of $\Delta\mathcal{L}$ as depicted in Fig. 1 for a subset of the test set. Reduction of this curvature-based loss term would decrease the training-loss curvature as well as the test-loss curvature, under the assumption that two curvatures are interlocked. By combining Eq. (6) and Eq. (8), Eq. (9) can be equivalently expressed as

The second order term has the batch-gradient projected Hessian component with a coefficient given by batch statistics $\mathcal{L}_{\mathcal{B}}$, $\hat{\mathcal{L}_{\mathcal{B}}^{\prime}}$, and $\mathcal{H}_{\mathcal{B}}$ evaluated at $(\phi,\theta_{0})$ besides the overall scalar $\gamma$. In this way, the range of perturbation is determined in a data-driven manner. As described earlier, $\gamma$ has an intuitive geometrical meaning, such that $\gamma=1$ corresponds to the minimum of the mini-batch loss and $\gamma=2$ corresponds to the opposite side of the quadratic mini-batch loss landscape.

It is worth mentioning that the proposed method often gives some order-of-magnitude large perturbations in the classifier parameter space compared to a typical SGD step. Nevertheless, the algorithm is surprisingly stable. Suppose now that $\gamma=2$. If the mini-batch loss is quadratic as in Eq. (6), a parameter perturbation with $\gamma=2$ makes the mini-batch loss invariant. It means this mini-batch experiences no harm by this (possibly very large) perturbation. This prevents the classifier from becoming too adversarial to the rest of the samples.

In reality, some mini-batches show asymmetric loss landscapes, such that one side for $\xi\in(0,\xi_{\mathcal{B}}^{\star})$ is close to quadratic but the other side is almost constant. This type of loss landscapes can be detected in the linear search. When such a landscape is found during the training, one may discard the mini-batch and repeat the mini-batch sampling step followed by the linear search algorithm.

The result shows that PoF does not improve generalization for Fashion-MNIST. This indicates that flattening the loss landscape in the classifier parameter space could not further improve the performance, probably because lower layers suffers severer loss deterioration $\Delta\mathcal{L}$. PoF clearly has a limitation in such a case.

As is evident from Fig. 5, PoF clearly reduces the Hessian components. The peak values are roughly reduced by a factor of two for both training and test sets. As SAM is a strong baseline having a flatness-enhancing functionality, it is surprising that there is still a room for the network to improve flatness along certain directions as well as to improve generalization just by additional 10 epoch post-training.