\section{Introduction}

Over the past few years, deep reinforcement learning has gained much popularity as it has been shown to perform better than previous methods on domains with very large state-spaces.
In one of the earliest deep reinforcement learning papers (hereafter the DQN paper), \citet{mnih2015human} presented a method for learning to play Atari 2600 video games, using the Arcade Learning Environment (ALE)~\citep{bellemare13arcade}, from image and performance data alone using the same deep neural network architecture and hyper-parameters for all the games.
DQN outperformed previous reinforcement learning methods on nearly all of the games and recorded better than human performance on most.
As many researchers tackle reinforcement learning problems with deep reinforcement learning methods and propose alternative algorithms, the results of the DQN paper are often used as a benchmark to show improvement.
Thus, implementing the DQN algorithm is important for both replicating the results of the DQN paper for comparison and also building off the original algorithm.
One of the main contributions of the DQN paper was finding ways to improve stability in their artificial neural networks during training.
There are, however, a number of other areas in the implementation of this method that are crucial to its success, which were only mentioned briefly in the paper.

We implemented a Deep Q-Network (DQN) to play the Atari games and replicated the results of \citet{mnih2015human}.
Our implementation, available freely online,\footnote{\url{www.github.com/h2r/burlap_caffe}} runs around 4x faster than the original implementation.
Our implementation is also designed to be flexible to different neural net network architectures and problem domains outside of ALE.
In replicating these results, we found a few key insights into the process of implementing such a system.
In this paper, we highlight key techniques that are essential for good performance and replicating the results of \citet{mnih2015human}, including termination conditions and gradient descent optimization algorithms, as well as expected results of the algorithm, namely the fluctuating performance of the network.

\section{Related Work}{}

The Markov Decision Process (MDP) \citep{bellman1957markovian} is the typical formulation used for reinforcement learning problems.
An MDP is defined by a five-tuple $(\mathcal{S, A, T, R, E})$;
$\mathcal{S}$ is the agent's state-space;
$\mathcal{A}$ is the agent's action-space;
$\mathcal{T}(s, a, s')$ represents the transition dynamics, which returns the probability that taking action $a$ in state $s$ will result in the state $s'$;
$\mathcal{R}(s, a, s')$ is the reward function, which returns the reward received when transitioning to state $s'$ after taking action $a$ in state $s$;
and $\mathcal{E} \subset \mathcal{S}$ is the set of terminal states, which once reached prevent any future action or reward.
The goal of planning in an MDP is to find a policy $\pi : S \rightarrow A$, a mapping from states to actions, that maximizes the expected future discounted reward when the agent chooses actions according to $\pi$ in the environment. A policy that maximizes the expected future discounted reward is an optimal policy and is denoted by $\pi^*$.

A key concept related to MDPs is the Q-function, $Q^\pi : S \times A \rightarrow \mathbb{R}$, that defines the expected future discounted reward for taking action $a$ in state $s$ and then following policy $\pi$ thereafter. According to the Bellman equation, the Q-function for the optimal policy (denoted $Q^*$) can be recursively expressed as:
\begin{equation}
Q^*(s, a) = \sum_{s' \in S} T(s, a, s') \left [ R(s, a, s') + \gamma \max_{a'} Q^*(s', a') \right ]
\end{equation}
where $0 \leq \gamma \leq 1$ is the discount factor that defines how valuable near-term rewards are compared to long-term rewards.
Given $Q^*$, the optimal policy, $\pi^*$, can be trivially recovered by greedily selecting the action in the current state with the highest Q-value: $\pi^*(s) = \argmax_a Q^*(s, a)$. This property has led to a variety of learning algorithms that seek to directly estimate $Q^*$, and recover the optimal policy from it. Of particular note is Q-Learning~\citep{watkins1989learning}.

In Q-Learning, an agent begins with an arbitrary estimate ($Q_0$) of $Q^*$ and iteratively improves its estimate by taking arbitrary actions in the environment, observing the reward and next state, and updating its Q-function estimate according to
\begin{equation}
Q_{t+1}(s_t, a_t) \gets Q_t(s_t, a_t) + \alpha_t \left[ r_{t+1} + \gamma \max_{a'} Q_t(s_{t+1}, a') - Q_t(s_t, a_t) \right],
\end{equation}
where $s_t$, $a_t$, $r_t$ are the state, action, and reward at time step $t$, and $\alpha_t \in (0, 1]$ is a step size smoothing parameter.
Q-Learning is guaranteed to converge to $Q^*$ under the following conditions: the Q-function estimate is represented tabularly (that is, a value is associated with each unique state-action pair), the agent visits each state and action infinitely often, and $\alpha_t \rightarrow 0$ as $t \rightarrow \infty$.
When the state-space of a problem is large (or infinite), Q-learning's $Q^*$ estimate is often implemented with function approximation, rather than a tabular function, which allows generalization of experience.
However, estimation errors in the function approximation can cause Q-learning, and other ``off policy'' methods, to diverge~\citep{baird1995residual}, requiring careful use of function approximation.

\section{Deep Q-Learning}

\begin{algorithm}[t]
\begin{algorithmic}
\State Initialize replay memory $D$ to capacity $N$
\State Initialize action-value function $Q$ with random weights $\theta$
\State Initialize target action-value function $\hat Q$ with weights $\theta^{-} = \theta$
\For{episode 1, $M$}
Initialize sequence $s_1 = \{ x_1 \}$ and preprocessed sequence $\phi_1 = \phi(s_1)$
\For{$t = 1, T$}
\State With probability $\varepsilon$ select a random action $a_t$
\State otherwise select $a_t = \argmax_a Q(\phi(s_t), a; \theta)$
\State Execute action $a_t$ in the emulator and observe reward $r_t$ and image $x_{t+1}$
\State Set $s_{t+1} = s_t, a_t, x_{t+1}$ and preprocess $\phi_{t+1} = \phi(s_{t+1})$
\State Store experience $(\phi_t, a_t, r_t, \phi_{t+1})$ in $D$
\State Sample random minibatch of experiences $(\phi_j, a_j, r_j, \phi_{j+1})$ from $D$
\State Set $y_j = \begin{cases}
r_j & \text{if episode terminates at step $j+1$}\\
r_j + \gamma \max_{a'} \hat Q(\phi_{j+1}, a'; \theta^{-}) & \text{otherwise}
\end{cases}$
\State Perform a gradient descent step on $(y_j - Q(\phi_j, a_j ; \theta))^2$ with respect to the weights $\theta$
\State Every $C$ steps reset $\hat Q = Q$
\EndFor
\EndFor
\end{algorithmic}
\caption{Deep Q-learning with experience replay}
\label{alg:dqn}
\end{algorithm}

Deep Q-Learning (DQN)~\citep{mnih2015human} is a variation of the classic Q-Learning algorithm with 3 primary contributions: (1) a deep convolutional neural net architecture for Q-function approximation; (2) using mini-batches of random training data rather than single-step updates on the last experience; and (3) using older network parameters to estimate the Q-values of the next state.
Pseudocode for DQN, copied from \citet{mnih2015human}, is shown in Algorithm~\ref{alg:dqn}.
The deep convolutional architecture provides a general purpose mechanism to estimate Q-function values from a short history of image frames (in particular, the last 4 frames of experience). The latter two contributions concern how to keep the iterative Q-function estimation stable.

In supervised deep-learning work, performing gradient descent on mini-batches of data is often used as a means to efficiently train the network. In DQN, it plays an additional role.
Specifically, DQN keeps a large history of the most recent experiences, where each experience is a five-tuple $(s, a, s', r, T)$, corresponding to an agent taking action $a$ in state $s$, arriving in state $s'$ and receiving reward $r$; and $T$ is a boolean indicating if $s'$ is a terminal state.
After each step in the environment, the agent adds the experience to its memory.
After some small number of steps (the DQN paper used 4), the agent randomly samples a mini-batch from its memory on which to perform its Q-function updates.
Reusing previous experiences in updating a Q-function is known as {\em experience replay}~\citep{lin1992self}.
However, while experience replay in RL was typically used to accelerate the backup of rewards, DQN's approach of taking fully random samples from its memory to use in mini-batch updates helps decorrelate the samples from the environment that otherwise can cause bias in the function approximation estimate.

The final major contribution is using older, or ``stale,'' network parameters when estimating the Q-value for the next state in an experience and only updating the stale network parameters on discrete many-step intervals. This approach is useful to DQN, because it provides a stable training target for the network function to fit, and gives it reasonable time (in number of training samples) to do so. Consequently, the errors in the estimation are better controlled.

Although these contributions and overall algorithm are straightforward conceptually, there are number of important details to achieving the same level of performance reported by \citet{mnih2015human}, as well as important properties of the learning process that a designer should keep in mind. We describe these details next.

\subsection{Implementation Details}

Large systems, such as DQN, are often difficult to implement since original scientific publications are not always able to describe in detail every important parameter setting and software engineering solution.
Consequently, some important low-level details of the algorithm are not explicitly mentioned or fully clarified in the DQN paper.
Here we highlight some of these key additional implementation details, which are provided in the original DQN code.\footnote{\url{www.github.com/kuz/DeepMind-Atari-Deep-Q-Learner}}

Firstly, every episode is started with a random number of ``No-op'' low-level Atari actions (in contrast to the agent's actions which are repeated for $4$ frames) between $0$ and $30$ in order to offset which frames the agent sees, since the agent only sees every $4$ Atari frames.
Similarly, the $m$ frame history used as the input to the CNN is the last $m$ frames that the agent sees, not the last $m$ Atari frames.
Additionally, before any gradient descent steps, a random policy is run for $\num{50000}$ steps to fill in some experiences in order to avoid over-fitting to early experiences.

Another parameter worth noting is the network update frequency.
The original DQN implementation only chose to take a gradient descent step every $4$ environment steps of the algorithm as opposed to every step, as Algorithm \ref{alg:dqn} might suggest.
Not only does this greatly increase the training speed (since learning steps on the network are far more expensive than forward passes), it also causes the experience memory to more closely resemble the state distribution of the current policy (since 4 new frames are added to the memory between training steps as opposed to 1) and may prevent the network from over-fitting.  \stnote{Are there results for this?  Could there be?  (I know this is a lot of work and probably not worth the tiem but wanted to ask the question...)}

\subsection{The Fluctuating Performance of DQN}
A common belief for new users of DQN is that performance should fairly stably improve as more training time is given. Indeed, average Q-learning learning curves in tabular settings are typically fairly stable improvements and supervised deep-learning problems also tend have fairly steady average improvement as more data becomes available. However, it is not uncommon in DQN to have ``catastrophic forgetting'' in which the agent's performance can drastically drop after a period of learning. For example, in Breakout, the DQN agent may reach a point of averaging a high score over $400$, and then, after another large batch of learning, it might be averaging a score of only around $200$. The solution \citet{mnih2015human} propose to this problem is to simply save the network parameters that resulted in the best test performance.

One of the reasons this forgetting occurs is the inherent instability of approximating the Q-function over a large state-space using these Bellman updates.
One of the main contributions of \citet{mnih2015human} was fighting this instability using experience replay and stale network parameters, as mentioned above.
Additionally, \citet{mnih2015human} found that clipping the gradient of the error term to be between $-1.0$ and $1.0$ further improved the stability of the algorithm by not allowing any single mini-batch update to change the parameters drastically.
These additions, and others, to the DQN algorithm improve its stability significantly, but the network still experiences catastrophic forgetting.
\jmnote{I think I might suggest moving the gradient clipping out as its own unique thing to helping the stabilization. So whereas the main contributions to stabilization are the experience replay and stale network parameters, gradient clipping is also an unsung hero.}

Another reason this catastrophic forgetting occurs is that the algorithm is learning a proxy, the Q-values, for a policy instead of approximating the policy directly.
A side effect of this method of policy generation is a learning update could increase the accuracy of a Q-function approximator, while decreasing the performance of the resulting policy.
For example, say the true Q-value for some state, $s$, and actions, $a_1$ and $a_2$, are $Q^*(s, a_1) = 2$ and $Q^*(s, a_2) = 3$, so the optimal policy at state $s$ would be to choose action $a_2$.
Now say the Q-function approximator for these values using current parameters, $\theta$, estimates $\hat Q(s, a_1; \theta) = 0$ and $\hat Q(s, a_2; \theta) = 1$, so the policy chosen by this approximator will also be $a_2$
But, after some learning updates we arrive at a set of parameters $\theta'$, where $\hat Q(s, a_1; \theta) = 2$ and $\hat Q(s, a_2; \theta) = 1$.
These learning updates clearly decreased the error of the Q-function approximator, but now the agent will not choose the optimal action at state $s$. \jmnote{Doesn't the nature paper show something like the plots of error in Q-function over time showing that it's steadily improving? If so you may want to call that out as evidence that even though the Q-function improves, you might get a really bad result in performance.}

Furthermore Q-values for different actions of the same state can be very similar if any of these actions does not have a significant effect on near-term reward.
These small differences are the result of longer-term rewards and are therefore critical to the optimal policy.
The consequence of trying to learn an approximator for this type of function is that very small errors in the Q-values can result is very different policies, making it difficult to learn long-term policies.

As an example of this, we will consider Breakout.
Breakout is an Atari game where the player controls a paddle with the goal of bouncing a ball to destroy all the bricks on the screen without dropping the ball.
There is an optimal strategy which is to destroy the bricks on the side of the screen so that the ball can be bounced above the bricks.
When the ball is above the bricks, the Q-values are much higher than they are when the ball is below the bricks, so we would expect a policy which follows the true Q-values to quickly exploit this policy.
Every time your paddle bounces the ball, the direction does not affect short-term rewards as a brick will be broken every time you bounce the ball
But, the ball direction will affect the distant reward of bouncing the ball above the bricks by breaking the bricks on the side of the screen.
Thus, it is difficult for a Q-function approximator to learn this long-term optimal policy.

Figure \ref{fig:q_values} shows Q-values approximated by the best network and a network that performed poorly very late into training on the same inputs near the beginning of a Breakout game.
The first frame illustrates a scenario where any action could be made and the agent could still prevent the ball from falling a few actions into the future.
But the actions made before bouncing the ball also allow the agent to aim the ball.
The Q-values in this case are very similar for both networks, but the chosen actions are different.
In the second scenario, if the agent does not take the left action, the ball will be dropped, which is a terminal state.
In this case, the Q-values are much more distinct.
Thus, this fluctuating performance is to be expected while running this algorithm.

\begin{figure}
  \centering
  \captionsetup[subfigure]{labelformat=empty}
  \begin{subfigure}[b]{0.2\textwidth}
    \includegraphics[width=\textwidth]{breakout_ambiguous}
  \end{subfigure}
  \begin{subfigure}[b]{0.35\textwidth}
    \includegraphics[width=\textwidth]{q_values_best_ambiguous}
  \end{subfigure}
  \begin{subfigure}[b]{0.35\textwidth}
    \includegraphics[width=\textwidth]{q_values_worst_ambiguous}
  \end{subfigure}

  \begin{subfigure}[b]{0.2\textwidth}
    \includegraphics[width=\textwidth]{breakout_unambiguous}
    \caption{Example Frames}
  \end{subfigure}
  \begin{subfigure}[b]{0.35\textwidth}
    \includegraphics[width=\textwidth]{q_values_best_unambiguous}
    \caption{Best Network}
  \end{subfigure}
  \begin{subfigure}[b]{0.35\textwidth}
    \includegraphics[width=\textwidth]{q_values_worst_unambiguous}
    \caption{Worst Network}
  \end{subfigure}
  \caption{
  A comparison of calculated Q-values by the networks that received the best and worst performance during testing that had been trained for at least 30,000,000 steps.
  The lighter cross-hatched bar indicates the action with the highest Q-value.
  The top frame corresponds to a situation where the actions don't have a significant affect on the near-future reward, while the bottom one shows a situation where the left action must be made to not loose a life.
  The ``Release'' action releases the ball at the beginning of every round or does nothing (the same as ``No-op'') if the ball is already in play.
  }
  \label{fig:q_values}
\end{figure}

\section{Machine Learning Libraries}

Our implementation uses the Brown-UMBC Reinforcement Learning and Planning (BURLAP) Java code library \citep{burlap}.
This library makes it easy to define a factored representation of an MDP and offers many well-known learning and planning algorithms as well as the machinery for creating new ones.

For running and interacting with the Atari video games, we used the Arcade Learning Environment (ALE) \citep{bellemare13arcade}.
ALE is a framework that provides a simple way to retrieve the screen and reward data from the Atari games as well as interact with the game through single actions.
We used ALE's FIFO Interface to interact ALE through Java.

To run and train our convolutional neural net, we used the Berkeley's Caffe (Convolutional Architecture for Fast Feature Embedding) library \citep{jia2014caffe}.
Caffe is a fast deep learning framework for constructing and training neural network architectures.
To interact with Caffe through our Java library, we used the JavaCPP library provided by Bytedeco.\footnote{\url{www.github.com/bytedeco/javacpp}}

\section{Results}

To measure our performance against that of \citet{mnih2015human}, we followed the same evaluation procedure as their paper on three games: Pong, Breakout, and Seaquest.
We trained the agent for $\num{50000000}$ steps (each step is 4 Atari frames) and tested performance every $\num{250000}$ steps.
We saved the network parameters that resulted in the best test performance.
We then evaluated the trained agent with the best performing network parameters on 30 games with and $\varepsilon$-greedy policy where $\varepsilon = 0.05$.
Each game was also initialized with a random number of ``No-op'' low-level Atari actions between $0$ and $30$.
We then took the average score of those games.

The comparison of our results and those of the DQN paper on Pong, Breakout, and Seaquest are shown in Table \ref{results}.
Each training process took about 3 days for our implementation and about 10 and a half days for the original implementation on our setup.

The differences in performance stem from the differences in gradient descent optimization algorithm and learning rate.
These differences are covered in more detail in section \ref{RMS}.

\begin{table}[t]
  \captionsetup{skip=8pt}
  \caption{Comparison of average game scores obtained by our DQN implementation and the original DQN paper.}
  \label{results}
  \centering
  \begin{tabular}{lll}
    \toprule
    Game     & Our implementation    & The original implementation \\
    \midrule
    Pong     & $19.7 \ (\pm 1.1)$    & $18.9 \ (\pm 1.3)$          \\
    Breakout & $339.3 \ (\pm 86.1)$  & $401.2 \ (\pm 26.9)$        \\
    Seaquest & $6309 \ (\pm 1027)$   & $5286 \ (\pm 1310)$         \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Key Training Techniques}

While implementing our DQN, we found there were a couple methods that were only mentioned briefly in the DQN paper, but critical to the overall performance of the algorithm.
Here we present these methods and explain why they have such a strong impact on training the network.

\subsection{Termination on the Loss of Lives}

In most of the Atari games, there is a notion of ``lives'' for the player, which correspond to the number of times the player can fail (such as dropping the ball in Breakout or running into a shark in Seaquest) before the game is over.
To increase performance, \citet{mnih2015human} chose to count the loss of a life (in the games involving lives) as a terminal state in the MDP during training.
This termination condition was not mentioned in much detail in the DQN paper, but is essential for achieving their performance.

Figure \ref{fig:lives-vs-no-lives} illustrates the difference between training with and without counting losing lives as terminal states in both Breakout and Seaquest.
In Breakout, the average score of the learner that uses end of lives as terminal states increases much faster than the other learner.
However, around halfway through training, the other learner is achieving similar performance, but with much higher variance.
Seaquest is a much more complex game with many more moving sprites and longer episode lengths.
In this game the learner that uses lives as terminal states performs significantly better than the other learner throughout training.
These figures illustrate that this additional prior information greatly benefits early training and stability and, in the more complex games, significantly improves the overall performance.

\begin{figure}
  \centering

  \begin{subfigure}[b]{0.48\textwidth}
    \includegraphics[width=\textwidth]{lives-vs-no-lives}
    \caption{Breakout}
  \end{subfigure}
  \begin{subfigure}[b]{0.48\textwidth}
    \includegraphics[width=\textwidth]{seaquest_lives-vs-no-lives}
    \caption{Seaquest}
  \end{subfigure}

  \caption{The average training test score received for Breakout and Seaquest at each test set when using lost lives as terminal states and when using the end of a game as terminal states (epoch = 250,000 steps).}
  \label{fig:lives-vs-no-lives}
\end{figure}

A terminal state in an MDP, as mentioned above, signifies to the agent that no more reward can be obtained.
Almost all the Atari games give positive rewards (Pong is a notable exception where a reward of $-1$ is received when the enemy scores a point), and thus, this addition essentially informs the agent that losing a life should be avoided at all costs.
This additional information given to the agent does seem reasonable: many human players know that loosing a life in an Atari game is bad the first time they play and it is difficult to imagine situations where the optimal policy would be to lose a life.

There are, however, a few theoretical issues with enforcing this
constraint.  The first being that this process is no longer Markovian
as the initial state distribution depends on the current policy.  An
example of this is in Breakout: If the agent performed well and broke many bricks before
losing a life, the new initial state for the next life will have many
fewer bricks remaining than if the agent performed poorly and broke very few bricks in the
previous life.  The other issue is that this signal gives strong
additional information to the DQN, making it challenging to extend to
domains where such strong signals are not available (e.g., real-world
robotics or more open-ended video games.)

Although ALE stores the number of lives remaining for each game, it does not provide this information to all the interfaces.
To work around this limitation, we modified ALE's FIFO Interface to provide the number of lives remaining along with the screen, reward, and terminal state boolean.
Our fork that provides this data to the FIFO interface is available freely online.\footnote{\url{www.github.com/h2r/arcade-learning-environment}}

\subsection{Gradient Descent Optimization} \label{RMS}

One potential issue in using the hyper-parameters provided by \citet{mnih2015human} is that they are not using the same RMSProp definition that many deep learning libraries (such as Caffe) provide.
The RMSProp gradient descent optimization algorithm was originally proposed by Geoffrey Hinton.\footnote{\url{www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf}}
Hinton's RMSProp keeps a running average of the gradient with respect to each parameter.
The update rule for this running average can be written as:
\begin{equation}
MeanSquare(w, t) = \gamma \cdot MeanSquare(w, t-1) + (1 - \gamma) \cdot (\frac{\partial E}{\partial w}(t))^2
\end{equation}
Here, $w$ corresponds to a single network parameter, $\gamma$ is the decay factor, and $E$ is the loss.
The parameters are then updated by:
\begin{equation}
w_t = w_{t-1} - \frac{\alpha}{\sqrt{MeanSquare(w, t) + \varepsilon}} \cdot \frac{\partial E}{\partial w}(t)
\end{equation}
where $\alpha$ corresponds to the learning rate, and $\varepsilon$ is a small constant to avoid division by $0$.

Although \citet{mnih2015human} cite Hinton's RMSProp, they use a
slight variation on the algorithm.  The implementation of this can be
seen in their GitHub
repository\footnote{\url{www.github.com/kuz/DeepMind-Atari-Deep-Q-Learner}}
in the NeuralQLearner.lua file on lines 266-273.  \stnote{Nice!}  This variation adds
a momentum factor to the RMSProp algorithm that is updated as follows:
\begin{equation}
Momentum(w, t) = \eta \cdot Momentum(w, t-1) + (1 - \eta) \cdot \frac{\partial E}{\partial w}(t)
\end{equation}
Here, $\eta$ is the momentum decay factor.
The parameter update rule is then modified to:
\begin{equation}
w_t = w_{t-1} - \frac{\alpha}{\sqrt{MeanSquare(w, t) - (Momentum(w, t))^2 + \varepsilon}} \cdot \frac{\partial E}{\partial w}(t)
\end{equation}

To account for this change in change in optimization algorithm, we had to modify the learning rate to something much lower than that of the \citet{mnih2015human} implementation (we used $0.00005$ as opposed to their $0.00025$).
We did not choose to implement this variant of RMSProp as it was not trivial to implement with the Java-Caffe bindings and Hinton's version produced similar results.

\section{Speed Performance}

Our implementation is training a bit less than 4x faster than the original implementation written in Lua using Torch.
The setup on which we tested these implementations is using two NVIDIA GTX 980 TI graphics cards along with an Intel i7 processor.
Our implementation runs at around $985$ Atari frames per second (fps) during training and $1584$fps during testing, while the Lua implementation runs at $271$fps during training and $586$fps during testing on our hardware (note that the algorithm only looks at every 4 frames, and so only 1 fourth of this number of frames are processed by the algorithm per second).
We attribute a large portion of this performance increase to cuDNN.

The NVIDIA CUDA Deep Neural Network library (cuDNN) is a proprietary NVIDIA library for running forward and backward passes on common neural network layers optimized specifically for the NVIDIA GPUs.
For both Torch and Caffe, cuDNN is available, but not used by default.
We compiled Caffe using cuDNN for our experiments, while the Lua implementation did not use this feature in Torch.
For comparison, when using Caffe without cuDNN, our implementation runs at around $268$fps during training and $485$fps during testing, which is a bit slower than the Lua implementation.

Another area that significantly increased the speed performance of our implementation was preallocating memory before training, which was also done in the original DQN implementation.
Allocating large amounts of memory is an expensive operation, so preallocating memory for large vectors, such as the experience memory and mini-batch data, and reusing it at each iteration significantly decreases this overhead.

\section{Conclusion}
In this paper we have presented a few key areas in the implementation of the DQN proposed by \citet{mnih2015human} that were essential to the overall performance of the algorithm, but were not covered in great detail in the original paper, in order to make it easier for researchers to implement their own versions of this algorithm.
We also highlighted some of the difficulties in approximating a Q-function with a CNN in such large state-spaces, namely the catastrophic forgetting.
We also have our implementation available freely online\footnote{\url{www.github.com/h2r/burlap_caffe}} and encourage researchers to use this as a tool for implementing novel algorithms as well as comparing performance with that of \citet{mnih2015human}.

\subsubsection*{Acknowledgments}
This material is based upon work supported by the National Science Foundation under grant numbers IIS-1637614 and IIS-1426452, and DARPA under grant numbers W911NF-10-2-0016 and D15AP00102.

\medskip

{
\bibliographystyle{plainnat}


\section{Experiments} \label{sec:exp}

We conduct experiments on various kinds of tasks to demonstrate the effectiveness of RGP. We first examine the utility of  models trained by RGP. To this end, we apply RGP  on the wide ResNet \citep{zagoruyko2016wide}  and  the BERT \citep{devlin2018bert} models, which are representative models for  computer vision and natural language modeling. The results are presented in Section~\ref{subsec:exp_resnet} and~\ref{subsec:exp_bert}.  The source code of our implementation is publicly available\footnote{\url{https://github.com/dayu11/Differentially-Private-Deep-Learning}}.

\begin{table}
\small
\renewcommand{\arraystretch}{1.2}
\centering
    \caption{Validation accuracy (in \%) of WRN28-4 on vision tasks . } \label{tbl:tbl_resnet}

        \begin{tabular}{l|l|l}
        \hline
        \hline
            Method         & SVHN &  CIFAR10    \\
            \hline
              Full (N.P.) &      97.2        &      93.3     \\\cline{1-3}
             Linear (N.P.) &      41.1        &     39.8            \\\cline{1-3}
             RGP (N.P.)   &        97.1       &    91.2      \\\cline{1-3}
             PowerSGD (N.P.) &      97.1         &        91.9         \\\cline{1-3}
                DP-SGD ($\epsilon=8$)  &      91.6         &       55.9   \\\cline{1-3}
             DP-PowerSGD ($\epsilon=8$) &         91.9      &     57.1           \\\cline{1-3}
                RGP-random ($\epsilon=8$)  &    91.7           &   51.0      \\\cline{1-3}
             RGP ($\epsilon=8$)&     94.2         &    63.4     \\\hline \hline
        \end{tabular}
\end{table}

\begin{table}
\small
\centering
    \caption{Validation accuracy (in \%) of RGP on vision tasks with varying $\epsilon$. The model architecture is WRN28-4. Numbers in  brackets denote the improvements compared to DP-SGD. } \label{tbl:vision_vary_eps}

        \begin{tabular}{l|l|l|l}
        \hline
        \hline
            Dataset         & $\epsilon=2$ &  $\epsilon=4$ & $\epsilon=6$    \\\hline
             SVHN &    87.3  (+4.1)      &  89.7   (+3.4)    & 92.3  (+3.9)  \\\hline
             CIFAR10 &     44.0 (+6.6)        &   53.3 (+6.4)   &     59.6 (+7.9)      \\\hline \hline
        \end{tabular}
\end{table}

Moreover, we empirically  evaluate the privacy risk of the models via the success rate of \emph{membership inference (MI) attack} \citep{shokri2017membership,sablayrolles2019white,yu2021how}. The results are  presented in Section~\ref{subsec:exp_mi}.

\textbf{Implementation.}  The number of iterations for power method is $1$. We use an open-source tool of moments accountant to compute the privacy loss\footnote{\url{https://github.com/tensorflow/privacy}}. For a given setting of hyperparameters, we  set $\sigma$ to be the smallest value so that the privacy budget is allowable to run desired epochs. All experiments are run on a node with four Tesla V100 GPUs.

\textbf{Baselines.} We implement several baseline algorithms for comparison.  For differentially private learning, the first baseline is \emph{DP-SGD} in \citet{abadi2016deep} and the second one is RGP with gradient carriers consisting of random orthonormal vectors, referred to as \emph{RGP-random}. We also include several non-private baselines, i.e., \textbf{(\romannumeral 1)} \emph{Full (N.P.)}: training the full model, \textbf{(\romannumeral 2)} \emph{Linear (N.P.)}: training only the linear classification layer, \textbf{(\romannumeral 3)} \emph{RGP (N.P.)}: training the model with reparametrization  but without gradient clipping or adding noise.

We consider differentially private  \emph{PowerSGD} \citep{vogels2019powersgd} as another baseline for vision tasks. PowerSGD approximates full gradients with low-rank matrices to reduce the communication cost. It first aggregates the individual gradients and then runs power iterations to find approximations of the principle components of the averaged gradient. Hence for DP-powerSGD, it is necessary to first perturb the aggregated gradient  and then project it into low-rank subspace otherwise  the sensitivity is hard to track after projection. As a consequence,  DP-powerSGD needs to compute the individual gradients explicitly, which costs huge memory as DP-SGD does. In Section~\ref{subsec:exp_resnet}, we add a DP-powerSGD baseline  with the same setting as that of RGP.

Additionally, some ablation experiments are conducted to study the influence of the residual weight and reparametrization ranks, which are relegated to the Appendix~\ref{app:sec:add-exp}.

\subsection{Experiments on Vision Tasks}\label{subsec:exp_resnet}

\textbf{Model.} We use wide ResNet models \citep{zagoruyko2016wide} for the vision tasks. The architecture is WRN28-4 with $\sim$1.5M parameters. All batch normalization layers are replaced with group normalization layers to accommodate private learning.

\textbf{Tasks.} We use two vision  datasets: SVHN \citep{netzer2011reading}  and CIFAR10 \citep{cifar}. SVHN  contains images of $10$ digits and CIFAR10 contains images of 10 classes of real-world objects.

\textbf{Hyperparameters.} We follow the hyperparameters in \citet{zagoruyko2016wide} except using a mini-batch size 1000. This mini-batch size is larger than the default because the averaging effect of large mini-batch reduces the noise variance. The reparametrization rank $r$ is chosen from $\{1, 2, 4, 8, 16\}$. We choose the privacy parameter $\delta<\frac{1}{n}$, and set $\delta=10^{-6}$ for SVHN and $\delta=10^{-5}$ for CIFAR10. We repeat each experiment 3 times and  report the average.

\textbf{Results.} The prediction accuracy with $\epsilon=8$ is presented in Table~\ref{tbl:tbl_resnet}. We can see that RGP (N.P.) achieves comparable performance with training the full model (N.P.). When trained with DP,  RGP outperforms DP-SGD by a considerable margin while enjoying a much lower memory cost. We also compare RGP with DP-SGD using different privacy budgets ($\epsilon=2/4/6$) and report the results  Table~\ref{tbl:vision_vary_eps}.

\subsection{Experiments on the Downstream Tasks of BERT}\label{subsec:exp_bert}

\textbf{Model.} We use the BERT\textsubscript{BASE} model in \citet{devlin2018bert}, which is pre-trained on a massive corpus collected from the Web. The BERT\textsubscript{BASE} model has $\sim$110M parameters.

\textbf{Tasks.} We use four tasks from the General Language Understanding Evaluation (GLUE) benchmark \citep{wang2018glue}, including MNLI, QQP, QNLI, and SST-2. The other  tasks from GLUE are excluded because their datasets are of small sizes (<10K) while differentially private learning  requires  large amount of data \citep{tramer2021differentially}.

\textbf{Hyperparameters.} We follow the hyperparameters in \citet{devlin2018bert}
 except for the mini-batch size and training epochs. The reparametrization rank $r$ is chosen from $\{1, 2, 4, 8\}$. The mini-batch size is 500 for SST-2/QNLI and 1000 for QQP/MNLI. To construct an update with desired mini-batch size, we accumulate the gradients of multiple micro-batches.   We choose $\delta = 10^{-5}$ for QNLI/SST-2 and $\delta =10^{-6}$ for QQP/MNLI. The privacy parameter $\epsilon$ is chosen from $\{1, 2, 4, 6, 8\}$.  The number of training epochs is 50 for $\epsilon>2$ and $20$ for $\epsilon\leq 2$. We run all experiments 5 times with different random seeds and report the average.

\textbf{Results.}   The prediction accuracy of RGP and other baselines is presented in Table~\ref{tbl:tbl_bert}. The results with varying DP parameter $\epsilon$ is plotted in Figure~\ref{fig:fig_bert}.  When trained without privacy guarantee, RGP (N.P.) achieves  test accuracy comparable with fine-tuning the full model.  When trained with differential privacy, RGP achieves the best performance. Its accuracy loss compared to non-private baselines is within $5\%$. The performance of RGP-random is worse than that of RGP because the random subspace does not capture gradient information as effectively as the subspace of historical updates.  DP-SGD achieves the worst performance because high-dimensional noise overwhelms the useful signal in gradients.  We note that DP-SGD runs the lowest because it needs to compute and store 110M floating-point numbers for each individual gradient.

 \begin{figure*}
    \centering
  \includegraphics[width=0.9\linewidth]{imgs/bert_varying_eps.pdf}
  \caption{Prediction accuracy of BERT on downstream tasks with varying $\epsilon$. For MNLI, we plot the average score of two test datasets.  }
  \label{fig:fig_bert}
\end{figure*}

\begin{table}
\small
\centering
    \caption{Prediction accuracy of BERT on downstream tasks (in \%). For DP-SGD, RGP, and RGP-random, a same $\epsilon=8$ is used.}

        \begin{tabular}{l|l|l|l|l|l}
        \hline
        \hline
            Method         & MNLI &  QQP  & QNLI  &  SST-2  &  Avg.  \\
            \hline
              Full (N.P.) &     84.8/83.7         &   90.2       &    91.6   &  93.4   &  88.7 \\\cline{1-6}
             Linear (N.P.) &     51.9/50.8         &    73.2     &    63.0     &    82.1   &    64.2     \\\cline{1-6}
                RGP (N.P.)  &  83.6/83.2    &   89.3    &   91.3    &   92.9   & 88.1  \\\cline{1-6}
            DP-SGD\tablefootnote{As shown in \citet{li2021large}, DP-SGD performs better when large batchsizes and full precision are used.}  &     54.6/53.4           &     74.5    &     63.6    &   82.3      &  65.7 \\\cline{1-6}
            RGP-random  &      74.6/73.3        &    81.7   &      82.1 &    87.8    & 79.9 \\\cline{1-6}
               RGP\tablefootnote{The performance of RGP is also better in the above setup. More details are in  \url{https://github.com/dayu11/Differentially-Private-Deep-Learning}. }  &       79.1/78.0       &     84.8    &     86.2  &   91.5  &     83.9
            \\\hline \hline
        \end{tabular}
    \label{tbl:tbl_bert}
\end{table}

\subsection{Defense Against Membership Inference Attack}\label{subsec:exp_mi}

\textbf{Setup.} We use membership inference (MI) attack to empirically evaluate the privacy risk of models trained with/without RGP. Following the membership decision in \citet{sablayrolles2019white}, we predict a sample from the training data if its loss value is smaller than a chosen threshold.  To evaluate the MI success rate, we construct a \emph{MI dataset}, which consists of the same number of training and test samples. Specifically, the MI dataset contains the whole test set and a random subset of the training set. We further divide the MI dataset evenly into two subsets. One is used to find the optimal loss threshold and the other one is used to evaluate the final attack success rate.

\textbf{Results.}
The MI success rates  are presented in Table~\ref{tbl:mi_bert}. For MNLI, QQP, QNLI, and SST-2 datasets, we conduct MI attacks on fine-tuned BERT\textsubscript{BASE} models. For SVHN and CIFAR10 datasets, we conduct MI attacks on trained WRN28-4 models. The MI attack on the models trained with RGP ($\epsilon=8$) is no better than random guessing ($50\%$ success rate), which empirically demonstrate the effectiveness of RGP in protecting privacy. Moreover, interestingly, the models trained with low-rank reparametrization alone also achieve much lower MI success rate than the fully trained model, which indicates the benefit of low-rank reparametrization in terms of privacy protection.

\section*{Acknowledgements}

Jian Yin is supported by NSFC (U1711262,  U1711261, U1811264, U1811261, U1911203, U2001211), Guangdong Basic and Applied Basic Research Foundation (2019B1515130001),  Key R\&D Program of Guangdong Province (2018B010107005).

\newpage

\section{Introduction}

A recent line of works \citep{shokri2017membership,carlini2019secret,carlini2020extracting} have exposed the potential privacy risks of trained models, e.g., data extraction from language model. Theoretically, learning with  \emph{differential privacy} \citep{dwork2006calibrating} is guaranteed to prevent such information leakage because differential privacy imposes an upper bound on the influence of any individual sample. Empirically, differential privacy also makes learning more resistant to attacks \citep{rahman2018membership,bernau2019assessing, zhu2019deep, carlini2019secret, ma2019data,lecuyer2019certified}.

To learn with differential privacy, many  algorithms have been proposed under different settings over the past decade, e.g.,  \citet{chaudhuri2009privacy,song2013stochastic,agarwal2018cpsgd,wang02019differentially,wang2019differentially,yu2020gradient,phan2020scalable,vietri2020private}, to name a few. Among them, \emph{gradient perturbation} is a popular choice because of its simplicity and wide applicability \cite{abadi2016deep}. In terms of simplicity, gradient perturbation only makes two simple modifications to the standard learning process. It  first clips the gradients of individual samples, referred to as individual gradients, to bound the sensitivity and then perturbs the aggregated gradient with random noise. In terms of wide applicability, it does not assume the objective to be convex and hence applies to deep  neural networks.

  \begin{figure}
    \centering
  \includegraphics[width=1.0\linewidth]{imgs/repara.pdf}
  \caption{The proposed  reparametrization scheme.  The residual weight makes the reparametrized output the same as the normal output and  $\partial{\bm{L}}$, $\partial{\bm{R}}$ naturally connected with the normal gradient.    }
  \label{fig:repara}
\end{figure}

Despite its advantages, there are two challenges when applying gradient perturbation to cutting-edge deep models. First, one needs to compute and store individual gradients. Recent works \citep{dangel2019backpack,Opacus} have developed toolkits to compute individual gradients for a mini-batch of data through a single forward/backward pass, but storing individual gradients consumes a huge amount of memory as each individual gradient requires the same amount of memory as the model itself. Second, both theoretical and empirical utilities of gradient perturbation suffer from bad dependence on the model size \citep{bassily2014differentially, papernot2020tempered,tramer2021differentially} because the intensity of the added noise scales proportionally with the model size.

To tackle these challenges, we reparameterize each weight matrix ${\bm{W}}$ of a deep neural network with a pair of low-rank \emph{gradient carriers} $\{{\bm{L}},{\bm{R}}\}$ and a \emph{residual weight} $\tilde{{\bm{W}}}$, as illustrated in Figure~\ref{fig:repara}. With this reparametrization, the forward signal and the  backward signal propagate the same as before. We show that the gradients on ${\bm{L}}$ and ${\bm{R}}$ are naturally connected with the  gradient on ${\bm{W}}$.  Especially if the gradient carriers consist of orthonormal vectors,  we can construct a projection of the gradient of ${\bm{W}}$ from the gradients of ${\bm{L}}$ and ${\bm{R}}$ that are of low dimension.  In other words, we can compute the projection of the gradient without computing the gradient itself. This property could save a huge amount of memory in DP-SGD where a large batch of individual gradients are  computed and stored.  We note that this could be also useful in other problems involving statistics of individual gradients, e.g. computing the gradient variance \citep{zhao2015stochastic,balles2016coupling,mahsereci2017probabilistic,balles2018dissecting}, which is out of our scope.

Based on the above framework, we propose \emph{reparametrized gradient perturbation (RGP)} for differentially private learning.  Specifically, after the backward process, RGP clips and perturbs the gradients of ${\bm{L}}$ and ${\bm{R}}$, which gives a certain level of privacy guarantee. Then RGP uses the noisy gradients to construct an update for the original weight. We note that because the gradient-carrier matrices are of much smaller dimension than the original weight matrix, the total intensity of the added noises is significantly smaller, which helps us break the notorious dimensional dependence of the utility of differentially private learning.

The key of the reparameterization scheme is how well the gradient projection approximates the original gradient. We argue that the approximation is good if 1) the original gradient of ${\bm{W}}$ itself is indeed low-rank and 2) its principal subspace aligns with ${\bm{L}}$ and ${\bm{R}}$. The first condition is empirically verified by showing the gradient of each layer is of low stable rank when training deep neural networks, which has also been exploited for gradient compression in distributed optimization \citep{vogels2019powersgd}. The second condition is guaranteed if ${\bm{L}}$ and ${\bm{R}}$ consists of the principal singular vectors of the original gradient, which, however, violates the differential privacy. Instead, in RGP, we approximately compute a few of principal vectors of the historical updates that are already published and free to use because of the post-processing property of differential privacy, and use them as gradient carriers. We theoretically prove that the optimality of using the historical update substitution for linear regression and empirically verify its efficacy for deep neural networks.

With RGP, we can easily train large models with differential privacy and achieve good utility on both the vision and language modeling tasks. For example, we use RGP to train the BERT model \citep{devlin2018bert} on downstream language understanding tasks. We establish rigorous differential privacy guarantee for such large model with a modest drop in accuracy. With a privacy budget $\epsilon=8$, we achieve an average accuracy $83.9\%$ on downstream tasks, which is within $5\%$ loss compared to the non-private baseline. We also use \emph{membership inference attack} \citep{shokri2017membership,sablayrolles2019white} to evaluate the empirical privacy risks and demonstrate that the models trained with RGP are significantly more robust to membership inference attack than the non-private ones.
 Overall, our contribution can be summarized as follows.
\begin{enumerate}[itemsep=0mm]
\item We propose reparametrized gradient perturbation (RGP) that reduces the memory cost and improves the utility when applying DP on large models.

\item We give a detailed analysis on the property of RGP. We propose using the historical update to find the principal subspace and give theoretical arguments.

\item Empirically we are able to efficiently train BERT with differential privacy on downstream tasks, and achieve both good accuracy and privacy protection.
\end{enumerate}

\subsection{Notations}

We introduce some basic notations. Vectors and matrices are denoted with bold lowercase letters, e.g., ${\bm{v}}$, and bold capital letters, e.g., ${\bm{M}}$, respectively. Sets are denoted with double-struck capital letters, e.g., ${\mathbb{S}}$. We use $[n]$ to denote the set of positive numbers $\{1,...,n\}$. Some preliminaries on differential privacy are presented in Appendix \ref{app:sec:preliminary}.

\section{Two Properties of the Gradient Matrix} \label{sec:grad_property}

We show two properties of the gradients of modern deep neural networks to justify the design choices of Algorithm~\ref{alg:dp_lrk_repara}. The first property is that the gradient of each weight matrix is naturally low-rank, which motivates us to use low-rank reparameterization. The second property is that the gradient of a weight matrix along the optimization path could stay in the same subspace, which motivates us to use the historical updates to generate the gradient-carrier matrices.

\subsection{Gradient Matrix Is of Low Stable Rank}
\label{subsec:grad_is_lrk}

Recent works have used the low-rank approximation to compress the gradients and reduce the communication cost in distributed optimization \citep{yurtsever2017sketchy, wang2018atomo, karimireddy2019error, vogels2019powersgd}. These existing works set up a good motivation to exploit the low stable rank property of the gradients of weight matrices.

We further verify this low-rank property which may give a hint about how to set the reparameterization rank $r$ in practice. We empirically compute the stable rank ($\|\cdot\|_F^2/\|\cdot\|^2_{2}$) of the gradient of the weight matrices in a BERT model and a wide ResNet model. The dataset for the BERT model is SST-2 from the GLUE benchmark \citep{wang2018glue}. The dataset for the wide ResNet model is CIFAR-10 \cite{cifar}. The experimental setup can be found in Section~\ref{sec:exp}. We plot the gradient stable rank in Figure~\ref{fig:stbl_rank}.

 \begin{figure}
    \centering
  \includegraphics[width=0.8\linewidth]{imgs/grad_stable_rank.pdf}
  \caption{Gradient stable rank ($\|\cdot\|_F^2/\|\cdot\|^2_{2}$). For ResNet, we plot the gradient rank of the classification layer and the first residual block. For BERT, we plot the gradient rank of the first fully-connected block and the first attention block.}
  \label{fig:stbl_rank}
\end{figure}

As shown in Figure~\ref{fig:stbl_rank}, both the gradients of BERT and ResNet models are naturally of low stable rank over the training process. Hence, low-rank gradient-carrier matrices would have a small approximation error if  we find the right gradient subspace. In Section \ref{subsec:historical_grad}, we argue that historical update is a good choice to identify the gradient subspace.

\subsection{Historical Gradients Are Correlated}
\label{subsec:historical_grad}

Suppose that ${\bm{W}}_t$ is a weight matrix at step $t$, and $\partial {\bm{W}}_t$ is the gradient with a batch of data ${\mathbb{D}}$ with a $r$-SVD $\partial {\bm{W}}_t = {\bm{U}}_t \Sigma_t {\bm{V}}_t^T$. For another step ${t'}$ with $t'>t$ and the same data ${\mathbb{D}}$, we have ${\bm{W}}_{t'}, \partial {\bm{W}}_{t'}$ and a $r$-SVD: $\partial {\bm{W}}_{t'} = {\bm{U}}_{t'} \Sigma_{t'} {\bm{V}}_{t'}^T$. We can project $\partial {\bm{W}}_{t'}$ onto the principal subspace of $\partial {\bm{W}}_t$ or $\partial {\bm{W}}_{t'}$ and measure the  projection residual
\begin{flalign}
&\|({\bm{I}} - {\bm{U}}_t{\bm{U}}_t^T)\partial {\bm{W}}_{t'}({\bm{I}}-{\bm{V}}_t{\bm{V}}_t^T)\|_F/\|\partial {\bm{W}}_{t'}\|_F,\label{eq:proj_res}
\\
&\|({\bm{I}} - {\bm{U}}_{t'}{\bm{U}}_{t'}^T)\partial {\bm{W}}_{t'}({\bm{I}}-{\bm{V}}_{t'}{\bm{V}}_{t'}^T)\|_F/\|\partial {\bm{W}}_{t'}\|_F,\label{eq:self_proj_res}
\end{flalign}
where Eq~(\ref{eq:proj_res}) is the projection residual using historical gradient, referred to as \emph{historical projection residual}, and Eq~(\ref{eq:self_proj_res}) is the projection residual using current gradient, referred to as \emph{self projection residual}. A small difference between Eq~(\ref{eq:proj_res}) and~(\ref{eq:self_proj_res}) indicates that the principal subspace of the current gradient aligns with that of the historical gradient.

 \begin{figure}[t]
    \centering
  \includegraphics[width=0.8\linewidth]{imgs/proj_residual.pdf}
  \caption{Projection residual with reparametrization rank $8$.  We use a fixed mini-batch with $500$ samples. For ResNet, we use the input convolution layer.  For BERT, we use the second matrix of the FC layer in the first encoder block. The definition of historical/self projection residual is in Eq~(\ref{eq:proj_res}) and~(\ref{eq:self_proj_res}).   }
  \label{fig:proj_res}
\end{figure}

We empirically examine the projection residual of a BERT model and a wide ResNet model. The tasks are the same as in Section~\ref{subsec:grad_is_lrk}.   At the beginning of each epoch, we evaluate the projection residual between the current gradient and the gradient of the previous epoch. The results are plotted in Figure~\ref{fig:proj_res}. We can see that the difference between  Eq~(\ref{eq:proj_res}) and~(\ref{eq:self_proj_res}) is small for both models.

To understand why historical gradients are correlated, we next use a linear regression problem to rigorously show that the gradients over time could live in the same subspace. Suppose we have a set of observations $\{({\bm{x}}_i, {\bm{y}}_i)\}_{i=1}^n$, where ${\bm{x}}_i \in {\mathbb{R}}^d$ is the feature vector and ${\bm{y}}_i\in {\mathbb{R}}^{p}$ is the target vector for all $i \in [n]$. The least-squares problem is given by
\begin{flalign}
\argmin_{\bm{W}}  \frac{1}{n}\sum_{i=1}^n \|{\bm{y}}_i - {\bm{W}} {\bm{x}}_i\|^2. \label{eq:least-squares}
\end{flalign}

\begin{restatable}{proposition}{gradalign}\label{prop:grad-align}
For the least squares problem (\ref{eq:least-squares}), if the model is updated by gradient descent with step size $\eta$
\begin{flalign}
{\bm{W}}_{t+1} \leftarrow {\bm{W}}_t - \eta \cdot \partial{\bm{W}}_t, \label{eq:gd}
\end{flalign}
then the gradients $\{\partial {\bm{W}}_t\}_{t\ge 1}$ share the same range and null space. That is to say, if $\partial {\bm{W}}_1$ is rank $r$ and has $r$-SVD $\partial {\bm{W}}_1 = {\bm{U}}_1 \Sigma_1 {\bm{V}}_1^T$, then for all $t\ge 1$, we have
\begin{flalign}
({\bm{I}} - {\bm{U}}_1{\bm{U}}_1^T) \partial {\bm{W}}_t = 0,\;  \partial {\bm{W}}_t ({\bm{I}} - {\bm{V}}_1{\bm{V}}_1^T)= 0.
\end{flalign}
\end{restatable}

\begin{proof}
The proof is relegated to Appendix \ref{apd:subsec:proof_sec4}.
\end{proof}

Hence we can use the historical updates ${\bm{W}}_t-{\bm{W}}_0$ to identify  gradient row/column subspaces as in Algorithm \ref{alg:dp_lrk_repara}.

That indicates that for the weight matrix ${\bm{W}}\in {\mathbb{R}}^{p \times d}$, if the gradient turns out to be low-rank $r$ due to the data $\{{\bm{x}}_i, {\bm{y}}_i\}$, we can possibly first identify the intrinsic subspace which is of $r(p+d)$ dimension instead of the original $p\cdot d$ number of parameters. Then we can work within this subspace for differentially private empirical risk minimization. This can both reduce the effect of noise and save the memory cost of gradient perturbation due to the small intrinsic dimension.
We note that identifying the low-rank subspace can be done approximately as in the algorithm, or by using some auxiliary public data as in \citet{zhou2021bypassing, yu2021do}.

\begin{remark}\label{rem:lst-sqr}
Suppose that the least-squares objective $L({\bm{W}}):=\frac{1}{n}\sum_{i=1}^n \|{\bm{y}}_i - {\bm{W}} {\bm{x}}_i\|^2$ is $\beta$-smooth and the gradient subspace is rank $r$ and can be exactly identified. Let the optimizer of RGP be gradient descent and $\sigma$ be set as in Proposition~\ref{prop:privacy}. If $\eta=\frac{1}{\beta}$, $T=\frac{n\beta\epsilon}{\sqrt{p}}$, and $\bar{{\bm{W}}}=\frac{1}{T}\sum_{t=1}^{T}{\bm{W}}_t$, then
\[\mathbb{E}[L(\bar{{\bm{W}}})]-L({\bm{W}}_*)\leq \mathcal{O}\left(\frac{\sqrt{(p+d)r\log(1/\delta)}}{n\epsilon}\right),\]
where ${\bm{W}}_*$ is the optimal point,  ${\bm{W}}_{t}$ is the output of Algorithm~\ref{alg:dp_lrk_repara} at step $t$.
\end{remark}
The proof of Remark \ref{rem:lst-sqr} can be adapted from \cite{yu2020gradient}.  Although  the exact low-rank property of the gradient  cannot be  rigorously proved for deep neural network because of the co-adaption across layers, we have empirically verified that the gradient matrices are still of low stable rank and stay in roughly the  same subspace over iterations (see Figure \ref{fig:stbl_rank} \& \ref{fig:proj_res}). Our algorithm exploits this fact to reparameterize weight matrices, which achieves better utility and reduces the memory cost compared with DP-SGD.

\section{Related Work}

\begin{table}

\renewcommand{\arraystretch}{1.3}
\centering
    \caption{Success rates of membership inference attack  against fine-tuned BERT models (in \%). The closer to 50, the better.} \label{tbl:mi_bert}
    \begin{adjustbox}{max width=0.45\textwidth}
        \begin{tabular}{l|l|l|l|l|l|l}
        \hline
        \hline
            Method         & MNLI &  QQP  & QNLI  &  SST-2   & SVHN & CIFAR10  \\
            \hline
              Full (N.P.) &        60.3      &      56.1    &   55.8   &  57.7  & 56.4 & 58.1  \\\cline{1-7}
                RGP (N.P.)  &      52.3         &   51.5    &    51.8   &   52.6 & 52.8 & 53.3 \\\cline{1-7}
                RGP ($\epsilon=8$)  &      49.9         &   50.0      &   50.4   &   50.1 & 50.1 & 50.3  \\
           \hline \hline
        \end{tabular}
    \end{adjustbox}
\end{table}

Differentially private learning has a poor dimensional dependency, i.e., the utility degrades dramatically when the model dimension gets large.  In the high-dimensional setting, related works usually assume the sparse structure  \citep{thakurta2013differentially, talwar2015nearly, wang2019sparse, wang2019differentially, cai2019cost} or specific problem structure \cite{chen2020locally,zheng2020locally}. However, these assumptions or specific structures do not hold for the gradient of deep neural networks. Here we emphasize the difference from our low-rank assumption. For the sparsity assumption, the bases are canonical and not private while for the low-rank assumption, it is ``sparse'' under certain bases but the bases are unknown and private. Hence the previous algorithms for sparsity cannot apply here.

Very recently, several works \citep{zhou2020bypassing,kairouz2020dimension, yu2021do} exploit the  redundancy of gradients of samples and suggest projecting the gradients into a low dimensional subspace that is identified by some public data points or historical gradients, in order to reduce the  noise effect when training large models. However, they all require storing and clipping whole individual gradients and hence are hard to train extremely large models.  Our work is orthogonal with theirs, i.e., we exploit the low-rank property of the gradient of each weight matrix, which truly breaks the  barrier of applying DP in large models.

Another recent approach of training non-convex models with differential privacy is based on the  knowledge transfer of machine learning models \emph{Private Aggregation of Teacher Ensembles (PATE)} \citep{papernot2016semi, papernot2018scalable, jordon2019pate}. They first train independent teacher models on disjoint shards of private data and then tune a student model with privacy by distilling noisy predictions of teacher models on some public samples, whose performance suffers from the data splitting \cite{yu2021do}. It is not clear how to apply PATE to train large language models like BERT. In contrast, our algorithms do not require public data and can be used in different settings with little change.

The phenomenon that the gradients of deep models live on a very low dimensional manifold  has been widely observed \citep{gur2018gradient, vogels2019powersgd, gooneratne2020low, li2020hessian, martin2018implicit, li2018algorithmic}. People have also used this fact to compress the gradient with low-rank approximation  in the distributed optimization scenario \citep{yurtsever2017sketchy, wang2018atomo, karimireddy2019error, vogels2019powersgd}.

\section{Conclusion}
In this paper, we present the reparametrized gradient perturbation (RGP) for applying DP on large models. The key design of RGP exploits two  properties of gradients in  deep neural network, which are 1) the gradient of each weight matrix is of low stable rank, 2) the principal components of historical gradients align well with that of the current gradient.  We also justify the designs with both theoretical and empirical evidence. Thanks to RGP, we are able to train BERT on several downstream tasks with DP guarantee and achieve small accuracy loss.

\vspace{-1mm}
\section{A Reparametrization Scheme}\label{sec:lrk}

In this section, we introduce a reparametrization scheme for the neural network weight matrices so that computing and storing individual gradients are affordable for large models. Specifically,  during each forward/backward process, for a layer with weight matrix ${\bm{W}}\in {\mathbb{R}}^{p\times d}$,  we reparametrize it as follows (see  Figure~\ref{fig:repara} for an illustration),
\begin{flalign}
{\bm{W}} \rightarrow {\bm{L}} {\bm{R}} + \tilde{{\bm{W}}}.{stop\_gradient()}, \label{eq:repara}
\end{flalign}
where ${\bm{L}}\in{\mathbb{R}}^{p\times r}, {\bm{R}}\in{\mathbb{R}}^{r\times d}$ are two low-rank gradient carriers  with $r\ll p \text{ or } d$, $\tilde{{\bm{W}}} = {\bm{W}}-{\bm{L}}{\bm{R}}$ represents the residual weight and $.{stop\_gradient()}$ means that we do not collect the gradient on $\tilde{{\bm{W}}}$.
Hence, such reparametrization does not change the forward signal and the backward signal, but only changes the gradient computation. Now we obtain the gradients on ${\bm{L}}$ and ${\bm{R}}$. We then unveil the connection between the gradient on ${\bm{W}}$ and the gradients on ${\bm{L}}$ and ${\bm{R}}$.

\begin{restatable}{theorem}{gradlr}\label{thm:grad_lr}
For a layer with weight matrix ${\bm{W}}$, suppose that $\partial{\bm{W}}$ is the gradient computed by back-propagation with a mini-batch data ${\mathbb{D}}$. Given two matrices ${\bm{L}}, {\bm{R}}$, we reparametrize ${\bm{W}}$ as in Eq~(\ref{eq:repara}) and compute the gradients $\partial{\bm{L}}$ and  $\partial{\bm{R}}$  by running the forward and backward process with the same mini-batch  ${\mathbb{D}}$, then
\begin{flalign}
\partial {\bm{L}} = (\partial{\bm{W}}){\bm{R}}^{T},\;\; \partial {\bm{R}} = {\bm{L}}^{T}(\partial{\bm{W}}).
\end{flalign}

\end{restatable}

Based on the above understanding, we can construct an update for ${\bm{W}}$ by using $\partial{\bm{L}}$ and $\partial{\bm{R}}$.

\begin{restatable}{corollary}{corogradlr}\label{corollary:grad_lr}
If the columns of ${\bm{L}}$ and the rows of ${\bm{R}}$ are orthonormal, respectively, and we use
\begin{flalign}
\label{eq:grad_lrk}
(\partial {\bm{L}}) {\bm{R}} + {\bm{L}}(\partial{\bm{R}}) - {\bm{L}}\mL^{T}(\partial{\bm{L}}){\bm{R}},
\end{flalign}
as the update for ${\bm{W}}$, then the update is equivalent to projecting $\partial{\bm{W}}$ into the subspace of matrices whose row/column spaces are spanned by ${\bm{L}}$ and ${\bm{R}}$.
\end{restatable}

\begin{proof}
The proofs of Theorem~\ref{thm:grad_lr} and Corollary~\ref{corollary:grad_lr} are relegated to Appendix~\ref{apd:subsec:proof_sec2}.
\end{proof}

We note that if ${\bm{L}}$ and ${\bm{R}}$ consist of orthonormal bases, Corollary~\ref{corollary:grad_lr} states that we can obtain the projection of  $\partial{\bm{W}}$ without explicitly computing and storing $\partial{\bm{W}}$! The size of gradient on ${\bm{L}}$ or ${\bm{R}}$ is much smaller than the size of $\partial{\bm{W}}$ if the gradient carriers are chosen to be low-rank. Therefore, this reparametrization provides a convenient way to compute and store projected gradients of a large matrix. This is extremely beneficial for the scenarios where individual gradients $\{\partial_i {\bm{W}}\}_{i=1}^{m}$ are required, e.g., approximating the variance of gradients and controlling the gradient sensitivity.

It is natural to ask how to choose ${\bm{L}}$ and ${\bm{R}}$ so that the update in Corollary~\ref{corollary:grad_lr} contains the most information of $\partial{\bm{W}}$. Ideally, we can first compute the aggregated gradient $\partial{\bm{W}}$ and run \emph{singular value decomposition} (SVD) $\partial{\bm{W}}={\bm{U}}{\bm{\Sigma}}{\bm{V}}^{T}$. Then we can choose the top few columns of ${\bm{U}}$ and ${\bm{V}}$ to serve as the gradient carriers. In this case, the update in Corollary~\ref{corollary:grad_lr} is equivalent to approximating $\partial{\bm{W}}$ with its top-$r$ principal components.

However, in the context of differential privacy, we can not directly decompose $\partial {\bm{W}}$  as it is private. In the sequel, we give a practical reparametrization scheme for differentially private learning, where we use the historical update to find ${\bm{L}}$ and ${\bm{R}}$ and argue the optimality under certain conditions.

One may wonder why not just replace ${\bm{W}}$ with ${\bm{L}}$ and ${\bm{R}}$ instead of doing the reparametrization. We note that the forward and the backward process remain the same as before if doing the reparametrization, and the only change is the gradient computation of ${\bm{W}}$. In contrast, if using ${\bm{L}}$ and ${\bm{R}}$ to replace the weight ${\bm{W}}$, this would not only reduce the expressive power but also hurt the optimization as the width varies dramatically across layers and the forward/backward signals cannot propagate well by common initialization strategies \cite{glorot2010understanding, he2016deep}.

\subsection{Reparametrization for Convolutional Layers}
\label{sec:lrk_conv}
In the above, we have described how to reparametrize a weight matrix, which covers the usual fully-connected layer and the attention layer in language models. In this subsection, we show the reparametrization of convolutional layers.  Let ${\bm{x}}\in\mathbb{R}^{d\times w' \times h'}$ be the input feature maps of one sample and ${\bm{h}}\in\mathbb{R}^{p\times w \times h}$ be the output feature maps. We describe how to compute the elements at one spatial position ${\bm{h}}_{:,i,j}\in\mathbb{R}^{p}$ where $i\in [0,w]$ and $j\in [0,h]$.

Let ${\bm{W}}\in \mathbb{R}^{p\times d\times k\times k}$ be the convolution kernels and ${\bm{x}}^{(i,j)}\in\mathbb{R}^{d\times k\times k}$ be the features that we need to compute ${\bm{h}}_{:,i,j}$. The output feature ${\bm{h}}_{:,i,j}$ can be computed as
${\bm{h}}_{:,i,j}=\bar{\bm{W}} {\bm{x}}^{(i,j)}$,
where $\bar{\bm{W}}\in\mathbb{R}^{p\times dk^{2}}$ is obtained by flattening the channel and kernel dimensions. Hence, we can use the same way as in Eq~(\ref{eq:repara}) to  reparametrize $\bar{\bm{W}}$:
\begin{flalign}
{\bm{h}}_{:,i,j} = {\bm{L}}{\bm{R}}{\bm{x}}^{(i,j)} + (\bar{\bm{W}}-{\bm{L}}{\bm{R}}){\bm{x}}^{(i,j)}.
\end{flalign}

Specifically, the operation of ${\bm{R}}$ and ${\bm{L}}$ are implemented by  two consequent convolutional layers with kernel sizes $r\times d\times k\times k$ \ and $p\times r\times 1\times 1$, respectively, where $r$ is the reparametrization rank. The residual weight is implemented by a convolutional layer of the original kernel size.

\section{Private Deep Learning with Reparametrized Gradient Perturbation}
\label{sec:dp_learning_lrk}

The above reparametrization strategy can significantly reduce the gradient dimension, which could help us circumvent the difficulties of applying differential privacy on large machine learning models. In this section, we propose a procedure ``reparametrized gradient perturbation (RGP)'' to train large neural network models  with differential privacy. Specifically, Section \ref{subsec:dp_learning_lrk_algo} introduces the whole procedure of RGP, Section \ref{subsec:privacy_rgp} gives the privacy guarantee of RGP, and Section \ref{subsec:complexity} presents the complexity analysis.

\begin{algorithm}[tb]
   \caption{Reparametrized Gradient Perturbation (RGP)}
   \label{alg:dp_lrk_repara}
\begin{algorithmic}[1]
  \STATE {\bfseries Input:} NN with weight matrices $\{{\bm{W}}^{(l)}\}_{l=1}^{H}$,  steps $T$,  probability $q$, variance $\sigma^2$, clipping threshold $C$, warm-up steps $T_{\text{warm-up}}$, Algorithm \ref{alg:decompose_pi} input $\{r, K\}$.

  \STATE Randomly initialize the weights and obtain $\{{\bm{W}}^{(l)}_{0}\}_{l=1}^H$;
  \FOR{$t=1$ {\bfseries to} $T$}
    \medskip
        \STATE Sample a minibatch $\{{\bm{x}}_{i}\}_{i\in S_t}$ with probability $q$;
    \medskip

      \STATE For all $l\in [H]$, compute historical updates
      $$\Delta_t^{(l)} \leftarrow {\bm{W}}^{(l)}_t - {\bm{W}}^{(l)}_0 \cdot 1_{\{t>T_{\text{warm-up}}\}};$$
      and run Alg.~\ref{alg:decompose_pi} with $\{\Delta_t^{(l)}, r, K\}$ to get ${\bm{L}}_t^{(l)},{\bm{R}}_t^{(l)}$;

    \medskip
    \STATE \textsl{//Forward/backward process with reparametrization.}

    \STATE Run reparametrized forward process with Eq~(\ref{eq:repara});
    \STATE Run backward process and compute individual gradients $\{\partial_i{\bm{L}}_t^{(l)},\partial_i{\bm{R}}_t^{(l)}\}_{l\in[H], i\in S_t}$;
    \medskip
    \STATE \textsl{//Bound gradient sensitivity and add noise.}

    \STATE Clip individual gradients with $L_{2}$ norm threshold $C$;

        \FOR{$l=1$ {\bfseries to} $H$}

    \STATE Sum individual gradients and get  $\{\partial{\bm{L}}_t^{(l)},\partial{\bm{R}}_t^{(l)}\}$;
    \STATE Perturbation with Gaussian noise ${\bm{z}}_{L,t}^{(l)},{\bm{z}}_{R,t}^{(l)}$ whose elements are independently from $\mathcal{N}(0,\sigma^{2}C^{2})$:
       $$\tilde{\partial}{\bm{L}}_t^{(l)}  \leftarrow \partial{\bm{L}}_t^{(l)} + {\bm{z}}_{L,t}^{{(l)}}, \quad \tilde{\partial}{\bm{R}}_t^{(l)} \leftarrow \partial{\bm{R}}_t^{(l)}+{\bm{z}}_{R,t}^{(l)};$$
    \STATE Use $\tilde{\partial}{\bm{L}}_t^{(l)}$, $\tilde{\partial}{\bm{R}}_t^{(l)}$, and Eq~(\ref{eq:grad_lrk}) to  construct $\tilde{\partial} {\bm{W}}_t^{(l)}$;
    \STATE Use off-the-shelf optimizer to get ${\bm{W}}_{t+1}^{(l)}$;
    \ENDFOR

  \ENDFOR

\end{algorithmic}
\end{algorithm}

\subsection{Reparametrized Gradient Perturbation Algorithm}
\label{subsec:dp_learning_lrk_algo}

The pseudocode of RGP is presented in Algorithm~\ref{alg:dp_lrk_repara}. The RGP proceeds for all the layers and we ignore the layer index for simplicity in the following discussion. At each update, for a layer with weight matrix ${\bm{W}}$, RGP consists of four steps: 1) generate the gradient-carrier matrices ${\bm{L}}$ and ${\bm{R}}$, 2) run the reparametrized forward/backward process and obtain the individual gradients $\{\partial_i {\bm{L}}\}_{i=1}^{m}$ and $\{\partial_i {\bm{R}}\}_{i=1}^{m}$, 3) clip and perturb the gradients, 4) reconstruct an approximated gradient on the original weight matrix.

In the RGP procedure, \textbf{step 1)}, which is also the core challenge, is to choose ``good" gradient-carrier matrices so that the reconstructed gradient can approximate the original gradient as well as possible. First, this requires for a given rank $r$, the generated gradient-carrier matrices should align with the principal components of the original gradient well. Moreover, to reconstruct the gradient in step 4), it requires the gradient carriers have orthonormal columns/rows.

For the first requirement, we use  historical updates to find the gradient carriers. The historical update is not sensitive because of the post-processing property of differential privacy.  In Section~\ref{subsec:historical_grad}, we give both empirical and theoretical arguments to  demonstrate that the principal subspace of the current gradient aligns with that of the historical update.   In our implementation, we use a warm-up phase in which the decomposition is directly done on the weight. We approximate the principal components via the power method (Algorithm~\ref{alg:decompose_pi}) instead of  the time-consuming full SVD. For the second requirement, we apply  the Gram-Schmidt process to orthonormalize ${\bm{L}}$ and ${\bm{R}}$.

\begin{algorithm}[tb]
   \caption{Decomposition via Power Method.}
   \label{alg:decompose_pi}
\begin{algorithmic}
  \STATE {\bfseries Input:} Historical update $\Delta$, reparametrization rank $r$, number of iterations $K$.
  \STATE {\bfseries Output:} Gradient carriers ${\bm{L}}\in\mathbb{R}^{p\times r}$,  ${\bm{R}}\in\mathbb{R}^{r\times d}$.

  \medskip

    \STATE Initialize  ${\bm{R}}$ from standard Gaussian distribution.

    \FOR{$k=1$ {\bfseries to} $K$}
    \STATE ${\bm{L}} \leftarrow \Delta {\bm{R}}^{T}$
    \STATE Orthonormalize the columns of ${\bm{L}}$.
    \STATE ${\bm{R}}={\bm{L}}^{T}\Delta$
    \ENDFOR
    \STATE Orthonormalize the rows of ${\bm{R}}$.

    \STATE Return ${\bm{L}}$, ${\bm{R}}$
\end{algorithmic}
\end{algorithm}

\textbf{Step 2)} of RGP is the reparametrization and a round of forward/backward propagations, as presented in Section \ref{sec:lrk}.

\textbf{Step 3)} is for differential privacy guarantee. The individual gradients $\{\partial_i {\bm{L}}, \partial_i {\bm{R}}\}_{i=1}^{m}$ are first clipped by a pre-defined threshold so that the sensitivity is bounded. Then, Gaussian noise is added to the aggregated gradient to establish a differential privacy bound. The energy of added noise is proportional to the dimension, i.e., the rank $r$ of the carrier matrices. Hence, in order to make the noise energy small, it encourages us to use smaller rank $r$. However, smaller rank would increase the approximation error in the \textbf{step 1)}. In practice, we trade off these two factors to choose a proper $r$.

In \textbf{step 4)}, we use the noisy aggregated gradients of gradient-carrier matrices to reconstruct the gradients of original weights, as depicted in Corollary~\ref{corollary:grad_lr}. The reconstructed gradients can then be used by any off-the-shelf optimizer.

\subsection{Privacy Analysis of RGP}
\label{subsec:privacy_rgp}

The privacy bound of Algorithm~\ref{alg:dp_lrk_repara} is given by  Proposition~\ref{prop:privacy}. The derivation of Proposition~\ref{prop:privacy} is based on the \emph{moments accountant} that is proposed in \citet{abadi2016deep}.  Moments accountant has tighter composition bound than the strong composition theorem in \citet{algofound}.  Moments accountant first tracks the privacy budget spent at each update. Then, it composes the spent budget of all updates and cast the final privacy cost into the classic $(\epsilon,\delta)$-differential privacy.

\begin{restatable}[\citet{abadi2016deep}]{proposition}{privacy}\label{prop:privacy}
There exist constants $c_1$ and $c_2$ so that given running steps $T$, for any $\epsilon<c_{1}q^{2}T$, Algorithm~\ref{alg:dp_lrk_repara} is $\left(\epsilon,\delta\right)$-differentially private for any $\delta>0$ if we choose \[\sigma\geq c_2\frac{q\sqrt{Tlog\left(1/\delta\right)}}{\epsilon}.\]
\end{restatable}

\begin{proof}
The proof outline is relegated to Appendix~\ref{apd:subsec:proof_sec3}.
\end{proof}

  The value of $\sigma$ in Proposition~\ref{prop:privacy} is based on an asymptotic bound on the moments of the privacy loss random variable. In practice, one can use the numerical tools \citep{wang2019subsampled,mironov2019renyi} to compute a tighter bound. So far we have depicted the overall picture of RGP. We next analyze  the computational and memory costs of RGP and compare them with that of DP-SGD.

\subsection{Complexity Analysis of RGP}
\label{subsec:complexity}
For the simplicity of notations, we only give the costs of one fully connected layer at one update (including forward and backward) and assume that the weight matrix is square. The shape of weight matrix, size of minibatch, number of power iterations, and rank of reparametrization are denoted by $(d\times d)$, $m$, $K$, and $r$, respectively.

The computational overhead of RGP consists of three parts. The first part is induced by matrix multiplication of power iteration, whose complexity is $\mathcal{O}(Krd^{2})$. The second part is induced by the Gram–Schmidt process, whose complexity is $\mathcal{O}(Kr^{2}d)$. The third part of overhead is the computational cost induced by gradient carriers during the forward/backward process, which is on the order of $\mathcal{O}(mrd)$.

RGP uses much less memory than DP-SGD in the practice.   Although RGP needs some extra memory to store the activation produced by the gradient carriers, it has a significant advantage over DP-SGD on the memory cost of storing individual gradients, which is one of the main challenges of learning with differential privacy. For RGP, the memory cost of individual gradients only scales linearly with model width $d$ in contrast with $d^2$ for DP-SGD. We summarize the computational cost of one update and the memory cost of storing individual gradients in  Table~\ref{tbl:complexity}.

\begin{table}
    \caption{Computation and memory costs of RGP (Algorithm~\ref{alg:dp_lrk_repara}) and DP-SGD \citep{abadi2016deep}, where $m$ is the size of mini-batch, $d$ is the model width, $r$ is the reparametrization rank, and $K$ is the number of power iterations.}
\label{tbl:complexity}
\centering
\small
\renewcommand{\arraystretch}{1.85}
\begin{tabular}{ P{2.45cm}|P{1.15cm}|P{3.4cm} }
 \hline \hline
   \backslashbox{Cost}{Method}            & DP-SGD 	& RGP  		 \\
 \hline
Computational cost     &   $\mathcal{O}(md^{2})$     & $\mathcal{O}(md^{2}+Krd^2+Kr^{2}d)$			\\\hline
Memory cost    &  $\mathcal{O}(md^{2})$    & $\mathcal{O}(mrd)$ 			\\

 \hline
 \hline
\end{tabular}
\end{table}

The low-rank nature of gradient permits us to choose a small $r$ without destroying utility (see Section~\ref{subsec:grad_is_lrk}). In practice, we typically choose the rank $r$ smaller than $10$.  For the number of power iterations in Algorithm~\ref{alg:decompose_pi}, we find that setting $K=1$ is sufficient to get good performance. Hence, in practice, we always choose small $r$ and $K$ for efficiency while not hurting the performance.

\section{Preliminary on Differential Privacy} \label{app:sec:preliminary}

Differential privacy (DP) \cite{dwork2006calibrating,dwork2014algorithmic} is widely recognized as a gold standard of privacy protection due to its mathematical rigor. It controls the maximum influence that any individual sample can produce. The definition of $(\epsilon,\delta)$-DP is given in Definition~\ref{def:dp}.

\begin{definition}[$(\epsilon,\delta)$-DP]
\label{def:dp}
A randomized mechanism $\mathcal{M}$  guarantees $(\epsilon,\delta)$-differential privacy if for any two neighboring input datasets ${\mathbb{D}}\sim {\mathbb{D}}^{'}$ and for any subset of outputs ${\mathbb{S}}$ it holds that $\text{Pr}[\mathcal{M}({\mathbb{D}})\in {\mathbb{S}}]\leq e^{\epsilon}\text{Pr}[\mathcal{M}({\mathbb{D}}^{'})\in {\mathbb{S}}]+\delta$.
\end{definition}

Two datasets are said to be neighboring datasets if they only differ in a single sample. When being applied to learning problems, DP requires the  learned models on neighboring datasets have approximately indistinguishable distributions.

\section{Missing Proofs} \label{app:sec:proof}

\subsection{Missing Proofs in Section \ref{sec:lrk}}
\label{apd:subsec:proof_sec2}
\gradlr*
\begin{proof}
The proof is based on the chain rule of back-propagation. Since the reparametrization does not  change the forward and backward signals, we assume the layer inputs are ${\mathbb{D}}=\{{\bm{x}}_{i}\}_{i=1}^{m}$, the corresponding outputs are $\{{\bm{h}}_i\}_{i=1}^{m}$ with ${\bm{h}}_i = {\bm{W}} {\bm{x}}_i$ and the backward signals on the layer output are $\{\partial {\bm{h}}_i\}_{i=1}^{m}$. By back-propagation, we have
\begin{flalign*}
&\partial {\bm{W}} = \sum_{{\bm{x}}_{i}\in {\mathbb{D}}} (\partial {\bm{h}}_i) {\bm{x}}_i^T, \\
&\partial {\bm{L}} =\sum_{{\bm{x}}_{i}\in {\mathbb{D}}}\partial {\bm{h}}_i ({\bm{R}} {\bm{x}}_i)^T,\;\; \partial {\bm{R}} =\sum_{{\bm{x}}_{i}\in {\mathbb{D}}} ({\bm{L}}^{T}\partial {\bm{h}}_i) {\bm{x}}_i^T.
\end{flalign*}
Proof is completed by the multiplication associativity.
\end{proof}

\corogradlr*

\begin{proof}
If the columns of ${\bm{L}}$ and the rows of ${\bm{R}}$ are orthonormal, the projection of $\partial {\bm{W}}$ onto  ${\bm{L}}$ and ${\bm{R}}$ is defined as,
\begin{flalign}
{\bm{L}}\mL^T (\partial {\bm{W}}) + (\partial {\bm{W}}) {\bm{R}}^T{\bm{R}} - {\bm{L}}\mL^T (\partial {\bm{W}}){\bm{R}}^T{\bm{R}}.
\end{flalign}
Substituting the above formula with $\partial {\bm{L}} = (\partial {\bm{W}}) {\bm{R}}^T$ and $\partial {\bm{R}} = {\bm{L}}^T (\partial {\bm{W}})$ in Theorem \ref{thm:grad_lr}, completes the proof.
\end{proof}

\subsection{Missing Proofs in Section \ref{sec:dp_learning_lrk}}
\label{apd:subsec:proof_sec3}

\privacy*

\begin{proof}
Although RGP releases projected gradient instead of releasing the whole gradient as in \citet{abadi2016deep}, moments accountant is still applicable because it applies to vectorized  function output.

Moments accountant tracks a bound on the moments of the privacy loss random variable, which is built on the ratio of the probability density functions of the output distributions of two neighboring datasets.  \citet{abadi2016deep} show the log moments of the privacy loss random variable composes linearly. Therefore one can compute the overall privacy cost by adding the log moments at every update. When the training is done, moments accountant  casts the accumulated log moments into $(\epsilon,\delta)$-DP via tail bound. Detailed proof can be found in Appendix B of \citet{abadi2016deep}.

\end{proof}

\subsection{Missing Proofs in Section \ref{sec:grad_property}}
\label{apd:subsec:proof_sec4}
\gradalign*
\begin{proof}
We can compute the gradient at step $t$
\begin{flalign*}
\partial {\bm{W}}_t &= \frac{1}{n}\sum_{i=1}^n ({\bm{W}}_t {\bm{x}}_i - {\bm{y}}_i) {\bm{x}}_i^T.
\end{flalign*}
Given the gradient descent update \eqref{eq:gd}, we can compute the gradient at ${\bm{W}}_{t+1}$ as follows
\begin{flalign*}
\partial {\bm{W}}_{t+1}
&=  \frac{1}{n}\sum_{i=1}^n (({\bm{W}}_t - \eta \cdot\partial{\bm{W}}_t){\bm{x}}_i - {\bm{y}}_i) {\bm{x}}_i^T\\
&= \frac{1}{n}\sum_{i=1}^n ({\bm{W}}_t{\bm{x}}_i - {\bm{y}}_i) {\bm{x}}_i^T  - \eta\cdot \partial{\bm{W}}_t \sum_{i=1}^n {\bm{x}}_i{\bm{x}}_i^T \\
& = \partial {\bm{W}}_t \left ({\bm{I}} - \eta \sum_{i=1}^n {\bm{x}}_i{\bm{x}}_i^T\right ).
\end{flalign*}
Hence we have $\partial {\bm{W}}_t = \partial {\bm{W}}_0 \left ({\bm{I}} - \eta \sum_{i=1}^n {\bm{x}}_i{\bm{x}}_i^T\right )^t$. The $\partial {\bm{W}}_t$ lives in the same subspace for all $t\ge 1$ as they have the same row/column spaces.
\end{proof}

\section{Additional Experiments}\label{app:sec:add-exp}
 \begin{figure*} [t]
    \centering
  \includegraphics[width=0.9\linewidth]{imgs/show_influ_rank.pdf}
  \caption{Prediction accuracy of BERT on four downstream tasks (in \%) with difference choices of reparametrization rank. We plot the average score of two test datasets for MNLI. }
  \label{fig:fig_bert_rank}
\end{figure*}
We present some ablation studies in this section to verify the effect of residual weight and reparametrization rank.  In Section~\ref{subsec:apd_rank}, we try RGP with difference rank choices. In Section~\ref{subsec:apd_residual}, we give a variant of RGP that simply discards the residual weight.

\subsection{On the Influence of Different Rank Choices} \label{subsec:apd_rank}

We present the results (see Figure \ref{fig:fig_bert_rank}) with different choices of reparametrization rank.  We consider four algorithms. The first one is fine-tuning the full model that serves as the baseline. The second one is RGP (N.P.) that trains the model with reparametrization but without gradient clipping or adding noise. The third one is RGP (Algorithm~\ref{alg:dp_lrk_repara}) and the  last one is RGP-random, which uses random orthogonal vectors as gradient-carrier matrices. The privacy parameter $\epsilon$ is $8$ and other settings are the same as those in Section~\ref{sec:exp}. The results are plotted in Figure~\ref{fig:fig_bert_rank}.  When the models are trained without noise, increasing the reparametrization rank makes the performance of RGP (N.P.) approach the performance of baseline. When the models are trained with privacy guarantee, increasing the rank sometimes decreases the performance because a larger rank induces more trainable parameters and hence higher noise dimension.

\subsection{On the Importance of Residual Weight}\label{subsec:apd_residual}

Recall that our reparametrization scheme reparametrizes the weight matrix as follows:

\begin{flalign}
{\bm{W}} \rightarrow {\bm{L}} {\bm{R}} + \tilde{{\bm{W}}}.{stop\_gradient()}. \label{eq:apd_repara}
\end{flalign}

We have shown that the residual weight $\tilde{{\bm{W}}}$ keeps the forward/backward signals unchanged and makes the gradients of ${\bm{L}}$ and ${\bm{R}}$ naturally connected with the original gradient. To empirically examine the effect of $\tilde{{\bm{W}}}$, we test the following scheme:

\begin{flalign}
{\bm{W}} \rightarrow {\bm{L}} {\bm{R}}. \label{eq:apd_repara_nores}
\end{flalign}

 \begin{figure}
    \centering
  \includegraphics[width=0.9\linewidth]{imgs/show_influ_residual.pdf}
  \caption{Prediction accuracy of BERT on two downstream tasks. All methods are trained without privacy guarantee.   }
  \label{fig:fig_bert_residual}
\end{figure}

We still use the historical update to generate ${\bm{L}}$ and ${\bm{R}}$. Other settings are the same as those in Section~\ref{sec:exp}. The results on two downstream tasks of the BERT model are presented in Figure~\ref{fig:fig_bert_residual}. Without residual weight, the model learns almost nothing from the re-constructed update and the final accuracy is close to the accuracy at initialization.

\end{appendix}