Optimization on Manifolds with Riemannian Jacobian Regularization

We thank the reviewer for the constructive feedback and would like to address the concerns of the reviewer as follows:

• Regarding the theoretical contribution: We appreciate the opportunity to elaborate on the theoretical contribution of our work. In the main theorem, we establish a generalization inequality for Riemannian manifolds, which, to the best of our knowledge, has not been previously reported, as prior results are limited to the Euclidean setting. A key distinction of our approach is that we directly bound the population loss using the empirical loss, rather than the worst-case empirical loss. To better understand the role of each component, as discussed in lines 230–237, if we focus solely on regularizing the gradient norm, we obtain a Riemannian sharpness-aware algorithm similar to Riemannian-SAM. However, our theorem motivates us to also regularize the Hessian term, which further reduces the gap between the general loss and the empirical loss, offering an improvement over Riemannian-SAM and other baselines. While Jacobian regularization has been previously introduced in the Euclidean context, it has not been theoretically justified. The Hessian term in our theorem justifies the use of Jacobian regularization as a means to further reduce sharpness, a result that is empirically supported by our ablation studies. Based on the reviewer's suggestions, we have incorporated this explanation in the rebuttal revision.

• Regarding the notation of the Euclidean norm: We use the notation $\|\cdot\|$ instead of the standardized norm-2 notations $\|\cdot\|_2$ for shortness, in which we omitted the subscript for shortness. We have included a clarification in the rebuttal revision based on the reviewer's suggestion.

• Regarding the numbering of the Equations: In our theoretical development, only a few equations are referred to later in the derivation. Then, for readability purposes, we decided to label only the equations that are being referred to.

• Regarding the first inequality on page 19, line 974: In Equation (5), we have derived that $|L_D(\theta)-L_D(\theta_i)|\leq K\rho$ Therefore, we have $L_D(\theta)-L_D(\theta_i)\leq |L_D(\theta)-L_D(\theta_i)|\leq K\rho$ which gives the inequality at Line 974.

• Comparison with SOTA Methods: Our work focuses on establishing a theoretical framework for a novel Riemannian optimizer. To illustrate the benefits of sharpness-aware Jacobian regularization on manifolds, we compare our algorithm with foundational optimizers, including SGD, SAM, and their Riemannian counterparts. We deliberately chose not to include comparisons with SOTA methods that diverge significantly from these fundamental optimizers in methodology, as this could detract from our primary theoretical focus. However, we greatly appreciate your suggestion. Due to time constraints, we will prioritize exploring them in future studies.

• Regarding the Appendix A.2: The lemma that the appendix refers to is the Lemma 1 in the main text. According to the suggestion of the reviewer, we have fixed this issue in the rebuttal revision.

We thank the reviewer for their constructive feedback. We would like to address the concerns as follows:

• Theoretical Contribution: Firstly, we respectfully disagree with the comment of the reviewer that the main result relies on [1] and [2]. In particular, [1] focuses on improving model robustness with Jacobian regularization; while [2] shows the benefit of norm regularization to improving generalization, which is unrelated to the Jacobian. Our main theorem otherwise demonstrates the theoretical benefits of the Jacobian to regularization. Our analysis also builds upon and strengthens the generalization inequalities established in previous works on sharpness-aware minimization techniques like SAM [3] and FisherSAM [4]. To the best of our knowledge, this is the first generalization inequality for Riemannian manifolds. Moreover, a key novelty in this theoretical analysis is that the general loss is directly bounded by the empirical loss, instead of the worst-case empirical loss like prior results such as SAM [3] or FisherSAM [4]. Additionally, although Jacobian regularization techniques have been applied in the Euclidean setting in earlier works [1, 5], these methods lacked a theoretical foundation on the effect of Jacobian regularization on the generalization ability. Our work provides the first theoretical justification for this practice and extends it to the Riemannian context. Based on the reviewer's suggestion, we have incorporated this explanation in the theoretical development section of the rebuttal revision.

• Methodological Contribution: The key feature of our algorithm that enhances its performance is the regularization of both the Riemannian gradient and the Jacobian. As demonstrated in our main theorem, this dual regularization effectively narrows the gap between the population loss $L_D$ and the empirical loss $L_S$, outperforming sharpness-aware methods that focus solely on implicitly regularizing the gradient norm. This is why we compare our method to Riemannian-SAM, which applies sharpness-aware minimization on manifolds, and SAM, a sharpness-aware technique in the Euclidean setting. Our results show that the sharpness in our models is significantly lower than in Riemannian-SAM, as confirmed in the ablation study, which in turn correlates with improved generalization performance.

• Regarding the empirical evaluations: In response to the reviewer’s suggestion for more rigorous experiments on a broader range of datasets, we have conducted additional experiments on three more challenging datasets: Tiny-ImageNet, Caltech101, and CUB-200-2011. The results are presented below.

In Section 5, we have also conducted the ablation study to show the sensitivity to the key hyperparameters including $\epsilon$ and $\rho$. The results show that RJR is robust across reasonable ranges of these two hyperparameters.

We would like to thank the reviewer for the insightful feedback. We would like to address the concerns below:

• Regarding the inner product/norm in Theorem 1: The first norm, corresponding to the Riemannian gradient term, should be expressed as the Riemannian norm. Similarly, the norm mentioned in lines 982–985 also refers to the Riemannian norm. We have corrected these typos in the revised manuscript and appreciate the reviewer for bringing this to our attention. Thank you for highlighting this oversight.

• Regarding the experiments on additional datasets: In response to the reviewer's suggestion, we conducted additional experiments on three more challenging datasets: TinyImageNet, Caltech101, and CUB-200-2011. The results, obtained using ResNet34 in the supervised setting, are summarized below. These experiments further validate the effectiveness and robustness of our approach across diverse and complex datasets.

The result shows that our method still outperforms the baselines.

• Regarding the choice of datasets to compare with Riemannian-SAM: Since the source code for Riemannian-SAM was not released, we opted to conduct experiments under two different settings: supervised and unsupervised learning for image classification. To ensure the fairness and reliability of the comparisons, we maintained consistent experimental settings across all methods evaluated. This consistency eliminates potential biases and highlights the relative performance of each approach under identical conditions.

• Regarding the comparisons with other orthogonal tricks: In Appendix A.1.2, we present an ablation study to highlight the superiority of our method compared to an alternative orthogonality technique proposed in [1]. Specifically, we designed a toy experiment to benchmark our approach against the method in [1], which regularizes with the term $|W^\top W - I_d|^2_2$ to encourage orthogonality in $W$. As illustrated in Figure 4, our method significantly outperforms this orthogonality trick in the toy example, demonstrating the advantages of learning on manifolds and the effectiveness of our approach. This result underscores the value of our method in achieving improved performance through principled manifold-based learning.

• Comparision with Trivialization [2]: Firstly, we note that Jacobian Regularization can be applied alongside Trivialization. Specifically, this involves replacing the base optimizer with RJR, where the manifold of interest is simplified to Euclidean space. Following the reviewer’s suggestion, we conducted additional comparisons, incorporating Trivialization with both SAM and Jacobian Regularization as the base optimizers. The results of these experiments are summarized below.

The results indicate that Trivialization with Jacobian Regularization achieves a slight improvement over both methods. The key insight behind this improvement is that regularizing the Jacobian reduces the generalization gap more effectively than sharpness-aware methods, as supported by Theorem 1. Additionally, we observed that Trivialization has a notably faster runtime. Although time constraints prevented us from conducting a comprehensive evaluation of Trivialization, the reviewer’s comment highlights a promising avenue for future research. Specifically, developing a deeper theoretical analysis of Trivialization's generalization ability could provide valuable insights. We sincerely thank the reviewer for this thoughtful suggestion.

• Regarding other minor issues: We have addressed the minor issues highlighted by the reviewer, including typos and abbreviations. These corrections have been incorporated into the revised rebuttal. We appreciate the reviewer’s attention to detail and thank them for pointing out these issues.

References:

[1] Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X. Yu. Orthogonal convolutional neural networks, 2020.

[2] Lezcano-Casado, M. (2019). Trivializations for Gradient-Based Optimization on Manifolds. In Advances in Neural Information Processing Systems (pp. 9154–9164).

Thanks for the response.

• The baseline seems weak, it hardly can tell if the benefits come from the network itself or the optimizer. Convincing experiments should be conducted on strong backbone networks.
• For the computer vision task, a large dataset is needed to validate the proposed optimizer.
• I'm not clear why this work cannot be applied to Riemannian networks, where most of the parameters lie in the manifold.

Based on the above, I will keep the original score.

Firstly, we would like to thank the reviewer for their valuable feedback. We address the concerns raised as follows:

• Regarding the main theorem: As noted after the theorem statement, our result extends prior work on generalization, including that of Foret et al. (2021) [1]. There are two key distinctions in our approach. First, we generalize the analysis from the Euclidean setting to a Riemannian manifold. Second, as the reviewer highlighted, we directly relate the general loss $L_D$ to the empirical loss $L_S$, rather than using the worst-case empirical loss within a neighborhood (i.e., $\max_{|\theta' - \theta| < \epsilon} L_S(\theta')$). This formulation leads to a more straightforward optimization procedure. We also emphasize even though under the local smoothness assumption, the term $\max_{|\theta' - \theta| < \epsilon} L_S(\theta')$ can be related to $L_S(\theta)$, our theorem does not rely on this assumption, since the derivation of our theorem is not a direct corollary of Foret et al. (2021).

• Regarding model robustness: As shown in previous studies on sharpness-aware minimization, such as those by Foret et al. (2021) [1] and Kim et al. (2022) [2], a key benefit of sharpness-aware methods is their robustness to parameter perturbations. To support this claim for RJR, we conducted an additional ablation study, included in Appendix A.3 of the rebuttal revision, which assesses the robustness of RJR against adversarial parameter perturbations. The results show that RJR performs better than both Riemannian SAM and SGD in terms of model robustness.

• Regarding other minor issues: The last row of the first table now correctly represents the 95% confidence interval for RJR. Additionally, we have addressed other minor issues raised by the reviewer in the revision.

Once again, we would like to thank the reviewer for this valuable feedback.

References:

[1] Foret, P., Malkin, N., & Cohn, T. (2021). Sharpness-aware minimization for efficiently improving generalization. Proceedings of the 38th International Conference on Machine Learning (ICML 2021), 26, 301-311.

[2] Kim, M., Li, D., Hu, S. X., & Hospedales, T. (2022). Fisher SAM: Information geometry and sharpness aware minimization. Proceedings of the 39th International Conference on Machine Learning, 162, 11148–11161. PMLR. https://proceedings.mlr.press/v162/kim22f.html​