DiSK: Differentially Private Optimizer with Simplified Kalman Filter for Noise Reduction

## Abstract

Abstract Differential privacy (DP) offers a robust framework for safeguarding individual data privacy. To utilize DP in training modern machine learning models, differentially private optimizers have been widely used in recent years. A popular approach to privatize an optimizer is to clip the individual gradients and add sufficiently large noise to the clipped gradient. This approach led to the development of DP optimizers that have comparable performance with their non-private counterparts in fine-tuning tasks or in tasks with a small number of training parameters. However, a significant performance drop is observed when these optimizers are applied to large-scale training. This degradation stems from the substantial noise injection required to maintain DP, which disrupts the optimizer’s dynamics.
This paper introduces DiSK , a novel framework designed to significantly enhance the performance of DP optimizers. DiSK employs Kalman filtering, a technique drawn from control and signal processing, to effectively denoise privatized gradients and generate progressively refined gradient estimations. To ensure practicality for large-scale training, we simplify the Kalman filtering process, minimizing its memory and computational demands.
We establish theoretical privacy-utility trade-off guarantees for DiSK and demonstrate provable improvements over standard DP optimizers like DPSGD in terms of iteration complexity upper-bound.
Extensive experiments across diverse tasks, including vision tasks such as CIFAR-100 and ImageNet-1k and language fine-tuning tasks such as GLUE, E2E, and DART, validate the effectiveness of DiSK . The results showcase its ability to significantly improve the performance of DP optimizers, surpassing state-of-the-art results under the same privacy constraints on several benchmarks.

## 1 Introduction

Data privacy has become one of the major concerns in modern machine learning systems. Differential Privacy (DP), with its rigorous mathematical foundation, offers a powerful solution. DP provides a framework for developing training algorithms that safeguard the privacy of individuals’ data used to train machine learning models.
Among various algorithms, Differentially Private Stochastic Gradient Descent (DPSGD) (Abadi et al., ) and its variants (Li et al., ; Yu et al., ; Tang et al., ) have emerged as popular choices for training various models, including computer vision (De et al., ) and language models (Bu et al., ; ).
These algorithms inject noise into the training process to guarantee privacy. However, this noise injection often comes at a significant cost to model performance, limiting the widespread adoption of DP optimizers (Jayaraman & Evans, ). For example, (McMahan et al., ; De et al., ) observed that DP training led to a $15\%$ drop in model accuracy on the Reddit dataset and a $30\%$ drop on CIFAR-10 compared to non-private training. This challenge has limited the application of DP optimizers primarily to small models or parameter-efficient fine-tuning, as highlighted by (Li et al., ).

Many recent works aim to improve differentially private (DP) training performance. These include applying low-pass filters to separate gradients from noise  (Zhang et al., ), injecting correlated noise with algorithms like DP-FTLR (Koloskova et al., ; Choquette-Choo et al., ), using sharpness-aware minimization for flatter loss landscapes that are less sensitive to DP noise (Park et al., ), adaptive clipping (Andrew et al., ; Lin et al., ; Hu et al., ), and model structure optimization (Bu et al., ; Papernot et al., ; De et al., ). However, these methods often require extra memory, lack theoretical guarantees, or have limited applicability. Therefore, there is a strong demand to design a new approach that 1) is memory and computation efficient for implementation, 2) has a theoretical guarantee, and 3) is compatible with a wide range of existing DP optimization algorithms and models.
To meet this demand, we leverage the Kalman filter, a tool from control theory, to improve gradient estimates in DP optimization. We further simplify our algorithm for memory and computational efficiency, while maintaining theoretical grounding and broad compatibility with existing DP algorithms and models.

The tools in signal processing and control theory have been leveraged to design novel optimization algorithms. In the context of DP optimization, the error-feedback approach has been adopted to reduce the bias caused by the clipping operation (Zhang et al., ); the low-pass and high-pass filters have been used to separate the gradient from the DP noise (Zhang et al., ; Koloskova et al., ; Choquette-Choo et al., ); the low-pass spatial filter, including Gaussian and Laplace filters have been used to smooth the privatized model or gradient across dimensions (Wang et al., ; ; Liu et al., ). These methods effectively improve DP optimizers’ performance.

### 1.1 Contributions

Our approach centers on constructing a dynamic system where the gradient serves as its state. We treat the privatized gradient as a noisy observation of the true gradient within this system. By applying a Kalman filter, we obtain a more accurate estimate of the true gradient by leveraging the gradient dynamics and past estimates, thereby enhancing the performance of DP optimizers. To address the inherent inefficiency of the Kalman filter, we introduce a series of simplifications that reduce memory and computational overhead. Our main contributions can be summarized as follows:

- •
Algorithm Design: We introduce a novel Kalman filter-based approach designed to mitigate DP noise and enhance the performance of various DP optimizers.
- •
Algorithm Simplification: We simplify the Kalman filtering process to significantly reduce memory and computational overhead. This simplified approach, DiSK, requires only one additional forward step and two extra optimizer states.
- •
Theoretical Analysis: We provide theoretical analyses of DiSK, demonstrating that the algorithm is convergent. Moreover, we showed that, compared to DPSGD, DiSK improves the iteration complexity upper-bound by a constant factor.
- •
Numerical Results: Extensive experiments across various models, datasets, and optimizers demonstrate that DiSK significantly boosts DP training performance.
Specifically, under the same privacy budgets, DiSK exhibits substantial improvements in test accuracy for training-from-scratch scenarios: a notable increase from $32.4\%$ to $36.9\%$ on the ImageNet-1k dataset, a considerable rise from $63\%$ to $75\%$ on CIFAR-10, and a remarkable improvement from $21\%$ to $42\%$ on CIFAR-100.
Furthermore, in fine-tuning tasks, DiSK demonstrates remarkable improvements: an increase from $85\%$ to $89\%$ on CIFAR-100 and an improvement from $81\%$ to $86\%$ on the GLUE dataset. These results surpass state-of-the-art DP training performance under the same privacy guarantees.

## 2 Preliminaries

### 2.1 Problem Definition & Notations

Typical training procedures require solving the empirical risk minimization (ERM) problem

$$ $\displaystyle\min_{\mathbf{x}\in\mathbb{R}^{d}}\left(F(\mathbf{x})=\frac{1}{N} \sum_{\xi\in\mathcal{D}}f(\mathbf{x};\xi)\right),$ (1) $$

where $\mathbf{x}\in\mathbb{R}^{d}$ is the optimization variable, $\mathcal{D}=\{\xi_{1},\ldots,\xi_{N}\}$ is the training dataset with $\left\lvert\mathcal{D}\right\rvert=N$ samples, and $f(\cdot)$ denotes the (possibly non-convex) loss function parameterized by $\mathbf{x}$ and evaluated on sample $\xi$. For solving the above optimization problem, we rely on iterative procedures, such as SGD, where the parameters are updated over the iterations $t=1,2,\ldots$.
Throughout the paper, we use $(\cdot)_{t}$ to denote the variables at iteration $t$, and use $\mathbf{I}_{d}$ to denote the identity matrix of dimension $d$.

### 2.2 Differentially private optimization

Let us start by recalling the definition of $(\epsilon,\delta)$-Differential Privacy:

###### Definition 1 ( $(\epsilon,\delta)$ -DP (Dwork & Roth, 2014 ) )

A randomized mechanism $\mathcal{M}$ is said to be $(\epsilon,\delta)$-differentially private if for any two neighboring datasets $\mathcal{D},\mathcal{D}^{\prime}$ ($\mathcal{D},\mathcal{D}^{\prime}$ differ only by one sample) and for any measurable output set $\mathcal{S}$, it holds that $\mathrm{Pr}[\mathcal{M}(\mathcal{D})\in\mathcal{S}]\leq\mathrm{e}^{\epsilon}
\mathrm{Pr}[\mathcal{M}(\mathcal{D}^{\prime})\in\mathcal{S}]+\delta.$

A widely used approach to achieving differential privacy (DP) when solving ERM problem (1) is to employ differentially private stochastic gradient descent (Abadi et al., ) and its variants, such as DP-Adam and DP-Lora (Yu et al., ). To ensure DP, DPSGD leverages the Gaussian mechanism (Dwork & Roth, ; Abadi et al., ), injecting carefully calibrated Gaussian noise into the gradients at each iteration of the optimization process. This noise injection effectively masks the contribution of individual data points, thereby providing the desired privacy guarantee.

###### Definition 2 (Gaussian Mechanism (Dwork & Roth, 2014 ; Balle & Wang, 2018 ) )

Suppose an algorithm $\mathcal{A}:\mathcal{D}\rightarrow\mathbb{R}^{d}$ has $\ell_{2}$ sensitivity $\Delta_{\mathcal{A}}$, i.e.,
$\max_{\mathcal{D},\mathcal{D}^{\prime}}\left\lVert\mathcal{A}(\mathcal{D})-
\mathcal{A}(\mathcal{D}^{\prime})\right\rVert\leq\Delta_{\mathcal{A}}.$
Then, for any $\epsilon>0$ and $\delta\leq 0.05$, by adding a carefully chosen random Gaussian noise to the output of the algorithm, we can make the algorithm $(\epsilon,\delta)$-DP. More specifically,
the mechanism $M(x)=\mathcal{A}(x)+\mathbf{w},\mathrm{with}~{}\mathbf{w}\sim\mathcal{N}(0,
\sigma_{\textup{DP}}^{2}I_{d})$
and $\sigma_{\text{DP}}$ \xweditsatisfies $\Phi\left(\frac{\Delta_{\mathcal{A}}}{2\sigma_{\text{DP}}}-\frac{\epsilon
\sigma_{\text{DP}}}{\Delta_{\mathcal{A}}}\right)-\mathrm{e}^{\epsilon}\Phi
\left(-\frac{\Delta_{\mathcal{A}}}{2\sigma_{\text{DP}}}-\frac{\epsilon\sigma_{
\text{DP}}}{\Delta_{\mathcal{A}}}\right)\leq\delta$ is $(\epsilon,\delta)$-DP, where $\Phi(t)=\mathbb{P}[\mathcal{N}(0,1)\leq t]$ is the cumulative
density function of normal distribution.

Figure: Algorithm 1 DPSGD algorithm

The DPSGD algorithm, outlined in Algorithm 1, operates by first sampling a mini-batch $\mathcal{B}_{t}$ of size $B$ and computing the per-sample gradients at each iteration $t$. To guarantee differential privacy, it then applies the Gaussian mechanism, which involves clipping each per-sample gradient to bound its sensitivity to a maximum value $C$ and subsequently injecting Gaussian noise. The clipping operation
$\mathrm{clip}\left(\nabla f,C\right)=\min\left\{1,\frac{C}{\left\lVert\nabla f
\right\rVert}\right\}\cdot\nabla f,$
often implemented by scaling the gradient when its norm exceeds $C$, limits the influence of any single data point, while the added noise further masks individual contributions, ensuring the desired privacy level (Abadi et al., ).

###### Theorem 1 (Privacy Guarantee (Abadi et al., 2016 ) )

Given the number of samples $N$, the batch-size $B$, total number of iterations $T$ and clipping threshold $C$, there exist positive constants $u,v$, such that for any $\epsilon<\frac{uB^{2}T}{N^{2}}$ and $0<\delta$, by choosing $\sigma_{\textup{DP}}^{2}\geq v\frac{C^{2}T\ln(\frac{1}{\delta})}{N^{2}\epsilon
^{2}}$, Algorithm 1 is $(\epsilon,\delta)$-DP.

Theorem 1 implies that the variance of the DP noise injected into the gradients, $\operatorname{\mathbb{E}}[\left\lVert\mathbf{w}_{t}\right\rVert^{2}]=d\sigma^{
2}_{\mathrm{DP}}$, scales linearly with both the number of iterations $T$ and the number of parameters $d$. This presents a significant challenge in modern deep learning, where models size $d$ is large (e.g., $22M$-$632M$ for ViT (Dosovitskiy et al., ), $137M$-$1.6B$ for GPT-2 (Radford et al., )) and require extensive training (e.g., $300K$ for ViT (Dosovitskiy et al., ) and $250K$ for training Llama (Touvron et al., )). Consequently, the magnitude of the injected DP noise can become substantial, leading to a considerable degradation in model performance.

###### Definition 1 ( ( ϵ , δ ) italic-ϵ 𝛿 (\epsilon,\delta) ( italic_ϵ , italic_δ ) -DP (Dwork & Roth, 2014 ) )

###### Definition 2 (Gaussian Mechanism (Dwork & Roth, 2014 ; Balle & Wang, 2018 ) )

###### Theorem 1 (Privacy Guarantee (Abadi et al., 2016 ) )

### 2.3 Kalman filter

The Kalman filter is a powerful algorithm that provides estimates of unknown variables by iteratively incorporating a series of measurements over time (Welch et al., ). To illustrate its application, let us consider a linear dynamic system characterized by the System update and Observation equations:

$$ $\displaystyle\boldsymbol{\theta}_{t}$ $\displaystyle=\mathbf{A}\boldsymbol{\theta}_{t-1}+\mathbf{u}_{t}+\mathbf{v}_{t},$ (System update) (2) $\displaystyle\boldsymbol{\psi}_{t}$ $\displaystyle=\mathbf{C}\boldsymbol{\theta}_{t}+\mathbf{w}_{t},$ (Observation) $$

where $\mathbf{A}\in\mathbb{R}^{d_{\theta}\times d_{\theta}}$, $\mathbf{C}\in\mathbb{R}^{d_{\psi}\times d_{\theta}}$ are the transition and observation matrices; $\boldsymbol{\theta}_{t}\in\mathbb{R}^{d_{\theta}}$ is the unknown variable to be tracked/estimated; $\boldsymbol{\psi}_{t}\in\mathbb{R}^{d_{\psi}}$ denotes the observation of the system; $\mathbf{u}_{t}\in\mathbb{R}^{d_{\theta}}$ denotes the known input; and $\mathbf{v}_{t}\in\mathbb{R}^{d_{\theta}},\mathbf{w}_{t}\in\mathbb{R}^{d_{\psi}}$ are the process and observation noises that follow Gaussian distribution $\mathcal{N}(0,\Sigma_{\mathbf{v}}),\mathcal{N}(0,\Sigma_{\mathbf{w}})$, respectively. Then, the Kalman filter uses the following updates to track $\{\boldsymbol{\theta}_{t}\}$ (Welch et al., ):

$$ $\displaystyle\tilde{\boldsymbol{\theta}}_{t|t-1}$ $\displaystyle=\mathbf{A}\tilde{\boldsymbol{\theta}}_{t-1}+\mathbf{u}_{t}$ (Prediction) (3) $\displaystyle\mathbf{P}_{t|t-1}$ $\displaystyle=\mathbf{A}\mathbf{P}_{t-1}\mathbf{A}^{\top}+\Sigma_{\mathbf{v}}$ $\displaystyle\mathbf{K}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}\mathbf{C}^{\top}(\mathbf{C}\mathbf{P}_{t|t-1} \mathbf{C}^{\top}+\Sigma_{\mathbf{w}})^{-1}$ $\displaystyle\tilde{\boldsymbol{\theta}}_{t}$ $\displaystyle=(\mathbf{I}_{d_{\theta}}-\mathbf{K}_{t}\mathbf{C})\tilde{ \boldsymbol{\theta}}_{t|t-1}+\mathbf{K}_{t}\boldsymbol{\psi}_{t}$ (Correction) $\displaystyle\mathbf{P}_{t}$ $\displaystyle=(\mathbf{I}_{d_{\theta}}-\mathbf{K}_{t}\mathbf{C})\mathbf{P}_{t| t-1}.$ $$

The filter first predicts the state $\tilde{\boldsymbol{\theta}}_{t|t-1}$ by the system dynamics, and compute the filter gain $\mathbf{K}_{t}\in\mathbb{R}^{d_{\theta}\times d_{\psi}}$ based on the covariance matrix $\mathbf{P}_{t}\in\mathbb{R}^{d_{\theta}\times d_{\theta}}$, and corrects the prediction with system observation $\boldsymbol{\psi}_{t}$ to obtain $\tilde{\boldsymbol{\theta}}_{t}$. The Kalman filter makes use of both the noisy observation and the prior knowledge of the system dynamics to obtain an accurate estimation of the state $\boldsymbol{\theta}_{t}$.

In this work, we treat the privatized gradients in the DP optimization as noisy observations of the true underlying gradients. By constructing gradient dynamics using Taylor expansion, we establish a framework for applying the Kalman filter to refine these noisy observations and obtain more accurate gradient estimates. To ensure practical applicability for large-scale models, we simplify the Kalman filter, significantly reducing its memory and computational footprint. These improved gradient estimates ultimately lead to enhanced performance in differentially private optimizers.

### 2.4 Related Works

Optimization with filters and controllers: The use of filters and controllers in designing and analyzing optimization algorithms has a rich history. Researchers have leveraged high-pass and low-pass filters to enhance gradient estimation in zeroth-order optimization (Chen et al., ), employed PID controllers for both centralized and distributed optimization (Wang et al., ), and analyzed optimizers through the lens of control theory, treating them as dynamic systems (Lessard et al., ; Hu & Lessard, ; Muehlebach & Jordan, ; Mohammadi et al., ; Badithela & Seiler, ; Cyrus et al., ; Scherer et al., ; Zhang et al., ).

Kalman filter for optimization: The Kalman filter has been utilized in convex optimization for reducing stochastic gradient noise (Bittner & Pronzato, ; Vuckovic, ). Vuckovic () uses the dynamics of the optimization variable and the gradient to construct the Kalman filter to analyze and improve the performance of momentum methods. Bittner & Pronzato () uses gradient and Hessian as its states to construct the dynamic system for SGD to construct a stopping rule. However, these approaches, with their direct application of the Kalman filter, incur prohibitively high computational and memory costs, ranging from $\mathcal{O}(d^{3})$ to $\mathcal{O}(d^{6})$, rendering them impractical for training large-scale machine learning models.

Improving DP optimization: Numerous techniques have been proposed to enhance DP optimization by mitigating the impact of DP noise. These include adaptive gradient clipping methods that dynamically adjust clipping thresholds (Andrew et al., ; Bu et al., ), parameter-efficient training strategies employing adapters, low-rank weights, or quantization (Yu et al., ; Luo et al., ; Yu et al., ), and the design of specialized model architectures less susceptible to noise perturbations (De et al., ; Papernot et al., ; Wang et al., ). Furthermore, drawing inspiration from signal processing, researchers have explored the use of colored high-frequency DP noise to separate it from the gradient (Koloskova et al., ) and the application of low-pass filters to extract the gradient signal from noisy observations (Zhang et al., ).

## 3 Algorithm design

This section introduces the general Noise Reduction for Differentially Private Optimizers with Simplified Kalman Filter (DiSK) framework. This approach leverages the inherent dynamics of the gradient and employs Kalman filtering to obtain denoised estimates of the true gradients from their noisy, privatized counterparts. To enhance its practicality for modern deep learning, Section 3.2 details a simplified version of the Kalman filter updates, designed for memory and computational efficiency.

### 3.1 Gradient dynamic and Kalman filter

To explain our idea of using the Kalman filter for denoising, let us start by first establishing a dynamic system for the gradients, comprising a “system update” equation and an “observation” equation: The system update of the gradient dynamics can be derived by the Taylor expansion and quantifying the change of the gradient at iteration $t$:

$$ $\displaystyle\nabla F(\mathbf{x}_{t})$ $\displaystyle=\nabla F(\mathbf{x}_{t-1})+\nabla^{2}F(\mathbf{x}_{t-1})^{\top}( \mathbf{x}_{t}-\mathbf{x}_{t-1})+R(\mathbf{x}_{t})$ (4) $\displaystyle=\nabla F(\mathbf{x}_{t-1})+\mathbf{H}_{t}(\mathbf{x}_{t}-\mathbf {x}_{t-1})+\mathbf{v}_{t},$ $$

where $\mathbf{H}_{t}:=\nabla^{2}F(\mathbf{x}_{t-1})\in\mathbb{R}^{d\times d}$ and $\mathbf{v}_{t}=\frac{1}{2}\int_{0}^{1}\nabla^{3}F(\mathbf{x}_{t-1})(z\mathbf{x
}_{t}+(1-z)\mathbf{x}_{t-1})^{\otimes 2}\mathrm{d}z$ is the remainder and $(\cdot)^{\otimes}$ denotes the tensor vector product.
The observation of the system is defined as the privatized gradient $\mathbf{g}_{t}$, which is a stochastic mapping of the true gradient:

$$ $\mathbf{g}_{t}=\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\mathrm{clip}\left( \nabla f(\mathbf{x}_{t},\xi),C\right)+\mathbf{w}_{t}=\mathbf{C}_{t}\nabla F( \mathbf{x}_{t})+\mathbf{w}^{\prime}_{t},$ (5) $$

where $\mathbf{w}^{\prime}_{t}$ is the observation noise containing the DP noise and sub-sampling noise. $\mathbf{C}_{t}$ is the observation matrix. If the clipping operation is inactive and $\mathcal{B}=\mathcal{D}$, i.e., the full batch gradient is used, then $\mathbf{C}_{t}=\mathbf{I}_{d}$. Otherwise, $\mathbf{C}_{t}$ depends on the clipping factor and the mini-batch $\mathcal{B}$.
Combining the system update (4) and the observation (5), we can model the system dynamic of the gradient as:

$$ $\displaystyle\nabla F(\mathbf{x}_{t})$ $\displaystyle=\nabla F(\mathbf{x}_{t-1})+\mathbf{H}_{t}(\mathbf{x}_{t}-\mathbf {x}_{t-1})+\mathbf{v}_{t},\qquad$ (System update) (6) $\displaystyle\mathbf{g}_{t}$ $\displaystyle=\mathbf{C}_{t}\nabla F(\mathbf{x}_{t})+\mathbf{w}^{\prime}_{t}.\qquad$ (Observation) $$

Let us compare our dynamic system with the general one in (2):
In our dynamic system, the gradient $\nabla F(\mathbf{x}_{t})$ plays the role of the state $\boldsymbol{\theta}_{t},$. Other parameters in (2) must change to the followings to obtain our gradient dynamic: $\mathbf{A}=\mathbf{I}_{d},\mathbf{C}=\mathbf{C}_{t}$, the input $\mathbf{u}_{t}=\mathbf{H}_{t}(\mathbf{x}_{t}-\mathbf{x}_{t-1})$, and the observation $\boldsymbol{\psi}_{t}=\mathbf{g}_{t}$. With the above mapping, we can apply the Kalman filter that combines the system update and the observation of the gradient to improve the overall estimation quality of the actual gradient beyond only using the observation $\mathbf{g}_{t}$. However, there are two key challenges when applying Kalman filter to (6): 1) the input $\mathbf{H}_{t}(\mathbf{x}_{t}-\mathbf{x}_{t-1})$ is hard to obtain as computing the Hessian matrix $\mathbf{H}_{t}$ is challenging for large models, and we can only approximate $\mathbf{H}_{t}(\mathbf{x}_{t}-\mathbf{x}_{t-1})$; 2) the observation matrix $\mathbf{C}_{t}$ is an unknown time-varying random matrix. Due to these differences, the traditional Kalman filter (Welch et al., ) cannot be directly applied. Instead, we apply the Kalman filter with noisy input and multiplicative observation noise proposed in Wu et al. () to our system (6), leading to the update rules:

$$ $\displaystyle\tilde{\mathbf{g}}_{t|t-1}$ $\displaystyle=\tilde{\mathbf{g}}_{t-1}+{\color[rgb]{1,0,1}\tilde{\mathbf{H}}_{ t}}(\mathbf{x}_{t}-\mathbf{x}_{t-1})$ (Prediction) $\displaystyle\mathbf{P}_{t|t-1}$ $\displaystyle=\mathbf{P}_{t-1}+{\color[rgb]{1,0,1}\Sigma_{\mathbf{H}}}+\Sigma_ {\mathbf{v}}$ $\displaystyle\mathbf{K}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}{\color[rgb]{1,0,1}\operatorname{\mathbb{E}}[ \mathbf{C}_{t}]^{\top}}\left(\Sigma_{\mathbf{w}}+{\color[rgb]{1,0,1} \operatorname{\mathbb{E}}[\mathbf{C}_{t}]}\left({\color[rgb]{1,0,1}\Sigma_{ \mathbf{C}}\mathbf{S}_{t}}+\mathbf{P}_{t|t-1}\right){\color[rgb]{1,0,1} \operatorname{\mathbb{E}}[\mathbf{C}_{t}]^{\top}-\Sigma_{\mathbf{H}}}\right)^{ -1}$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=\tilde{\mathbf{g}}_{t|t-1}+\mathbf{K}_{t}(\mathbf{g}_{t}-{\color [rgb]{1,0,1}\operatorname{\mathbb{E}}[\mathbf{C}_{t}]}\tilde{\mathbf{g}}_{t|t- 1})$ (Correction) $\displaystyle\mathbf{P}_{t}$ $\displaystyle=(\mathbf{I}-\mathbf{K}_{t}{\color[rgb]{1,0,1}\operatorname{ \mathbb{E}}[\mathbf{C}_{t}]})\mathbf{P}_{t|t-1}$ $\displaystyle{\color[rgb]{1,0,1}\mathbf{S}_{t}}$ $\displaystyle{\color[rgb]{1,0,1}=\operatorname{\mathbb{E}}[\tilde{\mathbf{g}}_ {t}\tilde{\mathbf{g}}_{t}^{\top}]},$ $$

where $\mathbf{P}_{t}$ denotes the covariance matrix of $\tilde{\mathbf{g}}_{t},$ $\Sigma_{\mathbf{H}},\Sigma_{\mathbf{w}},\Sigma_{\mathbf{C}},\Sigma_{\mathbf{v}}$ denote the covariance matrices of the random variables $\mathbf{H}_{t},\mathbf{w}_{t},\mathbf{C}_{t},\mathbf{v}_{t}$, respectively, and $\tilde{\mathbf{H}}_{t}$ is an instantiation/observation of the unknown Hessian matrix $\mathbf{H}_{t}$. The difference between the above Kalman filter and the original Kalman filter (3) is highlighted in magenta color. The variance of the Hessian $\tilde{\mathbf{H}}_{t}$ plays a role in updating $\mathbf{P}_{t|t-1},\mathbf{K}_{t}$, and the expectation $\operatorname{\mathbb{E}}[\mathbf{C}_{t}]$ of the random observation matrix $\mathbf{C}_{t}$ is used for the updates and its variance $\Sigma_{\mathbf{C}}$ also appears in the update of $\mathbf{K}_{t}$.
Compared to the noisy gradient $\mathbf{g}_{t}$, the output of the Kalman filter $\tilde{\mathbf{g}}_{t}$ has a smaller variance, resulting in improved performance when used in DP optimizers.

The resulting optimizer with the Kalman filter is given in Algorithm 2. The hyper-parameters of the Kalman filter are the expected observation matrix $\operatorname{\mathbb{E}}[\mathbf{C}]$ and the variances of the noises in the system, i.e., $\sigma^{2}_{\mathbf{w}},\Sigma_{\mathbf{C}},\Sigma_{\mathbf{H}}$. Here, we dropped subscript $t$ for $\mathbf{C}_{t}$ for simplicity in our modeling.

Figure: Algorithm 2 Optimizer with Kalman Filter

### 3.2 Algorithm simplification

While Algorithm 2 provides a general framework for applying Kalman filtering to various optimizers, it faces significant challenges in terms of computational and memory demands. Specifically, accurately computing the Hessian matrix ($\mathbf{H}$) for large models under DP constraints is infeasible, the matrix inversion step introduces a cubic computational complexity ($\mathcal{O}(d^{3})$), and storing the covariance matrix ($\mathbf{P}_{t}$) at each iteration requires quadratic memory ($\mathcal{O}(d^{2})$). To address these limitations and arrive at a practical and memory-efficient algorithm, we propose the following simplifications:

Constant $\mathbf{C}_{t}$: We simplify the random matrix $\mathbf{C}_{t}$ to an identity matrix $\mathbf{I}_{d}$, and simply model the randomness in the observation as additive noise only. This simplification is achieved in two steps corresponding to the two sources of randomness in $\mathbf{C}_{t}$. First, we remove the impact of clipping in $\mathbf{C}_{t}$. This is justified when 1) the clipping threshold $C$ is large enough so that clipping is inactive; or when 2) the clipped gradient $\nabla F_{C}(\mathbf{x})$ has zero curl, so that our method optimizes $F_{C}(\mathbf{x})$, where

$$ $F_{C}(\mathbf{x})=\int^{1}_{0}\nabla F_{C}(z\mathbf{x})^{\top}\mathbf{x} \mathrm{d}z,\;\nabla F_{C}(\mathbf{x})=\frac{1}{N}\sum_{\xi\in\mathcal{D}} \mathrm{clip}\left(\nabla f(\mathbf{x};\xi),C\right).$ (7) $$

Under this assumption, the minibatch clipped gradient $\frac{1}{B}\sum_{\xi\in\mathcal{B}}\mathrm{clip}\left(\nabla f(\mathbf{x},\xi)
,C\right)$ is an unbiased estimation of $\nabla F_{C}(\mathbf{x}).$ Thus, $\operatorname{\mathbb{E}}[\mathbf{C}_{t}]=\mathbf{I}_{d}$. Second, the randomness of sub-sampling in $\mathbf{C}_{t}$ is removed by assuming that the sub-sampled mini-batch gradient only causes additive noise, i.e., $\frac{1}{B}\sum_{\xi\in\mathcal{B}}\mathrm{clip}\left(\nabla f(\mathbf{x},\xi)
,C\right)=\nabla F_{C}(\mathbf{x})+\mathbf{w}_{\mathrm{SGD}},$ so that $\mathbf{C}=\mathbf{I}_{d}$ and $\Sigma_{\mathbf{C}}=0.$

Hessian estimation: We apply the following “trick” to bypass the explicit computation of the Hessian matrix, $\mathbf{H}_{t}$. Instead, we approximate the Hessian-vector product, $\mathbf{H}_{t}(\mathbf{d}_{t-1})$, using a finite difference method (Pearlmutter, ):

$$ $\mathbf{H}_{t}\mathbf{d}_{t-1}=\frac{\nabla F(\mathbf{x}_{t}+\gamma\mathbf{d}_ {t-1})-\nabla F(\mathbf{x}_{t})}{\gamma}+\mathcal{O}(\gamma)\approx\frac{1}{B} \sum_{\xi\in\mathcal{B}}\frac{\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1}; \xi)-\nabla f(\mathbf{x}_{t};\xi)}{\gamma},$ $$

This estimation eliminates the need for expensive Hessian computations and significantly reduces both memory and computational complexity.

Simplification of $\Sigma_{H}$: Few matrix computation steps in Algorithm 2 are extremely time-consuming and memory inefficient.
Specifically, the matrix inversion in Line 8 of Algorithm 2 incurs a cubic computational cost ($\mathcal{O}(d^{3})$), which becomes prohibitive for large models with billions of parameters. Therefore, we simplify the algorithm by assuming that the covariance matrix in the Kalman filter is a time-invariant diagonal matrix, i.e., $\Sigma_{H}=\sigma_{H}^{2}\mathbf{I}_{d}$. Under this assumption, the matrices $\mathbf{P}_{t}$ and $\mathbf{K}_{t}$ become $p_{t}\mathbf{I}_{d},k_{t}\mathbf{I}_{d}$, for some scalars $p_{t},k_{t}$. Thus, we reduce all matrix computations to efficient scalar-vector multiplications. Furthermore, the memory complexity for storing these matrices is reduced to $\mathcal{O}(1)$, making our algorithm a viable option for DP training of large-scale models.

Filter gain simplification: With the above simplification, the updates of $p_{t},k_{t}$ simplifies to $p_{t}=\frac{(\sigma^{2}_{\mathbf{w}}-\sigma^{2}_{H})(p_{t-1}+\sigma^{2}_{
\mathbf{v}}+\sigma^{2}_{H})}{p_{t-1}+\sigma^{2}_{\mathbf{v}}+\sigma^{2}_{
\mathbf{w}}}$, $k_{t}=\frac{p_{t-1}+\sigma^{2}_{\mathbf{v}}+\sigma^{2}_{H}}{p_{t-1}+\sigma^{2}
_{\mathbf{v}}+\sigma^{2}_{\mathbf{w}}}.$ Therefore,
$k_{t}$ converges to its stable value with a linear rate, i.e., $\left\lVert k_{t}-k_{\infty}\right\rVert=\mathcal{O}(c_{k}^{t}),$ where $c_{k}=\frac{2\sigma^{2}_{\mathbf{w}}+3\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}-
\sqrt{(\sigma_{\mathbf{v}}^{2}+\sigma^{2}_{H})(4\sigma^{2}_{\mathbf{w}}+\sigma
^{2}_{\mathbf{v}}-3\sigma^{2}_{H})}}{2\sigma^{2}_{\mathbf{w}}+3\sigma^{2}_{H}+
\sigma^{2}_{\mathbf{v}}+\sqrt{(\sigma_{\mathbf{v}}^{2}+\sigma^{2}_{H})(4\sigma
^{2}_{\mathbf{w}}+\sigma^{2}_{\mathbf{v}}-3\sigma^{2}_{H})}}\in(0,1)$ (see derivations in Appendix A.3). So we can use constant $k_{t}=\kappa,\forall\;t$ to further simplify the algorithm and avoid iteratively updating $p_{t}$ and recomputing $k_{t}$ for each step.

With the above simplifications, the complex Algorithm 2 simplifies to DiSK in Algorithm 3. For a detailed walkthrough of this simplification process, please refer to Appendix A.3. This simplified algorithm computes and privatizes the linear combination of the gradient evaluated at two points, $\mathbf{x}_{t},\mathbf{x}_{t}+\gamma\mathbf{d}_{t}$. Then, this privatized combination undergoes exponential weighted averaging to obtain $\tilde{\mathbf{g}}_{t}$, which serves as the input to the base optimizer (Line 7 of Algorithm 3).

Figure: Algorithm 3 DiSK: Differentially private optimizer with Simplified Kalman filter

### 3.3 Additional discussion

Memory and computation cost: Algorithm 3 requires one additional forward step to compute $f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)$. This means that the algorithm has at most twice the computational cost of the baseline DP-SGD algorithm. Moreover, Algorithm 3 only requires two additional states to store: $\tilde{\mathbf{g}}_{t}$ and $\mathbf{d}_{t}$. Compared to DPSGD, which requires storing the model and the gradient, DiSK has at most twice the memory cost; and compared to DPAdam, which requires storing the first- and second-order moments, the algorithm has at most $1.5\times$ the memory cost.

Connection to NAG and \xweditSTORM: Remarkably, Algorithm 3
has an implicit connection to the (unified) Nesterov accelerated gradient (NAG) method (Shen et al., ; Sutskever et al., ) \xweditand the Stochastic Recursive Momentum (STORM) algorithm (Cutkosky & Orabona, ). Specifically, by letting the OptimizerUpdate be SGD, assuming clipping is inactive, and setting $\mathbf{w}_{t}=0$ (i.e., without privatizing the gradient) in Algorithm 3, we can make a clear connection:
On one hand, DiSK reduces to NAG by choosing $\gamma=\frac{1-\kappa}{\kappa},\eta=\kappa,1-\kappa=\mu$, and $B=N$. On the other hand, the update of DiSK matches STORM by choosing $\gamma=-1,\kappa=\alpha$, and $B=1$. A detailed derivation and discussion is provided in Appendix A.4.
This observation reveals an intriguing connection between NAG (designed for acceleration) and STORM (focused on variance reduction), unifying them within the framework of DiSK, a Kalman filtering-based algorithm.

Connection to DOPPLER: DOPPLER (Zhang et al., ) and DiSK both use a filter to separate the gradient signal from the DP noise. If $\mathbf{g}_{t}$ only evaluates the gradient at one point instead of using the linear combination of gradients at two points in Algorithm 3, Line 5, DiSK becomes DOPPLER with a first-order filter. The key difference is that DOPPLER assumes an underlying low-frequency dynamic of the gradient and applies a time-invariant low-pass filter. Designing the optimal low-pass filter relies on the prior knowledge of the gradient frequency spectrum, which is hard to obtain in practice, and implementing such a high-order low-pass filter results in a large memory overhead. While in DiSK, we do not assume the frequency property of the gradient signal. Instead, we incorporate the gradient dynamics into the filtering procedure and use the Kalman filter, a predictive filtering approach, to reduce the impact of DP noise.

## 4 Theoretical analysis

This section includes theoretical analyses of Algorithm 3. Our study establishes the convergence, provable noise reduction, and the privacy-utility trade-off of the algorithm.
To facilitate our analysis, we make the following assumptions:

###### A 1 (Smoothness)

$f(\cdot,\xi)$ is $L$-smooth for any $\xi$, i.e.,
$\left\lVert\nabla f(\mathbf{x};\xi)-\nabla f(\mathbf{y};\xi)\right\rVert\leq L
\left\lVert\mathbf{x}-\mathbf{y}\right\rVert,\;\forall\xi\in\mathcal{D},\;
\forall\mathbf{x},\mathbf{y}\in\mathbb{R}^{d}.$

###### A 2 (Bounded Variance)

The per-sample gradient has bounded variance with
$\operatorname{\mathbb{E}}_{\xi\in\mathcal{D}}\left\lVert\nabla f(\mathbf{x};
\xi)-\nabla F(\mathbf{x})\right\rVert^{2}\leq\sigma^{2}_{\text{SGD}},\;\forall
\mathbf{x}\in\mathbb{R}^{d},$
where $\operatorname{\mathbb{E}}_{\xi\in\mathcal{D}}[\cdot]$ denotes the expectation taken on the randomness over $\xi$ that is uniformly sampled from dataset $\mathcal{D}$.

###### A 3 (Bounded Gradient)

Each per-sample gradient has a bounded norm, i.e.,
$\left\lVert\nabla f(\mathbf{x};\xi)\right\rVert\leq G,\;\forall\mathbf{x}\in
\mathbb{R}^{d},\forall\xi\in\mathcal{D}.$

Let us briefly comment on these assumptions: A1 and A2 are standard in non-convex optimization (Allen-Zhu & Hazan, ; Zaheer et al., ; Abadi et al., ); and A3 is commonly used in analyzing the convergence of DP algorithms (Abadi et al., ; Wang et al., ; Andrew et al., ) to avoid introducing the clipping bias. Since the impact of clipping is not the major focus of this paper, we follow this tradition and use A3 to simplify our theoretical analysis.

###### A 1 (Smoothness)

###### A 2 (Bounded Variance)

###### A 3 (Bounded Gradient)

### 4.1 Convergence analysis

We provide the following convergence results for Algorithm 3, assuming $\sigma_{\mathrm{DP}}$ being a constant.

###### Theorem 2

Assume A1-A3 hold. Fix $\sigma_{\mathrm{DP}}^{2}$ and choose $C\geq(1+\frac{2(1-\kappa)}{\kappa})G$, $\kappa,\eta$ satisfy

$$ $\eta<\frac{1+\kappa}{2L(1+2(1-\kappa)^{2}\beta L(2+\left\lvert 1+\gamma\right \rvert C_{\gamma}))},\;\kappa>1-\frac{1}{\sqrt{1+4\eta^{2}L^{2}+\xwedit{\left \lvert 1+\gamma\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)}}},$ $$

and run Algorithm 3 for $T$ iterations. Then,

$$ $\displaystyle\frac{1}{T}\sum^{T}_{t=0}\operatorname{\mathbb{E}}\left\lVert \nabla F(\mathbf{x}_{t})\right\rVert^{2}$ $\displaystyle\leq\frac{2(F(\mathbf{x}_{0})+\beta\left\lVert\nabla F(\mathbf{x} _{0})\right\rVert^{2}-F^{\star})}{C_{1}\eta T}$ $\displaystyle\quad+\frac{2(\beta+\eta^{2}L)\kappa^{2}}{C_{1}\eta}\left(\frac{ \xwedit{(2+\left\lvert 1+\gamma\right\rvert)}\sigma^{2}_{\mathrm{SGD}}}{B}+d \sigma^{2}_{\mathrm{DP}}\right),$ (8) $$

where $C_{\gamma}=1+\frac{4(2+1/\kappa+\left\lvert 1+\gamma\right\rvert)}{\gamma^{2}}$, $C_{1}=(1+\kappa-2\eta L)-4(\beta+\eta^{2}L)(1-\kappa)^{2}L^{2}\eta\xwedit{
\left(2+\left\lvert 1+\gamma\right\rvert C_{\gamma}\right)}>0,$
and $\beta\geq\frac{\eta(1-\kappa)/2+\eta^{2}L(1-\kappa)^{2}(1+4\eta^{2}L^{2}+
\xwedit{\left\lvert 1+\gamma\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}
\right)})}{1-(1-\kappa)^{2}(1+4\eta^{2}L^{2}+\xwedit{\left\lvert 1+\gamma
\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)})}\geq 0$ are some non-negative constants.

The proof of Theorem 2 is relegated to Appendix B. Notice that when $\kappa=1$, we have $\beta=0,C_{1}=2(1-\eta L)$ and the convergence result recovers that of DPSGD (Ghadimi & Lan, ; Zhang et al., ).
With this theorem, we can choose specific parameters and obtain the following corollary:

###### Corollary 1

Under the conditions of Theorem 2, choose $\gamma=-1,$

$$ $\displaystyle\eta$ $\displaystyle=\min\left\{\frac{1}{L(2+4/C_{\kappa}-C_{\kappa})},\frac{1}{C_{ \kappa}L}\sqrt{\frac{2C_{\kappa}L(F(\mathbf{x}_{0})-F^{\star})+\left\lVert \nabla F(\mathbf{x}_{0})\right\rVert^{2}}{T(2\sigma^{2}_{\mathrm{SGD}}/B+d \sigma^{2}_{\mathrm{DP}})}}\right\},$ $\displaystyle\beta$ $\displaystyle=\frac{\eta(1-\kappa)/2+\eta^{2}L(1-\kappa)^{2}(1+4\eta^{2}L^{2}) }{1-(1-\kappa)^{2}(1+4\eta^{2}L^{2})}\leq\frac{1}{2C_{\kappa}L},$ $\displaystyle\kappa$ $\displaystyle=C_{\kappa}L\eta\leq 1,\;\;\textrm{and}\;\;C_{\kappa}=\min\left\{ \frac{\left\lVert\nabla F(\mathbf{x}_{0})\right\rVert^{2}}{2L(F(\mathbf{x}_{0} )-F^{\star})},1\right\},$ $$

then Algorithm 3 satisfies:

$$ $\frac{1}{T}\sum^{T}_{t=0}\operatorname{\mathbb{E}}[\left\lVert\nabla F(\mathbf {x}_{t})\right\rVert^{2}]\leq 4\sqrt{\frac{C_{\kappa}L(F(\mathbf{x}_{0})-F^{ \star}))(2\sigma^{2}_{\mathrm{SGD}}/B+d\sigma^{2}_{\mathrm{DP}})}{T}}=\mathcal {O}\left(\sqrt{\frac{d}{T}}\right).$ (9) $$

Convergence improvement: The above result implies that the order of the number of iterations $T$ needed for convergence of Algorithm 3 is the same as of DPSGD (Ghadimi & Lan, ; Zhang et al., ).
However, Algorithm 3 has a constant factor improvement in its convergence rate over DPSGD. This improvement is resulted from the presence of $C_{\kappa}\leq 1$ in the numerator in the RHS of (9). More specifically, if $C_{\kappa}<1,$ i.e., $\left\lVert\nabla F(\mathbf{x}_{0})\right\rVert^{2}<2L(F(\mathbf{x}_{0})-F^{
\star}),$ then the convergence bound reduces by a factor of $\sqrt{1/C_{\kappa}},$
and DiSK  has clear theoretical improvement over vanilla DPSGD. Case I: For ($\mu$-strongly) convex problems, it is guaranteed that $2\mu(F(\mathbf{x}_{0})-F^{\star})\leq\left\lVert\nabla F(\mathbf{x}_{0})\right
\rVert^{2}\leq 2L(F(\mathbf{x}_{0})-F^{\star}).$ Therefore, the factor $C_{\kappa}\in[\mu/L,1].$ Case II: When training highly non-convex deep learning models, the Lipschitz constant $L$ can be large (Herrera et al., ), and $2L(F(\mathbf{x}_{0})-F^{\star})$ can be much larger than $\left\lVert\nabla F(\mathbf{x}_{0})\right\rVert^{2}$, which results in a considerable algorithm performance improvement compared to vanilla DPSGD.

While the above corollary is for the case of $\gamma=-1$, we can obtain performance improvement for $\gamma\neq-1$ as well. The choice $\gamma=-1$ is optimized for the worst case (function satisfying our assumptions) based on our upper-bound in (2). However, it is possible that for functions satisfying additional assumptions, other choices of $\gamma$ lead to better convergence result. This fact is further explained after presenting the proof of the theorem in Appendix B.2.

As a last remark in this subsecion, notice that our iteration complexity improvement does not require any additional assumption on the problem. In contrast, existing works with convergence improvement require additional assumptions, e.g., on the correlation between the gradients (Zhang et al., ), or on the trace of Hessian (Choquette-Choo et al., ; Li et al., ).

###### Theorem 2

###### Corollary 1

### 4.2 Privacy-utility trade-off

To provide DP, the Subsampled Gaussian mechanism is used in Lines 4 and 5 of Algorithm 3.
Instead of $\nabla f(\mathbf{x}_{t};\xi)$, we treat $\frac{1-\kappa}{\kappa\gamma}\nabla f(\mathbf{x}_{t}-\gamma\mathbf{d}_{t-1};
\xi)+(1-\frac{1-\kappa}{\kappa\gamma})\nabla f(\mathbf{x}_{t};\xi)$ as the per-sample gradient and apply the Subsampled Gaussian mechanism to privatize it. Therefore, Algorithm 3 and DPSGD share a similar privacy guarantee, as the privacy proof directly follows the Subsampled Gaussian mechanism and the composition of $T$ iterations in Theorem 1. Specifically, by combining Theorem 1 and Theorem 2, we obtain the following privacy-utility trade-off:

###### Theorem 3

Assume A1-A3 holds. Run Algorithm 3 for $T=\frac{\sqrt{2}N\epsilon\sigma_{\mathrm{SGD}}}{C\sqrt{Bd\ln(1/\delta)}}$ iterations, and choose $\kappa,\beta$ according to the choices in Corollary 1, then

$$ $\frac{1}{T}\sum^{T}_{t=0}\operatorname{\mathbb{E}}\left\lVert\nabla F(\mathbf{ x}_{t})\right\rVert^{2}\leq\frac{8C\sigma_{\mathrm{SGD}}\sqrt{C_{\kappa}L(F( \mathbf{x}_{0})-F^{\star})d\ln(1/\delta)}}{\sqrt{B}N\epsilon}=\mathcal{O}\left (\frac{\sqrt{d\ln(1/\delta)}}{N\epsilon}\right)$ (10) $$

Similar to our convergence result, compared to vanilla DPSGD, the privacy-utility trade-off of Algorithm 3 reduces by a constant factor of $\sqrt{1/C_{\kappa}}$. To the best of our knowledge, this is the first theoretical result on the utility improvement of a DPSGD-type algorithm without any additional assumptions on the problem.

###### Theorem 3

## 5 Numerical experiments

We perform extensive pre-training and fine-tuning experiments on various image classification (CV) and natural language processing (NLP) tasks using different base algorithms, privacy budgets, and models. The implementation details of the experiments are given in Appendix C.1. The link to the code is [https://anonymous.4open.science/r/KalmanDP-BEDB](https://anonymous.4open.science/r/KalmanDP-BEDB).

### 5.1 Experiment settings

Dataset: We train the models on one synthetic dataset, four CV datasets, including MNIST (Deng, ), CIFAR-10/CIFAR-100 (Krizhevsky et al., ), and ImageNet-1k (Deng et al., ), and three NLP dataset, including GLUE (Wang et al., ), E2E (Novikova et al., ), and DART (Nan et al., ).

Model: For the CV tasks, we use three different models, including a 5-layer CNN, WideResNet (WRN) (De et al., ), and ViT (Dosovitskiy et al., ), representing three typical CV model structures. For the NLP task, we use the RoBERTa (Liu et al., ) and the GPT-2 (Radford et al., ) models. For pre-training, the models are initialized with random weights, and for fine-tuning with ViT, RoBERTa, and GPT-2, we directly use the checkpoints on HuggingFace (Wolf et al., ).

Algorithm: We use the DP version of SGD and Adam for CV tasks, and AdamW for NLP tasks as baselines. Then, we apply DiSK to compare their performance. Additional results on LoRA (Li et al., ) are given in Appendix C.4. In the results, we use KF- to denote the privatized version of the algorithms with the DiSK framework.

Hyper-parameter choices: We tune the hyper-parameters using a grid search. Specifically, we conduct a grid search on the batch size $B$, total epochs $E=NT/B$, and step size $\eta$ for each given privacy budget $\epsilon$. For all experiments, we fix the privacy parameter $\delta=1/N^{1.1}$ to obtain a reasonable privacy notion. Detailed hyper-parameter choices are discussed in Appendix C.2.

### 5.2 Numerical results

CV tasks: We first train the CV models with randomly initialized weights on different image datasets. The results for 5-layer CNN on the MNIST dataset, 5-layer CNN on the CIFAR-10 dataset, and WRN-16-4 on the CIFAR-100 datasets with different privacy budgets are given in Figure 1. DiSK  significantly outperforms the base algorithm across all used privacy budgets.

The test accuracy curves during the training for 5-layer CNN on the MNIST dataset, WRN-16-4 on the CIFAR-100 dataset, and ViT-small on the ImageNet-1k dataset are given in Figure 2. The optimizer with DiSK  convergence faster than the base algorithm on all tasks and reaches a higher final accuracy at a given privacy budget. The test accuracy of CIFAR-100 achieves $41.8\%$, and ImageNet-1k achieves $36.4\%$, which outperforms the SOTA results that apply data augmentation under the same privacy budget ($40.6\%$ for CIFAR-100 (Bao et al., ) and $32.4\%$ for ImageNet-1k (De et al., )). Additional comparisons on different models, algorithms, and fine-tuning CIFAR-100 are given in Appendix C.3.

NLP tasks: We fine-tune a pre-trained RoBERTa-base model(^1^11[https://huggingface.co/FacebookAI/roberta-base](https://huggingface.co/FacebookAI/roberta-base)) on the GLUE datasets. The final test accuracy for is given in Table 1. We follow the same training script and hyper-parameter choices in the experiments as Bu et al. (). Compared with the base algorithm (DPAdamW), DiSK improves the final accuracy on all tasks by at least $3.8\%$ when $\epsilon=1$ and $1.3\%$ when $\epsilon=6.7$. Additional results for text generation tasks on the E2E ad DART datasets are given in Appendix C.4.

Figure: (a) MNIST
Refer to caption: extracted/5898046/figures/Mnist-epsilon.png

Figure: (a) MNIST, $\epsilon=0.5$
Refer to caption: extracted/5898046/figures/Mnist.png

**Table 1: Test accuracy of fine-tuning result on the GLUE dataset.**
| Task | ($\epsilon=\infty$) | $\epsilon=6.7$ | $\epsilon=1$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Non-DP | DP | KF-DP | KF-DPLora | DP | KF-DP | KF-DPLora |  |
| MNLI | 87.6 | 83.2 | 84.8 | 85.9 | 80.7 | 82.0 | 84.7 |
| QNLI | 92.8 | 87.5 | 88.9 | 90.5 | 86.0 | 88.7 | 90.3 |
| SST-2 | 94.8 | 91.5 | 92.8 | 93.1 | 91.4 | 91.5 | 92.9 |
| QQP | 91.9 | 85.8 | 88.5 | 89.0 | 84.2 | 86.9 | 87.8 |

Ablation study: We conduct ablation studies on the choice of the hyper-parameters of DiSK, specifically, how $\kappa,\gamma$ impact the algorithm performance. The results are presented in Appendix C.3.

Improvements over the SOTA: Table LABEL:tab:app:sota in Appendix C.5 summarizes the improvements of DiSK over the SOTA on different tasks.

## Reproducibility Statement

For the algorithm implementation, we provide a link to an anonymous downloadable source code in Section 5. We also discussed the data, model, and algorithms used in the experiment in Section 5 and in Appendix C. For our theoretical results provided in Section 4, we also include clear explanations of the assumptions in the section. The complete proof of the theorems can be found in Appendix B.

## Broader Impact

This paper presents work that aims to advance the field of Machine Learning, combining optimization with signal processing and control societies. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

## Appendix A Additional discussion

### A.1 Background on Kalman Filter

In this section, we provide an introduction and derivation of the Kalman filter. Kalman filter is introduced in (Kalman, ) and wildly used for control systems in accurately estimating system states with noisy observation and known system dynamics.
Kalman filter assumes the system is a linear (time-invariant) system:

$$ $\displaystyle\boldsymbol{\theta}_{t}$ $\displaystyle=\mathbf{A}\boldsymbol{\theta}_{t-1}+\mathbf{u}_{t}+\mathbf{v}_{t},$ (System update) $\displaystyle\boldsymbol{\psi}_{t}$ $\displaystyle=\mathbf{C}\boldsymbol{\theta}_{t}+\mathbf{w}_{t},$ (Observation) $$

where $\mathbf{A},\mathbf{C}$ are known matrix, and $\mathbf{v}_{t}\sim\mathcal{N}(0,\Sigma_{\mathbf{v}}),\mathbf{w}_{t}\sim
\mathcal{N}(,\Sigma_{\mathbf{w}})$ follows Gaussian distributions with known covariance matrices $\Sigma_{\mathbf{v}},\Sigma_{\mathbf{w}}$, and $\mathbf{u}_{t}$ is known input signal. The goal of the Kalman filter is to estimate $\boldsymbol{\theta}_{t}$ with the observation of $\boldsymbol{\psi}_{t}$ with the least mean-square error.

Derivation: First, we denote the estimation of $\boldsymbol{\theta}_{t}$ as $\tilde{\boldsymbol{\theta}}_{t}$, and the covariance of $\boldsymbol{\theta}_{t}-\tilde{\boldsymbol{\theta}}_{t}$ as

$$ $\mathbf{P}_{t}=\operatorname{\mathbb{E}}[(\boldsymbol{\theta}_{t}-\tilde{ \boldsymbol{\theta}}_{t})(\boldsymbol{\theta}_{t}-\tilde{\boldsymbol{\theta}}_ {t})^{\top}],$ $$

Then, the system dynamics give an unbiased prediction:

$$ $\tilde{\boldsymbol{\theta}}_{t|t-1}=\mathbf{A}\tilde{\boldsymbol{\theta}}_{t-1 }+\mathbf{u}_{t},$ (11) $$

and its covariance is

$$ $\mathbf{P}_{t|t-1}=\operatorname{\mathbb{E}}[(\boldsymbol{\theta}_{t}-\tilde{ \boldsymbol{\theta}}_{t|t-1})(\boldsymbol{\theta}_{t}-\tilde{\boldsymbol{ \theta}}_{t|t-1})^{\top}]=\mathbf{A}\mathbf{P}_{t-1}\mathbf{A}^{\top}+\Sigma_{ \mathbf{v}}.$ (12) $$

With $\tilde{\boldsymbol{\theta}}_{t|t-1}$, we have the prediction error $\Delta\boldsymbol{\psi}_{t}=\boldsymbol{\psi}_{t}-\mathbf{C}\tilde{\boldsymbol
{\psi}}_{t|t-1}.$ Since the system is linear, we would like to use the prediction error to correct the prediction:

$$ $\tilde{\boldsymbol{\theta}}_{t}=\tilde{\boldsymbol{\theta}}_{t|t-1}+\mathbf{K} _{t}\Delta\boldsymbol{\psi}_{t}=\mathbf{K}_{t}\boldsymbol{\psi}_{t}+(\mathbf{I }-\mathbf{K}_{t}\mathbf{C})\tilde{\boldsymbol{\theta}}_{t|t-1}.$ (13) $$

The goal of the Kalman filter is to minimize the mean-square error: $\min_{\mathbf{K}}\operatorname{\mathbb{E}}[\left\lVert\boldsymbol{\theta}_{t}-
\tilde{\boldsymbol{\theta}}_{t}\right\rVert^{2}],$ which is equivalent to minimizing $\mathrm{tr}(\mathbf{P}_{t}).$
From the definition of $\tilde{\boldsymbol{\theta}}_{t},$ we have:

$$ $\displaystyle\mathbf{P}_{t}$ $\displaystyle=\operatorname{\mathbb{E}}[(\boldsymbol{\theta}_{t}-\tilde{ \boldsymbol{\theta}}_{t})(\boldsymbol{\theta}_{t}-\tilde{\boldsymbol{\theta}}_ {t})^{\top}]$ $\displaystyle=\operatorname{\mathbb{E}}[(\boldsymbol{\theta}_{t}-\mathbf{K}_{t }\boldsymbol{\psi}_{t}+(\mathbf{I}-\mathbf{K}_{t}\mathbf{C})\tilde{\boldsymbol {\theta}}_{t|t-1})(\boldsymbol{\theta}_{t}-\mathbf{K}_{t}\boldsymbol{\psi}_{t} +(\mathbf{I}-\mathbf{K}_{t}\mathbf{C})\tilde{\boldsymbol{\theta}}_{t|t-1})^{ \top}]$ $\displaystyle=\operatorname{\mathbb{E}}[(\boldsymbol{\theta}_{t}-\tilde{ \boldsymbol{\theta}}_{t|t-1}-\mathbf{K}_{t}\mathbf{C}(\boldsymbol{\theta}_{t}- \tilde{\boldsymbol{\theta}}_{t|t-1})-\mathbf{K}_{t}\mathbf{w}_{t})(\boldsymbol {\theta}_{t}-\tilde{\boldsymbol{\theta}}_{t|t-1}-\mathbf{K}_{t}\mathbf{C}( \boldsymbol{\theta}_{t}-\tilde{\boldsymbol{\theta}}_{t|t-1})-\mathbf{K}_{t} \mathbf{w}_{t})^{\top}]$ $\displaystyle=\mathbf{P}_{t|t-1}-\mathbf{K}_{t}\mathbf{C}\mathbf{P}_{t|t-1}-( \mathbf{K}_{t}\mathbf{C}\mathbf{P}_{t|t-1})^{\top}+\mathbf{K}_{t}(\mathbf{C} \mathbf{P}_{t|t-1}\mathbf{C}^{\top}+\Sigma_{\mathbf{w}})\mathbf{K}_{t}^{\top}.$ $$

Taking partial derivative to the trace of $\mathbf{P}_{t}$ with respect to $\mathbf{K}_{t}$, and set it to zero, we have:

$$ $\frac{\partial\mathrm{tr}(\mathbf{P}_{t})}{\partial\mathbf{K}_{t}}=-2(\mathbf{ C}\mathbf{P}_{t|t-1})^{\top}+2(\mathbf{C}\mathbf{P}_{t|t-1}\mathbf{C}^{\top}+ \Sigma_{\mathbf{w}})\mathbf{K}_{t}^{\top}=0,$ $$

which gives

$$ $\mathbf{K}_{t}=\mathbf{P}_{t|t-1}\mathbf{C}^{\top}(\mathbf{C}\mathbf{P}_{t|t-1 }\mathbf{C}^{\top}+\Sigma_{\mathbf{w}})^{-1}.$ (14) $$

Substitute $\mathbf{K}_{t}$ back to $\mathbf{P}_{t},$ we can simplify

$$ $\displaystyle\mathbf{P}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}-\mathbf{K}_{t}\mathbf{C}\mathbf{P}_{t|t-1}-( \mathbf{K}_{t}\mathbf{C}\mathbf{P}_{t|t-1})^{\top}+\mathbf{K}_{t}(\mathbf{C} \mathbf{P}_{t|t-1}\mathbf{C}^{\top}+\Sigma_{\mathbf{w}})\mathbf{K}_{t}^{\top}$ (15) $\displaystyle=\mathbf{P}_{t|t-1}-\mathbf{K}_{t}\mathbf{C}\mathbf{P}_{t|t-1}=( \mathbf{I}-\mathbf{K}_{t}\mathbf{C})\mathbf{P}_{t|t-1}.$ $$

Combining (11)-(15) together, we have the update of the Kalman filter:

$$ $\displaystyle\tilde{\boldsymbol{\theta}}_{t|t-1}$ $\displaystyle=\mathbf{A}\tilde{\boldsymbol{\theta}}_{t-1}+\mathbf{u}_{t}$ $\displaystyle\mathbf{P}_{t|t-1}$ $\displaystyle=\mathbf{A}\mathbf{P}_{t-1}\mathbf{A}^{\top}+\Sigma_{\mathbf{v}}$ $\displaystyle\mathbf{K}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}\mathbf{C}^{\top}(\mathbf{C}\mathbf{P}_{t|t-1} \mathbf{C}^{\top}+\Sigma_{\mathbf{w}})^{-1}$ $\displaystyle\tilde{\boldsymbol{\theta}}_{t}$ $\displaystyle=(\mathbf{I}-\mathbf{K}_{t}\mathbf{C})\tilde{\boldsymbol{\theta}} _{t|t-1}+\mathbf{K}_{t}\boldsymbol{\psi}_{t}$ $\displaystyle\mathbf{P}_{t}$ $\displaystyle=(\mathbf{I}-\mathbf{K}_{t}\mathbf{C})\mathbf{P}_{t|t-1}.$ $$

Variants of Kalman filter: Kalman filter is designed for estimating the states following linear dynamics and achieves optimal performance, i.e., gives the smallest mean square error of the estimation when the system is linear (Welch et al., ). Extended Kalman filter (EKF) and unscented Kalman filter (UKF) are developed to deal with non-linear systems. EKF linearizes the non-linear system at each step and performs the Kalman filter on the linearized system (Ribeiro, ), while UKF takes the effect of the system non-linearity to the noise distribution into consideration and applies the unscented transform on the noise distribution and applies the Kalman filter on the system and the noise distribution (Wan & Van Der Merwe, ). Other extensions of the Kalman filter have been developed for special cases, including multiplicative noise and noisy input $\mathbf{u}_{t}$ (Wu et al., ).

### A.2 Other system dynamics for DPSGD

In this section, we want to discuss other possible formulations of the system dynamics for DPSGD to apply the Kalman filter.

Optimization variable and gradient version 1: Other than only using the gradients’ dynamics in the main paper, we can construct the system taking the dynamics of the optimization variable $\mathbf{x}$ into the system (Vuckovic, ):

$$ $\displaystyle\left[\begin{array}[]{c}\mathbf{x}_{t+1}\\ \nabla F(\mathbf{x}_{t+1})\end{array}\right]$ $\displaystyle=\left[\begin{array}[]{cc}\mathbf{I}&-\eta\mathbf{I}\\ 0&\mathbf{I}\end{array}\right]\left[\begin{array}[]{c}\mathbf{x}_{t}\\ \nabla F(\mathbf{x}_{t})\end{array}\right]+\left[\begin{array}[]{c}\eta\mathbf {w}_{t}\\ 0\end{array}\right],$ $\displaystyle\mathbf{y}_{t}$ $\displaystyle=\mathbf{x}_{t}.$ $$

However, this dynamic is inaccurate as the dynamics of the gradient are simplified to $\nabla F(\mathbf{x}_{t+1})=\nabla F(\mathbf{x}_{t})$. Although the update can further incorporate with the momentum methods, i.e., by adding a momentum $\mathbf{m}_{t}$ into the system, it fails to reveal the actual dynamic of the system.

Optimization variable and gradient version 2: Instead of treating the gradient as a constant, we can further utilize the Hessian to reveal the gradient dynamics:

$$ $\displaystyle\left[\begin{array}[]{c}\mathbf{x}_{t+1}\\ \mathbf{x}_{t}\\ \nabla F(\mathbf{x}_{t})\end{array}\right]$ $\displaystyle=\left[\begin{array}[]{ccc}\mathbf{I}-\eta\mathbf{H}&\eta\mathbf{ H}&-\eta\mathbf{I}\\ \mathbf{I}&0&0\\ \mathbf{H}&-\mathbf{H}&\mathbf{I}\end{array}\right]\left[\begin{array}[]{c} \mathbf{x}_{t}\\ \mathbf{x}_{t-1}\\ \nabla F(\mathbf{x}_{t-1})\end{array}\right]+\left[\begin{array}[]{c}\eta \mathbf{w}_{t}\\ 0\\ 0\end{array}\right],$ $\displaystyle\mathbf{y}_{t}$ $\displaystyle=\mathbf{x}_{t}.$ $$

This system is more accurate in evaluating the gradient at the cost of using an extra $\mathbf{x}_{t-1}$ state and a larger transition matrix. For non-linear problems $F(\mathbf{x})$, where $\mathbf{H}$ is not a constant, we can apply the extended Kalman filter and replace $\mathbf{H}$ with $\mathbf{H}_{t}$ that linearizes the problem at each time step.

Gradient and Hessian: In work (Bittner & Pronzato, ), the dynamics of the gradient and Hessian have been used to construct the system:

$$ $\displaystyle\left[\begin{array}[]{c}\nabla F(\mathbf{x}_{t+1})\\ \mathbf{h}_{t+1}\end{array}\right]$ $\displaystyle=\left[\begin{array}[]{cc}\mathbf{I}&\Delta\mathbf{X}_{t}\\ 0&\mathbf{I}\end{array}\right]\left[\begin{array}[]{c}\nabla F(\mathbf{x}_{t}) \\ \mathbf{h}_{t}\end{array}\right],$ $\displaystyle\mathbf{g}_{t}$ $\displaystyle=\nabla F(\mathbf{x}_{t})+\mathbf{w}_{t},$ $$

where $\mathbf{h}_{t}=[\mathbf{H}_{1,1},\dots,\mathbf{H}_{i,i+j},\dots,\mathbf{H}_{d,
d}]^{\top},$ with $j\in[0,\dots,d-i]$ represents the entries in the upper triangular part of the Hessian matrix. $\Delta\mathbf{X}_{t}$ is constructed such that $\Delta\mathbf{X}_{t}\mathbf{h}_{t}=\mathbf{H}_{t}(\mathbf{x}_{t+1}-\mathbf{x}_
{t}).$ The system treats the Hessian matrix as a constant matrix, and the transition matrix is of size $(d+d(d-1)/2)\times d+d(d-1)/2.$

Although there are different ways to construct the Kamlan filter for gradient noise reduction, the above systems are not implementable in practical deep-learning applications. Because the transition matrices of these systems are non-diagonal, the Kalman filters have non-diagonal gain $\mathbf{K}_{t}$ and $\mathbf{P}_{t}$. Therefore, the matrix inversion operation is unavoidable when Kalman filters are implemented based on these systems. The computation complexity for the matrix inversion can be $\mathcal{O}(d^{3})$ to $\mathcal{O}(d^{6}),$ and the memory consumption is $\mathcal{O}(d^{2})$ to $\mathcal{O}(d^{4})$ for storing the matrices of the Kalman filter.

### A.3 Algorithm Simplification

In this section, we explain how the simplification proceeds from Algorithm 2 to Algorithm 3. Recall that updates of Algorithm 2 is

$$ $\displaystyle\tilde{\mathbf{g}}_{t|t-1}$ $\displaystyle=\tilde{\mathbf{g}}_{t-1}+\tilde{\mathbf{H}}_{t}(\mathbf{x}_{t}- \mathbf{x}_{t-1})$ (Prediction) $\displaystyle\mathbf{P}_{t|t-1}$ $\displaystyle=\mathbf{P}_{t-1}+\Sigma_{\mathbf{H}}+\Sigma_{\mathbf{v}}$ $\displaystyle\mathbf{K}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}\operatorname{\mathbb{E}}[\mathbf{C}_{t}]^{ \top}\left(\Sigma_{\mathbf{w}}+\operatorname{\mathbb{E}}[\mathbf{C}_{t}]\left( \Sigma_{\mathbf{C}}\mathbf{S}_{t}+\mathbf{P}_{t|t-1}\right)\operatorname{ \mathbb{E}}[\mathbf{C}_{t}]^{\top}-\Sigma_{\mathbf{H}}\right)^{-1}$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=\tilde{\mathbf{g}}_{t|t-1}+\mathbf{K}_{t}(\mathbf{g}_{t}- \operatorname{\mathbb{E}}[\mathbf{C}_{t}]\tilde{\mathbf{g}}_{t|t-1})$ (Correction) $\displaystyle\mathbf{P}_{t}$ $\displaystyle=(\mathbf{I}-\mathbf{K}_{t}\operatorname{\mathbb{E}}[\mathbf{C}_{ t}])\mathbf{P}_{t|t-1}$ $\displaystyle\mathbf{S}_{t}$ $\displaystyle=\operatorname{\mathbb{E}}[\tilde{\mathbf{g}}_{t}\tilde{\mathbf{g }}_{t}^{\top}],$ $$

where $\mathbf{g}_{t}=\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\mathrm{clip}\left(
\nabla f(\mathbf{x}_{t};\xi),C\right)+\mathbf{w}_{t}.$
We apply four steps of simplification, including

- 1.
Replacing random $\mathbf{C}_{t}$ with constant $\mathbf{I}_{d}$.
- 2.
Use finite difference to estimate $\mathbf{H}_{t}\mathbf{d}_{t-1}$.
- 3.
Replace $\Sigma_{\mathbf{H}}$ with diagonal matrix $\sigma^{2}_{H}\mathbf{I}_{d}$, $\Sigma_{\mathbf{v}}$ with $\sigma^{2}_{\mathbf{v}}\mathbf{I}_{d}$, and $\Sigma_{\mathbf{w}}$ with $\sigma^{2}_{\mathbf{w}}\mathbf{I}_{d}$.
- 4.
Use fixed filter gain $\kappa$.

Step 1. By replacing $\mathbf{C}_{t}$ with $\mathbf{I}_{d}$, and $\Sigma_{\mathbf{C}}=0$, the update becomes

$$ $\displaystyle\tilde{\mathbf{g}}_{t|t-1}$ $\displaystyle=\tilde{\mathbf{g}}_{t-1}+\tilde{\mathbf{H}}_{t}(\mathbf{x}_{t}- \mathbf{x}_{t-1})$ (Prediction) $\displaystyle\mathbf{P}_{t|t-1}$ $\displaystyle=\mathbf{P}_{t-1}+\Sigma_{\mathbf{H}}+\Sigma_{\mathbf{v}}$ $\displaystyle\mathbf{K}_{t}$ $\displaystyle=\mathbf{P}_{t|t-1}\left(\Sigma_{\mathbf{w}}+\mathbf{P}_{t|t-1}- \Sigma_{\mathbf{H}}\right)^{-1}$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=\tilde{\mathbf{g}}_{t|t-1}+\mathbf{K}_{t}(\mathbf{g}_{t}-\tilde{ \mathbf{g}}_{t|t-1})$ (Correction) $\displaystyle\mathbf{P}_{t}$ $\displaystyle=(\mathbf{I}-\mathbf{K}_{t})\mathbf{P}_{t|t-1}.$ $$

Step 2. By using the finite difference to estimate $\mathbf{H}_{t}\mathbf{d}_{t-1}$, the prediction step becomes:

$$ $\tilde{\mathbf{g}}_{t|t-1}=\tilde{\mathbf{g}}_{t-1}+\frac{1}{B}\sum_{\xi\in \mathcal{B}_{t}}\frac{\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)- \nabla f(\mathbf{x}_{t};\xi)}{\gamma}\qquad\textit{(Prediction)}.$ $$

Step 3. By replacing $\Sigma$’s with diagonal matrix $\sigma^{2}\mathbf{I}_{d}$’s, the update becomes:

$$ $\displaystyle\tilde{\mathbf{g}}_{t|t-1}$ $\displaystyle=\tilde{\mathbf{g}}_{t-1}+\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}} \frac{\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)-\nabla f(\mathbf{x}_ {t};\xi)}{\gamma}$ (Prediction) $\displaystyle p_{t|t-1}$ $\displaystyle=p_{t-1}+\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}},\mathbf{P}_{t|t-1 }=p_{t|t-1}\mathbf{I}_{d}$ $\displaystyle k_{t}$ $\displaystyle=\frac{p_{t|t-1}}{p_{t|t-1}+\sigma^{2}_{\mathbf{w}}-\sigma^{2}_{H }}=\frac{p_{t-1}+\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}}{p_{t-1}+\sigma^{2}_{ \mathbf{w}}+\sigma^{2}_{\mathbf{v}}},\mathbf{K}_{t}=k_{t}\mathbf{I}_{d}$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=\tilde{\mathbf{g}}_{t|t-1}+k_{t}(\mathbf{g}_{t}-\tilde{\mathbf{g }}_{t|t-1})$ (Correction) $\displaystyle p_{t}$ $\displaystyle=(1-k_{t})p_{t|t-1}=\frac{(\sigma^{2}_{\mathbf{w}}-\sigma^{2}_{H} )(p_{t-1}+\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}})}{p_{t-1}+\sigma^{2}_{\mathbf {w}}+\sigma^{2}_{\mathbf{v}}}.$ $$

Step 4. As discussed in the main paper, the update of $k_{t},p_{t}$ becomes:

$$ $\displaystyle k_{t}$ $\displaystyle=\frac{p_{t-1}+\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}}{p_{t-1}+ \sigma^{2}_{\mathbf{w}}+\sigma^{2}_{\mathbf{v}}},$ $\displaystyle p_{t}$ $\displaystyle=\frac{(\sigma^{2}_{\mathbf{w}}-\sigma^{2}_{H})(p_{t-1}+\sigma^{2 }_{H}+\sigma^{2}_{\mathbf{v}})}{p_{t-1}+\sigma^{2}_{\mathbf{w}}+\sigma^{2}_{ \mathbf{v}}}.$ $$

Therefore, $p_{t}$ converges to $p_{\infty}=\frac{\sqrt{\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}}\sqrt{4\sigma^{2
}_{\mathbf{w}}-3\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}}-(\sigma^{2}_{H}+\sigma
^{2}_{\mathbf{v}})}{2}$, $k_{t}$ converges to $k_{\infty}=\frac{p_{\infty}+\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}}{p_{\infty}
+\sigma^{2}_{\mathbf{w}}+\sigma^{2}_{\mathbf{v}}},$ with rate $c_{k}=\frac{2\sigma^{2}_{\mathbf{w}}+3\sigma^{2}_{H}+\sigma^{2}_{\mathbf{v}}-
\sqrt{(\sigma_{\mathbf{v}}^{2}+\sigma^{2}_{H})(4\sigma^{2}_{\mathbf{w}}+\sigma
^{2}_{\mathbf{v}}-3\sigma^{2}_{H})}}{2\sigma^{2}_{\mathbf{w}}+3\sigma^{2}_{H}+
\sigma^{2}_{\mathbf{v}}+\sqrt{(\sigma_{\mathbf{v}}^{2}+\sigma^{2}_{H})(4\sigma
^{2}_{\mathbf{w}}+\sigma^{2}_{\mathbf{v}}-3\sigma^{2}_{H})}}.$

Figure: Figure 3: Convergence of $k_{t},p_{t}$.
Refer to caption: x1.png

Therefore, we define $\kappa=k_{\infty}$ and replace $k_{t}$ with $\kappa$, and the update becomes:

$$ $\displaystyle\tilde{\mathbf{g}}_{t|t-1}$ $\displaystyle=\tilde{\mathbf{g}}_{t-1}+\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}} \frac{\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)-\nabla f(\mathbf{x}_ {t};\xi)}{\gamma}$ (Prediction) $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=\tilde{\mathbf{g}}_{t|t-1}+\kappa(\mathbf{g}_{t}-\tilde{\mathbf{ g}}_{t|t-1})$ (Correction) $$

Rearrange the terms, we have:

$$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t}+\kappa\mathbf{g}_{t}+(1-\kappa) \frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\frac{\nabla f(\mathbf{x}_{t}+\gamma \mathbf{d}_{t-1};\xi)-\nabla f(\mathbf{x}_{t};\xi)}{\gamma}$ $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t}+\kappa\hat{\mathbf{g}}_{t}, \text{ with}$ $\displaystyle\hat{\mathbf{g}}_{t}$ $\displaystyle=\mathbf{g}_{t}+\frac{1-\kappa}{\kappa}\frac{1}{B}\sum_{\xi\in \mathcal{B}_{t}}\frac{\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)- \nabla f(\mathbf{x}_{t};\xi)}{\gamma}$ $\displaystyle=\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\left(\frac{1-\kappa}{ \kappa\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)+\left(1-\frac {1-\kappa}{\kappa\gamma}\right)\nabla f(\mathbf{x}_{t};\xi)\right)+\mathbf{w}_ {t},$ $$

where we substitute $\mathbf{g}_{t}=\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\mathrm{clip}\left(
\nabla f(\mathbf{x}_{t};\xi),C\right)+\mathbf{w}_{t}.$
Then, by privatizing $\hat{\mathbf{g}}_{t}$ and rename it as $\mathbf{g}_{t}$, we have:

$$ $\displaystyle\mathbf{g}_{t}$ $\displaystyle=\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\mathrm{clip}\left(\frac{ 1-\kappa}{\kappa\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)+ \left(1-\frac{1-\kappa}{\kappa\gamma}\right)\nabla f(\mathbf{x}_{t};\xi),C \right)+\mathbf{w}_{t}$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t}+\kappa\mathbf{g}_{t},$ $$

which is the update of Algorithm 3 (Lines 5 and 6).

### A.4 Connection between DiSK and NAG and STORM

Algorithm 3 has an implicit connection to the (unified) Nesterov Accelerated Gradient (NAG) method (Shen et al., ; Sutskever et al., ) \xweditand the Stochastic Recursive Momentum (STORM) algorithm (Cutkosky & Orabona, ). The update of (half-shifted) NAG can be written as (Sutskever et al., ):

$$ $\displaystyle\mathbf{m}_{t}$ $\displaystyle=\mu\mathbf{m}_{t-1}+\eta\nabla F(\mathbf{x}_{t}-\mu\mathbf{m}_{t -1})$ (16) $\displaystyle\mathbf{x}_{t+1}$ $\displaystyle=\mathbf{x}_{t}-\mathbf{m}_{t},$ $$

and the update of STORM writes (Cutkosky & Orabona, ):

$$ $\displaystyle\mathbf{m}_{t}$ $\displaystyle=(1-\alpha)\mathbf{m}_{t-1}+\alpha\nabla f(\mathbf{x}_{t},\xi_{t} )+(1-\alpha)(\nabla f(\mathbf{x}_{t},\xi_{t})-\nabla f(\mathbf{x}_{t-1},\xi_{t }))$ (17) $\displaystyle\mathbf{x}_{t+1}$ $\displaystyle=\mathbf{x}_{t}-\eta\mathbf{m}_{t},$ $$

In comparison, in Algorithm 3, \xweditby letting the update of OptimizerUpdate be SGD, and assuming clipping is inactive, $\mathbf{w}_{t}=0$, i.e., without DP, $\epsilon\rightarrow\infty$, the DiSK becomes:

$$ $\displaystyle\tilde{\mathbf{g}}_{t}$ $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t-1}+\frac{\kappa}{B}\sum_{\xi\in \mathcal{B}_{t}}\left(\frac{1-\kappa}{\kappa\gamma}\nabla f(\mathbf{x}_{t}+ \gamma\mathbf{d}_{t-1};\xi)+\big{(}1-\frac{1-\kappa}{\kappa\gamma}\big{)} \nabla f(\mathbf{x}_{t};\xi)\right)$ (18) $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t-1}+\frac{\kappa}{B}\sum_{\xi\in \mathcal{B}_{t}}\nabla f(\mathbf{x}_{t};\xi)+\frac{1-\kappa}{B\gamma}\sum_{\xi \in\mathcal{B}_{t}}\left(\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)- \nabla f(\mathbf{x}_{t};\xi)\right)$ $\displaystyle\mathbf{x}_{t+1}$ $\displaystyle=\mathbf{x}_{t}-\eta\tilde{\mathbf{g}}_{t}.$ $$

Comparing the above updates, we observe that (16) and (17) are two special cases of (18) with specific choices of the parameters. Specifically, by letting $\gamma=\frac{1-\kappa}{\kappa}>0,\eta=\kappa,1-\kappa=\mu$, $B=N$, we have $1-\frac{1-\kappa}{\kappa\gamma}=0$, and the update of $\tilde{\mathbf{g}}_{t}$ becomes $\tilde{\mathbf{g}}_{t}=\mu\tilde{\mathbf{g}}_{t-1}+\eta\nabla F(\mathbf{x}_{t}
-\frac{\mu}{\eta}\tilde{\mathbf{g}}_{t-1})$, and DiSK becomes NAG; on the other hand, by letting $\gamma=-1,\kappa=\alpha$, $B=1$, we have $\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1}=\mathbf{x}_{t}-(\mathbf{x}_{t}-\mathbf{x
}_{t-1})=\mathbf{x}_{t-1}$, and the update of $\tilde{\mathbf{g}}_{t}$ becomes $\tilde{\mathbf{g}}_{t}=(1-\alpha)\tilde{\mathbf{g}}_{t-1}+\alpha\nabla f(
\mathbf{x}_{t};\xi_{t})+(1-\alpha)\left(\nabla f(\mathbf{x}_{t};\xi)-\nabla f(
\mathbf{x}_{t-1};\xi)\right)$, which matches the update of STORM. From this discussion, we see that NAG and STORM, two algorithms for accelerated gradient and variance reduction, respectively, can be unified by DiSK, an algorithm based on Kalman filtering.

With this discussion, we observe that the key difference between NAG and STORM is the choice of $\gamma$. When $\gamma>0,$ the algorithm is close to NAG, which focuses on “exploring” and “accelerating”. When $\gamma<0$, the algorithm is close to STORM, which focuses on “exploiting” and “reducing noise”.

## Appendix B Proof for Section 4

In this section, we provide detailed proof for the results in Section 4.

We find the following relations useful:

$$ $\displaystyle\left\langle\mathbf{a},\mathbf{b}\right\rangle$ $\displaystyle\leq\frac{1}{2\alpha}\left\lVert\mathbf{a}\right\rVert^{2}+\frac{ \alpha}{2}\left\lVert\mathbf{b}\right\rVert^{2},$ (19) $\displaystyle\left\lVert\mathbf{a}+\mathbf{b}\right\rVert^{2}$ $\displaystyle\leq(1+\alpha)\left\lVert\mathbf{a}\right\rVert^{2}+(1+1/\alpha) \left\lVert\mathbf{b}\right\rVert^{2}.$ (20) $$

In the following sections, we use $\operatorname{\mathbb{E}}_{t}$ to denote the expectation conditioned on all the information before iteration $t$.
To prove Theorem 2, we first provide the following lemma to bound the difference between $\tilde{\mathbf{g}}_{t}$ and $\nabla F(\mathbf{x}_{t})$. Let us define $\Delta_{t}=\nabla F(\mathbf{x}_{t})-\tilde{\mathbf{g}}_{t}$. We have:
\xwedit

###### Lemma 1

Assume A1, A2 holds, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}\left\lVert\Delta_{t}\right\rVert^{2}$ $\displaystyle\leq(1-\kappa)^{2}\left(1+4\eta^{2}L^{2}+\left\lvert 1+\gamma \right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)\right)\left\lVert \Delta_{t-1}\right\rVert^{2}$ $\displaystyle\quad+2\eta^{2}L^{2}(1-\kappa)^{2}\left(2+\left\lvert 1+\gamma \right\rvert C_{\gamma}\right)\left\lVert\nabla F(\mathbf{x}_{t-1})\right \rVert^{2}$ $\displaystyle\quad+\kappa^{2}\left((2+\left\lvert 1+\gamma\right\rvert)\frac{ \sigma^{2}_{\text{SGD}}}{B}+d\sigma^{2}_{\text{DP}}\right),$ (21) $$

where we define $C_{\gamma}=\left(1+\frac{4(2+1/\kappa+\left\lvert 1+\gamma\right\rvert)}{
\gamma^{2}}\right).$

###### Lemma 1

### B.1 Proof of Lemma 1

By the update of $\tilde{\mathbf{g}}_{t}$, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}\left\lVert\Delta_{t}\right\rVert^{2}$ $\displaystyle=\operatorname{\mathbb{E}}_{t}\left\|\nabla F(\mathbf{x}_{t})-(1- \kappa)\tilde{\mathbf{g}}_{t-1}-\frac{\kappa}{B}\sum_{\xi\in\mathcal{B}_{t}} \nabla f(\mathbf{x}_{t},\xi)-\kappa\mathbf{w}_{t}\right.$ $\displaystyle\qquad\qquad\left.-\frac{1-\kappa}{\gamma B}\sum_{\xi\in\mathcal{ B}_{t}}\left(\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)-\nabla f( \mathbf{x}_{t};\xi)\right)\right\|^{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\operatorname{\mathbb{E}}_{t} \left\|(1-\kappa)(\nabla F(\mathbf{x}_{t})-\nabla F(\mathbf{x}_{t-1}))+(1- \kappa)(\nabla F(\mathbf{x}_{t-1})-\tilde{\mathbf{g}}_{t-1})\right.$ $\displaystyle\qquad\qquad+\kappa\left(\nabla F(\mathbf{x}_{t})-\frac{1}{B}\sum _{\xi\in\mathcal{B}_{t}}\nabla f(\mathbf{x}_{t},\xi)-\mathbf{w}_{t}\right)$ $\displaystyle\qquad\qquad-\frac{1-\kappa}{\gamma B}\sum_{\xi\in\mathcal{B}_{t} }\left(\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)-\nabla f(\mathbf{x} _{t};\xi)\right)\big{\|}^{2}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{=}}\operatorname{\mathbb{E}}_{t} \Bigg{\|}(1-\kappa)\underbrace{\left(\nabla F(\mathbf{x}_{t})-\nabla F(\mathbf {x}_{t-1})-\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\left(\nabla f(\mathbf{x}_{t };\xi)-\nabla f(\mathbf{x}_{t-1};\xi)\right)\right)}_{:=D_{1}}$ $\displaystyle\qquad\qquad+(1-\kappa)\Delta_{t-1}+\kappa\underbrace{\left( \nabla F(\mathbf{x}_{t})-\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\nabla f( \mathbf{x}_{t},\xi)\right)}_{:=D_{2}}-\kappa\mathbf{w}_{t}$ $\displaystyle\qquad\qquad-(1-\kappa)\underbrace{\frac{1}{B}\sum_{\xi\in \mathcal{B}_{t}}\left(\frac{1}{\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d} _{t-1};\xi)+\nabla f(\mathbf{x}_{t-1}\mathbf{d}_{t-1};\xi)-\frac{1+\gamma}{ \gamma}\nabla f(\mathbf{x}_{t};\xi)\right)}_{:=D_{3}}\Bigg{\|}^{2}$ $\displaystyle\stackrel{{\scriptstyle(c)}}{{=}}(1-\kappa)^{2}\operatorname{ \mathbb{E}}_{t}\left\lVert D_{1}\right\rVert^{2}+(1-\kappa)^{2}\left\lVert \Delta_{t-1}\right\rVert^{2}+\kappa^{2}\operatorname{\mathbb{E}}_{t}[\left \lVert D_{2}\right\rVert^{2}]+(1-\kappa)^{2}\operatorname{\mathbb{E}}_{t}\left \lVert D_{3}\right\rVert^{2}+\kappa^{2}\operatorname{\mathbb{E}}[\left\lVert \mathbf{w}_{t}\right\rVert^{2}]$ $\displaystyle\quad+2(1-\kappa)^{2}\left\langle\operatorname{\mathbb{E}}_{t}[D_ {1}],\Delta_{t-1}\right\rangle+2(1-\kappa)\kappa\left\langle\operatorname{ \mathbb{E}}_{t}[D_{2}],\Delta_{t-1}\right\rangle-2(1-\kappa)^{2}\left\langle \operatorname{\mathbb{E}}_{t}[D_{3}],\Delta_{t-1}\right\rangle$ $\displaystyle\quad+2\operatorname{\mathbb{E}}_{t}[\left\langle(1-\kappa)D_{1}, \kappa D_{2}\right\rangle]-2(1-\kappa)^{2}\operatorname{\mathbb{E}}_{t}[\left \langle D_{1},D_{3}\right\rangle]-2\operatorname{\mathbb{E}}_{t}[\left\langle \kappa D_{2},(1-\kappa)D_{3}\right\rangle],$ $\displaystyle\stackrel{{\scriptstyle(d)}}{{=}}(1-\kappa)^{2}(2+\left\lvert 1+ \gamma\right\rvert)\operatorname{\mathbb{E}}_{t}\left\lVert D_{1}\right\rVert^ {2}+(1-\kappa)^{2}(1+\kappa\left\lvert 1+\gamma\right\rvert)\left\lVert\Delta_ {t-1}\right\rVert^{2}+\kappa^{2}(2+\left\lvert 1+\gamma\right\rvert) \operatorname{\mathbb{E}}_{t}[\left\lVert D_{2}\right\rVert^{2}]$ $\displaystyle\quad+(1-\kappa)^{2}\left(1+\frac{2}{\left\lvert 1+\gamma\right \rvert}+\frac{1}{\kappa\left\lvert 1+\gamma\right\rvert}\right)\operatorname{ \mathbb{E}}_{t}\left\lVert D_{3}\right\rVert^{2}+\kappa^{2}d\sigma^{2}_{\text{ DP}},$ (22) $$

where $(a)$ we add and subtract $(1-\kappa)\nabla F(\mathbf{x}_{t-1})$ and rearrange the terms; $(b)$ adds and subtracts $\frac{1-\kappa}{B}\sum_{\xi\in\mathcal{B}_{t}}\left(\nabla f(\mathbf{x}_{t};
\xi)-\nabla f(\mathbf{x}_{t-1};\xi)\right)$; $(c)$ directly expands the square; and in $(d)$, we notice that $\operatorname{\mathbb{E}}_{t}[D_{1}]=0,\operatorname{\mathbb{E}}_{t}[D_{2}]=0$, so the first two inner products are zero, and we apply (19) to the other four inner products, with $\alpha=\kappa\left\lvert 1+\gamma\right\rvert,1,\left\lvert 1+\gamma\right\rvert$, respectively.
Next, we bound each term separately. For $\operatorname{\mathbb{E}}_{t}[\left\lVert D_{1}\right\rVert^{2}]$, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[\left\lVert D_{1}\right\rVert^{2}]$ $\displaystyle=\operatorname{\mathbb{E}}_{t}\left\lVert\nabla F(\mathbf{x}_{t}) -\nabla F(\mathbf{x}_{t-1})-\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\left( \nabla f(\mathbf{x}_{t};\xi)-\nabla f(\mathbf{x}_{t-1};\xi)\right)\right\rVert ^{2}$ (23) $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\operatorname{\mathbb{E}}_{t} \left\lVert\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\left(\nabla f(\mathbf{x}_{t };\xi)-\nabla f(\mathbf{x}_{t-1};\xi)\right)\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}L^{2}\eta^{2}\left\lVert \tilde{\mathbf{g}}_{t-1}\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}2L^{2}\eta^{2}(\left\lVert \Delta_{t-1}\right\rVert^{2}+\left\lVert\nabla F(\mathbf{x}_{t-1})\right\rVert ^{2}),$ $$

where $(a)$ uses the fact that $\operatorname{\mathbb{E}}\left\lVert X-\operatorname{\mathbb{E}}[X]\right
\rVert^{2}\leq\operatorname{\mathbb{E}}\left\lVert X\right\rVert^{2}$; $(b)$ applies A1 and $(c)$ adds and subtracts $\nabla F(\mathbf{x}_{t-1})$, and applies (20).
For $\operatorname{\mathbb{E}}_{t}[\left\lVert D_{2}\right\rVert^{2}]$, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[\left\lVert D_{2}\right\rVert^{2}]$ $\displaystyle=\operatorname{\mathbb{E}}_{t}\left\lVert\nabla F(\mathbf{x}_{t}) -\frac{1}{B}\sum_{\xi\in\mathcal{B}_{t}}\nabla f(\mathbf{x}_{t},\xi)\right \rVert^{2}\stackrel{{\scriptstyle A\ref{as:variance}}}{{\leq}}\frac{\sigma^{2} _{\text{SGD}}}{B}.$ (24) $$

For $\operatorname{\mathbb{E}}_{t}[\left\lVert D_{3}\right\rVert^{2}]$, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[\left\lVert D_{3}\right\rVert^{2}]$ $\displaystyle=\operatorname{\mathbb{E}}_{t}\left\lVert\frac{1}{B}\sum_{\xi\in \mathcal{B}_{t}}\left(\frac{1}{\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d} _{t-1};\xi)+\nabla f(\mathbf{x}_{t-1}\mathbf{d}_{t-1};\xi)-\frac{1+\gamma}{ \gamma}\nabla f(\mathbf{x}_{t};\xi)\right)\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{1}{B}\sum_{\xi\in \mathcal{B}_{t}}\Bigg{\|}\frac{1}{\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf {d}_{t-1};\xi)-\frac{1}{\gamma}\nabla f(\mathbf{x}_{t-1};\xi)$ $\displaystyle\qquad\qquad+\frac{1+\gamma}{\gamma}\nabla f(\mathbf{x}_{t-1} \mathbf{d}_{t-1};\xi)-\frac{1+\gamma}{\gamma}\nabla f(\mathbf{x}_{t};\xi)\Bigg {\|}^{2}$ $\displaystyle\stackrel{{\scriptstyle\eqref{eq:split}}}{{\leq}}2\left\lVert \frac{1}{\gamma}\nabla f(\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1};\xi)-\frac{1}{ \gamma}\nabla f(\mathbf{x}_{t-1};\xi)\right\rVert^{2}$ $\displaystyle\quad+2\left\lVert\frac{1+\gamma}{\gamma}\nabla f(\mathbf{x}_{t-1 };\xi)-\frac{1+\gamma}{\gamma}\nabla f(\mathbf{x}_{t};\xi)\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle A\ref{as:smooth}}}{{\leq}}\frac{2L^{2}}{ \gamma^{2}}\left\lVert\mathbf{x}_{t}+\gamma\mathbf{d}_{t-1}-\mathbf{x}_{t-1} \right\rVert^{2}+\frac{2L^{2}(1+\gamma)^{2}}{\gamma^{2}}\left\lVert\mathbf{x}_ {t-1}-\mathbf{x}_{t}\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{=}}\frac{4L^{2}(1+\gamma)^{2}\eta^{ 2}}{\gamma^{2}}\left\lVert\tilde{\mathbf{g}}_{t-1}\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}\frac{8L^{2}(1+\gamma)^{2} \eta^{2}}{\gamma^{2}}(\left\lVert\Delta_{t-1}\right\rVert^{2}+\left\lVert \nabla F(\mathbf{x}_{t-1})\right\rVert^{2}),$ (25) $$

where $(a)$ applies Jensens’ inequality to $\left\lVert\cdot\right\rVert^{2}$, and we add and subtract $\frac{1}{\gamma}\nabla f(\mathbf{x}_{t-1};\xi)$; $(b)$ apples (20); $(b)$ uses the fact that $\mathbf{d}_{t-1}=\mathbf{x}_{t}-\mathbf{x}_{t-1}=-\eta\tilde{\mathbf{g}}_{t-1}$; and $(c)$ adds and subtracts $\nabla F(\mathbf{x}_{t-1})$, and applies (20).
Plug in (23) – (B.1) to (B.1), we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}\left\lVert\Delta_{t}\right\rVert^{2}$ $\displaystyle\leq(1-\kappa)^{2}(2+\left\lvert 1+\gamma\right\rvert) \operatorname{\mathbb{E}}_{t}\left\lVert D_{1}\right\rVert^{2}+(1-\kappa)^{2}( 1+\kappa\left\lvert 1+\gamma\right\rvert)\left\lVert\Delta_{t-1}\right\rVert^{ 2}+\kappa^{2}d\sigma^{2}_{\text{DP}}$ $\displaystyle\quad+\kappa^{2}(2+\left\lvert 1+\gamma\right\rvert)\operatorname {\mathbb{E}}_{t}[\left\lVert D_{2}\right\rVert^{2}]+(1-\kappa)^{2}\left(1+ \frac{2}{\left\lvert 1+\gamma\right\rvert}+\frac{1}{\kappa\left\lvert 1+\gamma \right\rvert}\right)\operatorname{\mathbb{E}}_{t}\left\lVert D_{3}\right\rVert ^{2}$ $\displaystyle\leq(1-\kappa)^{2}\left(1+4\eta^{2}L^{2}+\left\lvert 1+\gamma \right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)\right)\left\lVert \Delta_{t-1}\right\rVert^{2}$ $\displaystyle\quad+2\eta^{2}L^{2}(1-\kappa)^{2}\left(2+\left\lvert 1+\gamma \right\rvert C_{\gamma}\right)\left\lVert\nabla F(\mathbf{x}_{t-1})\right \rVert^{2}$ $\displaystyle\quad+\kappa^{2}\left((2+\left\lvert 1+\gamma\right\rvert)\frac{ \sigma^{2}_{\text{SGD}}}{B}+d\sigma^{2}_{\text{DP}}\right),$ (26) $$

where we define $C_{\gamma}:=\left(1+\frac{4(2+1/\kappa+\left\lvert 1+\gamma\right\rvert)}{
\gamma^{2}}\right)$. This completes the proof of Lemma 1.

### B.2 Proof of Theorem 2

Then, we are ready to prove Theorem 2. By choosing $C\geq G(1+\frac{2(1-\kappa)}{\kappa\gamma}),$ the clipping is inactive. Recall that from the update rule, we have

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[\tilde{\mathbf{g}}_{t}]$ $\displaystyle=(1-\kappa)\tilde{\mathbf{g}}_{t-1}+\kappa\nabla F(\mathbf{x}_{t} )+(1-\kappa)(\nabla F(\mathbf{x}_{t})-\nabla F(\mathbf{x}_{t-1}))$ (27) $\displaystyle=\nabla F(\mathbf{x}_{t})-(1-\kappa)\Delta_{t-1}.$ $$

Then, by A1, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[F(\mathbf{x}_{t+1})]-$ $\displaystyle F(\mathbf{x}_{t})\leq-\eta\left\langle\nabla F(\mathbf{x}_{t}), \operatorname{\mathbb{E}}_{t}[\tilde{\mathbf{g}}_{t}]\right\rangle+\frac{\eta^ {2}L}{2}\operatorname{\mathbb{E}}_{t}\left\lVert\tilde{\mathbf{g}}_{t}\right \rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}-\eta\left\lVert\nabla F(\mathbf {x}_{t})\right\rVert^{2}+\eta(1-\kappa)\left\langle\nabla F(\mathbf{x}_{t}), \Delta_{t-1}\right\rangle+\frac{\eta^{2}L}{2}\operatorname{\mathbb{E}}_{t} \left\lVert\tilde{\mathbf{g}}_{t}\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}-\frac{\eta(1+\kappa-2\eta L) }{2}\left\lVert\nabla F(\mathbf{x}_{t})\right\rVert^{2}+\frac{\eta(1-\kappa)}{ 2}\left\lVert\Delta_{t-1}\right\rVert^{2}+\eta^{2}L\operatorname{\mathbb{E}}_{ t}\left\lVert\Delta_{t}\right\rVert^{2},$ (28) $$

where $(a)$ substitute (27); $(b)$ applies (19) to the second term and we add and subtract $\nabla F(\mathbf{x}_{t})$ to the last term and apply (20).

Define $\mathcal{L}_{t}:=F(\mathbf{x}_{t})+\beta\left\lVert\Delta_{t-1}\right\rVert^{2}$, we have:

$$ $\displaystyle\operatorname{\mathbb{E}}_{t}[\mathcal{L}_{t+1}]$ $\displaystyle-\mathcal{L}_{t}$ $\displaystyle\leq-\frac{\eta(1+\kappa-2\eta L)}{2}\left\lVert\nabla F(\mathbf{ x}_{t})\right\rVert^{2}-(\beta-\frac{\eta(1-\kappa)}{2})\left\lVert\Delta_{t-1 }\right\rVert^{2}+(\beta+\eta^{2}L)\operatorname{\mathbb{E}}_{t}\left\lVert \Delta_{t}\right\rVert^{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}-\frac{\eta(1+\kappa-2\eta L) }{2}\left\lVert\nabla F(\mathbf{x}_{t})\right\rVert^{2}-(\beta-\frac{\eta(1- \kappa)}{2})\left\lVert\Delta_{t-1}\right\rVert^{2}$ $\displaystyle\quad+(\beta+\eta^{2}L)(1-\kappa)^{2}(1+4L^{2}\eta^{2}+\xwedit{ \left\lvert 1+\gamma\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)}) \left\lVert\Delta_{t-1}\right\rVert^{2}$ $\displaystyle\quad+(\beta+\eta^{2}L)\kappa^{2}\left(\frac{\xwedit{(2+\left \lvert 1+\gamma\right\rvert)}\sigma^{2}_{\mathrm{SGD}}}{B}+d\sigma^{2}_{ \mathrm{DP}}\right)$ $\displaystyle\quad+2(\beta+\eta^{2}L)(1-\kappa)^{2}L^{2}\eta^{2}\xwedit{\left( 2+\left\lvert 1+\gamma\right\rvert C_{\gamma}\right)}\left\lVert\nabla F( \mathbf{x}_{t-1})\right\rVert^{2}$ $\displaystyle=-\frac{\eta(1+\kappa-2\eta L)}{2}\left\lVert\nabla F(\mathbf{x}_ {t})\right\rVert^{2}+2(\beta+\eta^{2}L)(1-\kappa)^{2}L^{2}\eta^{2}\xwedit{ \left(2+\left\lvert 1+\gamma\right\rvert C_{\gamma}\right)}\left\lVert\nabla F (\mathbf{x}_{t-1})\right\rVert^{2}$ $\displaystyle\quad-\left(\beta-\frac{\eta(1-\kappa)}{2}-(\beta+\eta^{2}L)(1- \kappa)^{2}(1+4L^{2}\eta^{2}+\xwedit{\left\lvert 1+\gamma\right\rvert\left( \kappa+2\eta^{2}L^{2}C_{\gamma}\right)})\right)\left\lVert\Delta_{t-1}\right \rVert^{2}$ $\displaystyle\quad+(\beta+\eta^{2}L)\kappa^{2}\left(\frac{\xwedit{(2+\left \lvert 1+\gamma\right\rvert)}\sigma^{2}_{\mathrm{SGD}}}{B}+d\sigma^{2}_{ \mathrm{DP}}\right),$ (29) $$

where $(a)$ applies Lemma 1 and $(b)$ rearrange the terms.
By choosing

$$ $\eta<\frac{1}{2L(1+2(1-\kappa)^{2}\beta L(2+\left\lvert 1+\gamma\right\rvert C _{\gamma}))},\;\kappa>1-\frac{1}{\sqrt{1+4\eta^{2}L^{2}+\xwedit{\left\lvert 1+ \gamma\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)}}},$ $$

$$ $\beta\geq\frac{\eta(1-\kappa)/2+\eta^{2}L(1-\kappa)^{2}(1+4\eta^{2}L^{2}+ \xwedit{\left\lvert 1+\gamma\right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma} \right)})}{1-(1-\kappa)^{2}(1+4\eta^{2}L^{2}+\xwedit{\left\lvert 1+\gamma \right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)})},$ $$

we have:

$$ $\displaystyle\frac{\eta(1+\kappa-2\eta L)}{2}-2(\beta+\eta^{2}L)(1-\kappa)^{2} L^{2}\eta^{2}\xwedit{\left(2+\left\lvert 1+\gamma\right\rvert C_{\gamma}\right )}>0,$ (30) $\displaystyle 1-(1-\kappa)^{2}(1+4\eta^{2}L^{2}+\xwedit{\left\lvert 1+\gamma \right\rvert\left(\kappa+2\eta^{2}L^{2}C_{\gamma}\right)})>0,$ $\displaystyle\beta-\frac{\eta(1-\kappa)}{2}-(\beta+\eta^{2}L)(1-\kappa)^{2}(1+ 4L^{2}\eta^{2}+\xwedit{\left\lvert 1+\gamma\right\rvert\left(\kappa+2\eta^{2}L ^{2}C_{\gamma}\right)})\geq 0.$ $$

Average from $t=0$ to $T-1$ and rearrange the terms, we have:

$$ $\displaystyle\frac{1}{T}\sum^{T}_{t=0}\operatorname{\mathbb{E}}\left\lVert \nabla F(\mathbf{x}_{t})\right\rVert^{2}\leq\frac{2(\mathcal{L}_{0}- \operatorname{\mathbb{E}}[\mathcal{L}_{T+1}])}{C_{1}\eta T}+\frac{2(\beta+\eta ^{2}L)\kappa^{2}}{C_{1}\eta}\left(\frac{\xwedit{(2+\left\lvert 1+\gamma\right \rvert)}\sigma^{2}_{\mathrm{SGD}}}{B}+d\sigma^{2}_{\mathrm{DP}}\right)$ $\displaystyle\quad\leq\frac{2(F(\mathbf{x}_{0})+\beta\left\lVert\nabla F( \mathbf{x}_{0})\right\rVert^{2}-F^{\star})}{C_{1}\eta T}+\frac{2(\beta+\eta^{2 }L)\kappa^{2}}{C_{1}\eta}\left(\frac{\xwedit{(2+\left\lvert 1+\gamma\right \rvert)}\sigma^{2}_{\mathrm{SGD}}}{B}+d\sigma^{2}_{\mathrm{DP}}\right),$ (31) $$

where we define $C_{1}:=(1+\kappa-2\eta L)-4(\beta+\eta^{2}L)(1-\kappa)^{2}L^{2}\eta\xwedit{
\left(2+\left\lvert 1+\gamma\right\rvert C_{\gamma}\right)},$ and in the last inequality we notice that $\mathcal{L}_{T+1}=F(\mathbf{x}_{T+1})+\beta\left\lVert\Delta_{T}\right\rVert^{
2}\geq F^{\star},$ and $\mathcal{L}_{0}=F(\mathbf{x}_{0})+\beta\left\lVert\nabla F(\mathbf{x}_{0})
\right\rVert^{2},$ as we initialize $\tilde{\mathbf{g}}_{0}=0.$ This completes the proof of Theorem 2.

On the choice of $\gamma=-1$. From the above proof, we see that in Appendix B.1, (B.1) $(c)$, we directly apply (19) to upper bound the cross-product terms by positive terms for the worst case, which results in the optimal choice of $\gamma=-1$. However, the inner products may be smaller than zero in some cases, making $\gamma=-1$ sub-optimal in practice.

## Appendix C Additional numerical results

### C.1 Experiment details

Coding: The code for the experiments will be provided online. We use PyTorch as the code base and the FastDP package (Bu et al., ) to privatize the optimizers. The Algorithm 3 is implemented as a PyTorch optimizer, which can be easily combined with any training scripts based on PyTorch. The modification is minimum:

The link to the full code of the experiments can be found at [https://anonymous.4open.science/r/KalmanDP-BEDB](https://anonymous.4open.science/r/KalmanDP-BEDB).

Hardware: All the experiments except the ImageNet-1k dataset are running with one Nvidia A40 (48GB memory) or one Nvidia V100 (32GB memory). The experiment on the ImageNet-1k dataset is running on one Nvidia H100 (80GB memory) GPU. The training time varies for different tasks depending on the data size and model size.

Training method: We use gradient accumulation to deal with the large batch size and use learning rate warm-up for $1/20$ of the training steps when training from randomly initialized weights. We also use the Cosine Annealing learning rate scheduler (Loshchilov & Hutter, ), which gradually decreases the learning rate.

### C.2 Choice of hyper-parameters

The main hyper-parameters in the algorithms are: epoch $E$, batch size $B$, step size $\eta$, clipping threshold $C$, Kalman filter parameters $\kappa$, and $\gamma$. In all experiments, we fix the clipping method as automatic clipping used in Bu et al. (), i.e., $\mathrm{clip}\left(\nabla f(\mathbf{x};\xi),C\right)=\nabla f(\mathbf{x};\xi)
\cdot\frac{C}{\left\lVert\nabla f(\mathbf{x};\xi)\right\rVert},$ and set $C=1$ for all experiments.

For each set of experiments, we conduct a grid search on the hyper-parameters $E,B,\eta$ and choose the optimal ones for the DP optimizer without DiSK. The search grids of each hyper-parameter are listed in Table 2;

**Table 2: Search grid of the CV pre-training experiments, the optimal hyper-parameters are in bold.**
|  | Search gird |  |  |
| --- | --- | --- | --- |
|  | MNIST | CIFAR | ImageNet |
| $E$ | $\{1,{\bf 2},3\}\times 20$ | $\{1,{\bf 2},3,4\}\times 40$ | $\{3,{\bf 4}\}\times 40$ |
| $B$ | $\{2,{\bf 5}\}\times 10^{3}$ | $\{0.5,1,2,{\bf 5}\}\times 10^{3}$ | $\{{\bf 5},10\}\times 10^{3}$ |
| $\eta$ | $\{3,{\bf 2.5},1,0.3,0.1\}\times 10^{-1}$ | $\{1,3,{\bf 5},7,10\}\times 10^{-3}$ | $\{10,3,1,{\bf 0.3},0.1\}\times 10^{-3}$ |
| $\kappa$ | $\{{\bf 0.7}\}$ | $\{9.9,9,8,{\bf 7},6,5\}\times 10^{-1}$ | $\{{\bf 0.7}\}$ |
| $\gamma$ | $\{{\bf 0.5}\}$ | $\{0.2,0.3,{\bf 0.5},1,2,3,\frac{1-\kappa}{\kappa}\}$ | $\{{\bf 0.5}\}$ |

Then, we fix $E,B.\eta$ and conduct the ablation study on $\kappa,\gamma$ as shown in Figure 8.

### C.3 Additional experiments on CV tasks

Training different models from scratch: We additionally train the WRN-16-4 and the ViT-small models on the CIFAR-10 with randomly initialized weights for different privacy budgets, and the test accuracies during the training are shown in Figure 4 for $\epsilon=4$. From the results, we can see that DiSK consistently outperforms the base optimizer.

Figure: (a) 5-layer CNN
Refer to caption: extracted/5898046/figures/Cifar-10-eps4.png

Fine-tuning on CIFAR-100: Besides training from scratch, we also compare the performance of fine-tuning a pre-trained ViT-small model on the CIFAR-100 dataset. The results for different $\epsilon$ are shown in Figure 5. For fine-tuning on the complex CIFAR-100 dataset, DiSK still improves the performance compared with DPAdam, and has less performance drop under large DP noise.

Figure: Figure 5: Fine-tuning ViT-small on CIFAR-100 with different $\epsilon$.
Refer to caption: extracted/5898046/figures/Cifar-100-ft.png

Comparison with existing methods and SOTA: We conduct comparisons with existing approaches for improving DP training performance. In Figure 6(a), we train a linear regression model with synthetic data and compared the performance of Noisy GD, Noisy GD with DOPPLER (NoisyLP), and with DiSK (NoisyKF). We inject Gaussian noise with different variances into the gradient and compare the final performance. We observe that DiSK has the lowest regression loss under all noise levels, indicating that the Kalman filter performs better in noise reduction than the Low-pass filter. In Figure 6(b), we compare the test accuracy of different methods, including DOPPLER (Zhang et al., ) and DP-FTRL (Kairouz et al., ; Choquette-Choo et al., ) on the CIFAR-10 dataset training the WRN from scratch. We observe that DiSK  significantly outperforms the SOTA algorithms on all privacy budgets.

Figure: (a) Synthetic data
Refer to caption: x2.png

Comparison with SOTA on ImageNet: Additionally, we conduct experiments to compare the performance of DiSK with SOTA result training from scratch on the ImageNet-1k dataset in De et al. (). The result is shown in Figure 7. We note that the reported test accuracy in the De et al. () is $32.4\%$ for privacy level $\epsilon=8$. However, we noticed that their accounting procedure has been improved in their GitHub repository and hence we re-ran their experiment using the new accountant to obtain our baseline. We kept all their options on, including group normalization, larger batch size, weight standardization, augmentation multiplicity, and model exponential moving average. In addition, we also follow their normalized clipping strategy, i.e.,

$$ $\mathrm{clip}\left(\nabla f,C\right)=\frac{\nabla f}{C}\min\left\{\frac{C}{ \left\lVert\nabla f\right\rVert},1\right\}.$ $$

Their updated code results in $36.35\%$ validation accuracy and the test accuracy of $33.56\%$ (as the SOTA result). In comparison, adding DiSK on top of their method results in a validation accuracy of $40\%$ ($3.7\%$ improvement) and a test accuracy of $36.89\%$ ($3.3\%$ improvement). This sets a new SOTA result for differentially private training on the ImageNet-1k dataset.

Figure: Figure 7: Comparison with the SOTA approach in De et al. () on the ImageNet-1k dataset.
Refer to caption: x3.png

Ablation study: We conduct ablation studies on the choice of the hyper-parameters of DiSK, specifically, how $\kappa,\gamma$ impact the algorithm performance. In Figure 8, we plot the accuracy on different combinations of $(\kappa,\gamma)$, and $(\kappa,\epsilon)$. We observe a clear trend of performance change for different combinations of the parameters, and there is an optimal choice of $\kappa,\gamma$ for different $\epsilon$’s.

Figure: (a) $(\kappa,\gamma),\epsilon=8$
Refer to caption: x4.png

### C.4 Additional experiments on NLP tasks

In this section, we provide additional results for NLP tasks.

Parameter-efficient fine-tuning on GLUE. We fine-tune a RoBERTa-base and a RoBERTa-large model from the Huggingface checkpoints(^2^22[https://huggingface.co/FacebookAI/roberta-base](https://huggingface.co/FacebookAI/roberta-base),[https://huggingface.co/FacebookAI/roberta-large](https://huggingface.co/FacebookAI/roberta-large)) on the GLUE dataset. We follow the same training scripts in Bu et al. () on the hyper-parameter choices of $\eta,B,E$ on the tasks and use rank $r=16$ for LoRA. We choose $\kappa=0.7,\gamma=0.5$ for DiSK. The results are listed in Table 3. With privacy budget $\epsilon=1,6.7$, DPLoRA with DiSK significantly outperforms SOTA results with vanilla DPLoRA on all tasks.

**Table 3: Test accuracy of fine-tuning result on the GLUE dataset.**
|  | $\epsilon=1$ | $\epsilon=6.7$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Algorithm | MNLI | QNLI | SST2 | QQP | MNLI | QNLI | SST2 | QQP |
| RoBERTa-base |  |  |  |  |  |  |  |  |
| AdamW ($\epsilon=\infty$) | 87.6 | 92.8 | 94.8 | 91.9 | 87.6 | 92.8 | 94.8 | 91.9 |
| Lora ($\epsilon=\infty$) | 87.5 | 93.3 | 95.1 | 90.8 | 87.5 | 93.3 | 95.1 | 90.8 |
| DPLora | 81.1 | 85.5 | 90.9 | 83.9 | 83.5 | 87.4 | 91.5 | 85.7 |
| KF-DPLora | 84.7 | 90.3 | 92.9 | 87.8 | 85.9 | 90.5 | 93.1 | 89.0 |
| RoBERTa-large |  |  |  |  |  |  |  |  |
| AdamW ($\epsilon=\infty$) | 90.3 | 94.7 | 96.4 | 92.2 | 90.3 | 94.7 | 96.4 | 92.2 |
| Lora ($\epsilon=\infty$) | 90.6 | 94.9 | 96.2 | 91.6 | 90.6 | 94.9 | 96.2 | 91.6 |
| DPLora | 85.6 | 89.5 | 90.9 | 85.1 | 87.8 | 90.8 | 94.3 | 87.4 |
| KF-DPLora | 87.9 | 92.5 | 95.2 | 88.2 | 89.4 | 92.6 | 95.4 | 89.6 |

Fine-tuning GPT-2 on text generation tasks. We fine-tune a GPT-2-small model with $137M$ parameters from the Huggingface checkpoints(^3^33[https://huggingface.co/openai-community/gpt2](https://huggingface.co/openai-community/gpt2)) on two text generation datasets, E2E and DART. We follow the same training scripts in Li et al. () on the hyper-parameter choices of $\eta,B,E$ on the tasks and choose $\kappa=0.7,\gamma=0.5$ for DiSK. The results on different metrics for the E2E dataset are given in Table 4, and the results for the DART dataset are in Table 5. With privacy budget $\epsilon=3,8$, DPAdamW with DiSK significantly outperforms SOTA results with vanilla DPAdamW on all metrics.

**Table 4: Performance of fine-tuning gpt-2 on the E2E dataset. (All metrics are higher the better)**
| Algorithm | BLEU ($\%$) | ROUGE-L ($\%$) | METEOR | NIST | CIDEr |
| --- | --- | --- | --- | --- | --- |
| AdamW ($\epsilon=\infty$) | 69.46 | 71.36 | 0.461 | 8.780 | 2.422 |
| DPAdamW ($\epsilon=3$) | 61.52 | 65.87 | 0.417 | 7.071 | 2.167 |
| KF-DPAdamW ($\epsilon=3$) | 68.35 | 70.23 | 0.456 | 8.636 | 2.399 |
| DPAdamW ($\epsilon=8$) | 64.99 | 67.34 | 0.425 | 8.387 | 2.192 |
| KF-DPAdamW ($\epsilon=8$) | 68.73 | 70.58 | 0.460 | 8.697 | 2.463 |

**Table 5: Performance of fine-tuning gpt-2 on the DART dataset. Val. Perp. stands for validation perplexity. (All metrics except Val. Perp. are higher the better)**
| Algorithm | Val. Perp. $\downarrow$ | BLEU ($\%$) | ROUGE-L ($\%$) | METEOR | NIST | CIDEr |
| --- | --- | --- | --- | --- | --- | --- |
| AdamW ($\epsilon=\infty$) | 0.921 | 44.56 | 58.66 | 0.379 | 8.733 | 2.773 |
| DPAdamW ($\epsilon=3$) | 1.427 | 33.96 | 52.38 | 0.310 | 6.090 | 1.864 |
| KF-DPAdamW ($\epsilon=3$) | 1.149 | 41.01 | 57.53 | 0.359 | 7.949 | 2.553 |
| DPAdamW ($\epsilon=8$) | 1.362 | 35.30 | 54.58 | 0.320 | 6.365 | 1.995 |
| KF-DPAdamW ($\epsilon=8$) | 1.102 | 42.12 | 58.11 | 0.364 | 8.111 | 2.628 |

### C.5 Improvement over SOTA

In the following Table LABEL:tab:app:sota, we provide a list of the experiment settings where our algorithm outperforms SOTA results. The paper that provides the SOTA results is cited for each setting, or the same as the line above. We observe that DiSK provides new SOTA for CV datasets, including CIFAR-10, CIFAR-100, and ImageNet-1k, and NLP datasets, including GLUE (MNLI, QNLI, QQP, SST-2), DART, and E2E.

**Table 6: Comparison with SOTA, PT=pre-training, FT=fine-tuning.**
| Dataset | TASK | Model | $\epsilon$ | Ours ($\%$) | SOTA ($\%$) |
| --- | --- | --- | --- | --- | --- |
| CIFAR-10 | PT | CNN | 0.5 | 59.7 | N/A |
| CIFAR-10 | PT | CNN | 2 | 68.8 | 67.2 (Tramer & Boneh, ) |
| CIFAR-100 | PT | WRN | 0.5 | 14.7 | N/A |
| CIFAR-100 | PT | WRN | 1 | 22.7 | 14.1 (Bao et al., ) |
| CIFAR-100 | PT | WRN | 2 | 30.0 | 21.5 |
| CIFAR-100 | PT | WRN | 4 | 37.1 | 33.3 |
| CIFAR-100 | PT | WRN | 8 | 42.0 | 40.6 |
| CIFAR-100 | FT | ViT | 0.5 | 83.49 | 78.3 (Mehta et al., ) |
| CIFAR-100 | FT | ViT | 1 | 85.36 | 81.8 (Bao et al., ) |
| CIFAR-100 | FT | ViT | 2 | 86.77 | 83.5 |
| CIFAR-100 | FT | ViT | 4 | 87.56 | 84.5 |
| CIFAR-100 | FT | ViT | 8 | 88.49 | 84.6 |
| ImageNet-1k | PT | ViT | 8 | 40.00 | 36.35 (De et al., ) |
| MNLI | FT | RoBERTa-base | 1 | 84.7 | 83.2 ($\epsilon=3$) (Bu et al., ) |
| QNLI | FT | RoBERTa-base | 1 | 90.3 | 87.4 ($\epsilon=3$) |
| QQP | FT | RoBERTa-base | 1 | 87.8 | 85.8 ($\epsilon=3$) |
| SST-2 | FT | RoBERTa-base | 1 | 92.9 | 92.3 ($\epsilon=3$) |
| MNLI | FT | RoBERTa-base | 6.7 | 85.9 | 83.8 ($\epsilon=8$) (Bu et al., ) |
| QNLI | FT | RoBERTa-base | 6.7 | 90.5 | 87.9 ($\epsilon=8$) |
| QQP | FT | RoBERTa-base | 6.7 | 89.0 | 86.6 ($\epsilon=8$) |
| SST-2 | FT | RoBERTa-base | 6.7 | 93.1 | 93.0 ($\epsilon=8$) (Li et al., ) |
| MNLI | FT | RoBERTa-large | 1 | 87.9 | 86.8 (Yu et al., ) |
| QNLI | FT | RoBERTa-large | 1 | 92.5 | 88.0 |
| QQP | FT | RoBERTa-large | 1 | 88.2 | 85.2 |
| SST-2 | FT | RoBERTa-large | 1 | 95.2 | 93.1 |
| MNLI | FT | RoBERTa-large | 6.7 | 89.4 | 89.0 (Yu et al., ) |
| QNLI | FT | RoBERTa-large | 6.7 | 92.6 | 92.5 |
| QQP | FT | RoBERTa-large | 6.7 | 89.6 | 88.4 |
| SST-2 | FT | RoBERTa-large | 6.7 | 95.4 | 95.3 |
| E2E (BLEU) | FT | GPT-2 | 3 | 68.35 | 61.52 (Li et al., ) |
| E2E (ROUGE-L) | FT | GPT-2 | 3 | 70.23 | 65.87 (Bu et al., ) |
| E2E (BLEU) | FT | GPT-2 | 8 | 68.73 | 63.60 (Bu et al., ) |
| E2E (ROUGE-L) | FT | GPT-2 | 8 | 70.58 | 67.53 (Li et al., ) |
| DART (BLEU) | FT | GPT-2 | 3 | 41.01 | 32.33 (Li et al., ) |
| DART (ROUGE-L) | FT | GPT-2 | 3 | 57.53 | 52.06 |
| DART (BLEU) | FT | GPT-2 | 8 | 42.12 | 35.06 |
| DART (ROUGE-L) | FT | GPT-2 | 8 | 58.11 | 54.57 |

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: I’m very impressed with these results and somewhat surprised about why they work so much better than existing methods. In developing new methods, or building on top of this, what parts of this method can be extended - what are they key things that differentiate this from existing methods that don’t work as well. This may seem a bit naive, but my question is effectively - based on how big of an improvement this algorithm produces, what should next steps focus on?

Comment: I did not go through the utility proofs in great detail and could have missed something there.
soundness: 3
strengths: Empirical results are very strong and very significant. Computational costs (e.g. time/memory/parallelizability) of the method are quite good too. You mention this in theory, it may be good to make even a small empirical statement about this as well.
summary: This paper introduces DiSK, a novel framework for enhancing differentially private (DP) optimizers by applying Kalman filtering to denoise privatized gradients, producing refined gradient estimates with reduced noise interference. By simplifying Kalman filtering for large-scale tasks, DiSK maintains efficiency and improves privacy-utility trade-offs, achieving significant performance gains over standard DP optimizers like DPSGD. Experimental results on vision and language tasks demonstrate DiSK's superiority under equivalent privacy constraints, achieving state-of-the-art performance across multiple benchmarks.
weaknesses: Why don’t people use Kalman filters for non-DP optimization? I don’t see why the same tricks you employed to not incur d^6 computational costs are not applicable to other non-private settings. 
This feels like an important consideration to understand for contextualizing the empirical claims made in this work

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. In Algorithm 2, are $\mathbb{E}[C]$ and $H_t$ data-dependent? If so, does Algorithm 2 satisfy differential privacy?

2. In Figure 2 (b), both the proposed algorithm KL-DPAdam and the baseline DP-Adam seem to achieve lower test accuracy than the SOTA results in De et al. 2022 Figure 1 (a). Does this contradict the claim that the proposed algorithm enables improvement compared to SOTA?
soundness: 3
strengths: - The paper proposed an interesting way of viewing gradient update in differentially private training as a linear dynamics system. The proposed algorithm provides a novel and computationally efficient way of (approximately) incorporating the Kalman filter into DP training for denoising gradient.
- Extensive numerical experiments that validate the effectiveness of the proposed method.
summary: The paper proposes to use the Kalman filter for denoising gradient computation in differentially private optimization. Specifically, the Kalman filter views the gradient vector as the state and derives the system update equation based on the Taylor expansion of the gradient. The observation is then written as Gaussian perturbation of the gradient. However, naive incorporation of the Kalman filter involves computing the Hessian matrix, which is computationally infeasible for large models under DP constraints. To address the computational challenge, the authors proposed to use a simplified Kalman filter with noisy input, multiplicative observation noise, and diagonal covariance matrix. Theoretical analysis shows that the proposed algorithm satisfies $O\left(\sqrt{d/T}\right)$ convergence rate for smooth loss function with bounded per-sample gradient and bounded per-sample gradient variance, which has a constant factor improvement over DP-SGD under certain assumptions. Experiments validate the improvement of the proposed algorithm compared to DP-SGD and DP-Adam over a variety of DP vision and language tasks.
weaknesses: - The reason for the gain of the proposed denoising algorithm is not perfectly clear from the presentation -- could the authors elaborate more on (1) why does the Kalman filter (as described in Algorithm 2) effectively denoise the gradient observation? (E.g., what is the minimized objective under denoising), (2) what are the assumptions for the simplifications made in Algorithm 3 and why are the assumptions meaningful?
- The authors mention that the algorithm performs better than all SOTA methods, including the DP-FTRL algorithm. However, DP-FTRL is also effectively a gradient denoising method that is proved optimal in terms of l2 error for releasing linear queries (gradient estimation in the context). Consequently, it is not very clear why the proposed Kalman filter method would be beneficial compared to another denoising method, and how much the improvement is.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: See weakness.
soundness: 3
strengths: - This paper effectively applies Kalman filtering from control theory to improve DP optimization algorithms, presenting an interesting connection between DP optimization and control theory. 
- The algorithm improves upon the traditional Kalman filter and requires only a constant factor more space and computation than DP-SGD. 
- The method is versatile and can be integrated with various DP optimization techniques, not just DP-SGD.
summary: This paper proposes using a Kalman filter to enhance the privacy-accuracy trade-off of DP optimizers. While traditional Kalman filtering introduces significant computational and memory demands, the authors provide a simplified approach that reduces these costs while retaining its advantages. Theoretically, they show that their method achieves better iteration complexity than DP-SGD by a constant factor. Experimentally, their method outperforms DP-SGD on various datasets.
weaknesses: - The privacy analysis of algorithm 3 is not clear to me. Is it possible to express $\sigma_{DP}$ explicitly in terms of other variables? 

- While the conditions in Theorem 2 and Figure 8 offer some guidance on hyperparameter selection, the interactions between different hyperparameters ($\gamma, \epsilon$ and $\kappa$) remain complex and unclear, making practical hyperparameter tuning challenging.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: My initial rating is a weak reject that I am willing to raise with good responses to these questions.

# Questions about ambiguity
- The way neighbors are defined has been left ambiguous. Are two datasets neighbors if adding/removing a record turns one into the other? Are two datasets neighbors if modifying one record turns one into another? The choice can result in significant differences in sensitivity and applicability of subsampling amplification results.

- What form of subsampling was used for the creation of minibatches. Simple random sample with a fixed size? Poisson sampling?

- What privacy accountant was used?

- Algorithms 1,2,3 should use parentheses to more clearly resolve the ambiguity of whether the noise term belongs
in the summation or outside.


# Experimental questions
- Did you match neighbor definitions in the comparisons to prior SOTA?

- Have you tried any of the privacy auditing frameworks (e.g., relatively recent NeurIPS papers) to check for possible errors? For example, depending on when you add noise and when you average over a batch, an error could easily appear to be SOTA and this is not uncommon in DP-SGD.

- The algorithms is 2 changes away from vanilla DP-SGD. The paper has hyperparmaeter tuning studies that are named ablation studies, but it does not have true ablation studies. There are 4 possible comparisons that can shed light on what is responsible for most of the apparent performance improvements. (1) vanilla dp-SGD (k=1, gamma=0); (2) weighted average before clipping, but no momentum on the gradient after clipping; (3) only momentum in the gradient and no pre-clipping averaging (gamma=0), (4) DiSK. 
    - Corollary 1 implies that gamma may be the more important parameter.

- The paper does not explain how hyperparameters were selected. Were they chosen based on validation data (good) or testing data (thus favoring any approach that has more hyperparameters).

- I tried to match up some of the SOTA numbers in the appendix as a spot check. 
    - Specifically, I looked at De et al. because they have a few prior SOTA numbers in the same place. I couldn't match their Imagenet numbers (32.4% top-1 accuracy to the 36.5 written in Appendix C5). 
   - I also noticed that some of the high CIFAR-10 training from scratch numbers in De et al. are completely missing from the Appendix.
   - It is not clear to me whether DiSK is appropriately compared to the prior SOTA. Do you match deltas? Do you match neighbor definitions and batch subsampling? Do you match privacy accountants?
   - Delta is a big difference maker, so you should put its numerical value(s) in every table.
   - Table 6 does not match imagenet for table 2.

# Other questions/comments
- What is that wacky green line in Figure 1?
- The paper uses C for too many different purposes. In some cases it is a clipping threshold, other times a constant in theoretical bounds, and other times a constant on properties of the data. This can get confusing even with subscripts.
- Line 121: where is the delta < 0.05 restriction coming from? If I am not mistaken, Definition 2 is all Balle & Wang and does not use the results from Dwork and Roth.
-Line 086-089: such comparisons are incomplete without specifying privacy parameters and adding citations.
soundness: 3
strengths: The paper reads very nicely. The authors take the reader through a tour of their intuition when explaining the algorithm. The theoretical results are also discussed nicely. Most papers simply present a Big-O result and switch topics. This paper explained the significance of the theorem conclusion.

I like the fact that the resulting algorithm is simple, increasing the chance of it being widely adopted.

The theoretical analysis is nice (theorem 2 and corollary 1) and shows that the provable convergence upper bound can be improved by a constant factor that depends on the model and data. It is not the same as claiming that faster convergence has been proved (Line 419, but that claim can be reworded to be accurate)

The model technique achieves state of the art in many settings (there is a small asterisk because I am not fully convinced yet, see below)..
summary: This paper proposes DiSK. DiSK is inspired by an application of a Kalman filter to the DP-SGD process. While the Kalman filter includes computationally intensive steps and requires knowledge of some unknown matrices, the simplification steps result in DP-SGD with two seemingly small changes. The first is that the clipped gradient is computed from an average of the gradient at two carefully chosen nearby points. The second is a momentum applied to the clipped gradient before being passed into an optimizer.

With these changes, the authors prove an upper bound on convergence rate. The provable constant in the upper bound for DiSK is smaller than the provable constant in the upper bound for vanilla SGD. Although this does not mean that the paper proved that DiSK has faster convergence (a claim that should be fixed), it does suggest that DiSK would converge faster. The analysis suggests that the pre-clipping averaging step may be the more important of the differences between DiSK and DP-SGD.

The paper also performs extensive experiments to compare against prior work to establish SOTA. It is not clear that the experiments match the deltas of the prior SOTA. I wasn't able to match the reported SOTA numbers to the papers that I checked, and some SOTA numbers are missing.
weaknesses: There is a lot of ambiguity in the paper. Critical details are not discussed to ensure comparisons are apples to apples (see below).

There are discrepancies in the experiments and comparisons to prior SOTA (see below)

Meta-review: This paper contributes to a growing body of work that aims to improve the performance (utility) of differentially private stochastic gradient descent (DP-SGD) by post-processing gradients. The main contribution is DiSK: a simplified Kalman Filter to denoise privatized gradients. The diSK is theoretically grounded and demonstrates favorable performance relative to competing state-of-the-art methods in various benchmarks.

The reviewers noted several positive points in the paper, including "strong and significant empirical results" and "Computational costs (e.g. time/memory/parallelizability) of the method are quite good too" (pdcW), the method's novel and computationally efficient way of (approximately) incorporating the Kalman filter into DP training (Nhmw), and the algorithm being practically feasible and theoretically grounded (PnV2).

The reviewers also noted several areas of improvement, including ambiguity in the experimental setups and comparisons to prior SOTA and discrepancies in experimental results (PnV2), clarity and presentation of the DiSK algorithm (Nhmw). They noted that the privacy analysis could be streamlined (j3TA).

Overall, this is a nice contribution of practical relevance that pushes the utility of DP-SGD forward -- a clear acceptance. I encourage the authors to add the promised changes in the final version of the paper.

Decision: Accept (Poster)