Paper: Kalman Filter for Online Classification of Non-Stationary Data (ICLR-2024)

Reviewer comment (Novelty):
This paper could benefit from a comparison with related works that utilize the Kalman filter for online continual learning and adaptation tasks, as seen in [1]. Such a comparison or discussion regarding the difference between the proposed method and [1] would enhance the paper's comprehensiveness.

Author response:
We would like to thank the reviewer for suggesting this paper. Another relevant paper is also based on Extended Kalman Filter (EKF) from Chang et al., which we discuss in our paper. The drawback of the paper you proposed and this other EKF work, consists in fact that these approaches scale poorly with problem dimensionality. EKF from Chang et al. has complexity $O(PC^2)$, where P is the number of parameters and C is the number of classes, so it is not  applicable in situations where C is large such as CLOC. Note that our method has complexity $O(mK+m^2)$ on top of online SGD complexity (or $O(PC+P^2)$  using Chang et al notation) so in practice runs as fast as online SGD which has $O(P K)$ cost.  Further, the paper you mentioned in [1], uses an ordinary EKF approach with additional modifications, which scales as $O(K^3+PK^2+P^2K)$. It is easy to see from the algorithm 2, where $H_t$ is the matrix of size $K \times P$ and $P_t$ is the matrix of size $P \times P$. Line 2 requires $H_t*P_{t-1}*H_t^T$ multiplication which is done in $O(KP^2+K^2 P)$. After that we need to invert the matrix of size $K \times K$ which takes $O(K^3)$. Therefore the total running complexity is $O(K^3+PK^2+P^2K)$. For large $K$ the term is dominated by $O(K^3+PK^2)$ and for large $P$ it is dominated by $O(P^2K)$. Moreover, the method requires us to store the matrix $P$ as well as to store (during computations) $H_t$, this results in $O(P^2+KP)$ storage complexity. Note that in Neural Networks settings, P is of the order of millions / billions. Our algorithm requires us to store $O(m^2+Km)$ which corresponds to the last layer covariance matrix and per-class mean vectors. Therefore, our algorithm is significantly cheaper than this variant, both in terms of computation time and memory. Even if we were to apply this approach to the last layer, we would still get a very large complexity of $O(K^3+mK^2+m^2 K)$, which would make this approach less practical than ours. Based on this thinking, comparing to the proposed approach might not be very informative, since we are also primarily interested in scalability.

Reviewer comment (Theory):
The paper assumes that the ground-truth model follows a linear predictor based on features, followed by adapting the last layer of the neural network using KF. The results in Table 1 suggest that the "Backbone finetuning" method outperforms "No backbone finetuning (Purely linear model)." This observation raises questions about the validity of the assumption regarding linear stochastic dynamics. It is essential to engage in an in-depth discussion regarding the correctness and practicality of these assumptions.

Author response:
We formulate the method assuming this linear model, but later on, we indeed, violate this assumption. Nevertheless, we still use the KF updates assuming the linear model and we see that in practice this works very well. That said, what we are doing is not completely unreasonable. If the model is fully linear, then the KF updates are EXACT, meaning that we have true posterior over the last layer at time t. If the model parameters change (thetas or gammas), the updates stop being exact, but become approximate and we in fact propagate the approximate posterior. Such a setting is well studied and is known as online Bayesian learning. Without explicitly saying that, this is in fact what is happening in our case when the model parameters change. We will add an explicit comment about it in the text.

Reviewer comment (Evaluation):
In the evaluation of CLOC in section 4.2.2, it's crucial to consider the real-time computational cost for each algorithm, as emphasized in [2]: In the evaluation of the proposed method could be a prudent step.

Author response:
This is an important point of making the model efficient real-time. However, the real-time computational costs are a function of algorithm complexity as well as the efficient implementation. Our algorithm has an additional  $O(m^2)$ computational overhead on top of online SGD, i.e. online SGD in the final layer has cost $O(Km)$ while our method has  $O(Km+m^2)$ where K is the number of classes and m is the dimensionality of the last layer. The term $O(m^2)$ comes from a matrix multiplication needed because Kalman Filter updates a $m \times m$ covariance matrix over the weights. This is a very small additional overhead on top of the online SGD, and in fact if $K$ is much larger than $m$ the term $O(Km)$ dominates  and the cost of online SGD and our method is roughly the same. We do not aim to provide an efficient implementation of our method for real-time applications, but this could be easily incorporated into an efficient real-time implementation of Online SGD.

Reviewer comment (Experiments):
It was not clear while reading the Related work section how using Kalman Filters in the representation space impacted the results compared to EKF on the full network as in Chang et al. I understood that the method presented here is more efficient, but how is the performance of the model in CL tasks impacted by only considering the last layer in the filter, especially when the representation network is trained from scratch?

Author response:
Our method (last-layer KF) and EKF are quite different. EKF linearises the model likelihood $p(y|x,\theta)$ wrt $\theta$ around the previous estimate $\theta_t$, and then uses the KF to update mean and variance for the distribution of $\theta$. These parameters represent all the parameters of the Neural Network (NN). In principle, the EKF approach could represent a richer non-stationarity structure because it models it in all the NN parameters. On the other hand, as a drawback, the approximation induced by the linearisation might be problematic. The second drawback is that the naive version of EKF scales $O(P^3)$ where P represents the total number of NN parameters. Chang et al. present a low-rank plus diagonal approximation for the precision matrix which renders the complexity of the algorithm to be $O(PC^2)$, where C is the number of classes. Therefore, even using this approximation becomes problematic when applying this technique to very large-scale classification problems. Moreover, EKF replaces gradient descent on the NN parameters by KF updates applied to the linearised model. This makes it hard to combine EKF with many optimization algorithms developed for NNs. Our method uses KF only in the last layer and then updates the representation using gradient descent on parameters through backpropagation. Since it is only the last layer, this method in principle could model a limited set of non-stationarities. The advantage of the method is that it is very simple, it does not require linearizing the system and it could be used together with many powerful algorithms for NN optimization. Moreover, as we show in the paper, our method allows us to learn the parameter controlling the drift (parameter gamma) in the last layer online, allowing the system to adapt to the non-stationarity in the data distribution. This makes our approach quite flexible. In many practical problems, it might be enough to model non-stationarity in the last layer only, but in order to conclude, an additional empirical study is required.
Finally, there exists a number of works which have shown that it is not necessary to reap the benefits of Bayesian Inference in NNs to be Bayesian wrt all parameters, and often just being Bayesian with the last layer works surprisingly well. Some papers are https://proceedings.mlr.press/v206/sharma23a.html,  https://proceedings.mlr.press/v139/daxberger21a.html, https://arxiv.org/abs/1502.05700

Reviewer comment (Experiments):
A further explanation of the backbone finetuning in the main text would be beneficial to understanding the Experiments section, since the impact of finetuning the representation neural network $\phi$ was often even greater than the impact of tuning the forgetting parameter $\gamma$, but the former did not appear in the list of contributions and was explained very quickly in the main text.

Author response:
Thank you for the suggestion, we will add more information in the experiments section about how we finetune the backbone. In fact, we already have some information about it in the text. In the main contributions in the introduction, we point out that finetuning backbone and showing that it works empirically, is one of the contributions. Moreover, we have Appendix C where we describe in more details how we finetune the backbone. We already have references to Appendix C in text –  in Section 2.1,(ii); in Algorithm 1. We also have Appendix F which contains additional information about finetuning backbone. The reason why we use Appendix is the space constraints. Following your suggestion, we will add a reference to Appendix C in the experiments section as well as a sentence about the impact of fine-tuning backbone.

Reviewer comment (Reproducibility):
A code implementation was not provided within the supplementary material.

Author response:
Unfortunately, we do not have an open-source code for this project. As an intermediate solution, we will write a more detailed pseudo-code in python format in Appendix, such that it is easier to reproduce.

Reviewer comment (Experiments):
Claims of efficiency in terms of computation were not supported by empirical results. The comparison to Chang et al. is supported by a theoretical analysis, but a comparison to the baselines Online SGD and ER++ would further support the claims of the method achieving improvements in accuracy metrics without sacrificing computational efficiency, especially when the backbone neural network is finetuned.

Author response:
Online SGD on the final layer costs $O(Km)$ (i.e. for computing analytically the gradients over last layer weights) while KF costs $O(Km+m^2)$,  i.e. O(Km)  for updating the posterior Gaussian mean and $O(m^2)$ for updating the covariance matrix, where K is the number of classes and m is the dimensionality of the last layer. So essentially KF has an additional $O(m^2)$ cost due to the need to update an $m \times m$ covariance matrix, and for large number of classes K when $O(Km)$ dominates there is no noticeable computational difference between the methods.  When finetuning the backbone then we need one backpropagation step for both methods and the cost of KF and Online SGD for this step is the same. Finally, we pay an additional cost when we finetune gamma, which requires taking the gradient over a single scalar parameter. This cost can be amortized together with backpropagation for updating the backbone. Therefore, the overall cost of our method is $O(Km+m^2+B)$, where B is the cost of backpropagation. ER++ pays an additional cost of backpropagation on the replay data as well as the additional memory cost required to store the data in the replay buffer. Given the backbone network our method has an additional memory overhead of $O(Km+m^2)$,  i.e. it requires storing a K x m matrix for the posterior means and an m x m matrix for storing the shared-across-classes covariance matrix for the last layer parameters. Online SGD just needs to store an extra K x m matrix for the last layer weights, so KF just has only an extra $O(m^2)$ storage overhead compared to Online SGD without replay. Therefore, overall our method has only a minor extra overhead compared to Online-SGD or ER++, and at the same time it has performance and data efficiency benefits. At prediction, the cost of our method is equivalent to Online SGD.

Reviewer comment (Presentation):
Some visualizations could be improved for better readability. For instance, Figure 2 has very thin lines in the legend and plot which makes it difficult to see what is happening.

Author response:
Thank you for this suggestion. We will improve the visibility of Figure 2.