LoKO: Low-Rank Kalman Optimizer for Online Fine-Tuning of Large Models

As mentioned in response to other reviewers, the adoption of a diagonal approximation for large-scale covariance and Hessian matrices is a widely accepted practice in machine learning due to its computational efficiency. For instance, [2] utilizes a diagonal Hessian approximation in the development of a second-order optimizer. Or, [3] demonstrates that the Adam algorithm effectively employs a diagonal approximation of the pure Newton’s method update. In addition, although the idea is interesting, the primary focus of this study is on the empirical investigation of LoKO's performance, rather than a detailed theoretical analysis of this approximation. However, we acknowledge the value of such analysis and consider it a promising direction for future work. One possible approach can be analyzing the evolution of the covariance matrix during Kalman iterations. By propagating the covariance matrix through the update equations, we should be able to somehow show that the off-diagonal elements converge to zero.

First of all, we sincerely appreciate your invaluable feedback. We have made every effort to address your concerns thoroughly, and below, we provide a detailed point-by-point response to each of your comments. Furthermore, all modifications have been clearly highlighted in the revised version of the paper that has been submitted.

As detailed in Section 5, our experiments cover a broad range of computer vision and language models applied to various datasets for classification tasks. The models used in our study are well-known and widely recognized within the classification domain. We believe that the current adopted models are appropriate for the task considered in our experiments. However, to address your concern more effectively, we would kindly ask if you could specify particular models that you believe should be investigated within the context of this paper.

Thank you for this insightful comment. As discussed in [4], AdamW optimizer benefits from the decoupled weight decay, which has contributed to its strong performance in language modeling tasks. On the other hand, the gradient-free Kalman algorithm used in LoKO does not incorporate a similar regularization mechanism, which may explain the slightly lower performance in certain cases compared to AdamW. We will clarify this point by adding the following discussion to Section 5.2 of the revised version of the paper:

“As can be seen in Table 2, LoKO demonstrates slightly lower performance than AdamW on certain language tasks. A justification for this observation is the fact that AdamW optimizer benefits from the decoupled weight decay, which has contributed to its strong performance in language modeling tasks [4]. On the other hand, the gradient-free Kalman algorithm used in LoKO does not incorporate a similar regularization mechanism, which may explain the slightly lower performance in certain cases compared to AdamW.”

We appreciate your comment and acknowledge the importance of studying dynamic adaptation. However, our work primarily focuses on the development and evaluation of the LoKO algorithm within the stationary dataset. As discussed in [1], dynamic adaptation typically is valid in the context of Online Continual Learning (OCL) with non-stationary data, which lies outside the primary focus of our current work. We believe that the current methodology is appropriate for the setting and domain considered in our experiments. Nonetheless, we recognize the importance of this issue and consider it an interesting avenue for future work.

[1] M. K. Titsias, A. Galashov, A. Rannen-Triki, R. Pascanu, Y. W. Teh, and J. Bornschein, “KALMAN FILTER FOR ONLINE CLASSIFICATION OF NON-STATIONARY DATA,” 2024.

[2] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma, “Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training,” Mar. 05, 2024, arXiv: arXiv:2305.14342. Accessed: Mar. 12, 2024. [Online]. Available: http://arxiv.org/abs/2305.14342

[3] I. Molybog et al., “A Theory on Adam Instability in Large-Scale Machine Learning,” Apr. 25, 2023, arXiv: arXiv:2304.09871. Accessed: Nov. 16, 2024. [Online]. Available: http://arxiv.org/abs/2304.09871

[4] S. Xie and Z. Li, “Implicit Bias of AdamW: $\ell_∞$-Norm Constrained Optimization,” in Proceedings of the 41st International Conference on Machine Learning, PMLR, Jul. 2024, pp. 54488–54510. Accessed: Nov. 17, 2024. [Online]. Available: https://proceedings.mlr.press/v235/xie24e.html

[5] Z. Han, C. Gao, J. Liu, J. Zhang, and S. Q. Zhang, “Parameter-Efficient Fine-Tuning for Large Models: A Comprehensive Survey,” Jul. 12, 2024, arXiv: arXiv:2403.14608. Accessed: Sep. 08, 2024. [Online]. Available: http://arxiv.org/abs/2403.14608

[6] R. Wightman et al., rwightman/pytorch-image-models: v0.8.10dev0 Release. (Feb. 07, 2023). Zenodo. doi: 10.5281/ZENODO.4414861.

[7] K. He, X. Zhang, S. Ren, and J. Sun, “Deep Residual Learning for Image Recognition,” Dec. 10, 2015, arXiv: arXiv:1512.03385. Accessed: Nov. 18, 2024. [Online]. Available: http://arxiv.org/abs/1512.03385

[8] M. Oquab et al., “DINOv2: Learning Robust Visual Features without Supervision,” Feb. 02, 2024, arXiv: arXiv:2304.07193. Accessed: Nov. 18, 2024. [Online]. Available: http://arxiv.org/abs/2304.07193

[9] Y. Liu et al., “RoBERTa: A Robustly Optimized BERT Pretraining Approach,” Jul. 26, 2019, arXiv: arXiv:1907.11692. Accessed: Nov. 18, 2024. [Online]. Available: http://arxiv.org/abs/1907.11692

[10] F. Heimes, “Extended Kalman filter neural network training: experimental results and algorithm improvements,” in SMC’98 Conference Proceedings. 1998 IEEE International Conference on Systems, Man, and Cybernetics (Cat. No.98CH36218), Oct. 1998, pp. 1639–1644 vol.2. doi: 10.1109/ICSMC.1998.728124.

[11] I. Rivals and L. Personnaz, “A recursive algorithm based on the extended Kalman filter for the training of feedforward neural models,” Neurocomputing, vol. 20, no. 1, pp. 279–294, Aug. 1998, doi: 10.1016/S0925-2312(98)00021-6.

As mentioned in response to other reviewers, initializing the covariance matrix requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for the covariance matrix can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. As long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix. As discussed in detail in Section 5.3 of the paper; to obtain the upper bounds for our case studies, we conducted a series of experiments across a wide range of initial values and tracked the average online accuracy. If the accuracy doesn’t reach to a specific threshold within a predetermined number of iterations, we consider the test to have diverged. The boundary values derived from this analysis have been reported in Table 4 of the paper, specifically for case studies where an upper bound for the initial values of covariance matrix is defined. To the best of our knowledge, no general criterion exists for this selection, and only a few studies have addressed the issue of initialization such as [10] and [11]. To clarify more, we will add the following discussion to the revised version of the paper:

“Initializing p ̂ requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for p ̂ can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. Our experiments demonstrate that as long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix. To obtain the upper bounds for our case studies, we conducted a series of experiments across a wide range of initial values and tracked the average online accuracy. If the accuracy doesn’t reach to a specific threshold within a predetermined number of iterations, we consider the test to have diverged. The boundary values derived from this analysis have been reported in Table 4, specifically for case studies where an upper bound for the initial values of covariance matrix is defined. To the best of our knowledge, no general criterion exists for this selection, and only a few studies have addressed the issue of initialization such as [10] and [11].”

As discussed in [5] and explained in the manuscript, the widespread use of deep neural networks, particularly large models, is largely driven by pre-training on extensive datasets followed by task-specific fine-tuning. The non-trainable parameters have been obtained from different pretrained models. For MNIST/DenseNet-121, we used a backbone pretrained on ImageNet via PyTorch Image Models (timm) [6]. For CIFAR-10/ResNet18 and CIFAR-100/ResNet50, we employed pretrained models from [7], which were trained on ImageNet as the backbone. For CIFAR-10/ViT-B16, CIFAR-100/ViT-B16, and ImageNet100/ViT-L16, we used DINOv2 [8], which was pretrained on ImageNet. And, for the language tasks we employed RoBERTa pretrained model from [9]. We will clarify your point by adding the following discussion to the revised version of the paper:

“We employ the following pre-trained models as backbones for our fine-tuning experiments:

o For MNIST/DenseNet-121, we used a backbone pretrained on ImageNet via PyTorch Image Models (timm) [6].

o For CIFAR-10/ResNet18 and CIFAR-100/ResNet50, we employed pretrained models from [7], which were trained on ImageNet as the backbone.

o For CIFAR-10/ViT-B16, CIFAR-100/ViT-B16, and ImageNet100/ViT-L16, we used DINOv2 [8], which was pretrained on ImageNet.

o For the language tasks we employed RoBERTa pretrained model from [9].”

Thank you for your constructive feedback. As discussed in Section 4.1, applying a diagonal approximation for the covariance matrix, combined with a low-rank decomposition approach, decrease the computational complexity from quadratic to linear in the number of trainable parameters. To clarify this further, we will include the following explanation in Appendix C.1 of the revised paper:

“Computational Analysis: Let n represent the number of parameters, n ̃ the number of trainable parameters (n ̃ ≪n), and m the number of model outputs. The computational complexities of each step in the LoKO algorithm compared to the vanilla Kalman filter are as follows: Prediction: The computational complexity of both LoKO and the vanilla Kalman filter for the Prediction step is O(1). Pre-Updating: For LoKO, estimating the observation noise covariance using Equation (7) has a complexity of O(m^2). When using Equation (12), the complexity increases to O(m^2 n ̃). In contrast, for the vanilla Kalman filter, the complexity for Equation (7) is O(m^2), while for Equation (12) it is O(m^2 n+n^2 m). Updating: The computational complexity for LoKO when calculating the Kalman gain (Equation (11a)) is O(m^3+m^2 n ̃), while for the vanilla Kalman filter (using Equation (3a)), it will be O(m^3+m^2 n+n^2 m). For updating parameters (Equation (11b)), LoKO has a complexity of O(mn ̃), compared to O(mn) for the vanilla Kalman filter. And finally, the complexity for covariance updating (Equation (11c)) in LoKO is O(n ̃), while for the vanilla Kalman filter (using Equation (3c)), it will be O(n^2 m). Thus, the total computational complexity in the worst-case scenario, for LoKO is O(m^3+m^2 n ̃), and for the vanilla Kalman filter will be O(m^3+m^2 n+n^2 m). In cases where the number of parameters is significantly larger than output size, the dominant term for LoKO is O(m^2 n ̃), while for the vanilla Kalman filter, it will be O(n^2 m). Therefore, LoKO reduces the computational complexity from quadratic to linear in the number of parameters as well as decreasing the number of trainable parameters (n ̃ ≪n).”

To evaluate the time efficiency of our LoKO algorithm, we compare the number of steps required for convergence, similar to the criterion used in [2]. Convergence is defined as the number of iterations at which the loss stays flat or decreases by less than a specified threshold. The time analysis will be presented as follows in the Appendix C.2 of the revised paper.

[1] S.-Y. Liu et al., “DoRA: Weight-Decomposed Low-Rank Adaptation,” Jul. 09, 2024, arXiv: arXiv:2402.09353. Accessed: Nov. 19, 2024. [Online]. Available: http://arxiv.org/abs/2402.09353

[2] Z. Han, C. Gao, J. Liu, J. Zhang, and S. Q. Zhang, “Parameter-Efficient Fine-Tuning for Large Models: A Comprehensive Survey,” Jul. 12, 2024, arXiv: arXiv:2403.14608. Accessed: Sep. 08, 2024. [Online]. Available: http://arxiv.org/abs/2403.14608

[3] G. I. Parisi, R. Kemker, J. L. Part, C. Kanan, and S. Wermter, “Continual lifelong learning with neural networks: A review,” Neural Netw., vol. 113, pp. 54–71, May 2019, doi: 10.1016/j.neunet.2019.01.012.

[4] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma, “Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training,” Mar. 05, 2024, arXiv: arXiv:2305.14342. Accessed: Mar. 12, 2024. [Online]. Available: http://arxiv.org/abs/2305.14342

Thank you for your constructive feedback. As discussed in Section 4.1, applying a diagonal approximation for the covariance matrix, combined with a low-rank decomposition approach, decrease the computational complexity from quadratic to linear in the number of trainable parameters. To clarify this further, we will include the following explanation in Appendix C.1 of the revised paper:

“Computational Analysis: Let n represent the number of parameters, n ̃ the number of trainable parameters (n ̃ ≪n), and m the number of model outputs. The computational complexities of each step in the LoKO algorithm compared to the vanilla Kalman filter are as follows: Prediction: The computational complexity of both LoKO and the vanilla Kalman filter for the Prediction step is O(1). Pre-Updating: For LoKO, estimating the observation noise covariance using Equation (7) has a complexity of O(m^2). When using Equation (12), the complexity increases to O(m^2 n ̃). In contrast, for the vanilla Kalman filter, the complexity for Equation (7) is O(m^2), while for Equation (12) it is O(m^2 n+n^2 m). Updating: The computational complexity for LoKO when calculating the Kalman gain (Equation (11a)) is O(m^3+m^2 n ̃), while for the vanilla Kalman filter (using Equation (3a)), it will be O(m^3+m^2 n+n^2 m). For updating parameters (Equation (11b)), LoKO has a complexity of O(mn ̃), compared to O(mn) for the vanilla Kalman filter. And finally, the complexity for covariance updating (Equation (11c)) in LoKO is O(n ̃), while for the vanilla Kalman filter (using Equation (3c)), it will be O(n^2 m). Thus, the total computational complexity in the worst-case scenario, for LoKO is O(m^3+m^2 n ̃), and for the vanilla Kalman filter will be O(m^3+m^2 n+n^2 m). In cases where the number of parameters is significantly larger than output size, the dominant term for LoKO is O(m^2 n ̃), while for the vanilla Kalman filter, it will be O(n^2 m). Therefore, LoKO reduces the computational complexity from quadratic to linear in the number of parameters as well as decreasing the number of trainable parameters (n ̃ ≪n).”

To evaluate the time efficiency of our LoKO algorithm, we compare the number of steps required for convergence, similar to the criterion used in [4]. Convergence is defined as the number of iterations at which the loss stays flat or decreases by less than a specified threshold. The time analysis is presented in Appendix C.2 of the revised paper.

As mentioned in response to other reviewers, initializing the covariance matrix requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for the covariance matrix can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. As long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix. For instance, we present evidence based on the results of CIFAR-100/ViT-B16 for four different covariance initialization values in Figure 7 (Appendix D.3 in the revised paper). As demonstrated, the final outcomes show minimal sensitivity to these initialization.

As discussed in detail in Section 5.3 of the paper; to obtain the upper bounds for our case studies, we conducted a series of experiments across a wide range of initial values and tracked the average online accuracy. If the accuracy doesn’t reach to a specific threshold within a predetermined number of iterations, we consider the test to have diverged. The boundary values derived from this analysis have been reported in Table 4 of the paper, specifically for case studies where an upper bound for the initial values of covariance matrix is defined. We will clarify this point by adding the following discussion to the revised version of the paper:

“Initializing p ̂ requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for p ̂ can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. Our experiments demonstrate that as long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix.”

We appreciate your comment and acknowledge the importance of studying on dynamic streaming type of data. However, our work primarily focuses on the development and evaluation of the LoKO algorithm within well-defined datasets. The type of investigation you mentioned falls under the broader category of Continual Learning [3], which is beyond the primary focus of our paper. Nevertheless, we agree that exploring LoKO in the context of dynamic streaming type of data is an interesting direction for future work, and we appreciate your suggestion.

As correctly pointed out in your comment, the primary limitation of full fine-tuning is the high computational cost, which can be prohibitive in many scenarios. Even with high GPU resources, based on discussion in [2], full fine-tuning can sometimes lead to catastrophic forgetting, which negatively impacts the model's ability to generalize effectively. LoKO, in contrast, aims to provide a more computationally efficient alternative while mitigating such issues.

Thank you for this valuable suggestion. We find your recommendation highly insightful. We have conducted additional experiments involving the DoRA [1] variant of LoRA, in combination with Kalman, AdamW, and AdaGrad. The results from these experiments are now included in the revised paper, with visualizations in Appendix D.1 (Figure 3, and Figure 4) and further analysis in Table 1, and Table 2, Section 5.2. As shown, the results are promising for the Kalman algorithm, demonstrating its compatibility and effectiveness when integrated with this Weight-Decomposed Low-Rank Adaptation technique.

First of all, we sincerely thank you for the invaluable feedback. Our utmost effort has been exerted in addressing your concerns, and we have subsequently provided a comprehensive point-by-point explanation in the following list of responses. Furthermore, all modifications have been clearly highlighted in the revised version of the paper that has been submitted.

Thanks. The additional results indeed look promising.

One question regarding to your reply. What do you mean about "weight-decomposed low-rank adaptation technique"? R u referring to the paper "[1] S.-Y. Liu et al., “DoRA: Weight-Decomposed Low-Rank Adaptation,” Jul. 09, 2024, arXiv: arXiv:2402.09353. Accessed: Nov. 19, 2024. [Online]. Available: http://arxiv.org/abs/2402.09353"?

The primary focus of this paper is on "online fine-tuning" scenarios, and we intentionally emphasize this context. While we acknowledge that offline settings with larger batch sizes are important, they fall outside the scope of this study.

Single epoch fine-tuning is a common practice in scenarios where large pretrained models are adapted to specific tasks or datasets particularly in large-scale datasets, as seen in [3], [4], [5]. However, we have performed multiple epochs of fine-tuning for smaller datasets (COLA and MRPC), as clearly mentioned in Section 5.2 of the paper.

Thank you for this valuable suggestion. We find your recommendation highly insightful. We have conducted additional experiments involving the DoRA [2] variant of LoRA, in combination with Kalman, AdamW, and AdaGrad. The results from these experiments are now included in the revised paper, with visualizations in Appendix D.1 (Figure 3, and Figure 4) and further analysis in Table 1, and Table 2, Section 5.2. As shown, the results are promising for the Kalman algorithm, demonstrating its compatibility and effectiveness when integrated with this Weight-Decomposed Low-Rank Adaptation technique.

To clarify, the primary focus of this study is on the empirical investigation of LoKO, rather than on theoretical analysis. There is no formal theoretical proof for LoRA (at least, we could not identify any in the literature). Given the challenges posed by the non-convex loss landscape, relatively few studies have tackled the theoretical analysis of optimizers such as Adam and SGD like [1]. Most works in this area evaluate the effectiveness of the algorithms primarily through empirical evidence. We acknowledge the value of a formal convergence analysis and consider it an interesting direction for future work. However, in the Appendix B, we have provided some theoretical insights as summarized below:

o Proposition 1: To establish the core contribution of this paper, we proved that the proposed LoKO formulation is mathematically equivalent to the vanilla Kalman filter algorithm when leveraging the low-rank decomposition technique in LoRA and applying a diagonal approximation of the covariance matrix. This proof relies on the properties of the transposed Khatri–Rao product, Einstein notation, and their connections to matrix multiplication.

o EMA technique for R matrix Estimation: We validated our EMA approach for estimating the noise covariance matrix (R) by employing the zero- and first-order terms of the Taylor series expansion around the most recent parameter updates, and expanding it based on the definition of the noise covariance.

We sincerely appreciate your invaluable feedback. We have made every effort to address your concerns and have provided a detailed point-by-point explanation in the following list of responses. Furthermore, all modifications have been clearly highlighted in the revised version of the paper that has been submitted.

[1] S. J. Reddi, S. Kale, and S. Kumar, “On the Convergence of Adam and Beyond,” Apr. 19, 2019, arXiv: arXiv:1904.09237. Accessed: Nov. 17, 2024. [Online]. Available: http://arxiv.org/abs/1904.09237

[2] S.-Y. Liu et al., “DoRA: Weight-Decomposed Low-Rank Adaptation,” Jul. 09, 2024, arXiv: arXiv:2402.09353. Accessed: Nov. 19, 2024. [Online]. Available: http://arxiv.org/abs/2402.09353

[3] A. Panda, B. Isik, X. Qi, S. Koyejo, T. Weissman, and P. Mittal, “Lottery Ticket Adaptation: Mitigating Destructive Interference in LLMs,” Jun. 25, 2024, arXiv: arXiv:2406.16797. Accessed: Nov. 17, 2024. [Online]. Available: http://arxiv.org/abs/2406.16797

[4] A. Komatsuzaki, “One Epoch Is All You Need,” Jun. 16, 2019, arXiv: arXiv:1906.06669. Accessed: Nov. 17, 2024. [Online]. Available: http://arxiv.org/abs/1906.06669

[5] S. Chen, M. Jiang, and Q. Zhao, “What Do Deep Saliency Models Learn about Visual Attention?,” Oct. 14, 2023, arXiv: arXiv:2310.09679. Accessed: Nov. 17, 2024. [Online]. Available: http://arxiv.org/abs/2310.09679

[6] S. Xie and Z. Li, “Implicit Bias of AdamW: $\ell_∞$-Norm Constrained Optimization,” in Proceedings of the 41st International Conference on Machine Learning, PMLR, Jul. 2024, pp. 54488–54510. Accessed: Nov. 17, 2024. [Online]. Available: https://proceedings.mlr.press/v235/xie24e.html

[7] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma, “Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training,” Mar. 05, 2024, arXiv: arXiv:2305.14342. Accessed: Mar. 12, 2024. [Online]. Available: http://arxiv.org/abs/2305.14342

Thank you for your constructive feedback. As discussed in Section 4.1, applying a diagonal approximation for the covariance matrix, combined with a low-rank decomposition approach, decrease the computational complexity from quadratic to linear in the number of trainable parameters. To clarify this further, we will include the following explanation in the Appendix C.1 of the revised paper:

“Computational Analysis: Let n represent the number of parameters, n ̃ the number of trainable parameters (n ̃ ≪n), and m the number of model outputs. The computational complexities of each step in the LoKO algorithm compared to the vanilla Kalman filter are as follows: Prediction: The computational complexity of both LoKO and the vanilla Kalman filter for the Prediction step is O(1). Pre-Updating: For LoKO, estimating the observation noise covariance using Equation (7) has a complexity of O(m^2). When using Equation (12), the complexity increases to O(m^2 n ̃). In contrast, for the vanilla Kalman filter, the complexity for Equation (7) is O(m^2), while for Equation (12) it is O(m^2 n+n^2 m). Updating: The computational complexity for LoKO when calculating the Kalman gain (Equation (11a)) is O(m^3+m^2 n ̃), while for the vanilla Kalman filter (using Equation (3a)), it will be O(m^3+m^2 n+n^2 m). For updating parameters (Equation (11b)), LoKO has a complexity of O(mn ̃), compared to O(mn) for the vanilla Kalman filter. And finally, the complexity for covariance updating (Equation (11c)) in LoKO is O(n ̃), while for the vanilla Kalman filter (using Equation (3c)), it will be O(n^2 m). Thus, the total computational complexity in the worst-case scenario, for LoKO is O(m^3+m^2 n ̃), and for the vanilla Kalman filter will be O(m^3+m^2 n+n^2 m). In cases where the number of parameters is significantly larger than output size, the dominant term for LoKO is O(m^2 n ̃), while for the vanilla Kalman filter, it will be O(n^2 m). Therefore, LoKO reduces the computational complexity from quadratic to linear in the number of parameters as well as decreasing the number of trainable parameters (n ̃ ≪n).”

To evaluate the time efficiency of our LoKO algorithm, we compare the number of steps required for convergence, similar to the criterion used in [7]. Convergence is defined as the number of iterations at which the loss stays flat or decreases by less than a specified threshold. The time analysis is presented in Appendix C.2 of the revised paper.

Thank you for this insightful comment. As discussed in [6], AdamW optimizer benefits from the decoupled weight decay, which has contributed to its strong performance in language modeling tasks. On the other hand, the gradient-free Kalman algorithm used in LoKO does not incorporate a similar regularization mechanism, which may explain the slightly lower performance in certain cases compared to AdamW. We will clarify this point by adding the following discussion to Section 5.2 of the revised version of the paper: “As can be seen in Table 2, LoKO demonstrates slightly lower performance than AdamW on certain language tasks. A justification for this observation is the fact that AdamW optimizer benefits from the decoupled weight decay, which has contributed to its strong performance in language modeling tasks [6]. On the other hand, the gradient-free Kalman algorithm used in LoKO does not incorporate a similar regularization mechanism, which may explain the slightly lower performance in certain cases compared to AdamW.”

First of all, we sincerely thank you for the invaluable feedback. We have carefully considered all your comments and made every effort to address your concerns comprehensively. All modifications have been clearly highlighted in the revised version of the paper that has been submitted. Below, we provide a detailed, point-by-point response to each of your comments:

[1] S. Murtuza and S. F. Chorian, “Node decoupled extended Kalman filter based learning algorithm for neural networks,” in Proceedings of 1994 9th IEEE International Symposium on Intelligent Control, Aug. 1994, pp. 364–369. doi: 10.1109/ISIC.1994.367790.

[2] P. G. Chang, M. Jones, and K. Murphy, “On diagonal approximations to the extended Kalman filter for online training of Bayesian neural networks”.

[3] P. G. Chang, G. Durán-Martín, A. Y. Shestopaloff, M. Jones, and K. Murphy, “Low-rank extended Kalman filtering for online learning of neural networks from streaming data,” Jun. 27, 2023, arXiv: arXiv:2305.19535. Accessed: May 31, 2024. [Online]. Available: http://arxiv.org/abs/2305.19535

[4] G. Welch, “An Introduction to the Kalman Filter,” 1997.

[5] H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma, “Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training,” Mar. 05, 2024, arXiv: arXiv:2305.14342. Accessed: Mar. 12, 2024. [Online]. Available: http://arxiv.org/abs/2305.14342

[6] I. Molybog et al., “A Theory on Adam Instability in Large-Scale Machine Learning,” Apr. 25, 2023, arXiv: arXiv:2304.09871. Accessed: Nov. 16, 2024. [Online]. Available: http://arxiv.org/abs/2304.09871

[7] J. Liu et al., “Towards Out-Of-Distribution Generalization: A Survey,” Jul. 27, 2023, arXiv: arXiv:2108.13624. doi: 10.48550/arXiv.2108.13624.

[8] S.-Y. Liu et al., “DoRA: Weight-Decomposed Low-Rank Adaptation,” Jul. 09, 2024, arXiv: arXiv:2402.09353. Accessed: Nov. 19, 2024. [Online]. Available: http://arxiv.org/abs/2402.09353

Thank you for this valuable suggestion. We find your recommendation highly insightful. We have conducted additional experiments involving the DoRA [8] variant of LoRA, in combination with Kalman, AdamW, and AdaGrad. The results from these experiments are now included in the revised paper, with visualizations in Appendix D.1 (Figure 3, and Figure 4) and further analysis in Table 1, and Table 2, Section 5.2. As shown, the results are promising for the Kalman algorithm, demonstrating its compatibility and effectiveness when integrated with this Weight-Decomposed Low-Rank Adaptation technique.

We appreciate your comment and acknowledge the importance of studying OOD generalization. However, our work primarily focuses on the development and evaluation of the LoKO algorithm within the Independent and Identically Distributed dataset. OOD generalization typically requires addressing issues in distributional shift scenarios [7], which lies outside the primary focus of our current work. Although we recognize that EMA-based noise covariance estimation may be sensitive to distribution shifts, we believe that the current methodology is appropriate for the setting and domain considered in our experiments. To address your concern more effectively, it would be helpful if you could clarify the specific types of OOD tasks you have in mind.

As mentioned in response to other reviewers, initializing the covariance matrix requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for the covariance matrix can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. As long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix. For instance, we present evidence based on the results of CIFAR-100/ViT-B16 for four different covariance initialization values in Figure 7 (Appendix D.3 in the revised paper). As demonstrated, the final outcomes show minimal sensitivity to the initialization.

As discussed in detail in Section 5.3 of the paper; to obtain the upper bounds for our case studies, we conducted a series of experiments across a wide range of initial values and tracked the average online accuracy. If the accuracy doesn’t reach to a specific threshold within a predetermined number of iterations, we consider the test to have diverged. The boundary values derived from this analysis have been reported in Table 4 of the paper, specifically for case studies where an upper bound for the initial values of covariance matrix is defined. We will clarify this point by adding the following discussion to the revised version of the paper: “Initializing p ̂ requires prior knowledge of the network weights and their associated uncertainties. In the absence of such knowledge, almost any initial values for p ̂ can be chosen, provided they are not too close to zero, and in some cases, they are not too far from an upper bound. Our experiments demonstrate that as long as the initialization falls within an acceptable range, the filter will eventually converge to the optimal values. Notably, the optimal values are not sensitive to the initial covariance matrix.”

The adoption of a diagonal approximation for large-scale covariance and Hessian matrices is a widely accepted practice in machine learning due to its computational efficiency. For instance, [5] utilizes a diagonal Hessian approximation in the development of a second-order optimizer. Or, [6] demonstrates that the Adam algorithm effectively employs a diagonal approximation of the pure Newton’s method update. Furthermore, it is well-known that even using a full covariance matrix does not guarantee convergence to a global optimum, given the highly non-convex nature of the loss landscape. Lastly, we note that text-to-image generation, while an important area of research, falls outside the scope of this paper.

We would like to clarify that the Gaussian assumption applies to each individual neural parameter, representing the uncertainty associated with each parameter, as discussed in [4]. It does not imply a single scalar Gaussian distribution over “all” neural parameters.

To clarify, the novelty is explicitly addressed in the paper at the sections of Introduction and Related Works, which is again summarized below:

o We employ the Low-Rank Adaptation (LoRA) method, introduced by Hu et al. (2021), to significantly reduce the number of trainable parameters, and scale up the Kalman algorithm, enhancing its compatibility for fine-tuning of large models.

o We adopt a diagonal approximation for the covariance matrix P to reduce the computational overhead from quadratic to linear and take advantage of the exponential moving average (EMA) for estimating the matrix R to improve performance without additional computational cost.

o We conduct various experiments to demonstrate LoKO’s performance in online fine-tuning tasks across computer vision and language modeling domains.

o To the best of our knowledge, this is the first successful attempt to fine-tune large models, including transformers with millions of parameters and complex architectures, using the Kalman filter algorithm.

We would like to emphasize that all previous related works and their relation to our work have been clearly explained and cited, including works with diagonal approximation: [1], [2], and [3] that is also mentioned in your comment. If there are any related or similar works that highlight areas of lower novelty in our research, we would greatly appreciate it if the reviewer could kindly provide specific references.