SepNorm: Generalization of Lion and Normalised Gradient Methods

1. Limited scope: The theoretical analysis is restricted to quadratic functions, which may not capture the complexity of real-world neural networks.

2. Performance in CV tasks: OrtSepNorm shows suboptimal results in certain CV tasks (e.g., ResNet architectures), suggesting that the method’s advantages may be more task-dependent.

3. Batch size: The noise reduction in SepNorm relative to Lion is noted, yet it would be beneficial to have a more comprehensive analysis of batch size dependency across various architectures.

4. Figures: The paper lacks figures comparing the optimizers’ performance over time (epochs), which would provide insights into convergence speed and stability differences across tasks.

5. No GitHub repository: At this day, there is no open-source implementation provided.

While the paper provides hyperparameters used in experiments (in Appendix), it does not detail how these were selected (e.g., through grid search or prior literature) nor offer specific recommendations for practitioners aiming to apply these optimizers. Providing more details on the tuning process would enhance the reproducibility of the results. Moreover, additional experiments on smaller, well-known benchmarks like CIFAR-10/ CIFAR10-100 would provide clearer visualization of the optimizer's advantage and allow for direct performance comparisons, enhancing reproducibility and practical insight.

• The writing needs significant improvement. There are numerous grammatical errors throughout the manuscript that affect readability. The background on overparameterization (Lines 31-32 and 86-89) have tenuous connections to the proposed optimizer. In fact, overparameterization is typically absent in the training of large language models (LLMs); for example, Llama3-8b is trained with approximately 2000 tokens per parameter. Additionally, the references are not well cited or discussed. SepNorm builds on signSGD, yet the original signSGD paper and closely related works are not cited. At Line 100, the reference to classical Nesterov’s Accelerated Gradient (NAG) to underscore SGD is misleading, as NAG is not directly related to SGD. Furthermore, the numbering of Theorems in the main text is inconsistent with their proofs in the appendix.

• SepNorm appears to function more as an engineering trick. The only modification from Block Normalized Gradient (BNG) is the multiplication by the square root of the block size. In my previous experiments with training a ViT using both BNG and SepNorm on CIFAR-10, I found that, with finely tuned learning rates for both optimizers, there was no significant difference in performance.

• The theoretical analyses are based on unrealistic assumptions. For instance, Theorems 3.1, 4.1, and 4.2 require $\langle \nabla F(w), w \rangle = 0 , \forall w$, which is impractical in real DNN training. Additionally, the analyses in Section 5 rely on an overly simplistic quadratic function, and Theorem 5.3 even requires $A$ to be diagonal or to have identical eigenvalues. These unrealistic assumptions provide limited insights into the behavior of SepNorm in practical DNN training. It would be more beneficial to provide a theoretical convergence analysis for a new optimizer under milder assumptions.

• The experimental setup is also questionable. The baseline BNG is missing from all comparison experiments. Moreover, the experimental settings lack justification; for example, the learning rates for AdamW and Lion when training ViT in Table 6 are set to unusually low values of 6.25e-5 and 6.25e-6, which are significantly lower than those in the original Lion paper. The batch size for Lion is set to 256, much smaller than the optimal 4096. Additionally, the test accuracy in Table 1 is notably inferior to that reported in the Lion paper. It is also somewhat unusual that the model selected for the language task is T5, rather than the more popular GPT-like decode-only Transformers.

Minor Issues

• In Theorem 4.3, there is no definition of $d$ provided anywhere.
• The proof for Theorem 4.3 is omitted, and it may not be straightforward to obtain.

1. The theory behind Lion (as described in Section 3) contains lots of hand-waving arguments that are not rigorous. For example, in line 160-line 161, the authors explained the possible benefit of Lion by the following statement: "Due to the vanishing gradient problem, some weights may receive small gradient components, especially those corresponding to first layers. " However, there is no evidence showing this is indeed the case or not. Similar hand-waving arguments also appear in line 197-line 208.

2. How is Theorem 3.1 related to the theory of Lion? I do not think there is a direct relationship. More explanations are needed.

3. In Section 4, the SepNorm seems to be a layer-wise (or group-wise) version of Normalized GD and Lion. However, it is unclear why Theorem 4.2 and Theorem 4.3 demonstrate the superior performance guarantees of the proposed methods. Also the proof can be trivially obtained by the property of scale-invariant networks.

4. Theorem 5.4 provides limited insights compared with normalized GD. How does the sharpness obtained by SepNorm imply the advantage over Lion and normalized GD? It is unclear to me. In addition, the proof of Theorem 5.4 seems to be a straightforward extension of [Arora et al. 2022].

5. In Section 6, there are missing details about how to select parameter groups for the proposed algorithms. In addition, the authors need to perform additional ablation studies to show the benefit of the proposed algorithms. For example, one baseline could be layer-wise optimizers such as LAMB (https://arxiv.org/pdf/1904.00962), which also uses group-wise (or layer-wise) learning rate. In addition, the experimental benefits of the proposed methods are marginal compared with existing optimizers such as AdamW.

It remains unclear to me the exact contributions of this work. If it is from the experimental side, it seems that the gain is marginal (e.g. Table 3). From the theoretical viewpoint, the analysis more or less follows that of SIGN and Normalized GD. It also lacks comparisons with previous theory works in section 5. In addition to this, the quadratic model perhaps is too simple to capture the underlying training dynamics.