algorithmic

# Momentum-Based Policy Gradient with Second-Order Information

Saber Salehkaleybar<EMAIL_ADDRESS>
School of Computer and Communication Sciences
EPFL Sadegh Khorasani<EMAIL_ADDRESS>
School of Computer and Communication Sciences
EPFL Negar Kiyavash<EMAIL_ADDRESS>
College of Management of Technology
EPFL Niao He<EMAIL_ADDRESS>
Department of Computer Science
ETH Zurich Patrick Thiran<EMAIL_ADDRESS>
School of Computer and Communication Sciences
EPFL

###### Abstract

Variance-reduced gradient estimators for policy gradient methods have been one
of the main focus of research in the reinforcement learning in recent years as
they allow acceleration of the estimation process. We propose a variance-
reduced policy-gradient method, called SHARP, which incorporates second-order
information into stochastic gradient descent (SGD) using momentum with a time-
varying learning rate. SHARP algorithm is parameter-free, achieving
$\epsilon$-approximate first-order stationary point with $O(\epsilon^{-3})$
number of trajectories, while using a batch size of $O(1)$ at each iteration.
Unlike most previous work, our proposed algorithm does not require importance
sampling which can compromise the advantage of variance reduction process.
Moreover, the variance of estimation error decays with the fast rate of
$O(1/t^{2/3})$ where $t$ is the number of iterations. Our extensive
experimental evaluations show the effectiveness of the proposed algorithm on
various control tasks and its advantage over the state of the art in practice.

##  1 Introduction

Reinforcement Learning (RL) has achieved remarkable success in solving various
complex tasks in games (Silver et al., 2017), autonomous driving (Shalev-
Shwartz et al., 2016), and robot manipulation (Deisenroth et al., 2013), among
other fields. In RL setting, an agent tries to learn the best actions by
interacting with the environment and evaluating its performance based on
reward signals. More specifically, in Markov Decision Processes (MDPs), the
mathematical formalism for RL, after taking an action, the state changes
according to a transition probability model and a reward signal is received
based on the action taken and the current state. The main goal of the learner
is to find a policy that maps the state space to the action space, maximizing
the expected cumulative rewards as the objective function.

Policy gradient methods (Sutton et al., 2000) are often used for obtaining
good policies in MDPs, especially for high-dimensional continuous action
space. In policy gradient methods, the policy is parameterized by an unknown
parameter $\theta$ and it is directly optimized using the stochastic first-
order gradient of cumulative rewards as it is infeasible to compute the
gradient exactly. REINFORCE (Williams, 1992), PGT (Sutton et al., 2000), and
GPOMDP (Baxter & Bartlett, 2001) are some classical methods that update the
policy by applying a stochastic gradient ascent step. These methods generally
require a large number of trajectories due to the large variance of gradient
estimates, stemming from randomness of transitions over trajectories.

In the RL literature, several methods have been proposed to reduce the
variance in policy gradient methods. For instance, Sutton et al. (2000)
proposed to consider a baseline in order to reduce variance of gradient
estimation. Konda & Tsitsiklis (2000) presented an actor-critic algorithm that
estimates the value function and uses it to mitigate the effect of large
variance. Schulman et al. (2015b) proposed GAE to control both bias and
variance by exploiting a temporal difference relation for the advantage
function approximation. More recent work such as TRPO (Schulman et al., 2015a)
considers a Kullback-Leibler (KL) divergence penalty term in order to ensure
that the updated policy remains close to the current policy or PPO (Schulman
et al., 2017) that uses clipped surrogate objective function to achieve the
same goal. In practice, it has been shown that these algorithms have better
performance compared with vanilla policy gradient method.

Most stochastic gradient based policy methods need $O(\epsilon^{-4})$
trajectories in order to achieve $\epsilon$-approximate first-order stationary
point ($\epsilon$-FOSP) of the objective function $J(\theta)$, i.e.,
$\mathbb{E}[\|\nabla J(\theta)\|]\leq\epsilon$ (Ghadimi & Lan, 2013; Shani et
al., 2020). In recent years, there have been several attempts to reduce the
variance of policy gradient by adapting variance reduction techniques proposed
previously in supervised learning context (a list of previous work is given in
Section 4). These methods can achieve sample complexity of $O(\epsilon^{-3})$
in RL setting and this rate is optimal in stochastic optimization under some
mild assumptions on the objective function and stochastic gradients (Arjevani
et al., 2020). In supervised learning problems, the objective function is
oblivious, in the sense that the randomness that selects the loss function
does not depend on the parameters that are to be optimized. On the other hand,
in RL setting, the distribution over trajectories is non-stationary and
changes over time as the parameters of policy are updated. To resolve this
issue, most previous work utilized importance sampling techniques, which may
degrade the effectiveness of the variance reduction process (Yang et al.,
2019). Moreover, to analyze the convergence rate of these methods, a strong
assumption on the variance of importance sampling weights is assumed, which
may not hold in RL setting. Most importantly, these methods often need huge
batch sizes, which is highly undesirable in practice.

In this paper, we propose Stochastic Hessian Aided Recursive Policy gradient
(SHARP) algorithm, which incorporates second-order information into SGD with
momentum. Our main contributions are summarized as follows:

  * •

Under some common regularity assumptions on the parameterized policy, SHARP
reaches $\epsilon$-FOSP with a sample complexity of $O(\epsilon^{-3})$.
Moreover, our algorithm does not use importance sampling techniques. As a
result, we can relax the strong additional assumptions on importance sampling
weights customary in the literature.

  * •

The batch size of SHARP is $O(1)$ and it does not require checkpoints, thanks
to the use of a second-order term in the updates and time-varying learning
rate and momentum weight.

  * •

SHARP is parameter-free in the sense that the initial learning rate and
momentum weight do not depend on the parameters of the problem. Moreover, the
variance of the estimation error decays with the rate of $O(1/t^{2/3})$, where
$t$ is the number of iterations.

  * •

Our experimental results show that SHARP outperforms the state of the art on
various control tasks, with remarkable performance in more complex
environments.

The rest of this paper is organized as follows: In Section 2, we define the
problem and provide some notations and background on variance reduction
methods in supervised learning. In Section 3, we describe the proposed
algorithm and analyze its convergence rate. In Section 4, we give a summary of
previous work and discuss how our proposed algorithm differs from them. In
Section 5, we evaluate the performance of the proposed algorithm against the
related work experimentally. Finally, we conclude the paper in Section 6.

##  2 Preliminaries

###  2.1 Notations and problem definition

Consider a discrete-time MDP
$\mathcal{M}=\\{\mathcal{S},\mathcal{A},P,R,\gamma,\rho\\}$ that models how an
agent interacts with a given environment. $\mathcal{S}$ and $\mathcal{A}$ are
state space and action space, respectively. $P(s^{\prime}|s,a)$ denotes the
probability of transiting to state $s^{\prime}$ from $s$ after taking action
$a$. The reward function $R$ returns reward $r(s,a)$ when action $a$ is taken
in state $s$. Parameter $\gamma\in(0,1)$ denotes the discount factor and
$\rho$ is the distribution of starting state. The actions are chosen according
to policy $\pi$ where $\pi(a|s)$ is the probability of taking action $a$ for a
given state $s$. Here, we assume that the policy is parameterized with a
vector $\theta\in\mathbb{R}^{d}$ and use shorthand notation $\pi_{\theta}$ for
$\pi_{\theta}(a|s)$.

For a given time horizon $H$, according to policy $\pi_{\theta}$, the agent
observes a sequence of state-action pairs
$\tau=(s_{0},a_{0},\cdots,s_{H-1},a_{H-1})$ called a trajectory. The
probability of observing a trajectory $\tau$ for a given policy $\pi_{\theta}$
is:

$p(\tau|\pi_{\theta})=\rho(s_{0})\prod_{h=0}^{H-1}P(s_{h+1}|s_{h},a_{h})\pi_{\theta}(a_{h}|s_{h}).$
(1)

The discounted cumulative reward for a trajectory $\tau$ is defined as
$R(\tau):=\sum_{h=0}^{H-1}\gamma^{h}r(s_{h},a_{h})$ and the expected return
for a policy $\pi_{\theta}$ is:

$J(\theta):=\mathbb{E}_{\tau\sim\pi_{\theta}}[R(\tau)].$ (2)

The main goal in policy-based RL is to find
$\theta^{*}=\arg\max_{\theta}J(\theta)$. As in many applications, $J(\theta)$
is non-convex and we settle instead for obtaining $\epsilon$-FOSP,
$\hat{\theta}$, such that $\mathbb{E}[\|\nabla
J(\hat{\theta})\|]\leq\epsilon$. It can be shown that:

$\nabla
J(\theta)=\mathbb{E}\Bigg{[}\sum_{h=0}^{H-1}\Psi_{h}(\tau)\nabla\log\pi_{\theta}(a_{h}|s_{h})\Bigg{]},$
(3)

where $\Psi_{h}(\tau)=\sum_{t=h}^{H-1}\gamma^{t}r(s_{t},a_{t})$. Therefore,
for any trajectory $\tau$,
$g(\tau;\theta):=\sum_{h=0}^{H-1}\Psi_{h}(\tau)\nabla\log\pi_{\theta}(a_{h}|s_{h})$
is an unbiased estimator of $\nabla J(\theta)$. The vanilla policy gradient
updates $\theta$ as follows:

$\theta\leftarrow\theta+\eta g(\tau;\theta),$ (4)

where $\eta$ is the learning rate.

The Hessian matrix of $J(\theta)$ can be written as follows (Shen et al.,
2019):

$\nabla^{2}J(\theta)=\mathbb{E}[\nabla\Phi(\theta;\tau)\nabla\log
p(\tau|\pi_{\theta})^{T}+\nabla^{2}\Phi(\theta;\tau)],$ (5)

where
$\Phi(\theta;\tau)=\sum_{h=0}^{H-1}\sum_{t=h}^{H-1}\gamma^{t}r(s_{t},a_{t})\log\pi_{\theta}(a_{h}|s_{h})$.
For a given trajectory $\tau$,
$B(\tau;\theta):=\nabla\Phi(\theta;\tau)\nabla\log
p(\tau|\pi_{\theta})^{T}+\nabla^{2}\Phi(\theta;\tau)$ is an unbiased estimator
of the Hessian matrix.

###  2.2 Variance reduced methods for gradient estimation

Variance reduced methods for estimating the gradient vector were originally
proposed for the stochastic optimization setting:

$\min_{\theta\in\mathbb{R}^{d}}\mathbb{E}_{z\sim p(z)}[f(\theta,z)],$ (6)

where a sample $z$ is drawn from distribution $p(z)$ and $f(.,z)$ is commonly
assumed to be smooth and non-convex function of $\theta$. This setting is
mainly considered in supervised learning context where $\theta$ corresponds to
the parameters of the training model and $z=(x,y)$ is the training sample,
with $x$ the feature vector of the sample and $y$ the corresponding label. In
this setting, the distribution $p(z)$ is invariant with respect to parameter
$\theta$.

The common approach for reducing the variance of gradient estimation is to
reuse past gradient vectors. The pseudo-code for this general framework for
variance reduction is given in Algorithm 1. After every pre-determined number
of iterations $Q$, there is a checkpoint to obtain an unbiased estimate of the
gradient, denoted by $h_{t}$, at the current parameter $\theta_{t}$ by taking
a batch of samples $\mathcal{B}_{check}$. Between any two consecutive
checkpoints, the gradient at the parameter $\theta_{t}$ is estimated according
to equation 8 by taking a batch of samples $\mathcal{B}$ drawn from $p(z)$.
The above framework appeared in several previous variance reduction methods in
stochastic optimization such as SARAH (Nguyen et al., 2017) and SPIDER (Fang
et al., 2018). Zhang (2021) discusses how to choose the size of batches and
the parameters $Q$ and $\eta$. In fact, there is a trade-off between $\eta$
and $|\mathcal{B}|$. If a small batch size is used, then $\eta$ is also
required to be small. The two extremes are SpiderBoost (Wang et al., 2019)
($|\mathcal{B}|=O(\epsilon^{-1}),\eta=O(1)$) and SARAH (Nguyen et al., 2017)
($|\mathcal{B}|=O(1)$, $\eta=O(\epsilon)$). Very recently, Li et al. (2021)
proposed PAGE, where in each iteration $t$, either a batch of samples is taken
with probability $p_{t}$ to update the gradient or the previous estimate of
the gradient is used with a small adjustment, with probability $1-p_{t}$.

Algorithm 1 Common framework in variance reduction methods

1: for $t=0,\cdots,T-1$ do

2: $\displaystyle h_{t}=$
$\displaystyle\frac{1}{|\mathcal{B}_{check}|}\sum_{z\in\mathcal{B}_{check}}\nabla
f(\theta_{t},z)$ if $t\equiv 0\pmod{Q}$, (7) $\displaystyle h_{t}=$
$\displaystyle h_{t-1}+\frac{1}{|\mathcal{B}|}\sum_{z\in\mathcal{B}}\nabla
f(\theta_{t},z)-\nabla f(\theta_{t-1},z),$ otherwise. (8)

3: $\theta_{t+1}\leftarrow\theta_{t}-\eta h_{t}$

4: end for

5: Return $\theta_{t}$ with $t$ chosen randomly from $\\{0,\cdots,T-1\\}$

In the context of RL, a sample $z$ corresponds to a trajectory $\tau$. Unlike
supervised learning, the distribution of these trajectories depends on the
parameters of policy generating them. Therefore, in the second term in the sum
in equation 8, namely $\nabla f(\theta_{t-1},z)$, $z$ (or trajectory $\tau$ in
RL context) is generated according to policy $\pi_{\theta_{t}}$ while
$\theta_{t-1}$ is the parameter of the policy at the previous iteration. In RL
setting, importance sampling technique is commonly used to account for the
distribution shift as follows:

$h_{t}=h_{t-1}+\frac{1}{|\mathcal{B}|}\sum_{\tau\in\mathcal{B}}g(\theta_{t};\tau)-w(\tau|\theta_{t},\theta_{t-1})g(\theta_{t-1};\tau),$
(9)

with the weights
$w(\tau|\theta_{t},\theta_{t-1})=\prod_{h=0}^{H-1}\frac{\pi_{\theta_{t-1}}(a_{h}|s_{h})}{\pi_{\theta_{t}}(a_{h}|s_{h})}$.

As we shall see in Section 4, nearly all variance reduction approaches in RL
employing the general framework of Algorithm 1, use an importance sampling
technique. This could significantly degrade the performance of the approach as
the gradient estimates depend heavily on these weights (Yang et al., 2019).
Besides, these variance reduction methods often need giant batch sizes at
checkpoints, which is not practical in RL setting. Finally, the hyper-
parameters of these approaches must be selected carefully as they often use
non-adaptive learning rates. To resolve the issue of requiring huge batch-
sizes, in the context of stochastic optimization, a variance reduction method
called STORM (Cutkosky & Orabona, 2019) was proposed with the following update
rule:

$\displaystyle h_{t}=(1-\alpha_{t})h_{t-1}+$ $\displaystyle\alpha_{t}\nabla
f(\theta_{t},z_{t})+(1-\alpha_{t})(\nabla f(\theta_{t},z_{t})-\nabla
f(\theta_{t-1},z_{t}))$
$\displaystyle\theta_{t+1}\leftarrow\theta_{t}-\eta_{t}h_{t},$ (10)

where $z_{t}$ is the sample drawn at iteration $t$ and $\alpha_{t}$ and
$\eta_{t}$ are the adaptive momentum weight and learning rate, respectively.
Compared with SGD with momentum, the main difference in STORM is the
correction term $\nabla f(\theta_{t},z_{t})-\nabla f(\theta_{t-1},z_{t})$ in
equation 2.2. Cutkosky & Orabona (2019) showed that by adaptively updating
$\alpha_{t}$ and $\eta_{t}$ based on the norm of stochastic gradient in
previous iterations, STORM can achieve the same convergence rate as previous
methods without requiring checkpoints nor a huge batch size.  Later, a
parameter-free version, called STORM+ (Levy et al., 2021), has been introduced
using new adaptive learning rate and momentum weight. However, to adapt these
methods in RL setting, we still need to use importance sampling techniques
because of the term $\nabla f(\theta_{t-1},z_{t})$. Recently, Tran & Cutkosky
(2022) showed that the correction term can be replaced with a second-order
term $\nabla^{2}f(\theta_{t},z_{t})(\theta_{t}-\theta_{t-1})$ by considering
additional assumption that objective function is second-order smooth. Besides,
the above Hessian vector product can be computed in $O(Hd)$ (similar to the
computational complexity of obtaining the gradient vector) by executing
Pearlmutter’s algorithm (Pearlmutter, 1994).

##  3 The SHARP Algorithm

In this section, we propose the SHARP algorithm, which incorporates second-
order information into SGD with momentum and provide a convergence guarantee.
SHARP algorithm is presented in Algorithm 2. At each iteration $t$, we draw
sample $b_{t}$ from a uniform distribution in the interval $[0,1]$ (line 5)
and then obtain $\theta_{t}^{b}$ as the linear combination of $\theta_{t-1}$
and $\theta_{t}$ with coefficients $1-b_{t}$ and $b_{t}$ (line 6). In line 7,
we sample trajectories $\tau_{t}$ and $\tau_{t}^{b}$ according to policies
$\pi_{\theta_{t}}$ and $\pi_{\theta_{t}^{b}}$, respectively. Afterwards, we
update the momentum weight $\alpha_{t}$ and the learning rate $\eta_{t}$ (line
8) and then compute the estimate of gradient at time $t$, i.e., $v_{t}$, using
the Hessian vector product
$B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1})$ and stochastic
gradient $g(\tau_{t};\theta_{t})$ (line 9). Finally we update $\theta_{t}$
based on a normalized version of $v_{t}$ in line 10.

Algorithm 2 The SHARP algorithm

Input: Initial point $\theta_{0}$, parameters $\alpha_{0},\eta_{0}$, and
number of iterations $T$

1: Sample trajectory $\tau_{0}$ with policy $\pi_{\theta_{0}}$

2: $v_{0}\leftarrow g(\tau_{0};\theta_{0})$

3: $\theta_{1}\leftarrow\theta_{0}+\eta_{0}\frac{v_{0}}{\|v_{0}\|}$

4: for $t=1,\cdots,T-1$ do

5: Sample $b_{t}\sim U(0,1)$

6: $\theta^{b}_{t}\leftarrow b_{t}\theta_{t}+(1-b_{t})\theta_{t-1}$

7: Sample trajectories $\tau_{t}$ and $\tau_{t}^{b}$ with policies
$\pi_{\theta_{t}}$ and $\pi_{\theta_{t}^{b}}$, respectively

8: $\eta_{t}\leftarrow\frac{\eta_{0}}{t^{2/3}}$,
$\alpha_{t}\leftarrow\frac{\alpha_{0}}{t^{2/3}}$

9:
$v_{t}\leftarrow(1-\alpha_{t})(v_{t-1}+B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1}))+\alpha_{t}g(\tau_{t};\theta_{t})$

10: $\theta_{t+1}\leftarrow\theta_{t}+\eta_{t}\frac{v_{t}}{\|v_{t}\|}$

11: end for

12: Return $\theta_{t}$ with $t$ chosen randomly from $\\{0,\cdots,T-1\\}$

###### Remark 1

By choosing a point uniformly at random on the line between $\theta_{t-1}$ and
$\theta_{t}$, we can ensure that
$B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1}))$ is an unbiased
estimate of $\nabla J(\theta_{t})-\nabla J(\theta_{t-1})$ (see equation 22 in
Appendix B). As we mentioned before, in the context of stochastic
optimization, Tran & Cutkosky (2022) used the second-order term
$\nabla^{2}f(\theta_{t},z_{t})(\theta_{t}-\theta_{t-1})$, which is biased as
the Hessian vector product is evaluated at the point $\theta_{t}$. As a
result, in order to provide the convergence guarantee, it is further assumed
that the objective function is second-order smooth in (Tran & Cutkosky, 2022).

###### Remark 2

To give an intuition why the second-order term is helpful in the update in
line 9, consider the following error term:

$\epsilon_{t}=v_{t}-\nabla J(\theta_{t}).$ (11)

We can rewrite the above error term as follows:

$\begin{split}\epsilon_{t}&=(1-\alpha_{t})(v_{t-1}-\nabla
J(\theta_{t})+B(\tau^{b}_{t};\theta^{b}_{t})(\theta_{t}-\theta_{t-1}))\\\
&\qquad\qquad+\alpha_{t}(g(\tau_{t};\theta_{t})-\nabla
J(\theta_{t})).\end{split}$ (12)

Now, for a moment, suppose that $\mathbb{E}[v_{t-1}]=\mathbb{E}[\nabla
J(\theta_{t-1})]$ (with total expectation on both sides). Then,

$\mathbb{E}[v_{t-1}-\nabla
J(\theta_{t})+B(\tau^{b}_{t};\theta^{b}_{t})(\theta_{t}-\theta_{t-1})]=0.$
(13)

As $v_{0}$ is an unbiased estimate of gradient at $\theta_{0}$, we can easily
show by induction that according to above equation,
$\mathbb{E}[v_{t}]=\mathbb{E}[\nabla J(\theta_{t})]$ for any $t\geq 0$.

In the next part, we provide a theoretical guarantee on the convergence rate
of SHARP algorithm.

###  3.1 Convergence Analysis

In this part, we analyze the convergence rate of Algorithm 2 under bounded
reward function and some regularity assumptions on the policy $\pi_{\theta}$.

######  Assumption 1 (Bounded reward)

For $\forall s\in\mathcal{S},\forall a\in\mathcal{A}$, $|R(s,a)|<R_{0}$ where
$R_{0}>0$ is some constant.

######  Assumption 2 (Parameterization regularity)

There exist constants $G,L>0$ such that for any $\theta\in\mathbb{R}^{d}$ and
for any $s\in\mathcal{S},a\in\mathcal{A}$:
(a) $\|\nabla\log\pi_{\theta}(a|s)\|\leq G$,
(b) $\|\nabla^{2}\log\pi_{\theta}(a|s)\|\leq L$.

Assumptions 1 and 2 are common in the RL literature (Papini et al., 2018; Shen
et al., 2019) to analyze the convergence of policy gradient methods. Under
these assumptions, the following upper bounds can be derived on
$\mathbb{E}[\|g(\tau;\theta)-\nabla J(\theta)\|^{2}]$ and
$\mathbb{E}[\|B(\tau;\theta)-\nabla^{2}J(\theta)\|^{2}]$.

######  Lemma 1 (Shen et al. (2019))

Under Assumptions 1 and 2:

$\begin{split}&\mathbb{E}[\|g(\tau;\theta)-\nabla
J(\theta)\|^{2}]\leq\sigma_{g}^{2}\\\
&\mathbb{E}[\|B(\tau;\theta)-\nabla^{2}J(\theta)\|^{2}]\leq\sigma_{B}^{2},\end{split}$
(14)

where $\sigma_{g}^{2}=\frac{G^{2}R_{0}^{2}}{(1-\gamma)^{4}}$ and
$\sigma_{B}^{2}=\frac{H^{2}G^{4}R_{0}^{2}+L^{2}R_{0}^{2}}{(1-\gamma)^{4}}$.

Based on these bounds, we can provide the following guarantee on the
convergence rate of SHARP algorithm. All proofs are provided in the appendix.

###### Theorem 1

Suppose that the initial momentum weight $\alpha_{0}\in(2/3,1]$ and initial
learning rate $\eta_{0}>0$. Under Assumptions 1 and 2, Algorithm 2 guarantees
that:

$\mathbb{E}\Bigg{[}\frac{1}{T}\sum_{t=1}^{T}\|\nabla
J(\theta_{t})\|\Bigg{]}\leq\frac{8\sqrt{C}+9C_{J}/\eta_{0}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}},$
(15)

where
$C=3\alpha_{0}(48\sigma_{B}^{2}\eta_{0}^{2}/\alpha_{0}+(6\alpha_{0}+1/\alpha_{0})\sigma_{g}^{2})/(3\alpha_{0}-2)$
and $C_{J}=R_{0}/(1-\gamma)$.

###### Corollary 1

The right hand side of equation 15 is dominated by the first term. If we set
$\eta_{0}$ in the order of $\sqrt{C_{J}/\sigma_{B}}$, then the number of
trajectories for achieving $\epsilon$-FOSP would be
$O(\frac{1}{(1-\gamma)^{2}\epsilon^{3}})$, where we assume that the time
horizon $H$ is set in the order of $1/(1-\gamma)$.

###### Remark 3

Along the iterations of the SHARP algorithm, it can be shown that the
following inequality holds for any $t\geq 1$ (see equation 20 in Appendix B):

$\mathbb{E}[\|\epsilon_{t}\|^{2}]\leq(1-\alpha_{t})\mathbb{E}[\|\epsilon_{t-1}\|^{2}]+O(\eta_{t}^{2}).$
(16)

Therefore, the variance of the estimation error decays with the rate of
$O(1/t^{2/3})$ (see Appendix C for the proof). To the best of our knowledge,
existing variance reduction methods only guarantee the decay of accumulative
variance. This appealing property of SHARP is largely due to the use of
unbiased Hessian-aided gradient estimator and normalized gradient descent.
Moreover, as a byproduct of these desirable properties, our convergence
analysis turns out to be more simple, compared to existing work (Cutkosky &
Orabona, 2019; Tran & Cutkosky, 2022). This could be of independent interest
for better theory of variance-reduced methods.

###### Remark 4

The SHARP algorithm is parameter-free in the sense that $\alpha_{0}$ and
$\eta_{0}$ are constants that do not depend on other parameters in the
problem. Therefore, for any choice of $2/3<\alpha_{0}\leq 1$ and $\eta_{0}>0$,
we can guarantee convergence to $\epsilon$-FOSP with the sample complexity of
$O(\epsilon^{-3})$. However, in practice, it is desirable to tune these
constants to have smaller constants in the numerators of convergence rates in
equation 15. For instance, $\sigma_{B}$ might be large in some RL settings and
we control the constant in the first term on the right hand side of equation
15 by tuning $\eta_{0}$. It is noteworthy that STORM+ is also parameter-free
but it requires adaptive learning rate and momentum weight that depend on
stochastic gradients in previous iterations.

###### Remark 5

Regarding the dependency on $\epsilon$, in the context of stochastic
optimization, Arjevani et al. (2020) have shown that under some mild
assumptions on the objective function and stochastic gradient, the rate of
$O(1/\epsilon^{3})$ is optimal in order to obtain $\epsilon$-FOSP, and cannot
be improved with stochastic $p$-th order methods for $p\geq 2$.

##  4 Related Work

In recent years, several variance-reduced methods have been proposed in order
to accelerate the existing PG methods. Papini et al. (2018) and Xu et al.
(2020) proposed SVRPG algorithm based on SVRG (Johnson & Zhang, 2013), with
sample complexity of $O(1/\epsilon^{4})$ and $O(1/\epsilon^{10/3})$,
respectively. This algorithm requires importance sampling techniques as well
as the following further assumption for guaranteeing convergence to
$\epsilon$-FOSP:

  * •

Bounded variance of importance sampling weights: For any trajectory $\tau$, it
is assumed that:

$Var\Bigg{(}\frac{p(\tau|\pi_{\theta_{1}})}{p(\tau|\pi_{\theta_{2}})}\Bigg{)}\leq
W,\qquad\forall\theta_{1},\theta_{2}\in\mathbb{R}^{d},$ (17)

where $W<\infty$ is a constant.

The above assumption is fairly strong as the importance sampling weight could
grow exponentially with horizon length $H$ (Zhang et al., 2021). In order to
remove importance sampling weights, Shen et al. (2019) proposed HAPG
algorithm, which uses second-order information and achieves better sample
complexity of $O(1/\epsilon^{3})$. However, it still needs checkpoints and
large batch sizes of $|\mathcal{B}|=O(1/\epsilon)$, and
$|\mathcal{B}_{check}|=O(1/\epsilon^{2})$.

In Table 1, we compare the main variance-reduced policy gradient methods
achieving $\epsilon$-FOSP in terms of sample complexity and batch
size111Please note that the sample complexity also depends on the other
parameters such as horizon length $H$ and discount factor $\gamma$. Here, we
just mention the dependency of sample complexity on $\epsilon$.. In this
table, after HAPG (Shen et al., 2019), all the proposed variance reduction
methods achieve a similar sample complexity. The orders of batch sizes are
also the same as in HAPG. Xu et al. (2019) proposed SRVR-PG, and used
stochastic path-integrated differential estimators for variance reduction.
This algorithm uses important sampling weights and the required batch sizes
are in the order of $|\mathcal{B}|=O(1/\sqrt{\epsilon})$ and
$|\mathcal{B}_{check}|=O(1/\epsilon)$. Later, Pham et al. (2020) proposed
HSPGA by adapting the SARAH estimator for reducing the variance of REINFORCE.
HSPGA still needs importance sampling weights, but the batch size is reduced
to $O(1)$. Huang et al. (2020) proposed three variants of momentum-based
policy gradient (called MBPG), which are based on the STORM algorithm
(Cutkosky & Orabona, 2019). Thus, the required batch size is in the order of
$O(1)$, similarly to STORM. However, it still needs to use importance sampling
weights. In (Zhang et al., 2020), similar update for the estimate of
stochastic gradient as the one in SHARP have been explored for Frank-Wolfe
type algorithms in the context of constrained optimization. Later, Zhang et
al. (2021) proposed TSIVR-PG with a gradient truncation mechanism in order to
resolve some of issues pertaining to the use of importance sampling weights.
In their convergence analysis, they are restricted to soft-max policy with
some specific assumptions on the parameterization functions. More recently,
two methods using mirror-descent algorithm based on Bregman divergence, called
VR-BGPO (Huang et al., 2022) and VR-MPO (Yang et al., 2022) have been
proposed, achieving $\epsilon$-FOSP if the mirror map in Bregman divergence is
the $l_{2}$-norm.  Very recently, based on PAGE (Li et al., 2021), Gargiani et
al. (2022) proposed the PAGE-PG algorithm which takes a batch of samples of
$O(\epsilon^{-2})$ for updating the parameters with probability $p_{t}$ or
reuse the previous estimate gradient with a small adjustment, with probability
$1-p_{t}$. The proposed algorithm requires importance sampling weights and
thus the additional assumption in equation 17 to guarantee convergence to
$\epsilon$-FOSP with a sample complexity of $O(\epsilon^{-3})$.

Method | SC | $|\mathcal{B}|$ | $|\mathcal{B}_{check}|$ | Checkpoint | IS | Further Assump.
---|---|---|---|---|---|---
SVRPG Xu et al. (2020) | $O(\frac{1}{\epsilon^{10/3}})$ | $O(\frac{1}{\epsilon^{4/3}})$ | $O(\frac{1}{\epsilon^{4/3}})$ | Needed | Needed | Assump. in equation 17
HAPG Shen et al. (2019) | $O(\frac{1}{\epsilon^{3}})$ | $O(\frac{1}{\epsilon})$ | $O(\frac{1}{\epsilon^{2}})$ | Needed | Not needed | -
SRVR-PG Xu et al. (2019) | $O(\frac{1}{\epsilon^{3}})$ | $O(\frac{1}{\sqrt{\epsilon}})$ | $O(\frac{1}{{\epsilon}})$ | Needed | Needed | Assump. in equation 17
HSPGA Pham et al. (2020) | $O(\frac{1}{\epsilon^{3}})$ | $O(1)$ | - | Not needed | Needed | Assump. in equation 17
MBPG Huang et al. (2020) | $\tilde{O}(\frac{1}{\epsilon^{3}})$ | $O(1)$ | - | Not needed | Needed | Assump. in equation 17
VRMPO Yang et al. (2022) | $O(\frac{1}{\epsilon^{3}})$ | $O(\frac{1}{\epsilon})$ | $O(\frac{1}{\epsilon^{2}})$ | Needed | Not needed | -
VR-BGPO Huang et al. (2022) | $O(\frac{1}{\epsilon^{3}})$ | $O(1)$ | - | Not Needed | Needed | Assump. in equation 17
PAGE-PG Gargiani et al. (2022) | $O(\frac{1}{\epsilon^{3}})$ | $O(1)$ | $O(\frac{1}{\epsilon^{2}})$ | Needed | Needed | Assump. in equation 17
This paper | $O(\frac{1}{\epsilon^{3}})$ | $O(1)$ | - | Not needed | Not needed | -
Table 1: Comparison of main variance-reduced policy gradient methods to
achieve $\epsilon$-FOSP based on sample complexity (SC), batch size
($|\mathcal{B}|$), batch size at checkpoints ($|\mathcal{B}_{check}|$), and
the need for checkpoints, importance sampling (IS), and additional
assumptions.

There exist few recent work on the global convergence of policy gradient
methods. For instance, Liu et al. (2020) showed global convergence of policy
gradient, natural policy gradient, and their variance reduced variants, in the
case of positive definite Fisher information matrix of the policy. Chung et
al. (2021) studied the impact of baselines on the learning dynamics of policy
gradient methods and showed that using a baseline minimizing the variance can
converge to a sub-optimal policy. Recently, Ding et al. (2022) studied the
soft-max and the Fisher non-degenerate policies, and showed that adding a
momentum term improves the global optimality sample complexities of vanilla PG
methods by $\tilde{O}(\epsilon^{-1.5})$ and $\tilde{O}(\epsilon^{-1})$,
respectively.

The aforementioned discussion for the main methods are summarized in Table 1.
For each method, we mention whether it needs checkpoints and importance
sampling weights.222To be more precise, although PAGE-PG, has no fixed
checkpoints, it takes a batch of $O(\epsilon^{-2})$ to get an unbiased
estimate of the gradient with probability $p_{t}$. Therefore, in this sense,
it requires checkpoints. All the discussed aforementioned methods require
Assumptions 1 and 2. In the last column, additional assumptions besides these
two are listed for each method.

Comparing the sample complexity of our algorithm with previous work, note that
all the algorithms (including ours) under SVRPG in Table 1 achieve the rate of
$O(1/\epsilon^{3})$ or $\tilde{O}(1/\epsilon^{3})$. Without any further
assumption, our proposed method is the only one that requires no checkpoints,
no importance sampling weights, and has a batch size of the order of $O(1)$.
As we will see in the next section, besides these algorithmic advantages, it
has remarkable performance compared to the state of the art in various control
tasks.

##  5 Experiments

In this section, we evaluate the performance of the proposed algorithm and
compare it with the previous work for control tasks in MuJoCo simulator
(Todorov et al., 2012) which is a physical engine, suitable for simulating
robotic tasks with good accuracy and speed in RL setting. We implemented SHARP
in the Garage library (garage contributors, 2019) as it allows for maintaining
and integrating it in future versions of Garage library for easier
dissemination. We utilized a Linux server with Intel Xeon CPU E5-2680 v3 (24
cores) operating at 2.50GHz with 377 GB DDR4 of memory, and Nvidia Titan X
Pascal GPU. The implementation of SHARP is available as supplementary
material.

We considered the following four control tasks with continuous action space:
Reacher, Walker, Humanoid, and Swimmer. In Reacher, there is an arm with two
degrees of freedom, aiming to reach a target point in the two-dimensional
plane. A higher rewards is attained if the arm gets closer to the target
point. In Walker, a humanoid walker tries to move forward in a two dimensional
space, i.e., it can only fall forward or backward. The state contains
velocities of different parts of body and joint angles and the actions
represent how to move legs and foot joints. The reward signal is based on the
current velocity of the agent. In Humanoid, a three-dimensional bipedal robot
is trained to walk forward as fast as possible, without falling over. The
state space is $376$-dimensional, containing the position and velocity of each
joint, the friction of the actuator, and contact forces. The action space is a
17-dimensional continuous space. Finally, in Swimmer, the agent is in a two-
dimensional pool and the goal is to move as fast as possible in the right
direction.

We compared SHARP algorithm with PG methods that provide theoretical
guarantees for converging to an approximate FOSP: PAGE-PG (Gargiani et al.,
2022), IS-MBPG (Huang et al., 2020) which is based on STROM, HAPG (Shen et
al., 2019) which does not require IS weights, and VR-BGPO (Huang et al., 2022)
which is a mirror descent based algorithm. We also considered REINFORCE
(Sutton et al., 2000) as a baseline algorithm. There are some other approaches
in the literature with theoretical guarantees such as VRMPO (Yang et al.,
2022), and STORM-PG (Yuan et al., 2020) but the official implementations are
not publicly available and our request to access the code from the authors
remained unanswered.

For each algorithm, we used the same set of Gaussian policies parameterized
with neural networks having two layers of 64 neurons each. Baselines and
environment settings (such as maximum trajectory horizon, and reward
intervals) were considered the same for all algorithms. We chose a maximum
horizon of 500 for Walker, Swimmer, and Humanoid and 50 for Reacher. More
details about experiments are provided in Appendix E.

Figure 1: Comparison of SHARP with other variance reduction methods on four
control tasks.

In the literature, it has been observed that most PG methods are quite
sensitive to parameter initialization or random seeds (Henderson et al.,
2018). Hence, it might be challenging in some cases to reproduce previous
results. Moreover, it is not clear how to compare methods in terms of
performance (e.g., the average episode return) and robustness (such as the
standard deviation (STD) of return) at the same time.

To resolve the above issues, we considered the perfromance-robustness (PR)
metric proposed in (Khorasani et al., 2023), capturing both the average return
and STD of return of an algorithm. In particular, for any algorithm $A$, after
observing $t$ number of state-actions pairs (which we call system probes), the
lower bound of the confidence interval 90 percent of average return over $n$
runs of the algorithm (denoted by $LCI_{A}(n,t)$) is computed. The PR metric
is defined by averaging $LCI_{A}(n,t)$ over all system probes $t=1,\cdots,T$
as follows:

$PR_{A}(n)=\dfrac{1}{T}\sum_{t=1}^{T}LCI_{A}(n,t),$ (18)

where $T$ is maximum number of system probes.

We used grid search to tune the hyper-parameters of all the considered
algorithms. For the algorithms excepts ours, the search space for each hyper-
parameter was chosen based on the one from the original papers. For each
configuration of the hyper-parameters, we ran each algorithm $A$, five times
and computed $PR_{A}(5)$. We selected the configuration which maximized
$PR_{A}(5)$ and then reported $PR_{A}(10)$ of each algorithm for the selected
configuration based on 10 different runs in Table 2. Our proposed method
achieved the highest PR in all environments compared with the other
algorithms.

Table 2: Comparison of SHARP with other variance-reduced methods in terms of
PR. In each environment, the highest PR is in bold.

| Reacher | Walker | Humanoid | Swimmer
---|---|---|---|---
HAPG | -19.51 | 104.296 | 161.12 | 70.62
IS-MBPG | -21.76 | 204.32 | 201.01 | 27.44
PAGE-PG | -17.39 | 247.58 | 356.58 | 19.26
REINFORCE | -20.10 | 79.98 | 156.21 | 20.28
VR-BGPO | -15.15 | 320.51 | 409.67 | 38.24
SHARP (our algorithm) | -10.34 | 325.98 | 422.98 | 83.45

We considered the confidence interval of the performance (average return) to
show the sensitivity of an algorithm to random seeds. As can be seen in Figure
1, SHARP not only achieves the highest average return after 10 million system
probes but also has a relatively small confidence interval.

##  6 Conclusion

We proposed a variance-reduced policy-gradient method, which incorporates
second-order information, i.e., Hessian vector product, into SGD with
momentum. The Hessian vector product can be computed with an efficiency that
is similar to that of obtaining the gradient vector. Therefore, the
computational complexity of the proposed algorithm per iteration remains in
the same order as first-order methods. More importantly, using the second-
order correction term enables us to obtain an estimate of the gradient,
completely bypassing importance sampling. Moreover, the batch size is $O(1)$
and there is no need for checkpoints, which makes the algorithm appealing in
practice. Under some regularity assumptions on the parameterized policy, we
showed that it achieves $\epsilon$-FOSP with sample complexity of
$O(\epsilon^{-3})$. SHARP is parameter-free in the sense that the initial
learning rate and momentum weight do not depend on the parameters of the
problem. Experimental results show its remarkable performance in various
control tasks, especially in some complex environments, such as HalfCheetah,
compared to the state of the art.

## References

  * Arjevani et al. (2020) Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Ayush Sekhari, and Karthik Sridharan.  Second-order information in non-convex stochastic optimization: Power and limitations.  In _Conference on Learning Theory_ , pp. 242–299. PMLR, 2020.
  * Baxter & Bartlett (2001) Jonathan Baxter and Peter L Bartlett.  Infinite-horizon policy-gradient estimation.  _Journal of Artificial Intelligence Research_ , 15:319–350, 2001.
  * Chung et al. (2021) Wesley Chung, Valentin Thomas, Marlos C Machado, and Nicolas Le Roux.  Beyond variance reduction: Understanding the true impact of baselines on policy optimization.  In _International Conference on Machine Learning_ , pp. 1999–2009. PMLR, 2021.
  * Cutkosky & Orabona (2019) Ashok Cutkosky and Francesco Orabona.  Momentum-based variance reduction in non-convex sgd.  _Advances in neural information processing systems_ , 32, 2019.
  * Deisenroth et al. (2013) Marc Peter Deisenroth, Gerhard Neumann, Jan Peters, et al.  A survey on policy search for robotics.  _Foundations and trends in Robotics_ , 2(1-2):388–403, 2013.
  * Ding et al. (2022) Yuhao Ding, Junzi Zhang, and Javad Lavaei.  On the global optimum convergence of momentum-based policy gradient.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1910–1934. PMLR, 2022.
  * Fang et al. (2018) Cong Fang, Chris Junchi Li, Zhouchen Lin, and Tong Zhang.  Spider: Near-optimal non-convex optimization via stochastic path integrated differential estimator.  _arXiv preprint arXiv:1807.01695_ , 2018.
  * garage contributors (2019) The garage contributors.  Garage: A toolkit for reproducible reinforcement learning research.  https://github.com/rlworkgroup/garage, 2019.
  * Gargiani et al. (2022) Matilde Gargiani, Andrea Zanelli, Andrea Martinelli, Tyler Summers, and John Lygeros.  Page-pg: A simple and loopless variance-reduced policy gradient method with probabilistic gradient estimation.  In _International Conference on Machine Learning_ , 2022.
  * Ghadimi & Lan (2013) Saeed Ghadimi and Guanghui Lan.  Stochastic first-and zeroth-order methods for nonconvex stochastic programming.  _SIAM Journal on Optimization_ , 23(4):2341–2368, 2013.
  * Henderson et al. (2018) Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger.  Deep reinforcement learning that matters.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Huang et al. (2020) Feihu Huang, Shangqian Gao, Jian Pei, and Heng Huang.  Momentum-based policy gradient methods.  In _International Conference on Machine Learning_ , pp. 4422–4433. PMLR, 2020.
  * Huang et al. (2022) Feihu Huang, Shangqian Gao, and Heng Huang.  Bregman gradient policy optimization.  In _International Conference on Learning Representations_ , 2022.
  * Johnson & Zhang (2013) Rie Johnson and Tong Zhang.  Accelerating stochastic gradient descent using predictive variance reduction.  _Advances in neural information processing systems_ , 26:315–323, 2013.
  * Khorasani et al. (2023) Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Niao He, and Matthias Grossglauser.  Efficiently escaping saddle points for non-convex policy optimization.  _arXiv preprint arXiv:2311.08914_ , 2023.
  * Konda & Tsitsiklis (2000) Vijay R Konda and John N Tsitsiklis.  Actor-critic algorithms.  In _Advances in neural information processing systems_ , pp. 1008–1014, 2000.
  * Levy et al. (2021) Kfir Levy, Ali Kavis, and Volkan Cevher.  Storm+: Fully adaptive sgd with momentum for nonconvex optimization.  In _35th Conference on Neural Information Processing Systems (NeurIPS 2021)_ , number CONF, 2021.
  * Li et al. (2021) Zhize Li, Hongyan Bao, Xiangliang Zhang, and Peter Richtárik.  Page: A simple and optimal probabilistic gradient estimator for nonconvex optimization.  In _International Conference on Machine Learning_ , pp. 6286–6295. PMLR, 2021.
  * Liu et al. (2020) Yanli Liu, Kaiqing Zhang, Tamer Basar, and Wotao Yin.  An improved analysis of (variance-reduced) policy gradient and natural policy gradient methods.  In _NeurIPS_ , 2020.
  * Nguyen et al. (2017) Lam M Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáč.  Sarah: A novel method for machine learning problems using stochastic recursive gradient.  In _International Conference on Machine Learning_ , pp. 2613–2621. PMLR, 2017.
  * Papini et al. (2018) Matteo Papini, Damiano Binaghi, Giuseppe Canonaco, Matteo Pirotta, and Marcello Restelli.  Stochastic variance-reduced policy gradient.  In _International conference on machine learning_ , pp. 4026–4035. PMLR, 2018.
  * Pearlmutter (1994) Barak A Pearlmutter.  Fast exact multiplication by the hessian.  _Neural computation_ , 6(1):147–160, 1994.
  * Pham et al. (2020) Nhan Pham, Lam Nguyen, Dzung Phan, Phuong Ha Nguyen, Marten Dijk, and Quoc Tran-Dinh.  A hybrid stochastic policy gradient algorithm for reinforcement learning.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 374–385. PMLR, 2020.
  * Schulman et al. (2015a) John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.  Trust region policy optimization.  In _International conference on machine learning_ , pp. 1889–1897. PMLR, 2015a.
  * Schulman et al. (2015b) John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel.  High-dimensional continuous control using generalized advantage estimation.  _arXiv preprint arXiv:1506.02438_ , 2015b.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Shalev-Shwartz et al. (2016) Shai Shalev-Shwartz, Shaked Shammah, and Amnon Shashua.  Safe, multi-agent, reinforcement learning for autonomous driving.  _arXiv preprint arXiv:1610.03295_ , 2016.
  * Shani et al. (2020) Lior Shani, Yonathan Efroni, and Shie Mannor.  Adaptive trust region policy optimization: Global convergence and faster rates for regularized mdps.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pp. 5668–5675, 2020.
  * Shen et al. (2019) Zebang Shen, Alejandro Ribeiro, Hamed Hassani, Hui Qian, and Chao Mi.  Hessian aided policy gradient.  In _International conference on machine learning_ , pp. 5729–5738. PMLR, 2019.
  * Silver et al. (2017) David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al.  Mastering the game of go without human knowledge.  _nature_ , 550(7676):354–359, 2017.
  * Sutton et al. (2000) Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour.  Policy gradient methods for reinforcement learning with function approximation.  In _Advances in neural information processing systems_ , pp. 1057–1063, 2000.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , pp. 5026–5033. IEEE, 2012.
  * Tran & Cutkosky (2022) Hoang Tran and Ashok Cutkosky.  Better sgd using second-order momentum.  _Advances in Neural Information Processing Systems_ , 35:3530–3541, 2022.
  * Wang et al. (2019) Zhe Wang, Kaiyi Ji, Yi Zhou, Yingbin Liang, and Vahid Tarokh.  Spiderboost and momentum: Faster variance reduction algorithms.  _Advances in Neural Information Processing Systems_ , 32:2406–2416, 2019.
  * Williams (1992) Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ , 8(3):229–256, 1992.
  * Xu et al. (2019) Pan Xu, Felicia Gao, and Quanquan Gu.  Sample efficient policy gradient methods with recursive variance reduction.  _arXiv preprint arXiv:1909.08610_ , 2019.
  * Xu et al. (2020) Pan Xu, Felicia Gao, and Quanquan Gu.  An improved convergence analysis of stochastic variance-reduced policy gradient.  In _Uncertainty in Artificial Intelligence_ , pp. 541–551. PMLR, 2020.
  * Yang et al. (2019) Long Yang, Gang Zheng, Haotian Zhang, Yu Zhang, Qian Zheng, Jun Wen, and Gang Pan.  Policy optimization with stochastic mirror descent.  _arXiv preprint arXiv:1906.10462_ , 2019.
  * Yang et al. (2022) Long Yang, Yu Zhang, Gang Zheng, Qian Zheng, Pengfei Li, Jianhang Huang, and Gang Pan.  Policy optimization with stochastic mirror descent.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pp. 8823–8831, 2022.
  * Yuan et al. (2020) Huizhuo Yuan, Xiangru Lian, Ji Liu, and Yuren Zhou.  Stochastic recursive momentum for policy gradient methods.  _arXiv preprint arXiv:2003.04302_ , 2020.
  * Zhang et al. (2021) Junyu Zhang, Chengzhuo Ni, Zheng Yu, Csaba Szepesvari, and Mengdi Wang.  On the convergence and sample efficiency of variance-reduced policy gradient method.  _arXiv preprint arXiv:2102.08607_ , 2021.
  * Zhang (2021) Liang Zhang.  Variance reduction for non-convex stochastic optimization: General analysis and new applications.  Master’s thesis, ETH Zurich, 2021.
  * Zhang et al. (2020) Mingrui Zhang, Zebang Shen, Aryan Mokhtari, Hamed Hassani, and Amin Karbasi.  One sample stochastic frank-wolfe.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 4012–4023. PMLR, 2020.

##  Appendix A Appendix

##  Appendix B Proof of Theorem 1

Let us define $\epsilon_{t}:=v_{t}-\nabla J(\theta_{t})$. Then, based on the
update in line 9 of Algorithm 2, we have:

$\epsilon_{t}=(1-\alpha_{t})\epsilon_{t-1}+\alpha_{t}U_{t}+(1-\alpha_{t})W_{t},$
(19)

where $U_{t}=(g(\tau_{t};\theta_{t})-\nabla J(\theta_{t}))$ and
$W_{t}=B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1})-(\nabla
J(\theta_{t})-\nabla J(\theta_{t-1}))$. Let $\mathcal{H}_{t}$ be the history
up to time $t$, i.e.,
$\mathcal{H}_{t}:=\\{\theta_{0},\tau_{0},\tau_{1},b_{1},\tau_{1}^{b},\cdots,\tau_{t},b_{t},\tau_{t}^{b}\\}$.

By taking the square $l_{2}$ norm of both sides of the above equation,

$\begin{split}\|\epsilon_{t}\|^{2}&=(1-\alpha_{t})^{2}\|\epsilon_{t-1}\|^{2}+\alpha_{t}^{2}\|U_{t}\|^{2}+(1-\alpha_{t})^{2}\|W_{t}\|^{2}+\\\
&\qquad
2\alpha_{t}(1-\alpha_{t})\langle\epsilon_{t-1},U_{t}\rangle+2\alpha_{t}(1-\alpha_{t})\langle
U_{t},W_{t}\rangle+2(1-\alpha_{t})^{2}\langle\epsilon_{t-1},W_{t}\rangle\\\
&\leq(1-\alpha_{t})^{2}\|\epsilon_{t-1}\|^{2}+2\alpha_{t}^{2}\|U_{t}\|^{2}+2(1-\alpha_{t})^{2}\|W_{t}\|^{2}+\\\
&\qquad
2\alpha_{t}(1-\alpha_{t})\langle\epsilon_{t-1},U_{t}\rangle+2(1-\alpha_{t})^{2}\langle\epsilon_{t-1},W_{t}\rangle,\end{split}$
(20)

where we used Young’s inequality in the first inequality for the following
term: $2\alpha_{t}(1-\alpha_{t})\langle
U_{t},W_{t}\rangle\leq(1-\alpha_{t})^{2}\|W_{t}\|^{2}+\alpha_{t}^{2}\|U_{t}\|^{2}$.
Now, by taking expectations on both sides,

$\begin{split}\mathbb{E}[\|\epsilon_{t}\|^{2}]&\leq(1-\alpha_{t})^{2}\mathbb{E}[\|\epsilon_{t-1}\|^{2}]+2\alpha_{t}^{2}\mathbb{E}[\|U_{t}\|^{2}]+2(1-\alpha_{t})^{2}\mathbb{E}[\|W_{t}\|^{2}]+\\\
&\quad
2\alpha_{t}(1-\alpha_{t})\mathbb{E}[\langle\epsilon_{t-1},U_{t}\rangle]+2(1-\alpha_{t})^{2}\mathbb{E}[\langle\epsilon_{t-1},W_{t}\rangle]\\\
&\overset{(a)}{\leq}(1-\alpha_{t})\mathbb{E}[\|\epsilon_{t-1}\|^{2}]+2\alpha_{t}^{2}\mathbb{E}[\|U_{t}\|^{2}]+2\mathbb{E}[\|W_{t}\|^{2}]+\\\
&\quad
2\alpha_{t}(1-\alpha_{t})\mathbb{E}[\langle\epsilon_{t-1},U_{t}\rangle]+2(1-\alpha_{t})^{2}\mathbb{E}[\langle\epsilon_{t-1},W_{t}\rangle]\\\
&\overset{(b)}{\leq}(1-\alpha_{t})\mathbb{E}[\|\epsilon_{t-1}\|^{2}]+2\alpha_{t}^{2}\mathbb{E}[\|U_{t}\|^{2}]+2\mathbb{E}[\|W_{t}\|^{2}],\end{split}$
(21)

$(a)$ follows as $\alpha_{t}\leq 1$ for all $t\geq 0$.
(b) follows because the following two terms are zero. First,
$\mathbb{E}[\langle\epsilon_{t-1},U_{t}\rangle]=\mathbb{E}[\mathbb{E}[\langle\epsilon_{t-1},U_{t}\rangle|\mathcal{H}_{t-1}]]=0$
as $\epsilon_{t-1}$ is determined given $\mathcal{H}_{t-1}$ and
$\mathbb{E}[U_{t}|\mathcal{H}_{t-1}]=0$ since $g(\tau_{t};\theta_{t})$ is an
unbiased estimation of $\nabla J(\theta_{t})$. Second,
$\mathbb{E}[\langle\epsilon_{t-1},W_{t}\rangle]=\mathbb{E}[\mathbb{E}[\langle\epsilon_{t-1},W_{t}\rangle|\mathcal{H}_{t-1}]]=0$
because

$\begin{split}\mathbb{E}[W_{t}|\mathcal{H}_{t-1}]&=\mathbb{E}[B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1})|\mathcal{H}_{t-1}]-(\nabla
J(\theta_{t})-\nabla J(\theta_{t-1}))\\\
&=\mathbb{E}_{b_{t}}[\mathbb{E}_{\tau_{t}^{b}}[B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1})|\mathcal{H}_{t-1}]]-(\nabla
J(\theta_{t})-\nabla J(\theta_{t-1}))\\\
&\overset{(c)}{=}\mathbb{E}_{b_{t}}[\nabla^{2}J(\theta_{t}^{b})(\theta_{t}-\theta_{t-1})|\mathcal{H}_{t-1}]-(\nabla
J(\theta_{t})-\nabla J(\theta_{t-1}))\\\
&\overset{(d)}{=}\int_{0}^{1}\nabla^{2}J(b\theta_{t}+(1-b)\theta_{t-1})(\theta_{t}-\theta_{t-1})db-(\nabla
J(\theta_{t})-\nabla J(\theta_{t-1}))=0,\end{split}$ (22)

$(c)$ It is due to the fact that for a given $\theta_{t}^{b}$,
$B(\tau_{t}^{b};\theta_{t}^{b})$ is an unbiased estimation of
$\nabla^{2}J(\theta_{t}^{b})$.
$(d)$ The integral follows since
$\theta_{t}^{b}=b_{t}\theta_{t}+(1-b_{t})\theta_{t-1}$ with $b_{t}$ uniformly
distributed in $[0,1]$, and its value is $\nabla J(\theta_{t})-\nabla
J(\theta_{t-1})$.

Now, from equation 21, we have:

$\begin{split}\mathbb{E}[\|\epsilon_{t-1}\|^{2}]&\overset{(a)}{\leq}\frac{1}{\alpha_{t}}\left(\mathbb{E}[\|\epsilon_{t-1}\|^{2}]-\mathbb{E}[\|\epsilon_{t}\|^{2}]\right)+\frac{2}{\alpha_{t}}\mathbb{E}[\|W_{t}\|^{2}]+2\alpha_{t}\sigma_{g}^{2}\\\
&\overset{(b)}{\leq}\frac{1}{\alpha_{t}}\left(\mathbb{E}[\|\epsilon_{t-1}\|^{2}]-\mathbb{E}[\|\epsilon_{t}\|^{2}]\right)+8\sigma_{B}^{2}\frac{\eta_{t-1}^{2}}{\alpha_{t}}+2\alpha_{t}\sigma_{g}^{2},\end{split}$
(23)

$(a)$ We use the bound in Lemma 1, i.e.,
$\mathbb{E}[\|U_{t}\|^{2}]\leq\sigma_{g}^{2}$.
$(b)$ It has been shown that $\|B(\tau;\theta)\|\leq\sigma_{B}$ for any
trajectory $\tau$ and $\theta\in\mathbb{R}^{d}$ (Shen et al., 2019).
Therefore,
$\|\nabla^{2}J(\theta)\|=\|\mathbb{E}_{\tau}[B(\tau;\theta)]\|\leq\sigma_{B}$.
Hence, $\nabla J(\theta)$ is Lipschitz with constant $\sigma_{B}$ and we have:

$\begin{split}\|W_{t}\|&\leq\|B(\tau_{t}^{b};\theta_{t}^{b})(\theta_{t}-\theta_{t-1})\|+\|\nabla
J(\theta_{t})-\nabla J(\theta_{t-1})\|\\\
&\leq\|B(\tau_{t}^{b};\theta_{t}^{b})\|\|\theta_{t}-\theta_{t-1}\|+\sigma_{B}\|\theta_{t}-\theta_{t-1}\|\\\
&\leq 2\sigma_{B}\|\theta_{t}-\theta_{t-1}\|,\end{split}$ (24)

where the first inequality is due to the Lipschitzness of $\nabla J(\theta)$,
and the second inequality results from the bound on
$\|B(\tau_{t}^{b};\theta_{t}^{b})\|$.

Summing the both sides of equation 23 from $t=1$ to $t=T$, we have:

$\mathbb{E}\left[\sum_{t=1}^{T}\|\epsilon_{t-1}\|^{2}\right]\leq-\frac{\mathbb{E}[\|\epsilon_{T}\|^{2}]}{\alpha_{T}}+\frac{\mathbb{E}[\|\epsilon_{0}\|^{2}]}{\alpha_{1}}+\underbrace{\sum_{t=1}^{T-1}\left(\frac{1}{\alpha_{t+1}}-\frac{1}{\alpha_{t}}\right)\mathbb{E}[\|\epsilon_{t}\|^{2}]}_{\text{(I)}}+8\sigma_{B}^{2}\underbrace{\sum_{t=1}^{T}\frac{\eta_{t-1}^{2}}{\alpha_{t}}}_{\text{(II)}}+2\sigma_{g}^{2}\underbrace{\sum_{t=1}^{T}\alpha_{t}}_{\text{(III)}}$
(25)

First note that
$\mathbb{E}[\|\epsilon_{0}\|^{2}]/\alpha_{1}\leq\sigma_{g}^{2}/\alpha_{0}$.
Now, we bound the other terms in the right hand side of the above inequality:
(I): For the coefficient in the sum, we have:
$1/\alpha_{t+1}-1/\alpha_{t}=((t+1)^{2/3}-t^{2/3})/\alpha_{0}\leq
2t^{-1/3}/(3\alpha_{0})\leq 2/(3\alpha_{0})$ where we used the gradient
inequality for the concave function $f(z)=z^{2/3}$. Therefore, this term can
be bounded by:
$(2/(3\alpha_{0}))\sum_{t=1}^{T-1}\mathbb{E}[\|\epsilon_{t}\|^{2}]$.
(II):
$\sum_{t=1}^{T}\eta_{t-1}^{2}/\alpha_{t}=\eta_{0}^{2}/\alpha_{0}(1+\sum_{t=1}^{T-1}(t+1)^{2/3}/t^{4/3})\leq\eta_{0}^{2}/\alpha_{0}(1+2^{2/3}\sum_{t=1}^{T-1}t^{-2/3})\leq
6\eta_{0}^{2}T^{1/3}/\alpha_{0}$.
(III): $\sum_{t=1}^{T}\alpha_{t}=\alpha_{0}\sum_{t=1}^{T}t^{-2/3}\leq
3\alpha_{0}T^{1/3}$.

Plugging these bounds in equation 25, we get:

$\mathbb{E}\left[\sum_{t=0}^{T-1}\|\epsilon_{t}\|^{2}\right]\leq CT^{1/3},$
(26)

where
$C:=3\alpha_{0}((48\sigma_{B}^{2}\eta_{0}^{2}+\sigma_{g}^{2})/\alpha_{0}+6\alpha_{0}\sigma_{g}^{2})/(3\alpha_{0}-2)$.

The previous inequality yields that:

$\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\epsilon_{t}\|]\leq\frac{1}{T}\sum_{t=0}^{T-1}\sqrt{\mathbb{E}[\|\epsilon_{t}\|^{2}]}\leq\sqrt{\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\epsilon_{t}\|^{2}]}\leq\frac{\sqrt{C}}{T^{1/3}},$
(27)

where we used Jensen’s inequality in the second inequality above.

Now, if we average from $t=0$ to $t=T-1$ in equation 32 in Lemma 2 and then
take expectations, we have:

$\begin{split}\mathbb{E}\left[\frac{1}{T}\sum_{t=0}^{T-1}\|\nabla
J(\theta_{t})\|\right]&\leq\frac{8}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\epsilon_{t}\|]+\frac{3\sigma_{B}}{2T}\sum_{t=0}^{T-1}\eta_{t}+\frac{3}{T}\mathbb{E}\left[\sum_{t=0}^{T-1}\frac{J(\theta_{t+1})-J(\theta_{t})}{\eta_{t}}\right]\\\
&\overset{(a)}{\leq}\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\mathbb{E}\left[\frac{3}{T}\sum_{t=0}^{T-1}\frac{J(\theta_{t+1})-J(\theta_{t})}{\eta_{t}}\right]\\\
&\leq\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\mathbb{E}\left[\frac{3}{T}\left(\frac{J(\theta_{T})}{\eta_{T-1}}-\frac{J(\theta_{0})}{\eta_{0}}+\sum_{t=1}^{T-1}J(\theta_{t})\left(\frac{1}{\eta_{t-1}}-\frac{1}{\eta_{t}}\right)\right)\right]\\\
&\overset{(b)}{\leq}\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\mathbb{E}\left[\frac{3}{T}\left(\frac{C_{J}(T-1)^{2/3}}{\eta_{0}}+\frac{C_{J}}{\eta_{0}}+\sum_{t=1}^{T-1}|J(\theta_{t})|\left(\frac{1}{\eta_{t}}-\frac{1}{\eta_{t-1}}\right)\right)\right]\\\
&\leq\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\mathbb{E}\left[\frac{3}{T}\left(\frac{C_{J}(T-1)^{2/3}}{\eta_{0}}+\frac{C_{J}}{\eta_{0}}+\sum_{t=1}^{T-1}C_{J}\left(\frac{1}{\eta_{t}}-\frac{1}{\eta_{t-1}}\right)\right)\right]\\\
&\overset{(c)}{\leq}\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\mathbb{E}\left[\frac{3}{T}\left(\frac{C_{J}(T-1)^{2/3}}{\eta_{0}}+\frac{C_{J}}{\eta_{0}}+\frac{C_{J}(T-1)^{2/3}}{\eta_{0}}\right)\right]\\\
&\leq\frac{8\sqrt{C}}{T^{1/3}}+\frac{6\sigma_{B}\eta_{0}}{T^{2/3}}+\frac{9C_{J}}{\eta_{0}T^{1/3}},\end{split}$
(28)

$(a)$ We used the bound in equation 27 for the first term and
$\sum_{t=0}^{T-1}\eta_{t}\leq 4\eta_{0}T^{1/3}$.
$(b)$
$|J(\theta)|=|\mathbb{E}_{\tau\sim\pi_{\theta}}[R(\tau)]|=|\mathbb{E}_{\tau\sim\pi_{\theta}}[\sum_{h=0}^{H-1}\gamma^{h}r(s_{h},a_{h})]|\leq\mathbb{E}_{\tau\sim\pi_{\theta}}[\sum_{h=0}^{H-1}\gamma^{h}|r(s_{h},a_{h})|]\leq\mathbb{E}_{\tau\sim\pi_{\theta}}[R_{0}/(1-\gamma)]=R_{0}/(1-\gamma)$.
Hence, we have: $|J(\theta)|\leq C_{J}$ for all $\theta\in\mathbb{R}^{d}$
where $C_{J}:=R_{0}/(1-\gamma)$.
$(c)$ Due to
$\sum_{t=1}^{T-1}1/\eta_{t}-1/\eta_{t-1}=1/\eta_{T-1}-1/\eta_{0}\leq(T-1)^{2/3}/\eta_{0}$.

##  Appendix C Proof of Remark 3

According to equation 21, we have:

$\begin{split}\mathbb{E}[\|\epsilon_{t}\|^{2}]\leq(1-\alpha_{t})\mathbb{E}[\|\epsilon_{t-1}\|^{2}]+8\sigma_{B}^{2}\eta_{t-1}^{2}+2\alpha^{2}_{t}\sigma_{g}^{2}.\end{split}$
(29)

Let us define $Z_{t}:=\mathbb{E}[\|\epsilon_{t}\|^{2}]$. Then, we can rewrite
the above equation as follows:

$Z_{t}\leq(1-\alpha_{t})Z_{t-1}+\frac{C_{Z}}{t^{4/3}},$ (30)

where $C_{z}=8\times
2^{4/3}\eta_{0}^{2}\sigma_{B}^{2}+2\alpha_{0}^{2}\sigma_{g}^{2}$. Now, by
induction, we will show that: $Z_{t}\leq C^{\prime}_{Z}/t^{2/3}$ where
$C^{\prime}_{Z}=C_{Z}/(\alpha_{0}-2/3)$. It can be easily seen that for the
base case $Z_{1}$, this statement holds. Now, for the induction step, suppose
that $Z_{t-1}\leq C^{\prime}_{Z}/(t-1)^{2/3}$ for some $t\geq 2$. Then, from
above equation, we have:

$\begin{split}\frac{C_{Z}}{t^{4/3}}&\overset{(a)}{\leq}\frac{\alpha_{0}C^{\prime}_{Z}-2C^{\prime}_{Z}/3}{t^{2/3}(t-1)^{2/3}}\\\
\implies&\frac{2C^{\prime}_{z}/3}{(t-1)t^{2/3}}\leq\frac{\alpha_{0}C^{\prime}_{Z}}{t^{2/3}(t-1)^{2/3}}-\frac{C_{Z}}{t^{4/3}}\\\
\overset{(b)}{\implies}&C^{\prime}_{Z}\left(\frac{1}{(t-1)^{2/3}}-\frac{1}{t^{2/3}}\right)\leq\frac{\alpha_{0}C^{\prime}_{Z}}{t^{2/3}(t-1)^{2/3}}-\frac{C_{Z}}{t^{4/3}}\\\
\implies&\left(1-\frac{\alpha_{0}}{t^{2/3}}\right)\frac{C^{\prime}_{Z}}{(t-1)^{2/3}}+\frac{C_{Z}}{t^{4/3}}\leq\frac{C^{\prime}_{Z}}{t^{2/3}}\\\
\overset{(c)}{\implies}&(1-\alpha_{t})Z_{t-1}+\frac{C_{Z}}{t^{4/3}}\leq\frac{C^{\prime}_{Z}}{t^{2/3}}\\\
\overset{(d)}{\implies}&Z_{t}\leq\frac{C^{\prime}_{Z}}{t^{2/3}},\end{split}$
(31)

where $(a)$ is due to definition of $C^{\prime}_{Z}$, $(b)$ is based on using
gradient inequality for the concave function $f(z)=z^{2/3}$, i.e.,
$t^{2/3}-(t-1)^{2/3}\leq 2/3(t-1)^{-1/3}$, $(c)$ is according to induction
hypothesis, and $(d)$ is due to equation 30.

##  Appendix D Supplemental Lemma

###### Lemma 2

Suppose that $\theta_{t}$’s are generated by executing Algorithm 2. Let
$\epsilon_{t}:=v_{t}-\nabla J(\theta_{t})$. Then, at any time $t$, we have:

$\|\nabla J(\theta_{t})\|\leq
8\|\epsilon_{t}\|+\frac{3\sigma_{B}\eta_{t}}{2}+\frac{3}{\eta_{t}}(J(\theta_{t+1})-J(\theta_{t})).$
(32)

From $\sigma_{B}$-smoothness of $J(\theta_{t})$ (Shen et al., 2019), we have:

$\begin{split}J(\theta_{t+1})-J(\theta_{t})&\geq\langle\nabla
J(\theta_{t}),\theta_{t+1}-\theta_{t}\rangle-\frac{\sigma_{B}}{2}\|\theta_{t+1}-\theta_{t}\|^{2}\\\
&=\eta_{t}\left\langle\nabla
J(\theta_{t}),\frac{v_{t}}{\|v_{t}\|}\right\rangle-\frac{\sigma_{B}}{2}\eta_{t}^{2}.\end{split}$
(33)

Regarding the first term above, we consider two cases, whether $\|\nabla
J(\theta_{t})\|\geq 2\|\epsilon_{t}\|$ or not. For the former case, we have:

$\begin{split}\left\langle\nabla
J(\theta_{t}),\frac{v_{t}}{\|v_{t}\|}\right\rangle&=\frac{\|\nabla
J(\theta_{t})\|^{2}+\langle\nabla J(\theta_{t}),\epsilon_{t}\rangle}{\|\nabla
J(\theta_{t})+\epsilon_{t}\|}\\\ &\geq\frac{\|\nabla
J(\theta_{t})\|^{2}}{2\|\nabla J(\theta_{t})+\epsilon_{t}\|}\\\
&\geq\frac{\|\nabla J(\theta_{t})\|^{2}}{2(\|\nabla
J(\theta_{t})\|+\|\epsilon_{t}\|)}\\\ &\geq\frac{\|\nabla
J(\theta_{t})\|}{3}\\\ &\geq\frac{\|\nabla
J(\theta_{t})\|}{3}-\frac{8}{3}\|\epsilon_{t}\|,\end{split}$ (34)

where in the first inequality, we used the bound $\langle\nabla
J(\theta_{t}),\epsilon_{t}\rangle\geq-\|\nabla
J(\theta_{t})\|\|\epsilon_{t}\|\geq-\|\nabla J(\theta_{t})\|^{2}/2$. For the
latter case,

$\begin{split}\left\langle\nabla
J(\theta_{t}),\frac{v_{t}}{\|v_{t}\|}\right\rangle&\geq-\|\nabla
J(\theta_{t})\|\\\ &=\frac{\|\nabla J(\theta_{t})\|}{3}-\frac{4\|\nabla
J(\theta_{t})\|}{3}\\\ &\geq\frac{\|\nabla
J(\theta_{t})\|}{3}-\frac{8\|\epsilon_{t}\|}{3}.\end{split}$ (35)

Plugging the bound on $\left\langle\nabla
J(\theta_{t}),v_{t}/\|v_{t}\|\right\rangle$ in equation 33, we get:

$\begin{split}&J(\theta_{t+1})-J(\theta_{t})\geq\eta_{t}\left(\frac{\|\nabla
J(\theta_{t})\|}{3}-\frac{8\|\epsilon_{t}\|}{3}\right)-\frac{\sigma_{B}}{2}\eta_{t}^{2}\\\
&\implies\|\nabla J(\theta_{t})\|\leq
8\|\epsilon_{t}\|+\frac{3\sigma_{B}\eta_{t}}{2}+\frac{3}{\eta_{t}}(\nabla
J(\theta_{t+1})-\nabla J(\theta_{t})).\end{split}$ (36)

##  Appendix E Details of Experiments

We used the default implementation of linear feature baseline from Garage
library. The employed linear feature baseline is a linear regression model
that takes observations for each trajectory and extracts new features such as
different powers of their lengths from the observations. These extracted
features are concatenated to the observations and used to fit the parameters
of the regression with least square loss function. In the experiments, we used
linear baseline for all the environments and methods.

In the following table, we provide the fine-tuned parameters for each
algorithm. Batch sizes are considered the same for all algorithms. The
discount factor is also set to $0.99$ for all the runs.

| Reacher | Walker | Humanoid | Swimmer
---|---|---|---|---
Max horizon | $50$ | $500$ | $500$ | $500$
Neural network sizes | $64\times 64$ | $64\times 64$ | $64\times 64$ | $64\times 64$
Activation functions | Tanh | Tanh | Tanh | Tanh
IS-MBPG $lr$ | 0.9 | 0.3 | 0.75 | 0.3
IS-MBPG $c$ | 100 | 12 | 5 | 12
IS-MBPG $w$ | 200 | 20 | 2 | 25
REINFORCE step-size | 0.01 | 0.01 | 0.001 | 0.01
SHARP $\alpha_{0}$ | 1.5 | 5 | 5 | 3
SHARP $\eta_{0}$ | 0.1 | 1 | 0.6 | 0.5
PAGE-PG $p_{t}$ | 0.4 | 0.4 | 0.6 | 0.6
PAGE-PG step-size | 0.01 | 0.001 | 0.0005 | 0.001
HAPG step-size | 0.01 | 0.01 | 0.01 | 0.01
HAPG $Q$ | 5 | 10 | 10 | 10
VR-BGPO $lr$ | 0.8 | 0.75 | 0.8 | 0.75
VR-BGPO $c$ | 25 | 25 | 25 | 25
VR-BGPO $w$ | 1 | 1 | 1 | 1
VR-BGPO $lam$ | 0.0005 | 0.0005 | 0.0025 | 0.0005
Table 3: Selected hyper-parameters for different methods.


# TEM: High Utility Metric Differential Privacy on Text

Ricardo Silva Carvalho  Theodore Vasiloudis  Oluwaseyi Feyisetan

###### Abstract

Ensuring the privacy of users whose data are used to train Natural Language
Processing (NLP) models is necessary to build and maintain customer trust.
Differential Privacy (DP) has emerged as the most successful method to protect
the privacy of individuals. However, applying DP to the NLP domain comes with
unique challenges. The most successful previous methods use a generalization
of DP for metric spaces, and apply the privatization by adding noise to inputs
in the metric space of word embeddings. However, these methods assume that one
specific distance measure is being used, ignore the density of the space
around the input, and assume the embeddings used have been trained on non-
sensitive data. In this work we propose Truncated Exponential Mechanism (TEM),
a general method that allows the privatization of words using any distance
metric, on embeddings that can be trained on sensitive data. Our method makes
use of the exponential mechanism to turn the privatization step into a
_selection problem_. This allows the noise applied to be calibrated to the
density of the embedding space around the input, and makes domain adaptation
possible for the embeddings. In our experiments, we demonstrate that our
method significantly outperforms the state-of-the-art in terms of utility for
the same level of privacy, while providing more flexibility in the metric
selection.

Differential Privacy, NLP

##  1 Introduction

Nowadays, text data are being used as input for a wide variety of machine
learning tasks, from next word prediction in mobile keyboards [9], to critical
tasks like predicting patient health conditions from clinical records [21].
Researchers have demonstrated that simple exploratory analysis tasks [13] or
the use of models trained on sensitive data [19] may breach the privacy of the
individuals involved. Even though there has been research focused on specific
privacy-preserving tasks with textual data, such as language models [15, 1,
20], to the best of our knowledge only a more recent line of work [6, 7, 8]
has been focusing on providing quantifiable privacy guarantees over the text
itself.

In this work we aim to provide privacy guarantees to textual data using the
formal notion of differential privacy (DP) [5], one of the most widely adopted
privacy frameworks in industry and academia. Specifically, we apply a
generalization called metric-differential privacy [2], which allows analysts
to tailor solutions over general distance metrics. Previous work [6, 7, 8] in
this setting considered each input word by its vector representation and added
noise to provide privacy guarantees. Additionally, as a noisy vector is
unlikely to exactly represent a valid word, these methods returned a nearest
neighbor approximation after querying the representation space. However, these
works consider the representation space as non-sensitive, as they do not
account for privacy loss in the nearest neighbor search. Moreover, the noise
added to a vector does not take into account the density of the region that
the vector lies in, which can potentially reduce the utility of a DP
algorithm.

Our main contribution is the design of a new mechanism which we call the
Truncated Exponential Mechanism (TEM,) that satisfies metric-DP over textual
data, posing the task as a selection problem. Instead of perturbing a
representation vector, our method selects an output from a set of possible
candidates where words closer to the input word in the metric space have
higher probability of being selected. TEM works by adapting the probabilities
of words being selected to the regions of a given input word, adjusting the
noise injected for better utility, and allows the application of any formal
distance function as metric. The mechanism includes a formal construction with
a truncation step to initially select from high utility words with high
probability, providing computationally efficient word selection with a tunable
error parameter. Our experiments show that TEM obtains higher utility when
compared to the state-of-the-art, for the same level of privacy.

##  2 Related Work

Fernandes et al. [6] worked on text data via the “bag of words” representation
of documents, and applied the Earth Mover’s metric to obtain privatized bags,
performing individual word privatization in the context of metric differential
privacy. Following this context, the Madlib111We refer to this algorithm with
the name used in [3] mechanism [8] adds noise to embedding vectors of words,
working on in the Euclidean space and adding Laplacian noise to the embedding
vectors. After introducing noise, the mechanism outputs the word that is
closest to the noisy vector in the embedding space. The algorithm presented in
[7] is a follow-up to [8] although it appeared later. This mechanism works in
a hierarchical embedding space, where the embedding vector of an input word is
perturbed with noise from a hyperbolic distribution.

Feyisetan et al. [7] compare the hyperbolic mechanism to Madlib [8]. However,
since the two algorithms use different metric functions, the evaluation of
privacy via only matching the $\varepsilon$ parameter of differential privacy
can be improved. In this sense, Feyisetan et al. [7] compares the privacy of
the two mechanisms, looking at the probability of not changing a word after
noise injection, i.e. the probability that the mechanism returns the exact
same word used as input. Even though this notion can be intuitively seen as a
level of indistinguishability, it cannot guarantee a fair comparison between
mechanisms. The issue of comparing metric-DP mechanisms with different metric
functions thus remains an open problem. In this work, we only compare
mechanisms using the same metric function (Euclidean distance) to ensure a
fair comparison.

##  3 Preliminaries

Consider a user giving as input a word $w$ from a discrete fixed domain
$\mathcal{W}$. For any pair of inputs $w$ and $w^{\prime}$, we assume a
distance function $d:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}_{+}$,
in a given space of representation of these words. More specifically, we
consider a word embedding model $\phi:\mathcal{W}\rightarrow\mathbb{R}^{n}$
will be used to represent words, and the distance function can be a valid
metric applicable to the embedding vectors.

Our goal is to select a word from $\mathcal{W}$, based on a given input, such
that the privacy of a user, with respect to this word choice, is preserved.
From an attacker’s perspective, the output of an algorithm working over input
$w$ or $w^{\prime}$ will become more similar as these inputs become closer
with respect to $d(w,w^{\prime})$. Intuitively, words that are distant in
metric space will be more easily distinguishable, compared to words that are
close.

With that in mind, we will work on Metric-Differential Privacy [2], a privacy
standard defined for randomized algorithms with input from a domain
$\mathcal{W}$ that are equipped with a distance metric
$d:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}_{+}$ satisfying the
formal axioms of a metric. In this context, the privacy guarantees given by
metric-DP depend not only on the privacy parameter $\varepsilon$, but also the
distance metric $d$ used.

######  Definition 3.1.

(Metric Differential privacy [2]). Given a distance metric
$d:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}_{+}$, a randomized
mechanism $\mathcal{M}:\mathcal{W}\rightarrow\mathcal{Y}$ is $\varepsilon
d$-differentially private if for any $w,w^{\prime}\in\mathcal{W}$ and all
outputs $y\in\mathcal{Y}$ we have:

$\displaystyle\Pr[\mathcal{M}(w)=y]\leq e^{\varepsilon
d(w,w^{\prime})}\Pr[\mathcal{M}(w^{\prime})=y]$ (1)

For the Euclidean distance metric, as discussed on Section 2, the current
state-of-the-art is the Madlib mechanism, which adds Laplacian noise to a
given vector in order to obtain a private output.

Algorithm 1 \- Madlib: Word Privatization Mechanism for Metric Differential
Privacy

Input: Finite domain $\mathcal{W}$, input word $w\in\mathcal{W}$ and privacy
parameter $\varepsilon$.
Output: Privatized element.

1: Compute embedding $\phi_{w}=\phi(w)$

2: Perturb embedding to obtain $\hat{\phi}_{w}=\phi_{w}+N$ with noise density
$p_{N}(z)\propto\exp(-\varepsilon\left\lVert z\right\rVert)$

3: Return perturbed word
$\hat{w}=argmin_{y\in\mathcal{W}}||\phi(y)-\hat{\phi}_{w}||$

For a Euclidean metric
$d:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}_{+}$, Madlib provides
metric differential privacy.

######  Theorem 1.

For a Euclidean metric $d$, Algorithm 1 is $\varepsilon d$-differentially
private.

Next we describe our algorithm that satisfies metric-DP, giving formal proof
of its privacy guarantees.

##  4 Metric Truncated Exponential Mechanism

At its core, our algorithm uses the Exponential Mechanism (EM) [16], which is
often used for selection in the context of differential privacy[5].

Algorithm 2 \- TEM: Metric Truncated Exponential Mechanism

Input: Finite domain $\mathcal{W}$, input word $w\in\mathcal{W}$, truncation
threshold $\gamma$, metric
$d_{\mathcal{W}}:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}_{+}$, and
privacy parameter $\varepsilon$.
Output: Privatized element.

1: Given input $w$, obtain the set $L_{w}$ such that each word $w_{i}\in
L_{w}$ satisfies $d_{\mathcal{W}}(w,w_{i})\leq\gamma$

2: Set the score $f(w,w_{i})$ of each $w_{i}\in L_{w}$ as
$f(w,w_{i})=-d_{\mathcal{W}}(w,w_{i})$

3: Create a $\bot$ element with score
$f(w,\bot)=-\gamma+2\ln(|\mathcal{W}\setminus L_{w}|)/\varepsilon$

4: For each word $w_{i}\in L(x)\cup\bot$, add Gumbel noise with mean 0 and
scale $2/\varepsilon$ to score $f(w,w_{i})$

5: Select $\hat{w}$ as the element with maximum noisy score from
$L(x)\cup\bot$

6: if $\hat{w}=\perp$ then return random sample of $\mathcal{W}\setminus
L_{w}$

7: else return $\hat{w}$

Algorithm 2, denoted as TEM, is using a variant of the exponential mechanism
with Gumbel noise [4], and more specifically adapted to metric-DP for any
given distance metric. TEM starts by selecting from words closer to the input
by a distance of less or equal than a threshold $\gamma$, also including a
$\bot$ element to account for the words outside the $\gamma$ distance. Privacy
proof is deferred to Appendix A.1.

######  Theorem 2.

For any formal distance metric $d$, Algorithm 2 is $\varepsilon
d$-differentially private.

To also optimize for utility, next we give a theoretical proposition for
defining $\gamma$ in order to select elements within $\gamma$ distance from
the input with high probability. We defer proof to Appendix A.2.

######  Theorem 3.

For $\beta>0$ and input $w\in\mathcal{W}$, TEM outputs elements with distance
less or equal than $\gamma$ from $w$ with probability at least $1-\beta$ for
$\gamma\geq\frac{2}{\varepsilon}\cdot\ln\frac{(1-\beta)(|\mathcal{W}|-1)}{\beta}$.

The result above gives a guarantee of outputting words that are close to the
input $w$ with high probability for a given $\gamma$ distance threshold. Thus,
it is a theoretical way to choose $\gamma$ without looking at the data, i.e.
without incurring privacy loss, while getting utility guarantees with high
probability.

Below we highlight some of the advantages of TEM, specifically compared to the
state-of-the-art.

###  4.1 Detailed comparison with previous work

TEM works for any given distance function that satisfies the axioms of a
metric. In this sense, it has advantages for future use, when compared to
previous mechanisms that used fixed metrics, such as Euclidean [8] and
Hyperbolic [7], where changing the metric would need additional privacy
analysis.

Previous work [6, 8, 7] considered the text privacy preservation problem
mainly as a task of releasing a word embedding vector after perturbing with
some noise. This means they add noise to each of the dimensions of the vector
they aim to release, treating every word embedding vector the same way. In
practice, this leads to adding the same amount of noise for any word in the
embedding space, regardless of whether the word lies in a dense or sparse
region. In contrast, TEM preserves privacy by posing the task as a selection
problem, giving words closer to the input word a higher probability of being
selected. Therefore, TEM has a more dynamic behaviour, adjusting the noise to
the domain of selection. In practice, for a given $\varepsilon$, TEM will add
less noise to regions with high density, and more noise to regions with low
density, therefore offering better utility in high-density areas.

Moreover, previous work assumed the word embeddings used were trained on a
separate dataset, distinct from the data being privatized. TEM does not have
that requirement, providing the potential for further utility gains through
the use of domain adaptation, i.e. fine-tuning public pre-trained embeddings
on the target sensitive data [18, 10].

In terms of computational cost, the bottleneck of previous work resides in the
nearest neighbor search of the noisy embedding vector obtained from the input.
TEM starts by getting the elements within distance $\gamma$ of the input, but
instead of querying the nearest neighbor, it queries for neighbors within a
given range. In this sense, both methods can rely on fast approximate nearest
neighbors implementations that support both querying nearest and by range,
such as [11]. Nonetheless, TEM is the only mechanism that, for a fixed domain,
is able to pre-process and store the search results for a given $\gamma$.
After this step, the range search cost becomes constant. For the interested
reader, we include a rewriting of Algorithm 2 with a pre-processing step on
Appendix B.

Next we empirically compare our mechanism with the state of the art Madlib [8]
mechanism.

##  5 Experiments

For fair comparison to Madlib we use TEM with Euclidean distance on a fixed
embedding space from GloVe [17]. Experiments use the IMDB reviews dataset
[14]. More details deferred to Appendix C.

Utility: To evaluate the utility of the metric-DP mechanisms, we build
sentiment classification models on training data privatized by each mechanism
and the baseline trained on sensitive data, and compare the accuracy of the
trained models on a test dataset.

Privacy: As both mechanisms use the Euclidean distance metric, their privacy
guarantees are matched by using the same $\varepsilon$. Nonetheless, for
illustration we include the results of a Membership Inference Attack (MIA)
[19], which tries to infer the presence of observations used to train a given
model based only on black-box access. Lower attack score is better,
representing more privacy preservation.

(a) Utility Evaluation

(b) Empirical Privacy Evaluation

Figure 1: Comparison of mechanisms with 95% confidence interval over 5 trials
for various $\varepsilon$. Baseline is built with models trained on original
data. TEM used $\gamma$ from Theorem 3 with $\beta=0.001$.

From the results of Figure 1(a) above we see that for fixed privacy level, TEM
outperforms Madlib significantly. More specifically, we see from TEM’s results
that $\varepsilon\geq 4$ gives a formal level of privacy that is not
significant. However, since Madlib adds more noise than needed for different
regions of the embedding space, it still achieves some empirical privacy, as
observed by the MIA results on Figure 1(b). Nonetheless, such enforcement is
not a formal guarantee of privacy, which for metric-DP is bounded by
$\varepsilon d(w,w^{\prime})$, therefore in this context it is a loose
guarantee of privacy for the embedding space considered. In this sense, when
comparing TEM and Madlib for metric-DP formal levels of privacy, i.e. same
$\varepsilon$ and metric space, we can clearly see better utility for TEM.
Finally, as an example, if we look at $\varepsilon=2$, where both mechanisms
have $\text{AUC}=0.50$ for MIA, we see Madlib with average test accuracy of
$52\%$ and TEM with $75\%$, representing a relative utility improvement of
$42\%$.

##  6 Conclusion

We presented TEM, a mechanism for text privatization on metric differential
privacy with formal guarantees. Unlike the current state-of-the-art, our
method allows the safe use of sensitive embeddings and provides flexibility of
metric definition. In addition, TEM adapts the noise introduced to improve
utility. Finally, it gives the possibility of performing pre-processing steps
for enhanced computational efficiency. Our empirical evaluation demonstrates
that TEM obtains better utility than the current state-of-the-art for the same
formal privacy guarantees. As future work, we plan to perform domain
adaptation, in order to leverage embeddings trained on sensitive data to
improve utility. Including privatization of word context vectors is also a
possible enhancement for improved accuracy.

## References

  * Bo et al. [2019] Haohan Bo, Steven HH Ding, Benjamin Fung, and Farkhund Iqbal.  Er-ae: Differentially-private text generation for authorship anonymization.  _arXiv_ , pages arXiv–1907, 2019.
  * Chatzikokolakis et al. [2013] Konstantinos Chatzikokolakis, Miguel E Andrés, Nicolás Emilio Bordenabe, and Catuscia Palamidessi.  Broadening the scope of differential privacy using metrics.  In _International Symposium on Privacy Enhancing Technologies Symposium_ , pages 82–102. Springer, 2013.
  * Diethe [2020] Thomas Diethe.  Preserving privacy in analyses of textual data, 2020.  URL https://www.amazon.science/blog/preserving-privacy-in-analyses-of-textual-data.
  * Durfee and Rogers [2019] David Durfee and Ryan M Rogers.  Practical differentially private top-k selection with pay-what-you-get composition.  In _NeurIPS_ , pages 3527–3537, 2019.
  * Dwork et al. [2006] Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith.  Calibrating noise to sensitivity in private data analysis.  In _Theory of cryptography conference_ , pages 265–284. Springer, 2006.
  * Fernandes et al. [2018] Natasha Fernandes, Mark Dras, and Annabelle McIver.  Author obfuscation using generalised differential privacy.  _CoRR_ , abs/1805.08866, 2018.  URL http://arxiv.org/abs/1805.08866.
  * Feyisetan et al. [2019] Oluwaseyi Feyisetan, Tom Diethe, and Thomas Drake.  Leveraging hierarchical representations for preserving privacy and utility in text.  In _2019 IEEE International Conference on Data Mining (ICDM)_ , pages 210–219. IEEE, 2019.
  * Feyisetan et al. [2020] Oluwaseyi Feyisetan, Borja Balle, Thomas Drake, and Tom Diethe.  Privacy- and utility-preserving textual analysis via calibrated multivariate perturbations.  In _Proceedings of the 13th International Conference on Web Search and Data Mining_ , WSDM ’20, page 178–186, New York, NY, USA, 2020. Association for Computing Machinery.  ISBN 9781450368223.  doi: 10.1145/3336191.3371856.  URL https://doi.org/10.1145/3336191.3371856.
  * Hard et al. [2018] Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage.  Federated learning for mobile keyboard prediction.  _arXiv preprint arXiv:1811.03604_ , 2018.
  * Jaech et al. [2016] Aaron Jaech, Larry Heck, and Mari Ostendorf.  Domain adaptation of recurrent neural networks for natural language understanding.  _Interspeech 2016_ , pages 690–694, 2016.
  * Johnson et al. [2017] Jeff Johnson, Matthijs Douze, and Hervé Jégou.  Billion-scale similarity search with gpus.  _arXiv preprint arXiv:1702.08734_ , 2017.
  * Joulin et al. [2017] Armand Joulin, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov.  Bag of tricks for efficient text classification.  _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers_ , 2017.  doi: 10.18653/v1/e17-2068.  URL http://dx.doi.org/10.18653/V1/E17-2068.
  * Kantarcioǧlu et al. [2004] Murat Kantarcioǧlu, Jiashun Jin, and Chris Clifton.  When do data mining results violate privacy?  In _SIGKDD_ , pages 599–604. ACM, 2004.
  * Maas et al. [2011] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts.  Learning word vectors for sentiment analysis.  In _Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies_ , pages 142–150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.  URL http://www.aclweb.org/anthology/P11-1015.
  * McMahan et al. [2017] H. Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang.  Learning differentially private recurrent language models, 2017.
  * McSherry and Talwar [2007] Frank McSherry and Kunal Talwar.  Mechanism design via differential privacy.  In _FOCS_ , volume 7, pages 94–103, 2007.
  * Pennington et al. [2014] Jeffrey Pennington, Richard Socher, and Christopher D Manning.  Glove: Global vectors for word representation.  In _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , pages 1532–1543, 2014.
  * Plank and Moschitti [2013] Barbara Plank and Alessandro Moschitti.  Embedding semantic similarity in tree kernels for domain adaptation of relation extraction.  In _Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1498–1507, 2013\.
  * Shokri et al. [2017] Reza Shokri, Marco Stronati, Congzheng Song, and Vitaly Shmatikov.  Membership inference attacks against machine learning models.  In _2017 IEEE Symposium on Security and Privacy (SP)_ , pages 3–18. IEEE, 2017.
  * Vu et al. [2019] XS Vu, SN Tran, and L Jiang.  dpugc: Learn differentially private representation for user generated contents.  In _20th International Conference on Computational Linguistics and Intelligent Text Processing_ , pages 1–16, 2019.
  * Yao et al. [2019] Liang Yao, Chengsheng Mao, and Yuan Luo.  Clinical text classification with rule-based features and knowledge-guided convolutional neural networks.  _BMC medical informatics and decision making_ , 19(3):71, 2019.

##  Appendix A Omitted Proofs

In this section we include the omitted proofs.

###  A.1 Proof of Theorem 2

To prove the privacy of TEM, we need to first show that the sensitivity of the
score function is still $d_{\mathcal{W}}(w,w^{\prime})$ after the truncation.
So first we prove a Lemma giving this result on the sensitivity, and after
that we prove the privacy of our version of EM.

######  Lemma A.1.

The sensitivity $\Delta
f=\max\limits_{i\in\mathcal{W}}\max\limits_{w,w^{\prime}\in\mathcal{W}}|f(i,w)-f(i,w^{\prime})|\leq
d_{\mathcal{W}}(w,w^{\prime})$.

###### Proof.

For a given input $w$, let us denote $\mathcal{I}(w)$ as the domain elements
that have $d_{\mathcal{W}}(i,w)\leq\gamma$, therefore keeping their distances
on the score function, while the elements on
$\mathcal{W}\setminus\mathcal{I}(w)$ have distances fixed as $\gamma$.

In this context, there are four possible cases we need to analyze for a given
$i$ and any pair $w,w^{\prime}$:, which we dive deep into now.

Case 1: $i\in\mathcal{I}(w)$ and $i\in\mathcal{I}(w^{\prime})$. If $i$ is in
both $\mathcal{I}(w)$ and $\mathcal{I}(w^{\prime})$, then it is using its
original distance on the score for both $w$ and $w^{\prime}$. Thus we have:

$\displaystyle
f(i,w)-f(i,w^{\prime})=-d_{\mathcal{W}}(i,w)+d_{\mathcal{W}}(i,w^{\prime})\leq
d_{\mathcal{W}}(w,w^{\prime})$

with the last inequality being the use of the triangle inequality for the
distance metric $d_{\mathcal{W}}$.

Case 2: $i\in\mathcal{I}(w)$ and $i\notin\mathcal{I}(w^{\prime})$. If $i$ is
in $\mathcal{I}(w)$ but not in $\mathcal{I}(w^{\prime})$, then it is using its
original distance on the score for $w$ and $\gamma$ for $w^{\prime}$,
therefore:

$\displaystyle f(i,w)-f(i,w^{\prime})=-d_{\mathcal{W}}(i,w)+\gamma$

Since $i$ is not in $\mathcal{I}(w^{\prime})$, it means that
$d_{\mathcal{W}}(i,w^{\prime})>\gamma$, or equivalently
$\gamma<d_{\mathcal{W}}(i,w^{\prime})$, which replacing on the result above
gives:

$\displaystyle
f(i,w)-f(i,w^{\prime})=-d_{\mathcal{W}}(i,w)+d_{\mathcal{W}}(i,w^{\prime})\leq
d_{\mathcal{W}}(w,w^{\prime})$

where we use triangle inequality on the last step.

Case 3: $i\notin\mathcal{I}(w)$ and $i\in\mathcal{I}(w^{\prime})$. If $i$ is
not in $\mathcal{I}(w)$ but is in $\mathcal{I}(w^{\prime})$, then it is using
$\gamma$ as the distance on score for $w$ and the original distance on score
for $w^{\prime}$, which gives us:

$\displaystyle f(i,w)-f(i,w^{\prime})=-\gamma+d_{\mathcal{W}}(i,w^{\prime})$

Since $i$ is in $\mathcal{I}(w^{\prime})$, it means that
$d_{\mathcal{W}}(i,w^{\prime})\leq\gamma$, which replacing on the result above
shows:

$\displaystyle f(i,w)-f(i,w^{\prime})$
$\displaystyle=-\gamma+d_{\mathcal{W}}(i,w^{\prime})$
$\displaystyle\leq-\gamma+\gamma=0\leq d_{\mathcal{W}}(w,w^{\prime})$

Case 4: $i\notin\mathcal{I}(w)$ and $i\notin\mathcal{I}(w^{\prime})$. If $i$
is not in both $\mathcal{I}(w)$ and $\mathcal{I}(w^{\prime})$, then it is
using $\gamma$ as distance on score for both $w$ and $w^{\prime}$, giving:

$\displaystyle f(i,w)-f(i,w^{\prime})=-\gamma+\gamma=0\leq
d_{\mathcal{W}}(w,w^{\prime})$

Finally, we note that showing $f(i,w)-f(i,w^{\prime})\geq-
d_{\mathcal{W}}(w,w^{\prime})$ on the same cases above follows by symmetry.

∎

With the sensitivity result we can now show the privacy guarantee of the
mechanism.

Theorem 2. For any formal distance metric $d$, Algorithm 2 is $\varepsilon
d$-differentially private.

###### Proof.

For a given output $y\in\mathcal{W}$ and any pair of inputs
$w,w^{\prime}\in\mathcal{W}$ we have:

$\displaystyle\frac{\Pr[M(w)=y]}{\Pr[M(w^{\prime})=y]}=$
$\displaystyle\frac{\exp(\frac{\varepsilon}{2}\cdot
f(y,w))}{\exp(\frac{\varepsilon}{2}\cdot
f(y,w^{\prime}))}\cdot\frac{\sum\limits_{z^{\prime}\in\mathcal{I}(w^{\prime})}\exp(\frac{\varepsilon}{2}\cdot
f(z^{\prime},w^{\prime}))+\sum\limits_{w^{\prime}\in\mathcal{W}\setminus\mathcal{I}(w^{\prime})}\exp(\frac{\varepsilon}{2}\cdot\gamma)}{\sum\limits_{z\in\mathcal{I}(w)}\exp(\frac{\varepsilon}{2}\cdot
f(z,w))+\sum\limits_{w\in\mathcal{W}\setminus\mathcal{I}(w)}\exp(\frac{\varepsilon}{2}\cdot\gamma)}$
(2)

We note that on the second term above we have the same domain of elements,
which even though they do not match among summations, they still satisfy the
sensitivity on Lemma A.1, which gives us for the first term on the right-hand
side of Equation A.1:

$\displaystyle\frac{\exp(\frac{\varepsilon}{2}\cdot
f(y,w))}{\exp(\frac{\varepsilon}{2}\cdot f(y,w^{\prime}))}\leq$
$\displaystyle\frac{\exp(\frac{\varepsilon}{2}\cdot(f(y,w^{\prime})+d_{\mathcal{W}}(w,w^{\prime})))}{\exp(\frac{\varepsilon}{2}\cdot
f(y,w^{\prime}))}\leq\exp(\frac{\varepsilon}{2}\cdot
d_{\mathcal{W}}(w,w^{\prime}))$ (3)

And similarly for the second term on the right-hand side of Equation A.1:

$\displaystyle\frac{\sum\limits_{z^{\prime}\in\mathcal{I}(w^{\prime})}\exp(\frac{\varepsilon}{2}\cdot
f(z^{\prime},w^{\prime}))+\sum\limits_{w^{\prime}\in\mathcal{W}\setminus\mathcal{I}(w^{\prime})}\exp(\frac{\varepsilon}{2}\cdot\gamma)}{\sum\limits_{z\in\mathcal{I}(w)}\exp(\frac{\varepsilon}{2}\cdot
f(z,w))+\sum\limits_{w\in\mathcal{W}\setminus\mathcal{I}(w)}\exp(\frac{\varepsilon}{2}\cdot\gamma)}\leq\exp(\frac{\varepsilon}{2}d_{\mathcal{W}}(w,w^{\prime}))$

Therefore, multiplying the two inequalities above gives us for Equation A.1:

$\displaystyle\frac{\Pr[M(w)=y]}{\Pr[M(w^{\prime})=y]}\leq$
$\displaystyle\exp(\frac{\varepsilon}{2}d_{\mathcal{W}}(w,w^{\prime}))\cdot\exp(\frac{\varepsilon}{2}d_{\mathcal{W}}(w,w^{\prime}))=\exp(\varepsilon
d_{\mathcal{W}}(w,w^{\prime}))$ (4)

which proves the metric-DP guarantee of the mechanism.

The only difference in writing from TEM to what we used in the probabilities
here is that instead of directly picking from all the elements with distance
greater than $\gamma$ we use $\bot$ first. So now we show they are equivalent.

Let $L(w)$ be a list of elements where each element $i$ satisfies
$d_{\mathcal{W}}(i,w)\leq\gamma$ and let $\bar{L}(w)$ be the list of remaining
elements $\mathcal{W}\setminus L(w)$. For elements in $\bar{L}(w)$, we see
that on TEM they are selected randomly after $\bot$ is selected. Since $\bot$
has score $-\gamma+2\ln(|\bar{L}(w)|)/\varepsilon$, on the exponential
mechanism this is equivalent to:

$\displaystyle\exp(\varepsilon/2\cdot(-\gamma+2\ln(|\bar{L}(w)|)/\varepsilon))$
$\displaystyle=$
$\displaystyle\exp(\varepsilon/2\cdot-\gamma)\cdot\exp(\varepsilon/2\cdot(2\ln(|\bar{L}(w)|)/\varepsilon))$
$\displaystyle=$
$\displaystyle\exp(\varepsilon/2\cdot-\gamma)\cdot\exp(\ln(|\bar{L}(w)|)))$
$\displaystyle=$
$\displaystyle\exp(\varepsilon/2\cdot-\gamma)\cdot|\bar{L}(w)|$

This result is essentially the same as using $|\bar{L}(w)|$ elements with
score $-\gamma$ directly for selection since, after selecting $\bot$, the
elements in $\bar{L}(w)$ are selected randomly, giving each a probability of
selection proportional to $\exp(\varepsilon/2\cdot-\gamma)$.

∎

###  A.2 Proof of Theorem 3

The following proof for Theorem 3 only uses the fact that on an exponential
mechanism we have the probability of a given element proportional to the score
of the element.

Theorem 3. For $\beta>0$ and input $w\in\mathcal{W}$, TEM outputs elements
with distance less or equal than $\gamma$ from $w$ with probability at least
$1-\beta$ for
$\gamma\geq\frac{2}{\varepsilon}\cdot\ln\frac{(1-\beta)(|\mathcal{W}|-1)}{\beta}$.

###### Proof.

This theorem is equivalent to guaranteeing that we only output elements
outside the $\gamma$ distance of the input with probability at most $\beta$.
The worst case to guarantee this condition is when we only have the input word
inside the $\gamma$ distance, and all of the remaining $|\mathcal{W}|-1$ words
are outside. Thus we calculate the probability in this theorem for this worst
case to obtain the maximum guarantee.

$\displaystyle\frac{(|\mathcal{W}|-1)\cdot\exp(-\varepsilon/2\cdot\gamma)}{\exp(-\varepsilon/2\cdot
0)+(|\mathcal{W}|-1)\cdot\exp(-\varepsilon/2\cdot\gamma)}$
$\displaystyle\leq\beta$
$\displaystyle\frac{1}{1/\big{(}(|\mathcal{W}|-1)\cdot\exp(-\varepsilon/2\cdot\gamma)\big{)}+1}$
$\displaystyle\geq\beta$ $\displaystyle
1/\big{(}(|\mathcal{W}|-1)\cdot\exp(-\varepsilon/2\cdot\gamma)\big{)}+1$
$\displaystyle\leq 1/\beta$ $\displaystyle
1/\big{(}(|\mathcal{W}|-1)\cdot\exp(-\varepsilon/2\cdot\gamma)\big{)}$
$\displaystyle\geq(1-\beta)/\beta$
$\displaystyle(|\mathcal{W}|-1)\cdot(1-\beta)/\beta$
$\displaystyle\leq\exp(\varepsilon/2\cdot\gamma)$
$\displaystyle(2/\varepsilon)\cdot\ln((|\mathcal{W}|-1)\cdot(1-\beta)/\beta)$
$\displaystyle\leq\gamma$

∎

##  Appendix B Efficient Version

Here we precisely describe a version of our algorithm with exact search for
the elements within $\gamma$ distance of any input.

A first glance at TEM reveals it can be computationally prohibitive for large
domains, so below we develop a version of our mechanism that is
computationally efficient leveraging the properties of our algorithm.

For a fixed finite domain $\mathcal{W}$, given a truncation threshold
$\gamma$, we pre-compute, for each possible input, the list of elements that
satisfy $d_{\mathcal{W}}(i,w)\leq\gamma$, which makes the search for possible
candidates of a given input $O(1)$. Another simplification already included in
TEM is to use a version of the Exponential Mechanism that uses Gumbel noise
[4], which helps avoid dealing with probabilities using the exponential
function. Finally, we also point out that we can group all elements below the
truncation threshold as one element $\bot$ with the aggregated count, and then
if $\bot$ gets selected, randomly sample one of the aggregated elements. Below
we give such simplified algorithm and prove it is equivalent to Algorithm 2.

Algorithm 3 Metric Truncated Exponential Mechanism

Input: Finite domain $\mathcal{W}=\\{1,...,m\\}$ of elements, element index
$x\in\mathcal{W}$, truncation threshold $\gamma$, metric
$d_{\mathcal{W}}:\mathcal{W}\times\mathcal{W}\rightarrow\mathbb{R}$, and
privacy parameter $\varepsilon$.
Output: Element index.

1: Pre-processing:

2: for each $x\in\mathcal{W}$ do

3: Create a list $L(w)$ of elements where each element $i$ satisfies
$d_{\mathcal{W}}(i,w)\leq\gamma$ and let $\bar{L}(w)$ be the list of remaining
elements $\mathcal{W}\setminus L(w)$.

4: Define for each $x\in\mathcal{W}$ the score of a $\bot$ element as
$f_{\bot}(w)=-\gamma+2\ln(|\bar{L}(w)|)/\varepsilon$

5: end for

6: Selection:

7: Given input $x\in\mathcal{W}$, for every element in $L(w)\cup\bot$ add
noise from a Gumbel distribution with mean 0 and scale $2/\varepsilon$ to each
score $-d_{\mathcal{W}}(\cdot,x)$

8: Set $y$ as the element with maximum noisy score

9: if $y=\perp$ then

10: Return random sample of $\bar{L}(w)$

11: else

12: Return $y$

13: end if

We now formally prove Algorithm 3 is equivalent to Algorithm 2.

######  Lemma B.1.

Algorithm 3 and Algorithm 2 are equal in distribution.

###### Proof.

The only difference between the algorithms is that on Algorithm 3 we pre-
process $L(w)$ and $\bar{L}(w)$ ahead of time. Since for a fixed domain they
do not change and are independent for each input word, the algorithms are
equal in distribution.

∎

Since the two algorithms are equivalent, Algorithm 3 also satisfies the same
privacy guarantee.

######  Corollary B.1.

Algorithm 3 is $\varepsilon d_{\mathcal{W}}$-DP.

##  Appendix C Experiments details

In this section, we describe the exact settings and include all of the details
for the mechanisms used, to allow reproducibility.

We used the IMDB dataset [14], which gives two different files: training data
and testing data, each with 25.000 examples. For the baseline we trained a
model using dataset TR1 with 50% of the IMDB training data and tested it on
TE1 with 50% of the IMDB testing data. For the privatized utility, we trained
models on TR1 after privatization by each mechanism and tested them on
original TE1.

For MIA the models trained as described above were attacked, we denote a given
target model as T. For the shadow model, denoted as S, we trained a model on
dataset TE2 having the other 50% of the IMDB testing data. To train the attack
model, denoted as A, we used as features the output of TE1 and TE2 given by S,
where TE2 is labeled as ”in” and TE1 as ”out”. After training model A, we
evaluated the inference attack with TR1 and TR2 (having another 50% of the
IMDB training data) with features being the output of TR1 and TR2 obtained by
a given target model T, where ground-truth for TR1 is ”in” and for TR2 is
”out.

For embeddings we used GloVe [17] with 300 dimensions. The sentiment
classification models follow the FastText classifier [12], whereas the attack
model is an MLP with two-layers having 64 hidden nodes each, and ReLU
activations. Each model was trained for 20 epochs with batch size of 64 and
default pytorch parameters for Adam optimizer.