Generalization Aware Minimization

Thank you for your insightful review and for recognizing the potential applications of our theoretical framework.

Concerns on Quadratic Loss Function Assumption

We acknowledge that the assumption of quadratic loss functions limits the direct applicability of our theoretical results. However, we emphasize that any smooth loss function can be locally approximated by a quadratic function around a point of interest. This local approximation is a standard technique in optimization and allows us to extend our insights to the non-convex loss landscapes typical in deep learning.

In practice, GAM leverages this local quadratic approximation to adjust the optimization trajectory, improving generalization even when the global loss landscape is complex.

Experimental Evaluation

We agree that a broader experimental evaluation would strengthen our paper. In the limited time of the rebuttal period, we conducted additional experiments on the ImageNet dataset (trained for 50 epochs) using a ResNet-20-like architecture. GAM achieved a top-1 test accuracy of 5.12%, compared to 4.89% for SAM at $\gamma=0.01$ and 4.04% for standard SGD (SAM at $\gamma=0.1$: 3.77%, SAM at $\gamma=0.001$: 4.55%). These results suggest that GAM's advantages extend to larger-scale datasets and modern architectures.

Clarification on Theorem 1

We apologize for any confusion regarding Theorem 1. The expectation on the left hand side of equation (7) is taken over the randomness in the loss functions, not over $\theta$. Thus, the expected loss remains a function of $\theta$ in general. The first term on the right hand side captures this dependence on $\theta$ while the second term is independent of $\theta$.

SAM as a Special Case of GAM

Yes, SAM can be seen as a special case of GAM with a single perturbation step ($T=1$) and a fixed perturbation coefficient ($\gamma_1$). GAM generalizes this approach by allowing multiple perturbation steps and learning the perturbation coefficients during training. This flexibility enables GAM to capture higher-order information about the loss landscape.

We have added a clarification in the paper to make this relationship explicit.

Clarification of Equations (16) and (17)

You are correct that Equations (16) and (17) both represent the right-hand side of Equation (15), broken down into different components.

We have revised the text to ensure that the progression from Equation (15) to Equations (16) and (17) is clear and logical.

Results on Wide Network (WN)

We omitted results on WN for MNIST because smaller models already achieve near-perfect accuracy on this dataset, leaving little room for improvement. Training larger models like WN on MNIST provides limited additional insight, and we did not consider it an experimental priority given the computational cost of training WN. However, if you consider it important, we are willing to include these results in the revised paper.

Effect of Hyperparameters

Thank you for suggesting an analysis of how these parameters affect GAM's performance. In the limited time of the rebuttal period, we have swept over the following hyperparameters to investigate their impact (fixing a CNN model on CIFAR-10):

• Batch Size
• Perturbation Steps ($T$)
• Perturbation Step Size ($\epsilon$)

We have included our results in Appendix E Figure 5; in summary, we find that GAM can become unstable under large $T$, performs best at small batch sizes, and is relatively insensitive to $\epsilon$. We hypothesize that reducing the learning rate of $\gamma$ can be one way to mitigate the instability for large $T$.

Thank you for your review and for acknowledging the originality and meaningful perspective of our work.

Experiments on Larger Models

We agree that evaluating GAM on commonly used models is important for demonstrating its practical impact. In the limited time of the rebuttal period, we conducted additional experiments on the ImageNet dataset (trained for 50 epochs) using a ResNet-20-like architecture. GAM achieved a top-1 test accuracy of 5.12%, compared to 4.89% for SAM at $\gamma=0.01$ and 4.04% for standard SGD (SAM at $\gamma=0.1$: 3.77%, SAM at $\gamma=0.001$: 4.55%). These results suggest that GAM's advantages extend to larger-scale datasets and modern architectures.

Marginal Improvement

We understand your concern regarding the size of the performance gains. However, these improvements are consistent across all evaluated datasets and architectures. It's important to note that in the context of deep learning, especially on benchmark datasets, even small improvements can be considered significant due to the already high performance levels achieved by state-of-the-art models. Our results are in line with the improvements reported in the original SAM paper over standard SGD (which are often less than 1%). The consistent outperformance of SAM by GAM suggests that our method offers meaningful benefits in terms of generalization.

Computation Cost

We appreciate your important point regarding the computational overhead introduced by GAM.

To mitigate this cost, we propose updating the GAM perturbation coefficients $\gamma_t$ less frequently during training. Instead of optimizing $\gamma_t$ at every iteration, we can update them periodically (e.g., every few iterations). This approach significantly reduces the computational overhead while maintaining the performance benefits of GAM.

In the limited time of the rebuttal period, we have added additional experiments in Section 4 evaluating the computational cost of GAM relative to SGD and SAM on CIFAR-10 with a CNN model. We find that standard GAM is roughly $4 \times$ as costly in training time as SGD (relative to $1.3 \times$ for SAM), but this cost can be reduced to $3 \times$ when updating $\gamma_t$ periodically. Periodic updates hurt GAM's accuracy, but still enable it to outperform SAM and SGD's accuracies.

Writing Clarifications

Thank you for highlighting areas where the writing could be improved. We have made all the clarifications you have suggested in our revision and are happy to make any further revision as well. We apologize for any confusion caused and will ensure that the revised paper is clear and precise.

Effect of Perturbation Steps ($T$)

Thank you for suggesting an analysis of how $T$ affects GAM's performance. In the limited time of the rebuttal period, we have swept over the following hyperparameters to investigate their impact (fixing a CNN model on CIFAR-10):

• Batch Size
• Perturbation Steps ($T$)
• Perturbation Step Size ($\epsilon$)

We have included our results in Appendix E Figure 5; in summary, we find that GAM can become unstable under large $T$, performs best at small batch sizes, and is relatively insensitive to $\epsilon$. We hypothesize that reducing the learning rate of $\gamma$ can be one way to mitigate the instability for large $T$.

Definition of Discrepancy Function

The discrepancy function $\Delta$ used in our method is the negative dot product between the gradient at the perturbed parameters and the gradient on an auxiliary minibatch. This function encourages alignment between the perturbed gradient and the gradient that better approximates the test loss.

We have updated the text to clarify this definition.

Thank you for your follow-up comments and for giving us the opportunity to clarify and address your concerns.

Inconsistencies in Reported Computational Costs

We apologize for any confusion regarding the reported computational costs of GAM. Initially, we mentioned that GAM was approximately 15 times slower than SGD based on a preliminary estimate focused solely on the time required for computing each individual parameter update, excluding factors like data loading, preprocessing, and other overheads common to both GAM and SGD.

In our latest experiments, we conducted a more comprehensive and precise measurement of the total training time including all steps of the training process. Under these conditions, we found that standard GAM is roughly $4 \times$ as costly in training time as SGD. These measurements are averaged $5$ over separate trials, giving us confidence in the validity of this measurement.

Alternative Options for the Discrepancy Function and Ablation Study

Thank you for highlighting the importance of the discrepancy function in GAM. In our current implementation, we use the negative dot product as the discrepancy function. This choice encourages alignment between the gradient at the perturbed parameters and the gradient on the auxiliary minibatch, promoting updates that generalize better.

We agree that exploring alternative discrepancy functions could provide valuable insights into the effectiveness of GAM.

Potential alternatives include:

• Squared error between the gradients
• Cosine similarity to measure the angle between the gradients

While we have not yet conducted an ablation study on different discrepancy functions, we are happy to do so. Due to the limited time remaining in the rebuttal period, we may not be able to present these experiments in the rebuttal. However, we recognize their importance and will prioritize them in the revised version of the paper.

Experimental Results on Large Models and Datasets

We understand your desire to see GAM evaluated on more advanced models and larger datasets, including NLP tasks. During the limited rebuttal period, we were able to show preliminary experimental results on ImageNet (in addition to several other experiments); these experiments showed that GAM outperforms SAM and SGD in this setting as well.

We acknowledge that the improvement may seem modest in percentage terms. Still, we highlight that in the context of SAM-based optimization methods, improvements on the order of $1$% are actually quite significant: SAM's improvements over SGD are also relatively small in percentage terms.

As for NLP datasets, we appreciate your suggestion and recognize the importance of evaluating GAM in different domains. Had this recommendation been made in your initial review, we may have prioritized allocating time to conduct NLP experiments (in addition to the several other additional experiments we did conduct). We are happy to conduct an experiment on an NLP task, but will unfortunately likely be unable to share our results by the close of the rebuttal period.

Thank you for your constructive feedback and for recognizing the plausibility and interest of our theoretical analysis.

Explanation of the Relationship Between Test Loss and Training Loss

We appreciate your request for a more intuitive explanation. The key insight is that the training loss landscape accurately captures the directions of curvature (the eigenvectors of the Hessian matrix) but may distort the curvature magnitudes (the eigenvalues). This distortion occurs due to noise from the data sampling process.

By understanding that the expected test loss landscape is a rescaled version of the training loss landscape, we can design optimization strategies that adjust for this discrepancy. Specifically, we can transform the training loss gradient to approximate the test loss gradient, guiding the optimization process toward solutions that generalize better.

Figure 1 in the paper provides a visual illustration of this concept.

We are happy to provide any further clarification on this, especially since this is a key aspect of our paper.

Effect of Batch Size, Perturbation Steps, and $\epsilon$

Thank you for suggesting an analysis of how these parameters affect GAM's performance. In the limited time of the rebuttal period, we have swept over the following hyperparameters to investigate their impact (fixing a CNN model on CIFAR-10):

• Batch Size
• Perturbation Steps ($T$)
• Perturbation Step Size ($\epsilon$)

We have included our results in Appendix E Figure 5; in summary, we find that GAM can become unstable under large $T$, performs best at small batch sizes, and is relatively insensitive to $\epsilon$. We hypothesize that reducing the learning rate of $\gamma$ can be one way to mitigate the instability for large $T$.

Comparison with SAM Using Larger $\gamma$

We appreciate your point about comparing GAM to SAM with larger perturbation sizes. In the limited time of the rebuttal window, we have conducted additional experiments on CIFAR-10 using the CNN architecture and found the following results for SAM over $5$ seeds:

At $\gamma=0.3$, accuracy: 0.64534 $\pm$ 0.02078

At $\gamma=1.0$, accuracy: 0.44048 $\pm$ 0.01161

Both these results are worse than our existing SAM results, giving us confidence that we have chosen an appropriate range of $\gamma$.

We note that simply increasing $\gamma$ in SAM does not capture the nonlinear transformation of the curvature magnitudes that GAM learns through its adaptive perturbation strategy. As shown in Figure 3, GAM learns a more highly curved mapping between the training and test loss landscapes, which cannot be replicated by merely increasing $\gamma$ in SAM.

Experiments on Larger Datasets

We agree that demonstrating GAM's effectiveness on larger datasets and architectures strengthens our work. In the limited time of the rebuttal period, we conducted additional experiments on the ImageNet dataset (trained for 50 epochs) using a ResNet-20-like architecture. GAM achieved a top-1 test accuracy of 5.12%, compared to 4.89% for SAM at $\gamma=0.01$ and 4.04% for standard SGD (SAM at $\gamma=0.1$: 3.77%, SAM at $\gamma=0.001$: 4.55%). These results suggest that GAM's advantages extend to larger-scale datasets and modern architectures.

Clarification on Equation (15)

In Equation (15), we reorganize the terms of the expected test loss into components based on their dependence on $\theta$. Specifically:

• Quadratic Term: Represents the curvature of the loss landscape with respect to $\theta$.

• Linear Term: Represents the gradient component that depends linearly on $\theta$.

• Constant Term: Collects all terms independent of $\theta$, including those dependent only on $\theta^*$.

Explanation of Equation (7)

Equation (7) conveys that the expected test loss is also a quadratic like the training loss, and has the same eigenvectors as the training loss but with eigenvalues transformed by an arbitrary function $D$. This transformation reflects the rescaling of the curvature magnitudes between the training and test losses. We are happy to further clarify this since this equation is central to our submission.

Thank you for your positive feedback and for acknowledging the clarity and potential impact of our work.

Size of Performance Gains

We understand your observation regarding the relatively small performance gains of GAM over SAM. However, these improvements are consistent across all evaluated datasets and architectures. It's important to note that in the context of deep learning, especially on benchmark datasets, even small improvements can be considered significant due to the already high performance levels achieved by state-of-the-art models.

Our results are in line with the improvements reported in the original SAM paper over standard SGD (which are often less than 1%) The consistent outperformance of SAM by GAM suggests that our method offers meaningful benefits in terms of generalization.

Additional Modalities

We appreciate your suggestion to evaluate GAM on other modalities. In the limited time of the rebuttal period, we conducted additional experiments on the ImageNet dataset (trained for 50 epochs) using a ResNet-20-like architecture. GAM achieved a top-1 test accuracy of 5.12%, compared to 4.89% for SAM at $\gamma=0.01$ and 4.04% for standard SGD (SAM at $\gamma=0.1$: 3.77%, SAM at $\gamma=0.001$: 4.55%).

These results suggest that GAM's advantages extend to larger-scale datasets and modern architectures. We agree that exploring GAM's effectiveness on language-related tasks and out-of-distribution performance is valuable, and we plan to investigate these areas in future work.

Connection to CR-SAM

Thank you for bringing up CR-SAM. While both CR-SAM and GAM aim to improve generalization by modifying the SAM framework, there are key differences:

Motivation and Approach: CR-SAM introduces a curvature regularization term to penalize high curvature in the loss landscape, thereby encouraging flatter minima. In contrast, GAM is derived from a perspective that directly targets the expected test loss by transforming the training loss gradient through multiple perturbations.

Flexibility in Perturbations: GAM allows for an arbitrary number of perturbation steps and learns the perturbation sizes online during training. This flexibility enables GAM to capture higher-order information about the loss landscape. CR-SAM, on the other hand, typically employs a fixed perturbation strategy.

Thank you for your thoughtful review and for highlighting both the strengths and areas for improvement in our paper.

Concerns on Computational Costs

We appreciate your important point regarding the computational overhead introduced by GAM. You are correct that GAM requires multiple perturbation steps, which increases computational cost compared to SAM and standard SGD.

To mitigate this cost, we propose updating the GAM perturbation coefficients $\gamma_t$ less frequently during training. Instead of optimizing $\gamma_t$ at every iteration, we can update them periodically (e.g., every few iterations). This approach significantly reduces the computational overhead while maintaining the performance benefits of GAM.

In the limited time of the rebuttal period, we have added additional experiments in Section 4 evaluating the computational cost of GAM relative to SGD and SAM on CIFAR-10 with a CNN model. We find that standard GAM is roughly $4 \times$ as costly in training time as SGD (relative to $1.3 \times$ for SAM), but this cost can be reduced to $3 \times$ when updating $\gamma_t$ periodically. Periodic updates hurt GAM's accuracy, but still enable it to outperform SAM and SGD's accuracies.

Intuition for Assumptions and Practical Significance of Theorem 1

Thank you for requesting a clearer explanation of Theorem 1's assumptions and their practical implications. In Appendix C of the paper, we provide justifications for each assumption. To summarize:

Quadratic Loss Functions: We assume that both the true (test) and observed (training) loss functions are quadratic. This serves as a local approximation to the loss landscape around a minimum, which is a common practice in optimization theory.

Rotational and Location Invariance: We assume that the generative process of the loss functions is invariant under rotations and translations. This means that the statistical properties of the loss landscape are the same in all directions and locations, simplifying the analysis.

Practically, Theorem 1 implies that the training loss landscape accurately captures the curvature directions (eigenvectors) of the test loss landscape but distorts the curvature magnitudes (eigenvalues). In other words, while the orientation information is preserved, the scale information is distorted. This understanding allows us to design optimization algorithms that correct for this distortion, leading to improved generalization.

Extension to Non-Quadratic Losses

You raise an important point about the applicability of our theoretical results to non-quadratic loss functions. We highlight that any smooth loss function can be approximated by a quadratic function in a sufficiently small neighborhood around a point. Therefore, we expect Theorem 1 to hold locally for any smooth loss function.

In practice, this means that our insights are relevant for the complex, non-convex loss landscapes encountered in deep learning. By considering local quadratic approximations, GAM can be effectively applied to train neural networks.

Intuition for the Test Loss Being a Rescaled Version of the Training Loss

Intuitively, the training loss landscape provides accurate information about the directions in which the loss changes most rapidly (the curvature directions), but it may misrepresent how sharply the loss changes along those directions (the curvature magnitudes). This distortion occurs due to noise from the data sampling process.

By recognizing that the test loss is a rescaled version of the training loss, we can adjust our optimization process to compensate for the distorted scale information. This involves modifying the training loss landscape to better align with the test loss landscape, ultimately improving generalization performance.

We are happy to provide any further clarification on this, especially since this is a key aspect of our paper.