Kalman Filter for Online Classification of Non-Stationary Data

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. What are the real-time computational costs associated with each algorithm?

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.

We would like to thank reviewer VM8e for the feedback. Please find our detailed answer below.

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. How does the proposed method compare to related Kalman filter-based techniques for online continual learning and adaptation, such as KF in [1]?

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.

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. Does the assumption concerning linear stochastic dynamics align with the outcomes presented in Table 1, given that the "Backbone finetuning" method outperforms the "No backbone finetuning (Purely linear model)"?

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.

We would like to thank reviewer NkkH for feedback. Please find our detailed response below.

The improvement over ER++ is marginal.

We would like to highlight that we compared our method to ER++ in two scenarios: stationary CIFAR-100 and CLOC (highly non-stationary classification problem). We saw marginal improvement over ER++ in stationary CIFAR-100, see Table 1. In fact, we didn't see much performance difference across methods. In the text we argued that this is due to the fact that this problem is stationary, therefore there is not much benefit in modeling non-stationarity. After that, we compared our method to ER++ on CLOC, see Figure 3. From this Figure we can see that our method offers a significant improvement over ER++. We hope this clarifies that our method is indeed significantly better than ER++ in the non-stationary regime that we most care about.

The experiments do not test against the best possible Kalman filter (in terms of the system matrices), or some restarted version of the best possible Kalman filter.

We are not sure what the reviewer means by "best possible Kalman filter". Please note that our paper is not about finding the best possible Kalman Filter nor to do Kalman Filter (KF) research. Our contribution can be summarized as follows. We propose to certain type of KF on the last layer of the Neural Network, that crucially is fully scalable in the large scale classification application and it allows modeling non-stationarity through online learning of single scalar parameter $\gamma$. We propose a procedure to learn non-stationarity parameters online as well as to finetune / learn the Neural Network backbone. Finally, we provide extensive empirical evidence of the effectiveness of our approach in practice. In contrast to our scalable KF method, other more standard KF methods having more complex systems matrices will not be scalable in our application since typically will have quadratic cost wrt the number of outputs/classes; see also responses to reviewers BEuY and VM8e.

Have you considered extending the work of Kozdoba et al (https://doi.org/10.1609/aaai.v33i01.33014098) showing that Kalman filters can be approximated arbitrarily closely using ARMA models? This could reduce the runtime of O(m^2) to O(d) for recursion depth of d.

We would like to thank you for suggesting this paper. It is not yet clear for us how exactly to extend this work to our case. The setting of the suggested papers studies the special class of linear models - ARMA models. In our paper, we work with Neural Networks and use the Kalman Filter in the last layer. On top of that we learn the dynamics parameters (gamma) online and finetune Neural Network representation. This is not entirely clear for us how to connect our paper to this suggested paper. The scope of our paper is not about Kalman Filter research, but rather of using Kalman Filter tools in the large scale non-stationary classification problem.

How exactly do you envision to use "reparametrize the integral to be an expectation under the standard normal and then apply Monte Carlo"? Plausibly, the scaleability of Monte Carlo could be a constraint? Could you present the statistical performance and runtime of the precomputing and the actual application of the method parametrized by S, as in the O(KS) or O(Km + m^2)?

As we have explained in the paper, see equation 11, the Monte-Carlo approximation is done over 1 dimensional scalars (which are phi(x)^t w) per class. Therefore the MC operations are very cheap. In practice, we found that S = 50 worked very well in all the considered settings. In practice, we found little difference in performance using smaller values for S, and because these operations are very cheap, we are not compute bound and could use even such a high S as 50.

We would like to thank the reviewer reom for the feedback. Please find our detailed answer below.

The paper is primarily focused on learning in the case of having a pre-trained extractor of a representation. Training a representation extractor can require considerable offline computational expenses. Also, having a fixed representation extractor does not ensure plasticity of the weights of the extractor itself, which may be important in cases where there is sufficient heterogeneity in the data over time that this becomes an issue. It is unclear how much of an issue this can be in practice, and therefore a comment, or some justification, should be made on the implicit assumption that plasticity in the last layer is sufficient.

You are absolutely right, having non-stationarity modeled only in the last layer is more restrictive compared to modeling it in all the layers. If we use pre-trained frozen representation and the distribution shift is significant, our approach might fail to capture non-stationarity. The fact that we can update both, representation and the parameter gamma, allows us to be more flexible and capture larger classes of non-stationarity. However, we believe that this only partially mitigates the issue with plasticity. If we were to use this method in settings with extreme plasticity loss, then it is likely that our approach may not work well and more would need to be done with neural network representation. However, we are not framing our method as the solution to the plasticity problem but rather as a practical algorithm to learn in non-stationary classification domains, and in the scenarios we studied the method, it seems to work reasonably well. It would be interesting to study the performance of the method in settings with significant plasticity loss. That's said, we will add a comment about that in the text and we thank you for this useful suggestion.

As presented, the methodology does not allow for a simple incorporation of additional weight blocks into the Kalman filter update, as this breaks the linear Gaussian assumption that the method relies on. This can perhaps be circumvented with a linearization of the observation model. Have the authors considered extensions of their method in this direction? This would lead to the additional question of what weights to learn with SGD and which weights to learn with the Kalman filter (intuitively, the ones requiring the most plasticity), and how to select them.

This is a very interesting suggestion. We did think about linearisation of the observation model. In fact, the paper Chang et al. does exactly that. The drawback of their approach is that they lose scalability (their complexity scales as O(P C^2), where P is the number of NN parameters and C is the number of classes) as well as that it introduces additional approximation error due to linearisation. However, what you suggest goes beyond a simple Extended Kalman Filter (EKF) approach. In fact, what you suggest is to have some set of parameters to be used for EKF (if it's not only the last layer) and the other to be used with a normal SGD. It is unclear for us how to best select such subsets and whether it could be done in a tractable and automatic way, but this gives us a great food for thought. We will think about how to do that in the future work. Nevertheless, we will add a comment about it in the paper.

We would like to thank reviewer BEuY for the feedback. Please find our detailed response below.

[...] not clear while reading the Related work section how using Kalman Filters in the representation space impacted the results compared to EKF[...], especially when the representation network is trained from scratch?

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

further explanation of the backbone finetuning in the main text would be beneficial...

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.

A code implementation was not provided...

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.

Some visualizations could be improved for better readability...

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