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

## Abstract

Abstract Training large models with millions or even billions of parameters from scratch incurs substantial computational costs. Parameter Efficient Fine-Tuning (PEFT) methods, particularly Low-Rank Adaptation (LoRA), address this challenge by adapting only a reduced number of parameters to specific tasks with gradient-based optimizers.
In this paper, we cast PEFT as an optimal filtering/state estimation problem and present Lo w-Rank K alman O ptimizer ( LoKO ) to estimate the optimal trainable parameters in an online manner. We leverage the low-rank decomposition in LoRA to significantly reduce matrix sizes in Kalman iterations and further capitalize on a diagonal approximation of the covariance matrix to effectively decrease computational complexity from quadratic to linear in the number of trainable parameters.
Moreover, we discovered that the initialization of the covariance matrix within the Kalman algorithm and the accurate estimation of the observation noise covariance are the keys in this formulation, and we propose robust approaches that work well across a vast range of well-established computer vision and language models. Our results show that LoKO converges with fewer iterations and yields better performance models compared to commonly used optimizers with LoRA in both image classifications and language tasks.
Our study opens up the possibility of leveraging the Kalman filter as an effective optimizer for the online fine-tuning of large models.

## 1 Introduction

The widespread adoption of deep neural networks, particularly large models, across various fields is mainly driven by the pre-training on extensive datasets followed by task-specific fine-tuning (Han et al., ). In recent years, the concept of online fine-tuning has attracted significant attention across diverse domains, ranging from robotics
(Yang et al., ; Fang et al., ; Julian et al., ), reinforcement learning
(Zheng et al., ; Nakamoto et al., ), computer vision (Gao et al., ; Kang et al., ), and natural language processing (Fan et al., ). Online fine-tuning refers to the process of continually updating a pre-trained model’s parameters as new data in a temporal stream of data becomes available, typically during deployment or in real-time applications.
As online full fine-tuning requires substantial computational resources and can negatively impact the generalization capabilities(Han et al., ), Parameter-Efficient Fine-Tuning (PEFT) techniques freeze the majority of the model parameters and selectively update a smaller subset.
Among PEFT approaches, the Low-Rank Adaptation (LoRA) technique has recently been widely recognized due to its efficient adaptation with low computational overhead (Han et al., ).
Many extensions to LoRA have also been proposed to improve its learning capacity and training stability (Liu et al., ; Zhang et al., ; Renduchintala et al., ; Dettmers et al., ; Valipour et al., ).

In these PEFT approaches, stochastic gradient descent (SGD) has been the dominant method for the parameter optimization.
The adaptive first-order optimizers, such as Adam (Kingma & Ba, ) and its variants, have demonstrated superior performance compared to traditional SGD ones(Ruder, ), as evidenced by their widespread adoption in LoRA extensions. Despite the simplicity and efficacy of these gradient-based optimizers, they exclusively rely on first-order derivatives, which may result in (sub-)optimal convergence and inefficient optimization (Reddi et al., ).

In contrast, many previous works (Singhal & Wu, ; ; Puskorius & Feldkamp, ; Williams, ; Heimes, ; Rivals & Personnaz, ) showed that recursive Extended Kalman Filter (EKF) algorithm – a method for state estimation of nonlinear systems using a data stream introduced by Kalman & Bucy () – can optimize relatively small models with performance surpassing the gradient-based counterparts.
This achievement remained obscure until when Ollivier () theoretically demonstrated that EKF is effectively equivalent to online natural gradient descent. This offers a new perspective on advanced optimization: instead of relying on complex techniques, we can recursively infer the optimal parameters as state estimation.
However, despite the commendable performance of the EKF in training neural models,
its practicality has been hampered, especially with the advent of large models.
Pariticularly, the EKF algorithm involves several sequential operations, including Linearization, Prediction, and Update steps,
all of which entail significant computational overhead.
The size of the crucial covariance matrix in EKF can even grow quadratically with the number of model parameters involved in training.

In this paper, we leverage the low-rank decomposition technique from LoRA to reduce the number of trainable parameters in specific layers.
We show that the EKF algorithm can be particularly useful for fine-tuning as fine-tuning only involves a small portion of model parameters.
We introduce LoKO, a Kalman-based scalable training algorithm, as an alternative to advanced optimizers for online fine-tuning of large models.
Particularly, our contributions are:

- •
We employ the Low-Rank Adaptation (LoRA) method, introduced by Hu et al. (), to significantly reduce the number of trainable parameters, and scale up the Kalman algorithm, enhancing its compatibility for fine-tuning of large models.
- •
We adopt a diagonal approximation for the covariance matrix ${\bm{P}}$ to reduce the computational overhead from quadratic to linear and take advantage of the exponential moving average (EMA) for estimating the matrix ${\bm{R}}$ to improve performance without additional computational cost.
- •
We conduct various experiments to demonstrate LoKO’s performance in online fine-tuning tasks across computer vision and language modeling domains.
- •
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.

In summary, LoKO shows outstanding performance on computer vision and language modeling benchmarks: MNIST (LeCun et al., ), CIFAR-10/100 (Krizhevsky et al., ), ImageNet100 (Vinyals et al., ), SST-2 (Socher et al., ), COLA (Warstadt, ), MRPC (Dolan & Brockett, ).
This paper contributes to ongoing efforts to develop a more efficient and fast optimization algorithm for online fine-tuning of increasingly more complex large models.

## 2 Related work

#### Parameter-Efficient Fine-Tuning (PEFT):

Traditional full fine-tuning typically demands significant computational resources and may damage the model’s generalization ability, occasionally leading to catastrophic forgetting (Han et al., ). In contrast, Parameter-Efficient Fine-Tuning (PEFT) efficiently freezes a portion of the parameters while updating a reduced number of trainable ones to mitigate these issues. Three primary approaches for PEFT are commonly used: plug-in adapters, parameter freezing and model reparameterization.
Plug-in adapters refer to the techniques of introducing an extra trainable adapter module to the pre-trained model, as demonstrated in works such as (Chen et al., ; Gao et al., ).
Parameter freezing is to freeze selected model parameters and update only a targeted subset, like BitFit (Zaken et al., ), Child-Tuning (Xu et al., ), IA^3 (Liu et al., ), and FISH Mask (Sung et al., ).
Model reparameterization, as the name suggests, reparameterizes model parameters using typically the Low-Rank Adaptation (LoRA) technique,
which adds low-rank weight matrices as trainable parameters(Hu et al., ).
Multiple extensions to LoRA have been proposed to improve learning capacity and training stability.
For instance, DoRA (Liu et al., ) decomposes pre-trained weights into magnitude and direction components to minimize the number of trainable parameters more efficiently.
AdaLoRA(Zhang et al., ) utilizes singular value decomposition (SVD) to dynamically allocate the parameter budget based on importance scoring.
Tied-LoRA(Renduchintala et al., ) leverages weight tying and selective training to reduce the number of trainable parameters further.
To optimize the memory efficiency, QLoRA (Dettmers et al., ), QA-LoRA (Xu et al., ), and LoftQ (Li et al., ) address the issue of memory usage by quantization technique. DyLoRA (Valipour et al., ) is a dynamic search-free LoRA to avoid exhaustive search for the most optimal rank.
However, all these extensions utilize the Adam optimizer or its variants like AdamW in parameter optimization.

#### Kalman Filter for Optimizing Neural Networks:

The idea of using the Kalman filter as parameter optimization in deep learning comes from Singhal & Wu (), which showed that the process of training neural networks can be conceptualized as tackling a system identification challenge for a nonlinear dynamic system, and thus Extended Kalman Filter (EKF) can be used to train neural network parameters.
The superior performance of the Kalman-based training algorithm compared to the traditional backpropagation technique drew attention to exploring the relationship between these two classical algorithms (Ruck et al., ; Ollivier, ).
To make the Kalman filter more scalable, several studies have addressed the computational complexity associated with this training algorithm,
notably using matrix partitioning techniques (Shah & Palmieri, ; Shah et al., ; Puskorius & Feldkamp, ), and a low-dimensional (block-)diagonal approximation of the covariance matrix (Murtuza & Chorian, ).
Only recently,
Ollivier (; ) demonstrated that training with a Kalman filter is, in fact, equivalent to an online natural gradient descent.
This finding renewed interest in this training method once again.
Various studies have since addressed the scalability challenges of the EKF optimizer.
For example, Chang et al. () introduced a diagonal Gaussian approximation, while Chang et al. () proposed a low-rank plus diagonal decomposition of the posterior precision matrix.
Hennig et al. () developed a matrix-free iterative algorithm to further enhance efficiency.
In the context of factorization models, Gómez-Uribe & Karrer () introduced a decoupled EKF (DEKF).
Furthermore, EKF has been applied to several specialized areas, such as continual learning (Titsias et al., ), test-time adaptation (Schirmer et al., ), and reinforcement learning (Shashua & Mannor, ; ; ; Totaro & Jonsson, ; Shih & Liang, ).
Other notable work includes loss adaptivity in the Kalman optimization algorithm (Davtyan et al., ), the Bayesian online natural gradient (Jones et al., ), and approaches for handling nonstationary data in online learning (Jones et al., ; ).
However, the practical implementation of these methods remained infeasible for large models.

## 3 Preliminaries and Background

### 3.1 Extended Kalman Filter (EKF)

Consider a general state-space model:

$$ $\displaystyle{\bm{s}}_{k}$ $\displaystyle=f({\bm{x}}_{k},{\bm{s}}_{k-1})+{\bm{w}}_{k},$ (1a) $\displaystyle{\bm{y}}_{k}$ $\displaystyle=h({\bm{x}}_{k},{\bm{s}}_{k})+{\bm{v}}_{k},$ (1b) $$

where ${\bm{s}}_{k}\in\mathbb{R}^{n}$, ${\bm{y}}_{k}\in\mathbb{R}^{m}$, and ${\bm{x}}_{k}\in\mathbb{R}^{p}$ denote the states, measurement, and input vectors at time $k$, respectively.
In this framework, the function $f(\cdot)$ is called the process model or the state transition function, and $h(\cdot)$ is the measurement model or observation function.
In addition, ${\bm{w}}_{k}\sim{\mathcal{N}}(0,{\bm{Q}}_{k})$ and ${\bm{v}}_{k}\sim{\mathcal{N}}(0,{\bm{R}}_{k})$ represent process noise and observation noise, respectively.
These noises are assumed to have known distributions, typically white Gaussian noise with zero mean.
The problem of state estimation for a nonlinear system, as depicted by equation 1a, can be addressed through the well-established recursive EKF algorithm (Welch et al., ). In
the following, we detail the various steps involved in implementing the extended Kalman filter (EKF):

- •
Prediction:
$\displaystyle{\bm{s}}_{k|k-1}$
$\displaystyle=f({\bm{x}}_{k},{\bm{s}}_{k-1})$
(2a)
$\displaystyle{\bm{P}}_{k|k-1}$
$\displaystyle={\bm{F}}_{k}{\bm{P}}_{k-1}{\bm{F}}_{k}^{\top}+{\bm{Q}}_{k}$
(2b)
where ${\bm{s}}_{k|k-1}$ and ${\bm{P}}_{k|k-1}$ are predicted (or prior) states and covariance, respectively. ${\bm{F}}_{k}$ denotes the Jacobian matrix of the function $f(\cdot)$ with respect to states at time $k$.
- •
Updating:
$\displaystyle{\bm{K}}_{k}$
$\displaystyle={\bm{P}}_{k|k-1}{\bm{H}}_{k}^{\top}({\bm{H}}_{k}{\bm{P}}_{k|k-1}
{\bm{H}}_{k}^{\top}+{\bm{R}}_{k})^{-1}$
(3a)
$\displaystyle{\bm{s}}_{k}$
$\displaystyle={\bm{s}}_{k|k-1}+{\bm{K}}_{k}({\bm{y}}_{k}-h({\bm{x}}_{k},{\bm{s
}}_{k|k-1}))$
(3b)
$\displaystyle{\bm{P}}_{k}$
$\displaystyle={\bm{P}}_{k|k-1}-{\bm{K}}_{k}{\bm{H}}_{k}{\bm{P}}_{k|k-1}$
(3c)
where ${\bm{K}}_{k}$ is the Kalman gain, and ${\bm{H}}_{k}$ indicates the Jacobian matrix of the function $h(\cdot)$ with respect to the states at time $k$. Finally, ${\bm{s}}_{k}$ and ${\bm{P}}_{k}$ are updated (or posterior) states and covariance matrix.

### 3.2 Low-Rank Adaptation (LoRA)

Low-Rank Adaptation (LoRA) (Hu et al., ) is a technique designed to efficiently fine-tune large pre-trained neural networks by reducing the number of trainable parameters.
Instead of updating the entire pre-trained weight matrix, LoRA introduces a low-rank decomposition that captures the essential changes needed for fine-tuning. Consider a layer in a neural network with a pre-trained weight matrix ${\bm{W}}_{0}\in\mathbb{R}^{d\times q}$, where $d$ and $q$ represent the dimensions of the weight matrix.
The output of this layer can be expressed as:

$$ ${\bm{z}}={\bm{W}}_{0}{\bm{x}}+\Delta{\bm{W}}{\bm{x}}$ (4) $$

Here, ${\bm{x}}$ is the input vector, and $\Delta{\bm{W}}$ represents the adjustment to the weights during fine-tuning.
LoRA modifies this by introducing two smaller matrices, ${\bm{A}}\in\mathbb{R}^{r\times q}$ and ${\bm{B}}\in\mathbb{R}^{d\times r}$, where $r$ is the chosen rank with $r\ll\min(d,q)$. The update to the weight matrix is then expressed as:

$$ ${\bm{z}}={\bm{W}}_{0}{\bm{x}}+{\bm{B}}{\bm{A}}{\bm{x}}$ (5) $$

At the beginning of training, the matrix ${\bm{A}}$ is initialized with a random Gaussian distribution ${\mathcal{N}}(\mathbf{0},\sigma^{2}{\bm{I}})$, and matrix ${\bm{B}}$ is initialized to zero.
Using this low-rank decomposition, LoRA reduces the number of trainable parameters from $d\times q$ to $r\times(d+q)$, where $r$ is much smaller than $d$ and $q$. This technique can be applied on certain layers of a neural network model, $h({\bm{x}},{\bm{\theta}})$, resulting in a re-parameterized version, $h_{LoRA}({\bm{x}},\tilde{{\bm{\theta}}})$, with a reduced number of trainable parameters. This reduction enables efficient fine-tuning of large models, making it feasible to apply the Kalman filter.

## 4 Low-rank Kalman Optimizer

### 4.1 Kalman Formulation for LoRA

Consider a pre-trained model in which certain layers are scaled down using the LoRA method. The modified model is parameterized by a reduced set of trainable parameters $\tilde{{\bm{\theta}}}_{k}\in\mathbb{R}^{\tilde{n}}$, where $\tilde{n}\ll n$:

$$ $\hat{{\bm{y}}}_{k}=h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k}).$ (6) $$

Here, ${\bm{x}}_{k}\in\mathbb{R}^{p}$ denotes the model input and $\hat{{\bm{y}}}_{k}\in\mathbb{R}^{m}$ represents the predicted output in the $k^{th}$ observed data. Let us adopt ${\bm{y}}_{k}$ as the true output. The $\hat{{\bm{y}}}_{k}$ represents the predicted values for the regression tasks and the predicted probabilities for the classification tasks. In both scenarios, the predicted output $\hat{{\bm{y}}}_{k}$ can be interpreted as the mean parameter of a Gaussian distribution over the actual output ${\bm{y}}_{k}$. This relationship can be expressed as white noise as ${\bm{v}}_{k}\sim{\mathcal{N}}(\mathbf{0},{\bm{R}}_{k})$, where ${\bm{R}}_{k}=\mathrm{Cov}({\bm{y}}_{k}|\hat{{\bm{y}}}_{k})$. More broadly, $\hat{{\bm{y}}}_{k}$ serves as the mean parameter for an exponential family distribution over $T({\bm{y}}_{k})$, where $T(\cdot)$ denotes the sufficient statistics for the exponential family. In this case, ${\bm{R}}_{k}=\mathrm{Cov}(T({\bm{y}}_{k})|\hat{{\bm{y}}}_{k})$ denotes the covariance matrix of the exponential family distribution.

Several approximations for the matrix ${\bm{R}}_{k}$ have been proposed in the literature. One common approach is to approximate it as the identity matrix, ${\bm{R}}_{k}\approx{\bm{I}}$, as shown in (Puskorius & Feldkamp, ; Murtuza & Chorian, ). Other formulations include ${\bm{R}}_{k}={\bm{I}}\cdot e^{-k/50}$(Singhal & Wu, ; ), and a more recent approximation, ${\bm{R}}_{k}=\text{diag}(\hat{{\bm{y}}}_{k})-\hat{{\bm{y}}}_{k}\hat{{\bm{y}}}_
{k}^{\top}$(Ollivier, ; Chang et al., ).
To obtain a more precise approximation of ${\bm{R}}_{k}$, we employ an Exponential Moving Average (EMA) approach based on the definition of the covariance matrix for estimating the matrix ${\bm{R}}_{k}$.
We make the simplifying assumption that $\tilde{{\bm{\theta}}}_{k}\approx\tilde{{\bm{\theta}}}_{k|k-1}$, allowing us to compute the covariance matrix as follows:

$$ $\displaystyle{\bm{R}}_{k}$ $\displaystyle=\beta{\bm{R}}_{k-1}+(1-\beta)\hat{{\bm{R}}}_{k},$ (7a) $\displaystyle\text{where}\quad\hat{{\bm{R}}}_{k}$ $\displaystyle=\left({\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{ k|k-1})\right)\left({\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{ k|k-1})\right)^{\top},$ (7b) $$

and $\beta\in(0,1)$ is the forgetting factor. The proofs and detailed derivations can be found in Appendix B for more details.

LoKO estimates the (sub-)optimal values of LoRA matrices ${\bm{A}}$ and ${\bm{B}}$ in real-time data streams through the Kalman algorithm. To formulate the online fine-tuning problem within the Kalman filtering framework, we assume there are no feedback loops in machine learning model, for example, a feed-forward neural network and transformers. This assumption enables us to consider the trainable parameters that are represented as the state vector of the process model ($\tilde{{\bm{\theta}}}_{k}\equiv{\bm{s}}_{k}$). Furthermore, the process model (or transition function) can be modeled as an identity function $f({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k-1})=\tilde{{\bm{\theta}}}_{k-1}$ with no process noise (${\bm{w}}_{k}=0$), which provides the prediction of the states at the next time step as $\tilde{{\bm{\theta}}}_{k}=\tilde{{\bm{\theta}}}_{k-1}$. Finally, by defining ${\bm{y}}_{k}=h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k})+{\bm{v}}_{k}$ as the measurement model (or observation function), we can apply the recursive Kalman filter algorithm to estimate (sub-)optimal values of $\tilde{{\bm{\theta}}}_{k}$.

Although the low-rank decomposition by LoRA offers a significant reduction in parameter size compared to the original parameter space $n$, the size of the covariance matrix ${\bm{P}}$ scales quadratically with the number of trainable parameters $\tilde{n}$. For large models such as deep neural networks, characterized by high-dimensional trainable parameters, the computational cost of implementing the Kalman algorithm becomes prohibitively expensive due to its $\tilde{n}^{2}$ complexity.
To address this challenge, one strategy involves decoupling the update phase of the Kalman filtering algorithm into smaller partitions, as outlined by Puskorius & Feldkamp (), which, however, may be infeasible for large models with millions of trainable parameters.
The other approach is to approximate the covariance matrix ${\bm{P}}$ with low-dimensional matrices.
Our empirical findings demonstrate that as the fine-tuning algorithm progresses, the covariance matrix of the feed-forward neural network asymptotically approaches a (block-)diagonal configuration:

$$ $\mathbb{E}[\hat{{\bm{p}}}]=\text{diag}({\bm{P}}).$ (8) $$

Therefore, we adopt a diagonal approximation of the covariance matrix, denoted as $\hat{{\bm{p}}}$. This approximation significantly reduces both computational and storage costs.

###### Proposition 1 .

Leveraging the low-rank decomposition technique in LoRA and applying the diagonal approximation of covariance matrix, the steps of the Low-Rank Kalman Optimizer (LoKO) can be outlined below:

- •
Prediction:
$\displaystyle\tilde{{\bm{\theta}}}_{k|k-1}$
$\displaystyle=\tilde{{\bm{\theta}}}_{k-1}$
(9a)
$\displaystyle\hat{{\bm{p}}}_{k|k-1}$
$\displaystyle=\hat{{\bm{p}}}_{k-1}$
(9b)
- •
Pre-Updating:
$\displaystyle{\bm{R}}_{k}=\beta{\bm{R}}_{k-1}+(1-\beta)\hat{{\bm{R}}}_{k}$
(10a)
- •
Updating:
$\displaystyle{\bm{K}}_{k}$
$\displaystyle=\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^{\top}\left({\bm{H}}_{
k}(\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^{\top})+{\bm{R}}_{k}\right)^{-1}$
(11a)
$\displaystyle\tilde{{\bm{\theta}}}_{k}$
$\displaystyle=\tilde{{\bm{\theta}}}_{k|k-1}+{\bm{K}}_{k}({\bm{y}}_{k}-h_{LoRA}
({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k|k-1}))$
(11b)
$\displaystyle\left(\hat{{\bm{p}}}_{k}\right)^{i}$
$\displaystyle=\left(\hat{{\bm{p}}}_{k|k-1}\right)^{i}-\left({\bm{K}}_{k}\right
)^{i}_{j}\left({\bm{H}}_{k}\right)^{j}_{i}\left(\hat{{\bm{p}}}_{k|k-1}\right)^
{i}$
(11c)
where the symbol $\bullet$ represents the transposed Khatri–Rao product, which is essentially the row-by-row Kronecker product of the vector $\hat{{\bm{p}}}_{k|k-1}$ and matrix ${\bm{H}}_{k}^{\top}$. The equation 11c represents the diagonal update of the covariance matrix, expressed with Einstein notation. For more details, see Appendix B.

Note that the operations used in equation 11a, and equation 11c can be computed efficiently. For example, in PyTorch, they can be seamlessly implemented using the $*$, and $\textbf{einsum}()$ operators, streamlining the computational process.

Importantly, in the above formulation, the matrix inversion required in the Kalman gain equation 11a has a dimension of $m^{2}$, where $m$ represents the size of the model output. This implies that the computational cost of the matrix inversion remains constant regardless of the size of the model. Consequently, whether the model is small or large, the computational expense of this operation remains constant, offering a consistent performance characteristic across different model sizes. In contrast, many advanced optimizers such as the Natural Gradient Descent (NGD) necessitate the inversion of the full pre-condition matrix (like the Fisher information matrix in NGD), which is of high dimensionality.
Moreover, the computation of the Jacobian matrix ${\bm{H}}_{k}$ does not require individual backpropagation processes for each output component. Leveraging GPU capabilities allows for parallel computation and thus efficiently streamlines the process.

###### Proposition 1 .

### 4.2 Initialization of p & Approximation of R

#### Initialization of 𝒑 ^ 0 subscript ^ 𝒑 0 \hat{{\bm{p}}}_{0} over^ start_ARG bold_italic_p end_ARG start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT :

Our findings demonstrate that the initialization of $\hat{{\bm{p}}}_{0}$ plays a crucial role in the performance of the filter as a training algorithm.
High values of $\hat{{\bm{p}}}_{0}$ indicate high uncertainty or lack of confidence in the initial learning parameters,
which can result in the Kalman filter making large corrections and experiencing potential instability and divergence.
In contrast, initialization $\hat{{\bm{p}}}_{0}$ with very low values suggests high confidence in the initial learning parameters.
In this case, the filter will heavily trust the initial model parameters. If these initial parameters are inaccurate, this can lead to slow updates or even no updates at all, as the filter gives insufficient weight to new measurements. Proper initialization of $\hat{{\bm{p}}}_{0}$ balances these extremes, ensuring that the filter can adapt appropriately to new data while maintaining stability and accuracy throughout the training process.
Therefore, it is essential to establish both upper and lower bounds for $\hat{{\bm{p}}}_{0}$. To ensure the initialization of $\hat{p}_{0}=[p_{1},p_{2},...p_{i},...,p_{n}]$ yields a positive definite diagonal matrix, we examined two methods for initialization of $\hat{{\bm{p}}}_{0}$:

- •
Method 1: Setting a constant positive value: $p_{i}=c\quad\forall i$.
- •
Method 2: Assigning random positive values drawn from a uniform distribution: $p_{i}\sim U(0,\text{upper_bound})\quad\forall i$.

Our experiments show that precisely estimating the ${\bm{R}}_{k}$ matrix greatly influences the bounds of $\hat{{\bm{p}}}_{0}$. The more accurate the estimation of ${\bm{R}}_{k}$, the broader the acceptable range for the initial $\hat{{\bm{p}}}_{0}$.

#### Approximation of 𝑹 k subscript 𝑹 𝑘 {\bm{R}}_{k} bold_italic_R start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT :

In addition to the method described in equation 7, we propose an alternative method for approximating the observation noise covariance, ${\bm{R}}_{k}$, with enhanced accuracy and without additional computation cost. Specifically, we incorporate an additional term from the first-order Taylor to approximate changes more precisely:

$$ $\displaystyle{\bm{R}}_{k}$ $\displaystyle=\beta{\bm{R}}_{k-1}+(1-\beta)\hat{{\bm{R}}}_{k},$ (12a) $\displaystyle\text{where}\quad\hat{{\bm{R}}}_{k}$ $\displaystyle=\left({\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{ k|k-1})\right)\left({\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{ k|k-1})\right)^{\top}+{\bm{H}}_{k}(\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^{ \top})$ (12b) $$

with the forgetting factor of $\beta\in(0,1)$. Note that this method will not add extra computational cost since the operation of ${\bm{H}}_{k}(\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^{\top})$ will be part of the Kalman gain calculation in equation 11a. See Appendix B for more details.

## 5 Experiments and Analysis

### 5.1 Experiments Setup

We assess the performance of LoKO by implementing it in various well-established computer vision and language models.
Our computer vision experiments involve online fine-tuning for image classification on the MNIST dataset using DenseNet-121 with 7 million parameters (Huang et al., ), the CIFAR-10/100 dataset with ResNet-18 and ResNet-50 (He et al., ), and ViT-B16 (Dosovitskiy, ), which contain 12 million, 26 million, and 86 million parameters, respectively. Furthermore, we evaluated ImageNet100 using ViT-L16 (Dosovitskiy, ), 305 million parameters. For text classification tasks, we examine the SST-2, CoLA, and MRPC datasets using RoBERTa-base and RoBERTa-large (Liu, ), which have 125 million and 355 million parameters, respectively.

We use AdamW (Loshchilov & Hutter, ) and AdaGrad (Duchi et al., ) as baseline methods since they are widely used in LoRA.
We set the batch size to 1 so that the learning algorithms train the models in the context of online classification tasks at each time step.
We configure AdamW with a linear learning rate decay schedule (weight_decay=0.0001), starting from 1e-4,
which empirically yields the highest accuracy after a preliminary sweep of various learning rates and schedules. In addition, $\beta_{1}$ and $\beta_{2}$ for this optimizer are set to 0.9 and 0.999, respectively. Furthermore, the learning rate of AdaGrad is set to 0.001 with no weight decay.
The hyperparameters in our LoKO algorithm are set according to our discussion in 5.3.
The diagonally approximated covariance vector $\hat{{\bm{p}}}_{k}$ is initialized with random numbers sampled from a uniform distribution between 0 and 0.2. Furthermore, we use our second method for $\hat{{\bm{R}}}_{k}$ estimation, and the forgetting factor for the EMA is set to $\beta=0.95$. We also adopt the pre-trained models introduced by (Caron et al., ) as the backbone for our fine-tuning experiments with ViT. As shown in Table 1, low-rank decomposition is applied on query layers in transformer-based architectures and convolution layers in CNN networks.
All experiments are conducted on an A100 GPU platform three times and the average and standard deviation of them are reported.

**Table 1: Total and trainable parameters for each model utilized in the experiments.**
| Model | LoRA layer | Num. of total parameters | Num. of trainable parameters |
| --- | --- | --- | --- |
| DenseNet-121 | convolution | 7.5 M | 166 K |
| ResNet18 | convolution | 11 M | 150 K |
| ViT-B16 | query | 86 M | 155 K |
| ResNet50 | convolution | 24 M | 520 K |
| ViT-L16 | query | 305 M | 400 K |
| RoBERTa-base | query | 125 M | 666 K |
| RoBERTa-large | query | 355 M | 1248 K |

### 5.2 Main results

We use two primary evaluation metrics in our study.
First, we calculate the moving average of the loss with a window size of 1000.
Second, we measure the average online accuracy, as adopted by Cai et al. (), up to the current timestep $k$.
The average online accuracy is defined as $\text{acc}(k)=\frac{1}{k}\sum_{i=1}^{k}\mathbbm{1}_{A}({\bm{y}}_{i})$, where $\mathbbm{1}_{\{\cdot\}}(\cdot)$ is an indicator function, and $A$ is the set of Top1 or Top5 prediction. This metric is computed online during the training process and serves to evaluate the algorithm’s capacity to incorporate new knowledge in real-time.
We evaluate the convergence speed of LoKO by comparing the number of observations required to achieve equivalent loss or accuracy levels across different algorithms.
Figure 1, 2 illustrate LoKO’s performance in online fine-tuning for image and language classification, compared to LoRA with AdamW and AdaGrad, across various well-established models and datasets. The figures illustrate *L1 losses* for images and *cross-entropy losses* for texts, along with accuracy during training. These results demonstrate that LoKO consistently matches or surpasses the performance of alternative optimizers. Additionally, Tables 2 and 3 present the average and standard deviation of online accuracy for a single epoch of online fine-tuning. Note that for CoLA and MRPC datasets, we performed multiple epochs due to their smaller dataset sizes. These results demonstrate that across all models and datasets, LoKO consistently outperforms (and in some cases performs almost equivalently) the other methods by achieving the highest average online accuracy with LoKO. The lower standard deviation in LoKO’s results also highlights its more consistent and robust performance compared to other methods. Furthermore, it is important to emphasize that (sub-)optimal values of the trainable parameters can be attained across a range of LoKO hyperparameters, provided they fall within the permissible initialization interval. This contrasts with gradient-based optimizers, whose performance is highly sensitive to the selection of appropriate hyperparameters, particularly the learning rate.

Figure: Figure 1: Performance of LoKO (blue) compared to LoRA/AdamW (red) and LoRA/AdaGrad (green) for different computer vision datasets and models. The upper rows show the training loss, and the lower rows display the average online accuracy versus the number of observed data.
Refer to caption: x1.png

**Table 2: Results for MNIST, CIFAR-10/100, and ImageNet100 datasets across various neural network models. The table presents the average and standard deviation of online accuracy achieved during a single epoch of online fine-tuning with different methods, where higher values indicate better performance. The boldfaced numbers represent the best values achieved for each column.**
| Method | MNIST | CIFAR-10 | CIFAR-100 | ImageNet100 |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| DenseNet-121 | ResNet18 | ViT-B16 | ResNet50 | ViT-B16 | ViT-L16 |  |
| LoKO (ours) | $\bm{{98.17}_{\pm 0.56}}$ | $\bm{96.05_{\pm 0.12}}$ | $\bm{99.93_{\pm 0.05}}$ | $\bm{79.39_{\pm 0.12}}$ | $\bm{91.56_{\pm 0.40}}$ | $\bm{97.71_{\pm 0.22}}$ |
| LoRA/AdamW | $94.73_{\pm 0.65}$ | $65.98_{\pm 0.52}$ | $98.98_{\pm 0.49}$ | $60.39_{\pm 0.40}$ | $85.95_{\pm 0.40}$ | $95.13_{\pm 0.13}$ |
| LoRA/AdaGrad | $86.22_{\pm 4.10}$ | $61.83_{\pm 0.10}$ | $99.17_{\pm 0.20}$ | $40.62_{\pm 0.05}$ | $80.04_{\pm 0.87}$ | $93.91_{\pm 1.28}$ |

**Table 3: Results for SST-2, COLA, and MRPC datasets across various neural network models. The table presents the average and standard deviation of online accuracy achieved during online fine-tuning with different methods, where higher values indicate better performance. For SST-2, we conducted a single epoch of training. For CoLA and MRPC, we performed multiple epochs due to their smaller dataset sizes. The bolded numbers represent the maximum values for each column.**
| Method | SST-2 | COLA | MRPC |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $RoB_{base}$ | $RoB_{large}$ | $RoB_{base}$ | $RoB_{large}$ | $RoB_{base}$ | $RoB_{large}$ |  |
| LoKO (ours) | ${88.41}_{\pm 0.2}$ | $88.21_{\pm 0.1}$ | $83.67_{\pm 0.3}$ | $82.69_{\pm 0.4}$ | $89.17_{\pm 0.3}$ | $88.70_{\pm 0.6}$ |
| LoRA/AdamW | $\bm{89.55_{\pm 0.1}}$ | $\bm{90.66_{\pm 0.2}}$ | $\bm{84.41_{\pm 0.1}}$ | $\bm{86.19_{\pm 0.9}}$ | $\bm{90.95_{\pm 0.3}}$ | $\bm{93.04_{\pm 0.4}}$ |
| LoRA/AdaGrad | $87.79_{\pm 0.2}$ | $86.25_{\pm 1.1}$ | $78.31_{\pm 0.0}$ | $79.72_{\pm 0.3}$ | $81.98_{\pm 0.1}$ | $83.88_{\pm 1.1}$ |

Figure: Figure 2: Comparison of LoKO (blue) with LoRA/AdamW (red) and LoRA/AdaGrad (green) across various language models and datasets. The top row presents the training loss, while the bottom row illustrates the average online accuracy against the number of data points observed. This figure presents results for RoBERTa-base. For comprehensive results across all models, refer to Appendix D.
Refer to caption: extracted/5925019/Figure/NLP/train_loss_roberta-base_lora_sst2.png

### 5.3 Ablation of Initialization and Covariance Approximation

To develop a comprehensive understanding of the behavior of the initialization of $\hat{{\bm{p}}}_{0}$, we begin by evaluating its impact on training from scratch using the Kalman filter. The MNIST dataset is employed for experimentation, utilizing the LeNet-5 architecture due to its fast training speed and low computational demands. Our experiments combine two proposed methods for initializing $\hat{{\bm{p}}}_{0}$ with four widely used parameter initialization strategies: Xavier Uniform and Xavier Normal (Glorot & Bengio, ), as well as Kaiming Uniform and Kaiming Normal (He et al., ).
Figure 3 illustrates the average online accuracy of the MNIST dataset after 1000 iterations of training from scratch.
As shown, when $\hat{{\bm{p}}}_{0}$ values are set far from the lower and upper bounds, the algorithm diverges, resulting in low prediction accuracy. The results indicate that initializing $\hat{{\bm{p}}}_{0}$ with random values drawn from a uniform distribution (second method)
provides a robust and wide range of choices for $\hat{{\bm{p}}}_{0}$ without causing divergence. In Table 4, we further investigate the initialization of $\hat{{\bm{p}}}_{0}$ in the online fine-tuning of various combinations of Dataset-Model using LoKO.

Figure: Figure 3: Initialization of $\hat{{\bm{p}}}_{0}$ with two different techniques: (left) Setting a constant positive value, and (right) Assigning random positive values drawn from a uniform distribution.
Refer to caption: extracted/5925019/Figure/P0-C.png

**Table 4: The lower and upper bounds for two proposed initialization methods of $\hat{{\bm{p}}}_{0}$**
| Dataset - Model | method 1 | method 2 |  |  |
| --- | --- | --- | --- | --- |
| min | max | min | max |  |
| MNIST - DenseNet-121 | 0.0001 | ${0.11}_{\pm 0.01}$ | 0.0001 | ${0.32}_{\pm 0.01}$ |
| CIFAR10 - ViT-B16 | 0.001 | ${105}_{\pm 3}$ | 0.001 | ${165}_{\pm 3}$ |
| CIFAR100 - ViT-B16 | 0.001 | ${0.2}_{\pm 0.01}$ | 0.001 | ${0.59}_{\pm 0.01}$ |
| ImageNet100 - ViT-L16 | 0.001 | ${1.0}_{\pm 0.1}$ | 0.001 | ${2.0}_{\pm 0.1}$ |

Our proposed methods for approximating ${\bm{R}}_{k}$ are presented in Table 5 alongside other estimation methods from the literature. Additionally, we have included the accuracy results on the MNIST test dataset for comparison.

**Table 5: Different approaches in the approximation of the observation noise covariance matrix**
| Ref. | Equation | Accuracy |
| --- | --- | --- |
| MNIST |  |  |
| Puskorius & Feldkamp () | $R_{k}=I$ | 92.59 % |
| Murtuza & Chorian () |  |  |
| Singhal & Wu () | $R_{k}=Ie^{-k/50}$ | 93.65 % |
| Singhal & Wu () |  |  |
| Ollivier () | $R_{k}=diag(\hat{\mathbf{y}}_{k})-\hat{\mathbf{y}}_{k}\hat{\mathbf{y}}_{k}^{\top}$ | 93.89 % |
| Chang et al. () |  |  |
| LoKO (method 1) | $R_{k}=\beta R_{k-1}+(1-\beta)\hat{R}_{k}$, | 93.81 % |
| $\hat{R}_{k}=(\mathbf{y}_{k}-\hat{\mathbf{y}}_{k})(\mathbf{y}_{k}-\hat{\mathbf{<br>y}}_{k})^{\top}$ |  |  |
| LoKO (method 2) | $R_{k}=\beta R_{k-1}+(1-\beta)\hat{R}_{k}$, | 94.51 % |
| $\hat{R}_{k}=(\mathbf{y}_{k}-\hat{\mathbf{y}}_{k})(\mathbf{y}_{k}-\hat{\mathbf{<br>y}}_{k})^{\top}+H_{k}(\hat{p}_{k|k-1}\bullet H_{k}^{\top})$ |  |  |

### 5.4 Sensitivity analaysis of β 𝛽 \beta italic_β in EMA

To assess the impact of hyperparameter $\beta$ on LoKO’s performance, we conducted a sensitivity analysis varying $\beta$ values. Through this analysis, we measured the average online accuracy across different $\beta$ settings. The results illustrate that excessively high or low $\beta$ values notably compromise LoKO’s performance. Optimal performance lies within the range of 0.9 to 0.98 for $\beta$. See Appendix C for more details.

## 6 Conclusion

In this study, we present the Low-Rank Kalman Optimizer (LoKO), a scalable Kalman-based training algorithm. LoKO offers a comparative alternative to advanced gradient-based optimizers for online fine-tuning scenarios. By leveraging the low-rank adaptation technique, and a diagonal approximation of the covariance matrix and incorporating a novel observation noise estimation technique based on EMA, we reformulate the EKF algorithm accordingly.
Empirical evaluations on the benchmark of computer vision and language models, including MNIST, CIFAR-10/100, ImageNet100, SST-2, COLA, and MRPC demonstrate that LoKO consistently achieves superiority over (or at least comparability with) commonly used optimizers, showcasing its potential for efficient online fine-tuning of large models. Although this paper focuses on online fine-tuning in computer vision and language models, further evaluations are necessary across a broader range of domains, including reinforcement learning. Due to space constraints, we leave these explorations for future work.

## Technical Appendices

## Appendix A LoKO Algorithm

The LoKO online fine-tuning method is outlined in Algorithm 1, with our modification to the EKF algorithm highlighted in red for clarity.

Figure: Algorithm 1 LoKO

## Appendix B Proofs and Detailed Derivations

This section presents the proofs and detailed derivations of the proposed LoKO algorithm, offering a comprehensive understanding of its underlying principles and mechanisms.

#### Proof for Proposition 1 :

Here, we present the derivation of equation 11a and equation 11c from Proposition 1. These equations represent our proposed method for calculating the Kalman gain and updating the covariance matrix using a diagonal approximation.

Proof of equation 11a:
Here, we will show that  equation 11a is equivalent to equation 3a under the assumption of covariance diagonal approximation. For this aim, we need to show that ${\bm{P}}_{k|k-1}{\bm{H}}_{k}^{\top}=\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^
{\top}$.

###### Proof.

For simplicity, let’s drop the subscripts of $\hat{{\bm{p}}}_{k|k-1}$, and ${\bm{H}}_{k}^{\top}$. Now, define a diagonal matrix: ${\bm{P}}=diag([p_{1},p_{2},...,p_{n}])$, where the diagonal elements $[p_{1},p_{2},...,p_{n}]$ are the elements of the vector $\hat{{\bm{p}}}$, and its off-diagonal elements are zeros. Then, the $i^{th}$ row of ${\bm{P}}{\bm{H}}^{\top}$ can be represented as: $\left({\bm{P}}{\bm{H}}^{\top}\right)_{i}=p_{i}\left({\bm{H}}^{\top}\right)_{i}$, where $\left({\bm{H}}^{\top}\right)_{i}$ represents the $i^{th}$ row of matrix ${\bm{H}}^{\top}$. Now, let’s represent the $i^{th}$ row of the transposed Khatri–Rao product of the vector $\hat{{\bm{p}}}$ and matrix ${\bm{H}}^{\top}$: $\left(\hat{{\bm{p}}}\bullet{\bm{H}}^{\top}\right)_{i}=\left(\hat{{\bm{p}}}
\right)_{i}\otimes\left({\bm{H}}^{\top}\right)_{i}$. Since $\left(\hat{{\bm{p}}}\right)_{i}=p_{i}$, and $p_{i}$ is a scalar, the equality $p_{i}\otimes\left({\bm{H}}^{\top}\right)_{i}=p_{i}\left({\bm{H}}^{\top}\right)
_{i}$ holds true. Consequently, ${\bm{P}}{\bm{H}}^{\top}=\hat{{\bm{p}}}\bullet{\bm{H}}^{\top}$.
∎

Proof of equation 11c:
Here also, we will demonstrate that under the assumption of covariance diagonal approximation, the equation 11c is equivalent to equation 3c in vanilla EKF algorithm.

###### Proof.

Again, by dropping the subscripts, let’s define a diagonal matrix: ${\bm{P}}$ whose diagonal elements are the elements of the vector $\hat{{\bm{p}}}$. Then, let’s expand the equation expressed with Einstein notation: $\left({\bm{K}}\right)^{i}_{j}\left({\bm{H}}\right)^{j}_{i}\left(\hat{{\bm{p}}}
\right)^{i}=\left({\bm{K}}\right)^{i}_{j}\left({\bm{H}}\right)^{j}_{i}\left({
\bm{P}}\right)^{i}_{i}=\left({\bm{K}}{\bm{H}}{\bm{P}}\right)^{i}_{i}$, where $\left({\bm{K}}{\bm{H}}{\bm{P}}\right)^{i}_{i}$ represents the $i^{th}$ element of the main diagonal of the matrix ${\bm{K}}{\bm{H}}{\bm{P}}$. Consequently, the equation 11c represents the diagonal version of the covariance update in equation 3c.
∎

###### Proof.

###### Proof.

#### Proof for 𝑹 k subscript 𝑹 𝑘 {\bm{R}}_{k} bold_italic_R start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT Approximation

Here, we derive Equation equation 7 and Equation equation 12, which corresponds to two different methods for estimating the matrix ${\bm{R}}_{k}$.

Proof of equation 7 for method 1:

###### Proof.

We start by the definition of the observation noise covariance matrix ${\bm{v}}_{k}$: ${\bm{R}}_{k}=\mathbb{E}[{\bm{v}}_{k}{\bm{v}}^{T}_{k}]$. Substituting the noise ${\bm{v}}_{k}$ with ${\bm{v}}_{k}={\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k})$, we get:
${\bm{R}}_{k}=\mathbb{E}\left\{\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}
}_{k},{\bm{x}}_{k})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k
},{\bm{x}}_{k})\right)^{\top}\right\}$. Assuming $\tilde{{\bm{\theta}}}_{k}\approx\tilde{{\bm{\theta}}}_{k|k-1}$, we will have:
${\bm{R}}_{k}=\mathbb{E}\left\{\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}
}_{k|k-1},{\bm{x}}_{k})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}
}_{k|k-1},{\bm{x}}_{k})\right)^{\top}\right\}$. To compute the expectation outlined above in an empirical manner, we define:
$\hat{{\bm{R}}}_{k}=\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{
\bm{x}}_{k})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{
\bm{x}}_{k})\right)^{\top}$ as the impact of new data on ${\bm{R}}_{k}$, and employ an incremental averaging approach: ${\bm{R}}_{k}=\frac{k-1}{k}{\bm{R}}_{k-1}+\frac{1}{k}\hat{{\bm{R}}}_{k}$. This incremental averaging can be expressed as an EMA approach: ${\bm{R}}_{k}=\beta{\bm{R}}_{k-1}+(1-\beta)\hat{{\bm{R}}}_{k}$, with the forgetting factor of $\beta$.
∎

Proof of equation 12 for method 2:

###### Proof.

Similar to method 1, we start with the definition of the covariance matrix of the observation noise ${\bm{v}}_{k}$: ${\bm{R}}_{k}=\mathbb{E}[{\bm{v}}_{k}{\bm{v}}^{T}_{k}]$. Based on ${\bm{v}}_{k}={\bm{y}}_{k}-h_{LoRA}({\bm{x}}_{k},\tilde{{\bm{\theta}}}_{k})$, and substituting the noise definition, we get: ${\bm{R}}_{k}=\mathbb{E}\left\{\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}
}_{k},{\bm{x}}_{k})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k
},{\bm{x}}_{k})\right)^{\top}\right\}$. Since the value of $\tilde{{\bm{\theta}}}_{k}$ is not available before Updating step, let’s approximate $h_{LoRA}(\tilde{{\bm{\theta}}}_{k},{\bm{x}}_{k})$ using first-order Taylor series expansion around the last updated parameters $\tilde{{\bm{\theta}}}_{k|k-1}$, which is:
$h_{LoRA}(\tilde{{\bm{\theta}}}_{k},{\bm{x}}_{k})=h_{LoRA}(\tilde{{\bm{\theta}}
}_{k|k-1},{\bm{x}}_{k})+{\bm{H}}_{k}(\tilde{{\bm{\theta}}}_{k}-\tilde{{\bm{
\theta}}}_{k|k-1})$. Thus: ${\bm{R}}_{k}=$
$\mathbb{E}\left\{\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{
\bm{x}}_{k})-{\bm{H}}_{k}(\tilde{{\bm{\theta}}}_{k}-\tilde{{\bm{\theta}}}_{k|k
-1})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{\bm{x}}_
{k})-{\bm{H}}_{k}(\tilde{{\bm{\theta}}}_{k}-\tilde{{\bm{\theta}}}_{k|k-1})
\right)^{\top}\right\}$.
Expanding the above expression with dropping subscript and $(\tilde{{\bm{\theta}}}_{k|k-1},{\bm{x}}_{k})$ from $h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{\bm{x}}_{k})$, and using $\overline{{\bm{\theta}}}_{k}=(\tilde{{\bm{\theta}}}_{k}-\tilde{{\bm{\theta}}}_
{k|k-1})$ for simplicity will yield: ${\bm{R}}_{k}=$
$\mathbb{E}\left\{{\bm{y}}_{k}{\bm{y}}_{k}^{\top}-{\bm{y}}_{k}h^{\top}-{\bm{y}}
_{k}\overline{{\bm{\theta}}}_{k}^{\top}{\bm{H}}^{T}_{k}-h{\bm{y}}_{k}^{\top}+
hh^{\top}+h\overline{{\bm{\theta}}}_{k}^{\top}{\bm{H}}^{T}_{k}-{\bm{H}}_{k}
\overline{{\bm{\theta}}}_{k}{\bm{y}}_{k}^{\top}+{\bm{H}}_{k}\overline{{\bm{
\theta}}}_{k}h^{\top}+{\bm{H}}_{k}\overline{{\bm{\theta}}}_{k}\overline{{\bm{
\theta}}}_{k}^{\top}{\bm{H}}^{T}_{k}\right\}$.
Now, let’s take the expectation by considering:
$\mathbb{E}[\overline{{\bm{\theta}}}_{k}]=\mathbb{E}[\tilde{{\bm{\theta}}}_{k}-
\tilde{{\bm{\theta}}}_{k|k-1}]=0$, and
$\mathbb{E}[\overline{{\bm{\theta}}}_{k}\overline{{\bm{\theta}}}_{k}^{\top}]=
\mathbb{E}[(\tilde{{\bm{\theta}}}_{k}-\tilde{{\bm{\theta}}}_{k|k-1})(\tilde{{
\bm{\theta}}}_{k}-\tilde{{\bm{\theta}}}_{k|k-1})^{\top}]={\bm{P}}_{k|k-1}$.
Now, we have:
${\bm{R}}_{k}=\mathbb{E}\left\{{\bm{y}}_{k}{\bm{y}}_{k}^{\top}-{\bm{y}}_{k}h^{
\top}-h{\bm{y}}_{k}^{\top}+hh^{\top}+{\bm{H}}_{k}{\bm{P}}_{k|k-1}{\bm{H}}^{T}_
{k}\right\}$.
We already knew that ${\bm{H}}_{k}{\bm{P}}_{k|k-1}{\bm{H}}^{T}_{k}={\bm{H}}_{k}(\hat{{\bm{p}}}_{k|k-
1}\bullet{\bm{H}}_{k}^{\top})$. Thus:
${\bm{R}}_{k}=\mathbb{E}\left\{\left({\bm{y}}_{k}-h\right)\left({\bm{y}}_{k}-h
\right)^{\top}+{\bm{H}}_{k}(\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_{k}^{\top})\right\}$. Similar to the previous approach, we compute the above expectation in an empirical manner by defining: $\hat{{\bm{R}}}_{k}=\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{
\bm{x}}_{k})\right)\left({\bm{y}}_{k}-h_{LoRA}(\tilde{{\bm{\theta}}}_{k|k-1},{
\bm{x}}_{k})\right)^{\top}+{\bm{H}}_{k}(\hat{{\bm{p}}}_{k|k-1}\bullet{\bm{H}}_
{k}^{\top})$
∎

###### Proof.

###### Proof.

## Appendix C Additional Experiment Details

#### Diagonal Approximation of Covariance Matrix:

Our empirical observations reveal that throughout the training process, the covariance matrix of a feedforward neural network tends towards a (block-)diagonal structure asymptotically. Figure 4 illustrates the evolution of the covariance matrix ${\bm{P}}_{k}\in\mathbb{R}^{n\times n}$ in LeNet-5 utilizing the Kalman optimizer. Initially, the matrix exhibits a fully dense positive-definite form, which progressively transitions towards a (block-)diagonal configuration as the training algorithm advances.

Figure: Figure 4: Evolution of covariance matrix ${\bm{P}}_{k}\in\mathbb{R}^{n\times n}$ in LeNet-5 using Kalman optimizer. The matrix starts with a fully dense positive-definite matrix, and with the progress of the training algorithm, it gradually converges to a (block-)diagonal configuration.
Refer to caption: extracted/5925019/Figure/diag-a.png

#### Sensitivity to β 𝛽 \beta italic_β in EMA:

The sensitivity analysis concerning $\beta$ is depicted in Figure 5. It is evident that extreme values, either too high or too low, significantly undermine LoKO’s performance. The optimal range for $\beta$ lies between 0.9 and 0.98, as indicated by the plot.

Figure: Figure 5: Sensitivity to $\beta$
Refer to caption: extracted/5925019/Figure/beta.png

## Appendix D Additional Experiment Results

Figure: Figure 6: Comparison of LoKO (blue) with LoRA/AdamW (red) and LoRA/AdaGrad (green) across various language models and datasets. For each combination, the top row presents the training loss, while the bottom row illustrates the average online accuracy against the number of data points observed.
Refer to caption: extracted/5925019/Figure/NLP/train_loss_roberta-base_lora_sst2.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - What is the novelty of the work beyond application to LoRA?
- Is this applicable beyond LoRA to other methods?
- Does the diagonal approximation cause problems in harder tasks like text-to-image generation?
- Is OOD generalization adversely affected due to the noise approximation?
soundness: 3
strengths: LoKO shows that Kalman optimizers work similarly well with LoRA as they do with bigger weight matrices in a neural network with better and faster convergence and efficient online learning.
summary: Low-Rank Kalman Optimizer (LoKO) is a novel optimizer for Parameter-Efficient Fine-Tuning (PEFT) aimed specifically at LoRA. LoKO frames fine-tuning as an optimal filtering problem, where trainable parameters are estimated as state variables updated dynamically with new observations. By leveraging Kalman’s recursive estimation, LoKO achieves efficient convergence with fewer updates. Additionally, the optimizer incorporates a diagonal approximation for the covariance matrix, which reduces computational complexity from quadratic to linear in the number of parameters.
weaknesses: "Novelty" The novelty of the work seems hard to gauge and should be explicitly noted in the paper. The use of EKF as an alternative optimizing strategy has already been pursued and its efficacy for faster convergence has been demonstrated in previous works some of which are cited in the paper. The diagonal approximation (amongst other approximation attempts as in the cited, Chang et al.) for the covariance matrix to reduce computations has also been utilized in a number of previous works regarding Kalman filtering. The novelty in this work seems to be in the application to the finetuning of neural networks using a specific low rank adapter, LoRA. There are no new findings or theoretical insights proposed other than what is expected of training a subset of neural weights (LoRA) with a Kalman optimizer. Unfortunately, this does not seem to be adequate especially given the limited diversity of experiments (explained below) and applicability to a single type of low rank adapter.

The gaussian assumption for the neural parameters and noise are unsupported by any empirical results in the paper.

"Diagonal Approximation" LoKO relies on a diagonal approximation of the covariance matrix to reduce computational costs. While effective, this approximation limits the optimizer's ability to capture parameter correlations fully, which could lead to suboptimal updates for tasks with high-dimensional dependencies​. Experiments on complex tasks like text-to-image generation or others should be utilized to ablate this aspect of the method.

In addition, the sensitivity to initialization of the covariance matrix requires significant ablation across different architectures and tasks to investigate its failings and properties, and whether it varies across different modalities.

"Noise Approximation" The exponential moving average (EMA) approach for noise covariance estimation may be unstable, particularly when new data distributions diverge significantly from the initial training data. It is necessary to test LoKO in OOD generalization tasks to ablate this approximation strategy.

Given this work relates to optimization specific to PeFT methods, its necessary to show its efficacy on different low rank adapters to show similar results and properties can be observed from them. Comparisons with DoRA, X-LoRA, etc. would be great additions to the work.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please address the comments in the weakness part.
soundness: 3
strengths: Casting PEFT as a state estimation problem using the Kalman filter is a novel and insightful approach. The combination of LoRA with EKF addresses the scalability issues traditionally associated with Kalman filters in large-scale models. The diagonal approximation of the covariance matrix and EMA-based estimation significantly reduce computational complexity, making the approach practical for large models.

The paper provides a clear and well-founded mathematical formulation of the proposed method. The derivation of the simplified Kalman update equations is thorough, and practical implementation considerations are discussed. The thoughtful discussion on covariance matrix initialization and its impact on performance demonstrates a deep understanding of the method's nuances.

The method is evaluated across multiple domains, including image and text classification tasks, demonstrating its versatility. Thorough comparisons with mainstream optimizers (AdamW and AdaGrad) show that LoKO often achieves better convergence and performance. Detailed analyses on the effects of covariance initialization and the forgetting factor provide valuable insights into the method's behavior.
summary: This paper introduces LoKO, a novel optimization algorithm that integrates the Extended Kalman Filter (EKF) with Low-Rank Adaptation (LoRA) for the online fine-tuning of large pre-trained models. By reframing parameter-efficient fine-tuning (PEFT) as a state estimation problem, the authors leverage the EKF to estimate optimal trainable parameters in an online manner. The use of LoRA's low-rank decomposition reduces the number of trainable parameters, and a diagonal approximation of the covariance matrix further decreases computational complexity from quadratic to linear. An Exponential Moving Average (EMA) approach is proposed for estimating the observation noise covariance without additional computational cost. Empirical evaluations on various computer vision and natural language processing tasks demonstrate that LoKO achieves competitive or superior performance compared to traditional gradient-based optimizers like AdamW and AdaGrad.
weaknesses: The paper lacks a formal convergence analysis of the proposed algorithm. Without theoretical guarantees, it's unclear under what conditions LoKO is expected to perform reliably.
The justification for the diagonal approximation of the covariance matrix is primarily empirical. A theoretical analysis or conditions under which this approximation holds would strengthen the contribution.
The theoretical implications of using EMA for observation noise covariance estimation are not fully explored. Understanding its impact on convergence and stability is important.
The experiments do not include comparisons with other recent PEFT methods such as QLoRA or AdaLoRA. Including these would better position LoKO within the current landscape. Evaluations are conducted over a single epoch, which may not capture the long-term behavior and stability of the optimizer. The focus on online fine-tuning with batch size 1 raises questions about the method's applicability and performance in standard offline settings with larger batch sizes. In some language tasks, LoKO performs slightly worse than AdamW. Investigating these cases could reveal insights into the method's limitations.

While the paper claims reduced computational complexity, it does not provide empirical runtime measurements or resource usage comparisons to substantiate these claims. There is limited discussion on potential trade-offs between computational cost and optimization performance compared to other methods.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What's the limitation of the proposed approach? It would be great to have more throughly memory usage and efficiency analysis and experiments, given the computational efficiency is one of the advantages of the proposed approach. 

2. How to initialize the LoKO? how sensitive of the model is? It would be great to have more studies on those.

3. How does the proposed LoKO compare against the standard LoRA and other LoRA variants such as QLoRA?
soundness: 4
strengths: LoKO seems to be an interesting approach that successfully combine the EKF and the LORA, this provides a good new perspective to approach the online gradient-descent method.

This paper provides a well structured experiments across computer vision and NLP. The results in general are not state-of-the-art, but they are decent to show the promising results of the LoKO.

The proof to show that the proposed approach equation is equivalent to that in the vanilla EKF algorithm seems to be sound.

The paper is well structured and well written. The logical flow from problem definition through methodology to empirical validation helps readers grasp the paper’s objectives and contributions effectively.

The proposed LoKO seems to provide high accuracy with lower computational footprint. This seems to be interesting to some practical applications where the computational efficiency is important.
summary: Previous works show that the recursive Extended Kalman Filter (EKF) algorithm can optimize relatively small models with performance surpassing its gradient-based counterparts. However, it is challenging to optimize neural networks with EKF because the size of the crucial covariance matrix in EKF can grow quadratically with the number of model parameters in training. In this paper, the core idea is to have EKF work with the LORA so that only a few parameters are involved in the EKF optimization. The proposed approach LoKO shows promising results on computer vision and language modeling benchmarks, such as MNIST, CIFAR 10/100, ImageNet 100, SST-2, COLA, MRPC.
weaknesses: The evaluation is limited. This paper has done quite a few evaluation on the proposed LoKO method. However, i find none of those datasets are well suited for the LoKO. In general, the model full fine-tuning can do better than that from the LoRA fine-tuning. The LoKO would be suboptimal compared to the model full-fining as long as one can afford the GPU computes. The proposed experiments do not cover the dynamic streaming type of data, which would be able to show the strength of the proposed LoKO.

The proposed LoKO seems to suggest to require one to initialize the covariance matrix, and noise estimation parameters. However, it is unclear how sensitive of those parameters to the final model performance. There are no studies to reveal the sensitive of the proposed approach.

There are multiple LoRA and its variants for PEFT. The experiments do not cover the comparison between the proposed one against the other PEFT methods such as QLoRA, QA-LoRA, LoftQ. It would be interesting and useful to see how the proposed LoKO compared against the other LoRA and its variants.

The experiments do not cover the study on the memory efficiency and what's the limitations of the proposed LoKO. It is important to understand the limits and boundary of the LoKO.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: In Table 1, the authors provide the tested models and their corresponding LoRA layers along with the trainable parameters. How do the remaining non-trainable parameters be obtained?

In Table 4, the authors provide the diagonal initialization bounds for parameter P for different datasets and models. How are these bounds determined? Is there a more general selection criterion?
soundness: 3
strengths: This paper innovatively proposes using the Kalman Filter algorithm combined with LoRA for online fine-tuning of large-scale models, and
reduces the computational complexity of fine-tuning through approximate solution methods.

The experiments also demonstrate that the proposed method achieves performance comparable to existing gradient descent-based algorithms across multiple models and datasets.

Finally, the paper investigates the sensitivity of the final performance to the initialization range of hyperparameters through experiments, thereby proving the robustness of the proposed method.
summary: This paper proposes an innovative Kalman filter-based low-rank optimization algorithm (referred to as LoKO) aimed at online fine-tuning of large-scale models.

The algorithm initially employs the LoRA algorithm to effectively reduce the number of trainable parameters through low-rank decomposition.
Subsequently, it adopts a diagonal approximation of the covariance matrix P and integrates Exponential Moving Average (EMA) techniques to estimate the covariance matrix R of the observation noise.

The implementation of these strategies significantly lowers the computational complexity of LoKO. Moreover, experimental results across
multiple computer vision and language models demonstrate that its performance is on par with existing advanced methods.
weaknesses: The LoKO method assumes that the model's trainable parameters remain unchanged during the observation process, meaning that the generation process of the observed variables remains consistent throughout the entire observation period. This assumption is often difficult to satisfy in online scenarios and results in the loss of the model's dynamic adaptation capability.

Although the authors observed through experiments that the matrix tends to become diagonal during training, they did not provide a
corresponding analysis. 

In Table 3, the LoKO method consistently underperforms compared to the baseline methods LoRA/AdamW, but the authors did not provide any related analysis.

One significant contribution of this paper is the proposal of a method for online fine-tuning of large-scale models. However, the largest model tested in the experiments, RoBERTa-large, has only 335M parameters. Therefore, further testing on even larger models is needed to validate the effectiveness of their method.

Finally, the paper lacks a comparison of the runtime efficiency between LoKO and the baseline methods.

Meta-review: ## Summary

LoKO is a novel optimization algorithm that integrates the Extended Kalman Filter (EKF) with Low-Rank Adaptation (LoRA) for online fine-tuning of large pre-trained models. It uses LoRA's low-rank decomposition to reduce the number of trainable parameters and a diagonal approximation of the covariance matrix, reducing computational complexity. LoKO achieves competitive or superior performance compared to traditional gradient-based optimizers, showing promising results on computer vision and language modeling benchmarks.


## Strengths
* LoKO combines Kalman optimizers with LoRA for efficient online learning and faster convergence, whereas it addresses scalability issues in large-scale models.
* LoKO is evaluated across multiple domains on mainstream optimizers, including image and text classification tasks, and often achieves better convergence and performance.
* The method achieves performance comparable to existing gradient descent-based algorithms across multiple models and datasets.

## Weaknesses
* The diagonal approximation for the covariance matrix is used to reduce computational costs, but it limits the optimizer's ability to capture parameter correlations fully.
* The paper does not provide a formal convergence analysis of the proposed algorithm, and the theoretical implications of using EMA for observation noise covariance estimation are not fully explored.
* The paper does not include comparisons with other recent PEFT methods such as QLoRA or AdaLoRA, and lacks a comparison of the runtime efficiency between LoKO and the baseline methods.
* The focus on online fine-tuning with batch size 1 raises questions about the method's applicability and performance in standard offline settings with larger batch sizes.
* The paper claims reduced computational complexity but does not provide empirical runtime measurements or resource usage comparisons to substantiate these claims.
* The LoKO method assumes that the model's trainable parameters remain unchanged during the observation process, which is often difficult to satisfy in online scenarios.

## Conclusions
Based on the reviews, author feedback and rebuttals, the paper has been improved thanks to the effort of the authors by adding more experiments, and clarification of the text. However, there were some concerns that did not satisfy the reviewers. For instance, there are a few citations about low-rank kalman filtering like *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* and have a clear differentiation instead of just citing them briefly. Therefore, the paper should be improved before being accepted.

Decision: Reject