Summary
The paper proposes a second-order optimizer named MKOR based on K-FAC. Compared with K-FAC, MKOR uses a rank-1 update to construct Kronecker factors enabling the inverse of Kronecker factors to be efficiently computed via the Sherman-Morrison-Woodbury (SMW) identity. In addition, the rank-1 update of factors in MKOR also helps reduce the communication traffic in distributed training with data parallel for aggregating factors. Experiments on the BERT-Large-Uncased model with the GLUE dataset show that MKOR outperforms first-order LAMB and second-order KAISA.
Strengths
1. An implicit inverse computation method with the SMW identity in computing the inverse preconditioner factors and a formula to control the balance of exploitation between the first-order and second-order information.
2. A hybrid method of MKOR with the first-order method, MKOR-H, which allows switchable from MKOR to the first-order optimizer in the later of training.
Weaknesses
1. An important baseline, Eva [A], which is extremely similar to MKOR, is missing.
2. Though MKOR-H seems to achieve some improvement, when switching MKOR to the first-order optimizer is unclear.
3. The comparison between MKOR and K-FAC seems to be not that fair (e.g., Figure 4(b)).
[A] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.
Questions
1. The key idea of using rank-1 updates to replace the original Kronecker factors in K-FAC has been proposed in Eva [A]. In [A], the authors provide a better derivation of the preconditioner without storing the factor matrices so that they use Kronecker vectors (KVs) to calculate the inverse of KVs using the SMW formula, which is much better than MKOR as MKOR needs to store the matrix form of factors (i.e., $L_t^m$ and $R_t^m$). The similarity between MKOR and Eva and the advantage of Eva over MKOR make the originality of this work weak. It is suggested to clarify the differences between MKOR and Eva, and highlight some advantages of MKOR over Eva.
2. The main novel idea in MKOR is the balance between second-order and first-order information in updating $L_t^m$ and $R_t^m$ using Eq. (7) and (8) by controlling $\zeta$. However, how to choose $\zeta$ (or dynamically changing $\zeta$ to switch MKOR to SGD just like MKOR-H) is very important. The paper doesn’t provide such an analysis of when and how to use $\zeta$ though some results show that it can help improve convergence.
3. Similarly to the previous problem, MKOR-H is a combination between first-order optimizer and MKOR according to the loss during training. However, how to measure whether the loss is changed is unclear (e.g., results in Figure 2 haven’t such information either) and when to switch to the first-order optimizer is unclear.
4. What are the update frequencies of the second-order information for MKOR and K-FAC in Figure 2 and Table 3? And how do you choose such frequencies for comparison?
5. In the original paper of KAISA, the update frequency can be set to be very high (e.g., 100-500) for achieving good accuracy on ResNet models with CIFAR10 and ImageNet datasets, is it possible that MKOR could be better than KAISA in these kinds of scenarios?
[A] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Limitations
See Questions.