Explaining f-Divergence-Based Regularization via Local Curvature and Sharpness-Aware Minimization

Technical Summary: Explaining f-Divergence-Based Regularization via Local Curvature and Sharpness-Aware Minimization

Problem Statement
In the context of lifelong learning agents and general deep learning, achieving robust generalization requires controlling overfitting and local sensitivity. Two prominent approaches for this are divergence-based regularization (which penalizes discrepancies between predictive distributions on original and perturbed inputs) and Sharpness-Aware Minimization (SAM), which seeks parameters with flat loss landscapes. While both methods are motivated by robustness to perturbations, their theoretical relationship remains largely unexplored. Specifically, it is unclear whether the local geometry induced by a divergence regularizer can be formally related to the Hessian structure that SAM implicitly controls.

Methodology
The paper establishes a unified theoretical framework by analyzing the local second-order geometry of $f$-divergence regularization. The authors proceed through the following steps:

1. Local Quadratic Expansion: They derive a second-order Taylor expansion for $f$-divergences when the model's predictive distribution $P$ and a perturbed distribution $Q$ are close. They show that in this local regime, any $f$-divergence reduces to a quadratic penalty scaled by $\phi''(1)/2$.
2. Geometric Interpretation:
   • Parameter Space: When perturbations occur in the parameter space ($\theta \to \theta + \delta$), the regularizer induces a penalty proportional to $\delta^\top F_x(\theta) \delta$, where $F_x(\theta)$ is the Fisher information matrix.
   • Input Space: When perturbations occur in the input space ($x \to T(x)$), the regularizer induces a "pullback" quadratic form involving the gradient of the density with respect to the input.
3. Connection to SAM: The authors compare this divergence-induced penalty to SAM. SAM is shown to locally approximate a spectral curvature penalty proportional to the dominant Hessian eigenvalue ($\lambda_{\max}(H)$). The paper demonstrates that for negative log-likelihood objectives with exponential-family output distributions (e.g., cross-entropy with softmax), the Fisher matrix coincides with the Generalized Gauss-Newton (GGN) matrix, which captures the positive-semidefinite component of the Hessian. Thus, divergence-based regularization and SAM become locally comparable through a shared curvature geometry.
4. Controlled Empirical Testbed: To validate these theoretical claims, the authors utilize the asymmetric $\alpha$-skew Jensen-Shannon Divergence (JSD) family. The local curvature coefficient for this family scales as $\alpha(1-\alpha)$, reaching a maximum at the symmetric point $\alpha = 1/2$. This allows for a controlled experiment where the strength of curvature penalization is varied without altering the fundamental structure of the regularizer.

Key Contributions

• Unified Local Expansion: The paper derives a local second-order expansion for $f$-divergence regularization, demonstrating that it induces curvature-sensitive penalties determined by Fisher geometry in parameter space and a pullback metric in input space.
• Formal Link to SAM: It establishes that under standard probabilistic loss assumptions (specifically NLL with exponential families), the geometry induced by divergence regularization is locally comparable to the second-order interpretations of SAM via the Fisher/GGN/Hessian relationship.
• General Perturbation Framework: The analysis extends to input-space perturbations, showing that divergence-based regularization offers a more general perturbation framework than standard SAM (which is typically defined via parameter perturbations) while preserving the same local sensitivity interpretation.
• Curvature Coefficient Derivation: For the asymmetric $\alpha$-skew JSD family, the authors explicitly derive that the local curvature coefficient is proportional to $\alpha(1-\alpha)$, identifying the symmetric case ($\alpha=1/2$) as the regime of maximal curvature penalization.

Results
The empirical validation was conducted on four benchmark datasets (CIFAR-10, Fashion-MNIST, EMNIST, and Oxford-IIIT Pet) using an EfficientNet-B2 model with input-space perturbations (random masking).

• Performance Trends: Across all datasets, model performance (measured by accuracy and Negative Log-Likelihood) was consistently best near the symmetric regime $\alpha = 1/2$. This aligns with the theoretical prediction that the strongest local curvature penalization (maximized at $\alpha=1/2$) yields superior generalization.
• Loss Landscape Visualization: Visualizations of the loss landscape on CIFAR-10 using ResNet-18 revealed that models trained with divergence regularization converge to flatter local minima compared to the baseline. Furthermore, the model trained with the symmetric setting ($\alpha=0.5$) exhibited the flattest basin, corresponding to the highest curvature penalization, while the asymmetric setting ($\alpha=0.9$) showed intermediate flatness.
• Statistical Metrics: Quantitative analysis of the loss landscape statistics (average loss, trace of the Hessian squared, and maximum eigenvalue) confirmed that the symmetric $\alpha=0.5$ setting resulted in the lowest average loss and reduced curvature metrics compared to both the baseline and highly asymmetric settings.

Significance and Claims
The paper claims to clarify the relationship between divergence-based regularization and sharpness-aware minimization from a local geometric perspective. It posits that while the two methods are not identical objectives, they become locally aligned through a common curvature geometry when specific probabilistic assumptions hold.

The authors modestly frame their contribution as primarily analytical. They clarify when divergence-based regularization, Fisher geometry, and sharpness-aware objectives become locally comparable. The empirical study focuses on validating the broader local-curvature perspective within the specific context of the $\alpha$-skew JSD family and input perturbations, rather than claiming a direct experimental equivalence to SAM in all settings. The work suggests that the strength of the induced local curvature is a meaningful factor in the effectiveness of divergence-based regularization, with the symmetric regime offering optimal performance in the tested scenarios.