# Deep Reinforcement Learning for Radio Resource Allocation in NOMA-based
Remote State Estimation
††thanks: *Wanchun Liu is the corresponding author.

Gaoyang Pang, Wanchun Liu*, Yonghui Li, and Branka Vucetic  School of
Electrical and Information Engineering, The University of Sydney, Australia
Emails: {gaoyang.pang, wanchun.liu, yonghui.li<EMAIL_ADDRESS>
###### Abstract

Remote state estimation, where many sensors send their measurements of
distributed dynamic plants to a remote estimator over shared wireless
resources, is essential for mission-critical applications of Industry 4.0.
Most of the existing works on remote state estimation assumed orthogonal
multiple access and the proposed dynamic radio resource allocation algorithms
can only work for very small-scale settings. In this work, we consider a
remote estimation system with non-orthogonal multiple access. We formulate a
novel dynamic resource allocation problem for achieving the minimum overall
long-term average estimation mean-square error. Both the estimation quality
state and the channel quality state are taken into account for decision making
at each time. The problem has a large hybrid discrete and continuous action
space for joint channel assignment and power allocation. We propose a novel
action-space compression method and develop an advanced deep reinforcement
learning algorithm to solve the problem. Numerical results show that our
algorithm solves the resource allocation problem effectively, presents much
better scalability than the literature, and provides significant performance
gain compared to some benchmarks.

###### Index Terms:

Remote state estimation, radio resource allocation, NOMA, deep reinforcement
learning, task-oriented communications.

##  I Introduction

Wireless networked control systems (WNCSs), consisting of spatially
distributed plants, sensors, machines, actuators and controllers, play an
essential role in the era of Industry 4.0 [1]. In particular, remote state
estimators for monitoring dynamic plant status in a real-time manner are
critical in WNCSs to enable high-quality closed-loop control. In Industry 4.0,
massive wireless sensors are deployed for remote state estimation of spatially
distributed plants. Thus, it is essential to manage the limited wireless radio
resources for remote state estimation of many distributed plants.

Existing works on wireless resource allocation mainly focused on data-oriented
communications, and the design targets are transmission throughput, latency,
and reliability [2]. Advanced data-driven machine learning approaches, such as
supervised learning and reinforcement learning, have been adopted when
resource allocation problems cannot be solved effectively by conventional
model-based methods [3]. Different from data-oriented communications, the
resource allocation design in a remote estimation system should be task-
oriented as the goal is to minimize the long-term average remote estimation
mean-square error (MSE) of the dynamic plant states [2].

Significant efforts have been devoted to the wireless resource allocation of
remote estimation systems with orthogonal multiple access (OMA) (see [4, 5]
and the references therein), where each frequency channel (i.e., a subcarrier)
can only be assigned to a single sensor to avoid inter-user interference
completely. Non-orthogonal multiple access (NOMA) [6] allows simultaneous
transmission of multiple sensors’ packets at the same frequency channel and
offers higher transmission capacity than OMA. Each sensor packet can be
detected by processing received superimposed signals.

However, existing works on dynamic resource allocation in NOMA-based remote
estimation only considered the simple multi-sensor-single-channel setting, and
focused on either power allocation [7, 8, 9] or channel assignment problems
[10], rather than joint design ones. Moreover, the developed policy
optimization methods for resource allocation can only handle very small-scale
systems, e.g., a five-sensor-single-channel setting [8], due to the curse of
dimensionality in policy optimization. We aim to address these limitations of
NOMA-based remote estimation. The novel contributions of our work are
summarized below.

1) New system model. We investigate a NOMA-based remote estimation system with
multiple sensors and frequency channels. It requires a joint design of channel
assignment and power control of each sensor for achieving the optimal overall
remote estimation performance. Such a system has not been investigated before.

2) Novel problem formulation. We formulate the dynamic resource allocation
problem into a Markov decision process (MDP) problem. It takes into account
both estimation quality states and channel quality states for decision making.
To the best of our knowledge, such a problem has not been considered in the
literature of remote state estimation. The action space of the formulated MDP
is a multi-dimensional hybrid one consisting of discrete channel assignment
and continuous power allocation, while classical MDP solutions only work for
discrete action spaces. The new (enlarged) state space and the hybrid action
space make the decision-making problem difficult to solve, especially in large
systems.

3) Advanced dynamic resource allocation algorithm with large state and action
spaces. To handle the optimal resource allocation problem, we develop an
advanced data-driven deep reinforcement learning (DRL) algorithm that
generates low-dimensional continuous virtual actions, and propose a novel
action mapping scheme to map virtual actions into real hybrid actions in
resource allocation. Extensive simulation results illustrate that the proposed
DRL algorithm can effectively solve the dynamic resource allocation problem
with a much larger scale than the literature, and provides significant
performance gain compared with some benchmark policies, especially when the
system scale is large.

##  II System Model of Remote State Estimation

The studied remote estimation system consists of $N$ dynamic plants, each
monitored by a sensor, and a remote estimator, as shown in Fig. 1. The $N$
sensors send the measurement data to the remote estimator over $M$ frequency
channels (i.e., subcarriers), where $M<N$. Each wireless device is equipped
with a single antenna. The remote estimator applies a dynamic resource
allocation policy for channel assignment and power control of each sensor.

Figure 1: A NOMA-based $N$-sensor-$M$-channel remote estimation system.

###  II-A Local State Estimation

We consider discrete-time linear time-invariant (LTI) systems, where the model
of plant $n$ is given as [5, 11]

$\displaystyle\mathbf{x}_{n}(t+1)$
$\displaystyle=\mathbf{A}_{n}\mathbf{x}_{n}(t)+\mathbf{w}_{n}(t)$ (1)
$\displaystyle\mathbf{y}_{n}(t)$
$\displaystyle=\mathbf{C}_{n}\mathbf{x}_{n}(t)+\mathbf{v}_{n}(t)$

where $\mathbf{x}_{n}(t)\in\mathbb{R}^{l_{n}}$ is the plant state vector;
$\mathbf{A}_{n}\in\mathbb{R}^{l_{n}\times l_{n}}$ is the plant state
transition matrix; $\mathbf{y}_{n}(t)\in\mathbb{R}^{r_{n}}$ is the sensor
measurement vector; $\mathbf{C}_{n}\in\mathbb{R}^{r_{n}\times l_{n}}$ is the
measurement matrix; $\mathbf{w}_{n}(t)\in\mathbb{R}^{l_{n}}$ and
$\mathbf{v}_{n}(t)\in\mathbb{R}^{r_{n}}$ are the plant disturbance and sensing
measurement noise vectors, respectively. These noise vectors are independent
and identically distributed (i.i.d.) zero-mean Gaussian processes with
covariance matrices $\mathbf{W}_{n}$ and $\mathbf{V}_{n}$, respectively.

Due to the measurement distortion and noise in (1), each sensor has a local
estimator to pre-estimate the plant state $\mathbf{x}_{n}(t)$ based on the raw
measurement $\mathbf{y}_{n}(t)$ before sending it to the remote estimator. We
adopt the classical Kalman filter (KF) for generating the estimated state
$\mathbf{x}_{n}^{s}(t)$ as it is the optimal estimator of LTI systems in the
average estimation MSE [5]. The definition of local estimation error
covariance $\mathbf{P}_{n}^{s}(t)$ is given as

$\mathbf{P}_{n}^{s}(t)\triangleq\mathbb{E}\left[\left(\mathbf{x}_{n}^{s}(t)-\mathbf{x}_{n}(t)\right)\left(\mathbf{x}_{n}^{s}(t)-\mathbf{x}_{n}(t)\right)^{\mathrm{T}}\right].$
(2)

Since our work focuses on the remote state estimation, each local estimator is
assumed to be stable and operate in the steady state, i.e., the estimation
error covariance of the local KF is a constant
$\mathbf{P}_{n}^{s}(t)\triangleq{\bar{\mathbf{P}}}_{n},\forall
t\in\mathbb{N}$, where $\mathbb{N}$ is the set of positive integers[5, 11, 8,
4].

###  II-B Wireless Channel Model and Communications

We consider finite-state Markov block-fading channels [12]. The overall
channel power gain state matrix is denoted as $\mathbf{G}(t)$, where each
column vector
$\mathbf{g}_{n}(t)\triangleq(g_{n,1}(t),\ldots,g_{n,M}(t))^{\mathrm{T}}$
represents the channel power gain between sensor $n$ and the remote estimator
over $M$ channels. Each channel power gain has $H$ states, i.e.,
$g_{n,m}(t)\in\mathcal{G}\triangleq\left\\{h_{1},h_{2},\ldots,h_{H}\right\\}$.
The channel state vector
$\mathbf{g}_{n}(t)\in\mathcal{G}^{M}\triangleq\left\\{\tilde{\mathbf{g}}_{1},\tilde{\mathbf{g}}_{2},\ldots,\tilde{\mathbf{g}}_{H^{M}}\right\\}$
is modeled as a multi-state Markov chain. Since sensors are dislocated and
have different radio propagation environments, we assume that their Markov
channel states are independent. The remote estimator has the knowledge of
channel state information obtained via standard channel estimation techniques.
The channel state transition matrices are unknown to the remote estimator,
because the estimation of a multi-dimensional Markov chain model is
computationally intensive [13].

We adopt short-packet communications for sensor packet transmissions [14].
Given the packet length $l$ (i.e., the number of symbols per packet), the
number of data bits $b$, the signal-to-noise ratio (SNR) $\gamma_{n}$ of
sensor $n$, we have the Shannon capacity
$\mathcal{C}(\gamma_{n})=\log_{2}{(1+\gamma_{n})}$ and the channel dispersion
$\mathcal{V}(\gamma_{n})=(1-(1+\gamma_{n})^{-2})(\log_{2}{e})^{2}$. Then, the
decoding failure probability of sensor $n$’s packet can be approximated as
[15]

$\varepsilon\left(\gamma_{n}\right)\approx\mathcal{Q}\left(\frac{\mathcal{C}\left(\gamma_{n}\right)-\frac{b}{l}}{\sqrt{\frac{\mathcal{V}\left(\gamma_{n}\right)}{l}}}\right),$
(3)

where
$\mathcal{Q}(x)=(\frac{1}{\sqrt{2\pi}})\int_{x}^{\infty}{e^{-\frac{t^{2}}{2}}\text{d}t}$
is the Gaussian Q-function.

###  II-C Multiple-Access Scheme

In the NOMA scheme, each sensor takes at most one channel for transmission,
while each channel can be allocated to multiple sensors at the same time. At
each time slot, the remote estimator needs to determine both the sensor-to-
channel assignment and the sensor transmission power for interference
management. Let
$\mathbf{d}_{n}(t)\triangleq(d_{n,1}(t),d_{n,2}(t),\ldots,d_{n,M}(t))^{\mathrm{T}}\in\left\\{0,1\right\\}^{M}$
denotes the binary channel selection action of sensor $n$ at the $M$ channels,
where $\sum_{m=1}^{M}d_{n,m}(t)\leq 1$. Let
$\mathbf{p}^{\mathrm{tx}}(t)\triangleq(P_{1}^{\mathrm{tx}}(t),P_{2}^{\mathrm{tx}}(t),\ldots,P_{N}^{\mathrm{tx}}(t))^{\mathrm{T}}\in\mathbb{P}^{N}$
denotes the transmission power of $N$ sensors at time $t$, where
$\mathbb{P}\triangleq[0,P_{max}]$. Then, the received signal power of sensor
$n$ is
$P_{n}^{\mathrm{rx}}(t)=P_{n}^{\mathrm{tx}}(t)\left(\left(\mathbf{d}_{n}(t)\right)^{\mathrm{T}}\mathbf{g}_{n}(t)\right)$.

To decode sensor signals at the same channel, the remote estimator performs
successive interference cancellation (SIC) with the decreasing order of the
received sensor signal power [6], i.e., the strongest/weakest sensor signal is
decoded first/last. The first sensor packet is decoded by treating all other
sensor signals as interference. Once it is decoded successfully, the sensor
signal can be reconstructed perfectly and thus removed from the received
signal. Then, the second sensor packet will be decoded without the
interference of the first one. The decoder stops once a decoding failure
occurs or the last sensor packet has been successfully decoded. Assuming that
$P_{1}^{\mathrm{rx}}(t)\geq P_{2}^{\mathrm{rx}}(t)\geq\ldots\geq
P_{N}^{\mathrm{rx}}(t)$, the signal-to-interference-plus-noise ratio (SINR)
for decoding sensor $n$’s packet is

$\gamma_{n}(t)=\frac{P_{n}^{\mathrm{rx}}(t)}{\sum_{i=n+1}^{N}\left(\mathbf{d}_{n}(t)\right)^{\mathrm{T}}\mathbf{d}_{i}(t)P_{i}^{\mathrm{rx}}(t)+\sigma^{2}},$
(4)

where $\sigma^{2}$ is the receiving noise power. The decoding failure
probability of sensor $n$ can be obtained as

$\hat{\varepsilon}_{n}(t)=\\!\begin{cases}\\!\varepsilon\left(\gamma_{1}(t)\right),&n=1\\\
U_{n,1}+\sum_{k=2}^{n}\left(U_{n,k}\prod_{i=1}^{k-1}\left(1\\!-\\!U_{k,i}\right)\right),&n>1.\end{cases}$
(5)

where $U_{n,n}\\!=\\!\varepsilon\left(\gamma_{n}(t)\right)$ and
$U_{n,i}\\!=\\!\left(\mathbf{d}_{n}(t)\right)^{\mathrm{T}}\\!\mathbf{d}_{i}(t)\hat{\varepsilon}_{i}(t),\forall
i\\!\leq\\!n$.

###  II-D Remote State Estimation

Due to transmission scheduling and packet detection errors, sensor $n$’s
packet may not be received by the remote estimator at each time slot. Let
$\zeta_{n}(t)\\!=\\!1$ denote the successful packet detection of sensor $n$ at
$t$. To provide real-time state estimation of all plants, the remote estimator
employs a minimum mean-square error (MMSE) state estimation for each plant [5,
11]

$\hat{\mathbf{x}}_{n}(t)=\left\\{\begin{array}[]{cc}\mathbf{x}_{n}^{s}(t)&\text{
if }\zeta_{n}(t)=1\\\ \mathbf{A}_{n}\widehat{\mathbf{x}}_{n}(t-1)&\text{
otherwise.}\end{array}\right.$ (6)

Thus, the remote estimation error covariance is

$\displaystyle\mathbf{P}_{n}(t)$
$\displaystyle\triangleq\mathbb{E}\left[\left(\hat{\mathbf{x}}_{n}(t)-\mathbf{x}_{n}(t)\right)\left(\hat{\mathbf{x}}_{n}(t)-\mathbf{x}_{n}(t)\right)^{\mathrm{T}}\right]$
(7) $\displaystyle=\left\\{\begin{array}[]{cc}\mathbf{\bar{P}}_{n}&\text{ if
}\zeta_{n}(t)=1\\\
\mathbf{A}_{n}\mathbf{P}_{n}(t-1)\mathbf{A}_{n}^{\mathrm{T}}+\mathbf{W}_{n}&\text{
otherwise }\end{array}\right.$ (10)

where (10) is obtained by taking (6) and (2) into (7). Recall that
$\mathbf{\bar{P}}_{n}$ is the local estimation error covariance.

Now we define $\tau_{n}(t)$ as the age of information (AoI) of sensor $n$ at
time slot $t$, representing the time interval since the last successful
transmission of sensor $n$ to the remote estimator. Therefore, the AoI state
of sensor $n$ has the updating rule below

$\tau_{n}(t)=\left\\{\begin{array}[]{cc}1&\text{ if }\zeta_{n}(t-1)=1\\\
\tau_{n}(t-1)+1&\text{ otherwise. }\end{array}\right.$ (11)

A larger AoI indicates that the remote estimation is less accurate. Jointly
using (10) and (11), the remote estimation error covariance can be concisely
rewritten as a function of AoI as

$\mathbf{P}_{n}(t)=f_{n}^{\tau_{n}(t)}\left(\mathbf{\bar{P}}_{n}\right),$ (12)

where
$f_{n}^{1}(\mathbf{X})\\!=\\!\mathbf{A}_{n}\mathbf{X}\mathbf{A}_{n}^{\mathrm{T}}\\!+\\!\mathbf{W}_{n}$
and
$f_{n}^{\tau_{n}}(\mathbf{X})\\!=\\!f_{n}^{1}\\!\left(f_{n}^{\tau_{n}-1}(\mathbf{X})\right)$.

To quantify the remote estimation quality of sensor $n$ at time $t$, we define
the _estimation cost function_ , i.e., the sum estimation MSE of the plant
vector state, as

$\displaystyle J_{n}(t)$
$\displaystyle\triangleq\mathbb{E}\left[\left(\hat{\mathbf{x}}_{n}(t)-\mathbf{x}_{n}(t)\right)^{\mathrm{T}}\left(\hat{\mathbf{x}}_{n}(t)-\mathbf{x}_{n}(t)\right)\right]$
(13)
$\displaystyle=\operatorname{Tr}\left(\mathbf{P}_{n}(t)\right)=\operatorname{Tr}\left(f_{n}^{\tau_{n}(t)}\left(\mathbf{\bar{P}}_{n}\right)\right),$

where $\operatorname{Tr}(\cdot)$ is the matrix trace operator. Thus, the AoI
state and the local estimation error covariance jointly determine the remote
estimation quality. A smaller $J_{n}(t)$ indicates that the remote state
estimation is more accurate. As proved in our previous work [11], if the
spectral radius of plant $n$’s state transition matrix $\mathbf{A}_{n}$ is
greater than $1$, the increasing AoI gives rise to the exponential growth of
estimation cost.

##  III Problem Formulation

We aim to design a deterministic and stationary resource allocation policy
denoted as $\pi(\cdot)\in\Pi$ that generates channel allocation and power
control actions of all sensors at each time slot, for achieving the optimal
discounted long-term average estimation quality [5], i.e.,

$J^{*}=\min_{\pi(\cdot)\in\Pi}\lim_{T\rightarrow\infty}\mathbb{E}\left[\sum_{t=0}^{T-1}\sum_{n=1}^{N}\lambda^{t}J_{n}(t)\right]$
(14)

where $\lambda\in(0,1)$ is a discount factor and
$\mathbb{E}\left[\cdot\right]$ is the expectation operator. We formulate such
a sequential decision-making problem into an MDP.

###  III-A MDP Formulation

In general, the MDP takes both the channel quality state and the AoI state of
all sensors as inputs and generates the resource allocation action at each
time. The Markovian property holds directly due to the Markov channel modeling
and the AoI state updating rule.

State: Given the channel state matrix $\mathbf{G}(t)\in\mathcal{G}^{M\times
N}$, and the AoI state
$\bm{\tau}(t)=(\tau_{1}(t),\tau_{2}(t),\ldots,\tau_{N}(t))\in\mathbb{N}^{N}$,
the state of the MDP is defined as
$\mathbf{s}\left(t\right)\triangleq\left\\{\mathbf{G}(t),\bm{\tau}(t)\right\\}\in\mathcal{G}^{M\times
N}\times\mathbb{N}^{N}$.

Action: The hybrid action
$\mathbf{a}(t)\triangleq\left\\{\mathbf{D}(t),\mathbf{p}^{\mathrm{tx}}(t)\right\\}$
takes into account both the discrete channel allocation $\mathbf{D}(t)$ and
the continuous power control $\mathbf{p}^{\mathrm{tx}}(t)\in\mathbb{P}^{N}$.
The discrete action space is denoted as $\mathcal{A}$ with the cardinality of
$\left|\mathcal{A}\right|=(M+1)^{N}$, as each sensor has $M+1$ options for
channel selection. The hybrid action space is denoted as
$\mathcal{A}\times\mathbb{P}^{N}$.

Transition: The state-transition probability of MDP consists of the channel
state transition and the AoI state transition. Since the former does not
depend on the latter, the state-transition probability from state
$\mathbf{s}(t)$ to state $\mathbf{s}(t+1)$ under a particular action
$\mathbf{a}(t)$ can be written as
$\operatorname{Pr}[\mathbf{s}(t+1)|\mathbf{s}(t),\mathbf{a}(t)]=\operatorname{Pr}[\mathbf{G}(t+1)|\mathbf{G}(t)]\operatorname{Pr}[\bm{\tau}(t+1)|\mathbf{s}(t),\mathbf{a}(t)]$.
Recall that the knowledge of
$\operatorname{Pr}[\mathbf{G}(t+1)|\mathbf{G}(t)]$ is unavailable (Section
II-B).

Policy: The policy is a mapping between the state and the action as
$\mathbf{a}(t)=\pi(\mathbf{s}(t))$, where $\pi(\cdot)\in\Pi$.

Reward: We define the reward of the MDP as the negative sum estimation cost,
i.e., $r(t)=-J(t)=-\sum_{n=1}^{N}J_{n}(t)$. Thus, one needs to find a resource
allocation policy for maximizing the discounted long-term average reward
$\lim_{T\rightarrow\infty}\mathbb{E}\left[\sum_{t=0}^{T-1}\lambda^{t}r(t)\right]$.

###  III-B Challenges for Solving the MDP

In the absence of state transition probabilities, conventional model-based MDP
algorithms cannot work (e.g., value and policy iteration algorithms), and thus
data-driven reinforcement learning approaches are preferable. Due to the curse
of dimensionality introduced by the large state space $\mathcal{G}^{N\times
M}\times\mathbb{N}^{N}$, conventional data-driven reinforcement learning
algorithms, such as Q-learning, are not applicable. To deal with the large
state space issue, we resort to more advanced DRL algorithms adopting deep
neural networks (DNNs) for value function approximation, though there are
several challenges remaining.

1) Large and hybrid action space: The MDP has a large action space, even with
relatively small numbers of sensors and channels. For example, when
$N=10,M=5$, there are $\left(M+1\right)^{N}=60466176$ discrete actions for
channel assignment. For the continuous action part, a commonly adopted method
is action discretization. In the simplest 2-level quantization scenario, there
are $2^{N}=1024$ actions for power control and thus $1024\times
60466176\approx 6\times{10}^{10}$ discrete actions in total for joint channel
assignment and power control. However, deep Q-network (DQN), the most popular
DRL algorithm adopted for solving MDPs with discrete actions, cannot handle
such an MDP. The large action space entails large DNNs and makes the DQN
difficult to train, requiring a huge storage space and a tremendous amount of
computing power. Solving MDPs with large action spaces are challenging.

2) Training difficulties: In our MDP problem, the channel state is of high
dynamics and the AoI state is also a stochastic function of the channel state
and the action. In addition, as mentioned in Section II-D, the absolute value
of reward grows exponentially fast with respect to the AoI state, which means
a DRL agent needs to deal with a fairly large range of rewards, inducing a
highly fluctuated training process that is difficult to converge. This feature
is different from many existing works applying DRL in wireless communications
problems [3], where the reward/cost commonly grows up in linear or log scale
with the increasing state. Therefore, the highly stochastic states and the
large reward range can lead to unstable training process and make a DRL agent
difficult to converge to an optimal policy.

##  IV Deep Reinforcement Learning Algorithm

To solve the challenges above, first, we design low-dimensional continuous
virtual actions of the resource allocation problem and propose a novel scheme
that maps virtual action $\tilde{\mathbf{a}}(t)$ to real hybrid discrete and
continuous action $\mathbf{a}(t)$. Second, we choose a proper DRL framework
with a continuous action space, which provides a stable training process to
learn the optimal policy with virtual actions. The framework is illustrated in
Fig. 2.

Figure 2: The framework for generating resource allocation actions based on a
trained actor NN (Section IV-B) and the action mapping scheme (Section IV-A).

###  IV-A Virtual Action Design and Action Mapping Scheme

Considering the hybrid action feature of resource allocation, it is natural to
design a multi-dimensional continuous virtual action of each sensor to
represent the hybrid action. However, how to design a continuous virtual
action and map it to a discrete channel selection action for effective DRL is
highly non-trivial. A naive method is to adopt a scalar continuous virtual
action for each sensor and discretize it linearly into $M+1$ levels for
channel selection. Level 0 means no channel selection and Level $m>0$ denotes
the selection of channel $m$. At first glance, the method compresses the
$M$-size discrete action space by a one-dimensional continuous one. However,
the method implicitly converts the original non-Euclidean action space into
the Euclidean one. This potentially requires the DRL agent to learn a highly
discontinuous multi-level-multi-stage piecewise policy function. Ideally, a
slight change of state can make a significant change in action output. This
makes the DRL difficult to converge to an optimal one.

To solve the issue of discontinuous function approximation, we propose to
convert the decimal channel selection action (i.e., from $0$ to $M$) of each
sensor to a binary sequence with a length of $\lceil\log_{2}(M+1)\rceil$.
Thus, we design the virtual action for each sensor with
$\lceil\log_{2}(M+1)\rceil$ elements. Each element is a real number, and the
positive and negative values are mapped to ‘1’ and ‘0’, respectively, for real
action mapping. Now, the virtual action has $N\lceil\log_{2}(M+1)\rceil$
dimension for channel selection of all sensors. Each element of the virtual
action only handles the selection of two discrete actions. Changing of a
single virtual action entry from negative to positive represents the reduced
and increased likelihood of one action and the other, respectively. The
continuity can make the DRL easy to train. The original $(M+1)^{N}$-size
discrete action space is compressed by the
$N\lceil\log_{2}(M+1)\rceil$-dimension continuous one, which scales linearly
and logarithmically in $N$ and $M$, respectively. One dimensional virtual
action $\tilde{P}^{\mathrm{tx}}_{n}(t)\in\mathbb{R}$ is added for each
sensor’s transmit power control, and the mapping between the virtual to real
power control action is
$P^{\mathrm{tx}}_{n}(t)=P_{\max}(\operatorname{clamp_{-1}^{1}}[\tilde{P}^{\mathrm{tx}}_{n}(t)]+1)/2$,
where the operator $\operatorname{clamp_{-1}^{1}}[\cdot]$ is for truncating a
variable to values between $-1$ and $1$. Thus, the virtual action
$\tilde{\mathbf{a}}(t)$ has $N\lceil\log_{2}(M+1)\rceil+1$ elements in total.

###  IV-B Policy Optimization with Virtual Actions

There are many DRL frameworks with various features developed during the past
decade. To match the key problem features in Section III-B, we adopt the
proximal policy optimization (PPO) [16], a _policy-based_ DRL method that
learns a _stochastic policy_ via _on-policy learning_ , for solving the MDP.
The reasons include: 1) compared with off-policy DRL methods, on-policy ones
are more stable in highly stochastic environments due to the guaranteed
monotonic performance improvement based their policy updating mechanism [17];
2) existing research shows that stochastic policy based DRL algorithms lead to
smoother function approximations [18], and thus provide higher speed for
training convergence and better training stability than deterministic ones
when the training environment is of high dynamics [19]; 3) policy-based DRL
generates multi-dimensional continuous actions directly and can deal with MDPs
with continuous action spaces.

The PPO agent has a pair of actor and critic neural networks (NNs), where the
NN parameters, including weights and biases, are denoted by $\theta$ and
$\varphi$, respectively. Each of the NNs has an input layer, multiple hidden
layers, an output layer, and adopts a fully connected and feedforward
structure. The input layers of the actor and critic NNs receive the MDP state
$\mathbf{s}(t)$. All hidden layers adopt the activation function of rectified
linear unit (ReLU) for nonlinear function approximation.

We design the actor NN for generating a multi-dimensional continuous virtual
action $\tilde{\mathbf{a}}(t)$ with stochastic policy
$\tilde{\pi}(\tilde{\mathbf{a}}(t)|\mathbf{s}(t);\theta)$, which is a
probability density function of $\tilde{\mathbf{a}}(t)$ given the current
state $\mathbf{s}(t)$. The virtual action is then mapped to the real action
$\mathbf{a}(t)$ for radio resource allocation as discussed in Section IV-A to
obtain the next state $\mathbf{s}(t+1)$ and reward $r(t+1)$. The actor output
layer generates mean $\bm{\mu}(\mathbf{s}(t);\theta)$ and standard deviation
$\bm{\sigma}(\mathbf{s}(t);\theta)$ of $\tilde{\mathbf{a}}(t)$. We adopt tanh
and softplus as the activation functions for the mean and the standard
deviation outputs, respectively. The former is to bound the mean within
$[-1,1]$, while the latter is to guarantee a positive standard deviation. The
critic NN estimates the state-value function $V(\mathbf{s}(t);\varphi)$ given
the actor’s policy $\tilde{\pi}(\tilde{\mathbf{a}}(t)|\mathbf{s}(t);\theta)$,
i.e.,

$V(\mathbf{s}(t);\\!\varphi)\\!\approx\\!\mathbb{E}\\!\left[\sum_{k=0}^{\infty}\lambda^{k}r(t\\!+k)\bigg{|}\tilde{\pi}(\tilde{\mathbf{a}}(t\\!+k)|\mathbf{s}(t\\!+k);\\!\theta),\\!\forall
k\\!\geq\\!0\right]\\!\\!.$ (15)

Thus, the critic NN has a single output.

The training of the PPO agent alternates between experience generation and
policy update. In the following, we only present the key steps, and the
detailed algorithm can be found in [16].

1) Experience generation. By using current policy
$\tilde{\pi}(\cdot|\cdot;\theta_{\text{old}})$, the PPO agent samples data
$(\mathbf{s}(t),\tilde{\mathbf{a}}(t),r(t))$ with length of $T$ through
interacting with the environment. By leveraging the generated experience, the
advantage function $A(t)$ and the reward-to-go function $R(t)$ for each
$t=0,\dots,T-1$ can be calculated as

$A(t)=\sum_{k=t}^{T-1}(\lambda\alpha)^{k-t}(r(k)+\lambda
V(\mathbf{s}(k+1);\varphi)-V(\mathbf{s}(k);\varphi))$ (16)

and

$\displaystyle R(t)=r(t)+\lambda V(\mathbf{s}(t+1);\varphi)$ (17)

respectively, where $\alpha$ is a hyper-parameter named as the smoothing
factor.

2) Policy update. The PPO agent creates a mini-batch data set by randomly
sampling $B$ data from the generated $T$-length experience earlier, i.e.,
$\\{(\mathbf{s}(t_{i}),\tilde{\mathbf{a}}(t_{i}),A(t_{i}),R(t_{i}))\\}$ where
$t_{i}\in\\{t_{1},\dots,t_{B}\\}\subset\\{0,\dots,T-1\\}$.

The loss function for updating the critic NN is defined as

$L_{\text{C}}(\varphi)=\frac{1}{B}\sum_{i=1}^{B}\left(R(t_{i})-V\left(\mathbf{s}(t_{i});\varphi\right)\right)^{2},$
(18)

which is named as the temporal difference error, specifying difference of the
state-value estimations based on time steps $t_{i}$ and $t_{i}+1$. The loss
function for updating the actor NN has been elaborately designed for achieving
high training stability and is much more complex:

$\displaystyle L_{\text{A}}(\theta)\\!=$ (19)
$\displaystyle\frac{1}{B}\\!\\!\sum_{i=1}^{B}\\!\left(\min\\!\left\\{p(\mathbf{s}(t_{i});\theta)\\!A(t_{i}),c\left(\mathbf{s}(t_{i});\theta\right)\\!A(t_{i})\right\\}\\!+\\!w\hbar_{\theta_{\text{old
}}}\\!\\!\left(\\!\tilde{\mathbf{a}}(t_{i})\right)\right)$

where
$p(\mathbf{s}(t_{i});\theta)=\frac{\tilde{\pi}(\tilde{\mathbf{a}}(t_{i})|\mathbf{s}(t_{i});\theta)}{\tilde{\pi}\left(\tilde{\mathbf{a}}(t_{i})|\mathbf{s}(t_{i});\theta_{\text{old
}}\right)}$ is a density ratio, and
$c(\mathbf{s}(t_{i});\theta)=\max\left\\{\min\left\\{p(\mathbf{s}(t_{i});\theta),1+\omega\right\\},1-\omega\right\\}$
is a clip function with a hyper-parameter $\omega$.
$\hbar_{\theta_{\text{old}}}\\!\\!\left(\tilde{\mathbf{a}}(t_{i})\right)$ is
the entropy loss function of $\tilde{\mathbf{a}}(t_{i})$. Since
$\tilde{\mathbf{a}}(t_{i})$ follows a Gaussian distribution, its entropy loss
can be directly obtained based on its standard deviation
$\bm{\sigma}(\mathbf{s}(t_{i});\theta_{\text{old}})$. $w$ denotes entropy loss
weight factor. Then, the critic NN and the actor NN parameters $\varphi$ and
$\theta$ can be updated by optimizing the loss functions
$L_{\text{C}}(\varphi)$ and $L_{\text{A}}(\theta)$, respectively, using the
widely adopted Adam optimizer.

After training, the virtual action $\tilde{\mathbf{a}}(t)$ can be generated
deterministically based on the maximum likelihood method for online
deployment, i.e., $\tilde{\mathbf{a}}(t)=\bm{\mu}(\mathbf{s}(t);\theta)$.

##  V Numerical Experiments

###  V-A Experiment Setup

Our numerical experiments are implemented on a computing platform with two
Intel Xeon Gold 6256 CPUs @ 3.60 GHz and a 192 GB RAM. Each of the actor and
critic NNs of the PPO-based DRL agent has three hidden layers with sizes of
$\left\lceil 70K\right\rceil,\left\lceil 50K\right\rceil,\left\lceil
30K\right\rceil$, respectively, where $K=\sqrt{N/M}\log_{2}(M+1)$. The state
input of each NN has $N(M+1)$ dimensions, which is the same as the MDP. The
output size of the critic NN is $1$, while that of the actor NN is discussed
in Section IV-B. The dynamic system matrices $\mathbf{A}_{n}$ are randomly
generated by leveraging the method presented in [5], where the spectrum radius
is drawn uniformly from the range of $(1,1.3)$. The channel transition
matrices are generated randomly. Table I summarizes the details of the remote
estimation system parameters and the DRL parameters.

TABLE I: Summary of Experiment Setup Items | Value
---|---
Remote state estimation system parameters
Transmit power budget [dBm], $P_{max}$ | 23
Receiving noise power [dBm], $\sigma^{2}$ |  $-$60
Code rate [bps], $b/l$ | 2
Block length [symbols], $l$ | 200
Markov channel power gain states, $\mathcal{G}$ | $\\{10^{-8},10^{-7},\cdots,10^{-1}\\}$
Channel state transition matrix | Randomly generalized
Traning parameters
Episode number, $E$ | $\left\lceil 250\times\frac{N}{M}\times\sqrt{NM}\right\rceil$
Maximum time steps per episode, $T$ | 128
Learning rate of actor NN | 0.0001
Learning rate of critic NN | 0.001
Mini-batch size, $B$ | 128
Discount factor, $\lambda$ | 0.95
Learning agent parameters of this work (PPO)
The smoothing factor, $\alpha$ | 0.95
Entropy loss weight factor, $w$ | 0.01
Clip factor, $\omega$ | 0.2
Learning agent parameters of the benchmark (DQN)
Initial epsilon for exploring action space | 1
Epsilon decay rate | 0.999
Minimum epsilon | 0.01
Experience buffer length | $1000NM$

We adopt the DQN-based channel assignment of a remote estimation system with
OMA as a benchmark. As shown in [5], the DQN-based algorithm performs better
than the heuristic algorithms (e.g., the greedy and the round-robin policies).
Thus, we only need to compare our algorithm with the DQN-based one. In
addition, we use the naive action mapping based PPO (discussed in Section
IV-A) as a benchmark of our proposed novel action mapping scheme.

###  V-B Performance Evaluation

In Table II, we compare the remote estimation performance, i.e., the average
sum MSE of all plants, between the proposed DRL-based algorithm and the
benchmarks with various system scales and sensor-to-channel ratios (SCRs). The
remote estimation system with a smaller average estimation MSE has a better
performance. Average estimation MSEs in Table II are calculated by
$10000$-step simulations.

TABLE II: Performance Comparison of the Proposed Algorithm and the Benchmarks in Terms of Averaged Estimation MSE System scale | OMA | NOMA
---|---|---
($N,M$, SCR) | DQN | Naive action mapping | This work
(6, 3, 2) | 46.6243 | 39.0462 | 38.0663
(10, 5, 2) | $-$ | 65.0766 | 63.0768
(20, 10, 2) | $-$ | 154.9627 | 125.4118
(30, 15, 2) | $-$ | 297.5236 | 218.3914
(40, 20, 2) | $-$ | 397.5561 | 286.1228
(50, 25, 2) | $-$ | 486.9513 | 360.9214
(10, 4, 2.5) | $-$ | 78.8600 | 77.6075
(20, 8, 2.5) | $-$ | 172.9613 | 157.2169
(30, 12, 2.5) | $-$ | 318.8591 | 252.8929
(40, 16, 2.5) | $-$ | 404.8909 | 345.0570
(50, 20, 2.5) | $-$ | 563.4724 | 434.4692
(15, 5, 3) | $-$ | 159.5629 | 129.8431
(24, 8, 3) | $-$ | 275.5904 | 230.7575
(33, 11, 3) | $-$ | 438.7793 | 343.1833
(42, 14, 3) | $-$ | 578.4324 | 448.6448
(51, 17, 3) | $-$ | 765.4839 | 592.1996

We see that the DQN algorithm with OMA only works for the
$6$-sensor-$3$-channel setting and does not even converge for larger systems.
The proposed DRL algorithm with NOMA can scale up to $50$ sensors and $25$
channels. We also see that the proposed action mapping scheme-based algorithm
can provide a $25$% average estimation MSE reduction than the naive mapping
scheme-based one when the system is large (e.g., $N\geq 50$). The performance
gap increase with the growing system scale. We also see that the estimation
performance decreases with the increasing SCR as expected, due to the
decreasing amount of wireless resource.

##  VI Conclusion

We have proposed a practical remote estimation system with the NOMA scheme. We
have developed an advanced DRL algorithm for resource allocation with large
hybrid state and action spaces. Our experiments have showcased that the
proposed DRL algorithm is able to effectively address the resource allocation
problem and provide significant performance gain than the benchmarks. For
future work, we will investigate distributed DRL algorithms for resource
allocation of large-scale systems and compare them with the present
centralized allocation scheme.

## References

  * [1] P. Park, S. Coleri Ergen, C. Fischione, C. Lu, and K. H. Johansson, “Wireless network design for control systems: A survey,” _IEEE Commun. Surv. Tutor._ , vol. 20, no. 2, pp. 978–1013, May 2018.
  * [2] E. Uysal _et al._ , “Semantic communications in networked systems: A data significance perspective,” _arXiv preprint (accepted by IEEE Network)_ , Mar. 2021. [Online]. Available: https://doi.org/10.48550/arXiv.2103.05391
  * [3] A. Zappone, M. Di Renzo, and M. Debbah, “Wireless networks design in the era of deep learning: Model-based, AI-based, or both?” _IEEE Trans. Commun._ , vol. 67, no. 10, pp. 7331–7376, Jun. 2019.
  * [4] W. Liu, K. Huang, D. E. Quevedo, B. Vucetic, and Y. Li, “Deep reinforcement learning for wireless scheduling in distributed networked control,” _arXiv preprint_ , Sep. 2021. [Online]. Available: https://doi.org/10.48550/arXiv.2109.12562
  * [5] A. S. Leong, A. Ramaswamy, D. E. Quevedo, H. Karl, and L. Shi, “Deep reinforcement learning for wireless sensor scheduling in cyber-physical systems,” _Automatica_ , vol. 113, pp. 1–8, Mar. 2020. Art. no. 108759\.
  * [6] S. M. R. Islam, N. Avazov, O. A. Dobre, and K.-s. Kwak, “Power-domain non-orthogonal multiple access (NOMA) in 5G systems: Potentials and challenges,” _IEEE Commun. Surv. Tutor._ , vol. 19, no. 2, pp. 721–742, Oct. 2016.
  * [7] Y. Li, A. Mehr, and T. Chen, “Multi-sensor transmission power control for remote estimation through a SINR-based communication channel,” _Automatica_ , vol. 101, pp. 78–86, Mar. 2019.
  * [8] A. Forootani, R. Iervolino, M. Tipaldi, and S. Dey, “Transmission scheduling for multi-process multi-sensor remote estimation via approximate dynamic programming,” _Automatica_ , vol. 136, pp. 1–14, Feb, 2022. Art. no. 110061\.
  * [9] M. Pezzutto, L. Schenato, and S. Dey, “Transmission scheduling for remote estimation with multi-packet reception under multi-sensor interference,” _IFAC-PapersOnLine_ , vol. 53, no. 2, pp. 2628–2633, Apr. 2021.
  * [10] ——, “Transmission power allocation for remote estimation with multi-packet reception capabilities,” _Automatica_ , vol. 140, pp. 1–13, Jun. 2022. Art. no. 110257.
  * [11] W. Liu, D. E. Quevedo, Y. Li, K. H. Johansson, and B. Vucetic, “Remote state estimation with smart sensors over Markov fading channels,” _IEEE Trans. Autom. Control_ , early access, Jun. 2021, doi: 10.1109/TAC.2021.3090741.
  * [12] P. Sadeghi, R. A. Kennedy, P. B. Rapajic, and R. Shams, “Finite-state Markov modeling of fading channels - a survey of principles and applications,” _IEEE Signal Process. Mag._ , vol. 25, no. 5, pp. 57–80, Sep. 2008.
  * [13] Y. He _et al._ , “Deep-reinforcement-learning-based optimization for cache-enabled opportunistic interference alignment wireless networks,” _IEEE Trans. Veh. Technol._ , vol. 66, no. 11, pp. 10 433–10 445, Sep. 2017\.
  * [14] Y. Polyanskiy, H. V. Poor, and S. Verdu, “Channel coding rate in the finite blocklength regime,” _IEEE Trans. Inform. Theory_ , vol. 56, no. 5, pp. 2307–2359, May 2010.
  * [15] W. Liu, G. Nair, Y. Li, D. Nesic, B. Vucetic, and H. V. Poor, “On the latency, rate, and reliability tradeoff in wireless networked control systems for IIoT,” _IEEE Internet Things J._ , vol. 8, no. 2, pp. 723–733, Jan. 2021\.
  * [16] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” _arXiv preprint_ , Aug. 2017. [Online]. Available: https://doi.org/10.48550/arXiv.1707.06347
  * [17] T. Yang, L. Zhao, W. Li, and A. Y. Zomaya, “Reinforcement learning in sustainable energy and electric systems: A survey,” _Annu. Rev. Control_ , vol. 49, pp. 145–163, Apr. 2020.
  * [18] M. Morales, _Grokking deep reinforcement learning_ , 1st ed. USA: Manning Publications, 2020, ch. 11, pp. 342–343.
  * [19] Y. Xiao, J. Liu, J. Wu, and N. Ansari, “Leveraging deep reinforcement learning for traffic engineering: A survey,” _IEEE Commun. Surv. Tutor._ , vol. 23, no. 4, pp. 2064–2097, Aug. 2021.


$\displaystyle\leq\sqrt{\frac{3\log\left(\frac{3SAn}{\delta}\right)}{\pi_{\textup{min}}N_{t}(s)}}\sum_{s}d^{\pi_{E_{i}}}_{s^{\prime}}(s)\sum_{a:\pi_{E_{1}}(a|s)=0}\pi_{E_{i}}(a|s)\zeta(s,a)$
$\displaystyle=\sqrt{\frac{3\log\left(\frac{3SAn}{\delta}\right)}{\pi_{\textup{min}}tA}}\sum_{s}d^{\pi_{E_{i}}}_{s^{\prime}}(s)\sum_{a:\pi_{E_{1}}(a|s)=0}\pi_{E_{i}}(a|s)\zeta(s,a),$

where the third step follows from the fact that, since $N_{t}(s)\geq 1$,
$\hat{\pi}_{E_{i}}(a|s)=0$ for all $(s,a)$ such that $\pi_{E_{i}}(a|s)=0$, the
fifth one, instead, from Lemma 8, and the last one from the fact that, in US-
IRL-SE, $N_{t}(s)=\sum_{a}N_{t}(s,a)=tA$.171717Notice that Equation (56)
guarantees that we can apply Lemma 8. Therefore, we have obtained an element-
wise upper-bound on $\epsilon_{2}(\zeta)$ of the following form:

$\displaystyle\epsilon_{2}(\zeta)\leq\sqrt{\frac{3\log\left(\frac{3SAn}{\delta}\right)}{\pi_{\textup{min}}tA}}d^{\pi_{E_{i}}}\pi_{E_{i}}\bar{B}^{\pi_{E_{1}}}\zeta\coloneqq\rho_{t}d^{\pi_{E_{i}}}\pi_{E_{i}}\bar{B}^{\pi_{E_{1}}}\zeta.$
(68)

We are finally ready to define the set $\widetilde{\Psi}$. More specifically:

$\displaystyle\widetilde{\Psi}\coloneqq\left\\{\zeta\in\mathbb{R}^{S\times
A}_{\geq
0}:(1+\rho_{t})d^{\pi_{E_{i}}}\pi_{E_{i}}\bar{B}^{\pi_{E_{1}}}\zeta\leq\xi_{i}-\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)\right\\}.$
(69)

As one can verify, the definition of $\widetilde{\Psi}$ follows from upper-
bounding $\epsilon_{1}(\zeta)$ and $\epsilon_{2}(\zeta)$ with Equations (67)
and (68). As a direct consequence, whenever $\zeta\in\widetilde{\Psi}$, we
have that $\zeta$ belongs to $\widehat{\Psi}$ as well.

#### Step 3: Ensuring that the feasible region of $\widetilde{\Psi}$ is non-
empty

At this point, one might be tempted to directly study the projection of
$||\zeta-\textup{proj}_{\widetilde{\Psi}}\left(\zeta\right)||_{\infty}$.
Nevertheless, we notice that, for sufficiently large values of $\alpha_{t}$
and $\beta_{t}$, $\widetilde{\Psi}$ might be empty.

Sufficient conditions to guarantee that $\widetilde{\Psi}$ is not empty are
the following ones:

$\displaystyle\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\alpha_{t}\leq\frac{\xi_{i}}{2}$
(70)
$\displaystyle\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\beta_{t}\leq\frac{\xi_{i}}{2}.$
(71)

Using Lemma 12 in (Jonsson et al.,, 2020), it is possible to verify that
Equations (57) and (58) are sufficient conditions for Equations (70)-(71) to
hold.

#### Step 4: Picking $\tilde{\zeta}\in\widetilde{\Psi}$ to upper-bound
$||\zeta-\textup{proj}_{\widehat{\Psi}}\left(\zeta\right)||_{\infty}$

At this point, we are ready to conclude our proof. As we have previously
verified, the set $\widetilde{\Psi}$ is non-empty. We can now study
$||\zeta-\textup{proj}_{\widehat{\Psi}}\left(\zeta\right)||_{\infty}$. More
precisely, we have that:

$\displaystyle||\zeta-\textup{proj}_{\widehat{\Psi}}\left(\zeta\right)||_{\infty}\leq||\zeta-\textup{proj}_{\widetilde{\Psi}}\left(\zeta\right)||_{\infty}.$

To further upper bound this Equation, we notice that we can always pick, by
definition of the infinite norm projection, any
$\tilde{\zeta}\in\widetilde{\Psi}$. In other words, we have that:

$\displaystyle||\zeta-\textup{proj}_{\widehat{\Psi}}\left(\zeta\right)||_{\infty}\leq||\zeta-\textup{proj}_{\widetilde{\Psi}}\left(\zeta\right)||_{\infty}\leq||\zeta-\tilde{\zeta}||_{\infty}.$

More specifically, we choose $\tilde{\zeta}$ in the following way. If for all
$i\in\left\\{2,\dots,n+1\right\\}$, $\pi_{E_{i}}(a|s)=0$, then we pick
$\tilde{\zeta}(s,a)=\zeta(s,a)$;181818Notice that, if for all
$i\in\left\\{2,\dots,n+1\right\\}$, $\pi_{E_{i}}(a|s)=0$, $\zeta(s,a)$ does
not contribute to any of the linear constraints that introduced by Equation
(6). otherwise, we pick $\zeta(s,a)$ in the following way:

$\displaystyle\tilde{\zeta}(s,a)=\frac{1-\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}{1+\rho_{t}}\zeta(s,a).$
(72)

First of all, we verify that this choice of $\tilde{\zeta}$ belongs to
$\widetilde{\Psi}$. Plugging Equation (72) into the definition of
$\widetilde{\Psi}$ we obtain that:

$\displaystyle\left(1-\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)\right)\frac{1+\rho_{t}}{1+\rho_{t}}d^{\pi_{E_{i}}}\pi_{E_{i}}\bar{B}^{\pi_{E_{1}}}\zeta\leq\xi_{i}-\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right).$

However, since $\zeta\in\Psi$, we have that
$d^{\pi_{E_{i}}}\pi_{E_{i}}\bar{B}^{\pi_{E_{1}}}\zeta\leq\xi_{i}$, thus
leading to:

$\displaystyle\left(\xi_{i}-\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)\right)\leq\xi_{i}-\frac{2\sqrt{2}\gamma\xi_{i}}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right),$

which is always true, and, consequently $\tilde{\zeta}\in\widetilde{\Psi}$.

To conclude the proof, it remains to analyze
$||\zeta-\tilde{\zeta}||_{\infty}$. At this point, we notice that:

$\displaystyle||\zeta-\tilde{\zeta}||_{\infty}$
$\displaystyle=\max_{\begin{subarray}{c}(s,a):\pi_{E_{1}(a|s)}=0,\\\ \exists
i:\pi_{E_{i}}(a|s)>0\end{subarray}}\Big{|}\zeta(s,a)-\frac{1-\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}{1+\rho_{t}}\zeta(s,a)\Big{|},$

Indeed, for a state-action pair $(s,a)$ such that $\pi_{E_{1}}(a|s)=0$, we can
notice that any value $\zeta(s,a)$ will not affect the resulting reward
function.191919In other words, if $(s,a)$ is such that $\pi_{E_{1}}(a|s)=0$,
we can add the following constraint to any of the set we defined:
$\zeta(s,a)=0$. The resulting feasible reward set is left unchanged.
Furthermore, whenever there is no sub-optimal expert such that
$\pi_{E_{i}}(a|s)>0$, then $\tilde{\zeta}(s,a)=\zeta(s,a)$ holds by
definition. To conclude, with simple algebraic manipulations we can obtain the
following result:

$\displaystyle||\zeta-\tilde{\zeta}||_{\infty}$
$\displaystyle\leq\max_{\begin{subarray}{c}(s,a):\pi_{E_{1}(a|s)}=0,\\\
\exists
i:\pi_{E_{i}}(a|s)>0\end{subarray}}\Big{|}(1+\rho_{t})\zeta(s,a)-\left({1-\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\zeta(s,a)\Big{|}$
$\displaystyle=\max_{\begin{subarray}{c}(s,a):\pi_{E_{1}(a|s)}=0,\\\ \exists
i:\pi_{E_{i}}(a|s)>0\end{subarray}}\Big{|}\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\zeta(s,a)\Big{|}$
$\displaystyle=\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\max_{\begin{subarray}{c}(s,a):\pi_{E_{1}(a|s)}=0,\\\
\exists i:\pi_{E_{i}}(a|s)>0\end{subarray}}\Big{|}\zeta(s,a)\Big{|}$
$\displaystyle\leq\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\},$

where in the last step we have upper-bounded $\zeta(s,a)$ with an upper bound
that follows directly from Equation (8), thus concluding the proof. ∎

We can appreciate as Lemma 10 provides, under certain conditions on the time-
step $t$, a high-probability upper bound on the difference
$\inf_{\hat{r}\in\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}}||r-\hat{r}||_{\infty}$
for any choice of $r\in\mathcal{R}_{\bar{\mathfrak{B}}}$. Furthermore, as we
can see, the proof is fairly involved due to the necessity of upper bounding
the error term
$||\zeta-\textup{proj}_{\widehat{\Psi}}\left(\zeta\right)||_{\infty}$. At this
point, by deriving symmetric results (i.e., Lemma 9 and Lemma 10), it is
possible to derive an identical upper-bound for
$\inf_{{r}\in\mathcal{R}_{{\bar{\mathfrak{B}}}}}||r-\hat{r}||_{\infty}$ for
any choice of $\hat{r}\in\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}$. As a
consequence, Lemma 10 provides the following high-probability upper-bound on
the Hausdorff distance between $\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}$
and $\mathcal{R}_{{\bar{\mathfrak{B}}}}$:

$\displaystyle
H_{\infty}\left(\mathcal{R}_{{\bar{\mathfrak{B}}}},\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}\right)\leq\frac{2\sqrt{2}\gamma}{1-\gamma}\beta_{t}+\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}.$
(73)

See 6

###### Proof.

Let us start from Lemma 10 and consider $t$ such that the condition of Lemma
10 are satisfied. As previously discussed, due to Lemma 10, we have that, at
time $t$, the algorithm US-IRL-SE induces an estimated feasible reward set
such that the following holds with high-probability:

$\displaystyle
H_{\infty}\left(\mathcal{R}_{{\bar{\mathfrak{B}}}},\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}\right)\leq\frac{2\sqrt{2}\gamma}{1-\gamma}\beta_{t}+\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}.$

To conclude the proof, we need to find a sufficiently large $t$ such that the
following holds:

$\displaystyle
H_{\infty}\left(\mathcal{R}_{{\bar{\mathfrak{B}}}},\mathcal{R}_{\widehat{\bar{\mathfrak{B}}}}\right)\leq\frac{2\sqrt{2}\gamma}{1-\gamma}\beta_{t}+\left(\rho_{t}+{\frac{2\sqrt{2}\gamma}{1-\gamma}\left(\alpha_{t}+\beta_{t}\right)}\right)\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\leq\epsilon,$

for any $\epsilon\in(0,1)$. To this end, it is sufficient to find the smallest
$t$ that satisfies the following conditions:

$\displaystyle\rho_{t}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\leq\frac{\epsilon}{4}$
(74)
$\displaystyle\frac{2\sqrt{2}\gamma}{(1-\gamma)}\beta_{t}\leq\frac{\epsilon}{4}$
(75)
$\displaystyle\frac{2\sqrt{2}\gamma}{(1-\gamma)}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\beta_{t}\leq\frac{\epsilon}{4}$
(76)
$\displaystyle\frac{2\sqrt{2}\gamma}{(1-\gamma)}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\alpha_{t}\leq\frac{\epsilon}{4}$
(77)

At this point, let us focus on
$\rho_{t}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\leq\frac{\epsilon}{4}$.
By simple algebraic manipulations we obtain the following sufficient condition
for $t$:

$\displaystyle
tA\geq\frac{48q_{1}^{2}\log\left(\frac{3SAn}{\delta}\right)}{\pi_{\textup{min}}\epsilon^{2}}.$

For $\frac{2\sqrt{2}\gamma}{(1-\gamma)}\beta_{t}\leq\frac{\epsilon}{4}$,
instead, we obtain the following condition on $t$:202020This follows by
applying Lemma 12 in Jonsson et al., (2020)

$\displaystyle
t\geq\frac{128\gamma^{2}}{(1-\gamma)^{2}\epsilon^{2}}\left[\log\left(\frac{3SAn}{\delta}\right)+(S-1)\log\left(\frac{16384\gamma^{4}}{(1-\gamma)^{4}\epsilon^{4}}\left(\log\left(\frac{3SAn}{\delta}\right)+(S-1)\left(\sqrt{e}+\sqrt{\frac{1}{S-1}}\right)^{2}\right)\right)\right].$

Similarly, for
$\frac{2\sqrt{2}\gamma}{(1-\gamma)}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\beta_{t}\leq\frac{\epsilon}{4}$,
we need at least the following number of samples:

$\displaystyle
t\geq\frac{128q_{1}^{2}\gamma^{2}}{(1-\gamma)^{2}\epsilon^{2}}\left[\log\left(\frac{3SAn}{\delta}\right)+(S-1)\log\left(\frac{16384q_{1}^{4}\gamma^{4}}{(1-\gamma)^{4}\epsilon^{4}}\left(\log\left(\frac{3SAn}{\delta}\right)+(S-1)\left(\sqrt{e}+\sqrt{\frac{1}{S-1}}\right)^{2}\right)\right)\right]$

Finally, for
$\frac{2\sqrt{2}\gamma}{(1-\gamma)}\min\left\\{\pi_{\textup{min}}^{-1}\max_{i}\xi_{i},(1-\gamma)^{-1}\right\\}\alpha_{t}\leq\frac{\epsilon}{4}$,
we obtain:

$\displaystyle
tA\geq\frac{128q_{1}^{2}\gamma^{2}}{(1-\gamma)^{2}\epsilon^{2}}\left[\log\left(\frac{3SAn}{\delta}\right)+(A-1)\log\left(\frac{16384q_{1}^{2}\gamma^{4}}{(1-\gamma)^{4}\epsilon^{4}}\left(\log\left(\frac{3SAn}{\delta}\right)+(A-1)\left(\sqrt{e}+\sqrt{\frac{1}{A-1}}\right)^{2}\right)\right)\right].$

At this point, we notice that the total number of samples gathered from the
generative at step $t$ is given by $tSA$. Therefore, from the previous
equations, together with the conditions of Lemma 10, we obtian the following
sufficient condition to guarantee that US-IRL-SE is
$(\epsilon,\delta)$-correct.

$\displaystyle
tSA=\mathcal{O}\left(\max\left\\{\frac{q_{1}^{2}S\log\left(\frac{1}{\delta}\right)}{\pi_{\textup{min}}^{2}\epsilon^{2}},\frac{q_{2}^{2}SA\left(S+\log\left(\frac{1}{\delta}\right)\right)}{(1-\gamma)^{2}\epsilon^{2}},\frac{S\log\left(\frac{1}{\delta}\right)}{\pi_{\textup{min}}^{2}}\right\\}\right),$
(78)

where the first term is due to Equation (74), the second one comes from
Equations (75)-(76), and the last term arises from Equation (56), and is
independent w.r.t. the desired accuracy $\epsilon$.212121All the other
equations introduce negligible terms in the $\mathcal{O}(\cdot)$ notation.
Specifically, we notice that for sufficiently small values of $\epsilon$
(i.e., $\epsilon<q_{1}$), the last term is always dominated by the first one,
which concludes the proof.

∎

##  Appendix D Stochastic Optimal Expert

In this section, we discuss how to extend our analysis to the case in which
the optimal expert plays a stochastic policy. First of all, let us define
$\pi_{\textup{min},E_{1}}$ as the minimum probability with which the optimal
expert plays its actions.222222Notice that, in principle,
$\pi_{\textup{min},E_{1}}\neq\pi_{\textup{min}}$ that we defined for the sub-
optimal experts.

At this point, we notice that, for obtaining sample-complexity guarantees of
the US-IRL-SE algorithm, Lemma 9 implies that we are interested in learning
_which are_ the actions that are played with positive probability by the
optimal expert. In other words, for any non-zero vector
$x\in\mathbb{R}^{S\times A}$, we want the following to hold with high-
probability:

$\displaystyle\Big{|}\left(\bar{B}^{\pi_{E_{1}}}-\bar{B}^{\hat{\pi}_{E_{1}}}\right)x\Big{|}=0.$
(79)

To this end, one can apply Lemma D.3 of Metelli et al., (2023), which implies
that with a number of samples that is:

$\displaystyle\mathcal{O}\left(S\frac{\log\left(\frac{1}{\delta}\right)}{\log\left(1/(1-\pi_{\textup{min},E_{1}})\right)}\right),$
(80)

Equation (79) holds w.p. at least $1-\delta$. Once this condition is verified,
the rest of the proof of the complexity of US-IRL-SE follows identically to
the one of Theorem 6.232323Notice, indeed, that we can apply Lemma 10 as-is,
but introducing the further constraints that $t$ should be sufficiently large
so that Equation (79) holds. We notice that this introduces an additional
maximum in the result of the sample-complexity. However, it is also possible
to extend the proof of the lower bound of Metelli et al., (2023) (see Theorem
D.1 in Metelli et al., (2023)), to show that the dependency resulting from
(80) is unavoidable. In this sense, US-IRL-SE remains minimax optimal whenever
the performance of the sub-optimal experts are sufficiently close to the ones
of the optimal agent.

##  Appendix E Per-expert Complexity of IRL-SE

In this section, we provide a description of how to change the PAC learning
formalism to show that Equation (33) actually represents a lower bound to the
number of data that is necessary to gather from each of the different sub-
optimal experts.

Specifically, we define a learning algorithm for an IRL problem
$\bar{\mathfrak{B}}$ as a tuple $\mathfrak{A}=\left(\tau,\nu\right)$, where
$\tau$ is a stopping time that controls the end of the data acquisition phase,
and $\nu=\left(\nu_{t}\right)_{t\in\mathbb{N}}$ is a history-dependent
sampling strategy over
$\mathcal{S}\times\mathcal{A}\times\left(\mathcal{S}\right)^{n+1}$. More
precisely,
$\nu_{t}\in\Delta^{\mathcal{S}\times\mathcal{A}\times\left(\mathcal{S}\right)^{n+1}}_{\mathcal{D}_{t}}$,
where
$\mathcal{D}_{t}=\left(\left(\mathcal{S}\times\mathcal{A}\right)^{n+2}\times\mathcal{S}\right)^{t}$
. At each time step $t\in\mathbb{N}$, the algorithm selects, by means of
$\nu_{t}$:

  * (i)

a state-action pair $(S_{t},A_{t})$ and observes a sample $S^{\prime}_{t}\sim
p(\cdot|S_{t},A_{t})$ from the environment

  * (ii)

a state $S_{t}^{(i)}$ for each expert $i\in\left\\{1,\dots,n+1\right\\}$ and
observes a sample $A_{t}^{(i)}\sim\pi_{E_{i}}(\cdot|S_{t}^{(i)})$

The observed realizations are then used to update the sampling strategy
$\nu_{t}$, and the process goes on until the stopping rule is satisfied. At
the end of the data acquisition phase, the algorithm leverages the collected
data to output the estimate of the feasible reward set
${\mathcal{R}}_{\widehat{\bar{\mathfrak{B}}}_{\tau}}$ that is induced by the
resulting empirical IRL problem $\widehat{\bar{\mathfrak{B}}}_{\tau}$. Given
this formalism, we are interested in designing learning algorithms that, for
any desired accuracy $\epsilon\in(0,1)$ and any risk parameter
$\delta\in(0,1)$, guarantee that:

$\displaystyle\mathbb{P}_{\mathfrak{A},\bar{\mathfrak{B}}}\left(H_{\infty}(\mathcal{R}_{\bar{\mathfrak{B}}},{\mathcal{R}}_{\widehat{\bar{\mathfrak{B}}}_{\tau}})>\epsilon\right)\leq\delta.$
(81)

We refer to these algorithms as $(\epsilon,\delta)$-correct identification
strategies. For $(\epsilon,\delta)$-correct strategies, we define their sample
complexity as the sum of the number of _unitary_ samples gathered from the
generative model. In other words, let us denote with $N_{t}(s,a)$ the number
of samples gathered, at step $t$, from the environment, and let
$N_{t}^{(i)}(s)$ denotes the number of samples gathered at step $t$ from the
$i$-th expert at state $s$. Then, the sample complexity is given by:

$\displaystyle\sum_{(s,a)}N_{\tau}(s,a)+\sum_{i}\sum_{s}N^{(i)}_{\tau}(s).$

Given this learning formalism, it is straightforward to extend the results of
Theorem 5 to this setting. More specifically, Theorem 9 can be used to lower
bound $\mathbb{E}\left[\sum_{s}N^{(i)}_{\tau}(s)\right]$ for each sub-optimal
expert, thus showing a significant increase in the sample complexity (i.e.,
linear in the number of sub-optimal experts).

##  Appendix F Further details on US-IRL-SE

In this section, we provide further details on US-IRL-SE. More precisely, we
specify the exact expression of the parameter $m$ that is used to provide
$(\epsilon,\delta)$-correc

#### Expression of $m$

From the proof of Theorem 6, it is possible to derive an exact expression of
$m$ that can be used in US-IRL-SE. More specifically, since
$t=N_{t}(s,a)$,242424This is a direct consequence of the uniform sampling
strategy. it is sufficient to take $m$ as the minimum $t$ that satisfies
Equations (74)-(77) and Equations (56)-(58).

#### Computational Complexity of US-IRL-SE

At each iteration, US-IRL-SE will query the generative model in all state-
action pair $(s,a)\in\mathcal{S}\times\mathcal{A}$. Notice that the generative
model samples (i) the environment (ii) all the expert’s policy. Assuming a
unitary cost for generating each sample, a single query to the generative
model results in a computational complexity of $\mathcal{O}\left(n\right)$.
Therefore, the total computational complexity of US-IRL-SE is given
$\mathcal{O}\left(SAmn\right)$.


# Self-Distillation Mixup Training for Non-autoregressive Neural Machine
Translation

Jiaxin Guo1, Minghan Wang1, Daimeng Wei1, Hengchao Shang1, Yuxia Wang2,
Zongyao Li1, Zhengzhe Yu1, Zhanglin Wu1, Yimeng Chen1, Chang Su1,
Min Zhang1, Lizhi Lei1, Shimin Tao1, Hao Yang1
1Huawei Translation Services Center, Beijing, China
2The University of Melbourne, Melbourne, Australia
{guojiaxin1,wangminghan,weidaimeng,shanghengchao,
lizongyao,yuzhengzhe,wuzhanglin2,chenyimeng,suchang8,
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Recently, non-autoregressive (NAT) models predict outputs in parallel,
achieving substantial improvements in generation speed compared to
autoregressive (AT) models. While performing worse on raw data, most NAT
models are trained as student models on distilled data generated by AT teacher
models, which is known as sequence-level Knowledge Distillation. An effective
training strategy to improve the performance of AT models is Self-Distillation
Mixup (SDM) Training, which pre-trains a model on raw data, generates
distilled data by the pre-trained model itself and finally re-trains a model
on the combination of raw data and distilled data. In this work, we aim to
view SDM for NAT models, but find directly adopting SDM to NAT models gains no
improvements in terms of translation quality. Through careful analysis, we
observe the invalidation is correlated to Modeling Diversity and Confirmation
Bias between the AT teacher model and the NAT student models. Based on these
findings, we propose an enhanced strategy named SDMRT by adding two stages to
classic SDM: one is Pre-Rerank on self-distilled data, the other is Fine-Tune
on Filtered teacher-distilled data. Our results outperform baselines by
0.6$\sim$1.2 BLEU on multiple NAT models. As another bonus, for Iterative
Refinement NAT models, our methods can outperform baselines within half
iteration number, which means 2$\times$ acceleration.

##  1 Introduction

Autoregressive (AT) neural machine translation (Bahdanau et al., 2015) models
decode tokens one-by-one, which ensure the robustness of the intrinsic
language model but make the inference slow (§2.1). Recently, non-
autoregressive (NAT) models generate all outputs in parallel. It speeds up the
inference, but at the cost of breaking the dependency between adjacent tokens,
leading to worse performance than standard AT models (§2.2).

Knowledge distillation (Hinton et al., 2015) is a method that trains a student
model to perform better by learning from a stronger teacher model. This method
has been proved to be necessary for NAT models training in (Zhou et al.,
2020), and most NAT models are trained on distilled data generated by an AT
teacher model to achieve competitive performance. This method of distilled
data generation is called sequence-level knowledge distillation (Kim and Rush,
2016) (§2.3).

When the teacher model is the same as the student model, we call it Self-
Distillation. Using the combination of the original data and the self
distilled data to re-train the model itself is Self-Distillation Mixup (SDM)
Training. (Freitag et al., 2017) has proved that SDM can improve the
performance of AT models (§2.4).

In this work, we study a general training strategy to improve the NAT
performance. A promising attempt to address this issue is to improve the AT
teacher model by using some strategies such as SDM Training. The problem is
that, higher quality distilled data generated by a better AT teacher model may
not improve the performance of NAT student models generally (§3.1). Another
approach is to adopt SDM Training to NAT models directly. But the second
problem occurs that, classic SDM Training has no effect when NAT models are
trained on AT-distilled data (§3.2).

We propose an enhanced training progress named Self-Distillation Mixup
Training with Pre-Rerank and Fine-Tune (SDMRT). In the stage of distilled data
generation, we use Rerank Algorithm (Lee et al., 2021) to reduce the Modeling
Diversity between AT and NAT models. In the stage of model training, we Fine-
Tune the NAT models on the filtered data to reduce the Confirmation Bias (§4).

Our contributions are as follows:

  * •

We first apply SDM to NAT training and introduce an enhanced training strategy
named SDMRT to significantly improve translation quality of multiple NAT
models (§4.1; §5.2).

  * •

We analyse why classic SDM fails on NAT models trained on AT-distilled data.
Specifically, we theoretically demonstrate Modeling Diversity and Confirmation
Bias between the AT teacher model and the NAT student models (§4.2; §5.3;
§5.3).

  * •

We carefully study the effect of different Rerank strategies on self-distilled
data and Data-filtering methods on teacher-distilled data. Furthermore, we
find our SDMRT can be regarded as an acceleration method to reduce iterations
for Iterative Refinement NAT models (§5.3; §5.3; §5.3).

##  2 Background

###  2.1 Autoregressive NMT Models

Let $x$ denote the input sequence and $y$ denote the output sequence. Let
$\mathcal{D}$ denote the original parallel data:
$\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$. In order to model the joint
probability of the output sequence, autoregressive (AT) models usually
generate each output token $y_{t}$ conditioned on the previously generated
ones(Bahdanau et al., 2015). The formulation is as following:

$\displaystyle p(y|x)=\prod_{t=1}^{T}p(y_{t}|\textbf{y}_{<t},x)$ (1)

###  2.2 Non-autoregressive NMT Model

Non-autoregressive (NAT) models (Du et al., 2021; Gu and Kong, 2021; Saharia
et al., 2020) generate all outputs in parallel and break the dependency
between the output tokens. The basic formulation of a NAT model independently
factors the conditional distribution:

$\displaystyle p(y|x)=\prod_{t=1}^{T}p(y_{t}|x)$ (2)

According to the decoding strategy, NAT can be classified into two types,
fully non-autoregressive models and iterative refinement models:

#### Fully Non-autoregressive models

  * •

Vanilla NAT (Gu et al., 2018) Vanilla NAT (V-NAT) is the first NAT model. We
use the simplified version where the decoder’s inputs are directly copied from
the encoder without considering latent variables.

  * •

NAT-CRF (Sun et al., 2019) NAT-CRF designs an efficient approximation for
Conditional Random Fields (CRF) for decoder’s outputs.

#### Iterative Refinement (IR) NAT models

  * •

iNAT (Lee et al., 2018) iNAT extends the vanilla NAT by iteratively refining
the translation. At each refinement step, the outputs of the previous step
would be taken as inputs.

  * •

CMLM (Ghazvininejad et al., 2019) CMLM adopts a masked language model to
progressively generate the sequence from entirely masked inputs. At each
refinement step, the outputs of the previous step would be partially re-masked
and taken as inputs.

###  2.3 Knowledge Distillation

Knowledge distillation (Hinton et al., 2015) was originally proposed for
training a weaker student classifier on the targets predicted from a stronger
teacher model. A typical approach is using the label probabilities produced by
the teacher as “soft targets”
$q_{i}=\exp(z_{i}/\tau)\big{/}\sum_{j}\exp(z_{j}/\tau)$ for training the
student model, where $q_{i}$ and $z_{i}$ are the probability and the logit of
class $i$ respectively and $\tau$ is the temperature. Prior work has shown the
effectiveness of adopting knowledge distillation in adversarial defense,
neural network compression, and fast inference for speech synthesis.

#### Seq-Level KD

In the context of sequence generation, Kim and Rush (2016) extend knowledge
distillation to the sentence level using “hard targets” from a pretrained
large teacher model to train a small sequence generation model. More
precisely, the teacher distribution $q(t|x)$ is approximated by its mode:
$q(t|x)\approx\mathbbm{1}\\{t=\mathop{\arg\max}\limits_{t\in\mathcal{T}}q(t|x)\\}$
with the following objectives:

$\displaystyle\mathcal{L}_{seq-
KD}=-\mathbbm{E}_{x\sim\mathcal{D}}\sum_{t\in\mathbf{T}}q(t|x)\log p(t|x)$ (3)
$\displaystyle\
=-\mathbbm{E}_{x\sim\mathcal{D},\hat{y}=\mathop{\arg\max}\limits_{t\in\mathbf{T}}q(t|x)}[\log
p(t=\hat{y}|x)]$

where $t\in\mathcal{T}$ is the space of possible target sequences. This can
also be seen as a special case of standard distillation over the sentence
space when the temperature $\tau$ approaches 0, which is equivalent to taking
the $\arg\max$ over all feasible translations. While the “hard target”
$\hat{y}$ is the most likely translation predicted by the teacher, in practice
we use beam search as an approximation. As mentioned earlier, almost all the
existing literature trains NAT models using sequence-level knowledge
distillation from a pre-trained AT model to achieve competitive performance.

###  2.4 Self-Distillation Mixup Training

We clarify Self-Distillation Training and Self-Distillation Mixup Training in
this section. The overview of these training strategies is shown in Figure 1.

Figure 1: An overview of (a) Baseline Training, (b) Self-Distillation (SD)
Training and (c) Self-Distillation Mixup (SDM) Training for an AT model.

#### Self-Distillation

As shown in Figure 1 (b), we define Self-Distillation (SD) as training a
student model on the distilled data generated by the teacher model which is as
same as the student model.

#### Self-Distillation Mixup

The strategy of Mixup Training is transferred from Self-training (Arazo et
al., 2020). Self-training is an algorithm to train a model to fit pseudo-
labels predicted by another previously-learned model. It has been very
successful for learning with unlabeled data while mixup training on the
combinations of labeled data and pseudo-labeled data.

In NMT task, distilled data can be considered as pseudo-labels. As shown in
Figure 1 (c), we define Self-Distillation Mixup (SDM) as training the student
model on the combinations of original data and distilled data.

Prior research Freitag et al. (2017) has proved that SDM Training can improve
the translation quality. We re-implement the results of AT performance on
IWSLT14 EN$\rightarrow$DE task, shown in Table 1.

Setup | Paralled Data | BLEU |  $\Delta$BLEU
---|---|---|---
Baseline | Raw Data (160K) | 28.34 | -
SD | Distilled Data (160K) | 28.29 | -0.05
SDM | Raw Data + | 29.51 | +1.17*$\uparrow$
| Distilled Data (320K) |  |

Table 1: Results of AT performance under SD and SDM training strategies on
IWSLT14 EN$\rightarrow$DE task.

##  3 Problem Formulation

###  3.1 Motivations

NAT models relay more on the distilled data generated by AT models than the
raw data. But prior research (Zhou et al., 2020) has proved that, higher
quality distilled data does not necessarily improve the performance of NAT
models. V-NAT (Gu et al., 2018) achieves better performance on distilled data
generated by Transformer-small than Transformer-base/big. CMLM (Ghazvininejad
et al., 2019) only achieves a little bit improvement of 0.22 BLEU when trained
on distilled data by Transformer-big comparing with Transformer-base.

As mentioned in §2.4, a better AT model can be got by using SDM Training. We
study several NAT models’ performance under AT-baseline teacher model and AT-
SDM teacher model. We do this preliminary experiments on IWSLT14
EN$\rightarrow$DE task.

Teacher Model | Student Model
---|---
|  | NAT-CRF | CMLM
AT-Baseline | 28.34 | 22.41 | 26.78
AT-SDM | 29.51 | 21.83$\downarrow$ | 26.83 $\uparrow$

Table 2: Results of several NAT models’ performance under AT-baseline teacher
model and AT-SDM teacher model on IWSLT14 EN$\rightarrow$DE task.

As shown in Table 2, we reach a conclusion consistent with prior research
(Zhou et al., 2020) that, higher quality distilled data does not necessarily
improve the performance of NAT models. The little bit difference between our
experiments and theirs is that we generate distilled data by the same size AT
teacher models with the different training strategies, while they use the
different size teacher models with the same training strategies.

#### Problem 1:

SDM can improve the performance of the AT teacher models and lead to higher
quality distilled data, but it cannot further improve the performance of NAT
student models generally. We make some theoretical analysis in §4.2.

###  3.2 Straightforward but Failed Formulation

A straightforward way is to directly adopt SDM Training on NAT models. The
overview of the training strategies is shown in Figure 2. We do experiments of
all strategies on IWSLT14 EN$\rightarrow$DE task.

Figure 2: An overview of Baseline Training and SDM Training on (a) Raw Data
and (b) Distilled Data for a NAT model. Figure 3: An overview of Self-
Distillation Mixup Training with Pre-Rerank and Fine-Tune (SDMRT) Training on
distilled data for a NAT model.

Arch | Raw data | AT-distilled data
---|---|---
| baseline | SDM | baseline | SDM
AT | 28.34 | 29.51 $\uparrow$ | - | -
NAT-CRF | 18.41 | 19.61 $\uparrow$ | 22.41 | 21.93 $\downarrow$
CMLM | 25.04 | 25.97 $\uparrow$ | 26.78 | 26.67 $\downarrow$

Table 3: Results of several NAT models’ performance under classic SDM training
on IWSLT14 EN$\rightarrow$DE task.

In Table 3, it is clear that SDM Training can significantly improve the
translation quality of NAT models trained on raw data, which has the similar
conclusion with AT. But SDM Training gains no improvement when NAT models is
trained on distilled data. Such inconsistent results confuse us.

Arch | Test Set | D-Test Set
---|---|---
| baseline | SDM | baseline | SDM
NAT-CRF | 22.41 | 21.93 $\downarrow$ | 36.45 | 36.96 $\uparrow$
CMLM | 26.78 | 26.67 $\downarrow$ | 43.78 | 44.11 $\uparrow$

Table 4: Results of several NAT models’ performance on Test Set and D-Test Set
on IWSLT14 EN$\rightarrow$DE task.

The only difference is the different Train Set. One is original, and the other
is distilled. But the Test Set is the same one. So we build a Distilled Test
Set (D-Test Set), which keeps the source side of the test set and replace the
target side with the prediction by the AT teacher model. Furthermore, we re-
evaluate the performance of the NAT models on D-Test Set and the results in
Table 4 show SDM training worked on D-Test Set. So the key is the diversity
between the distilled data and the original data. It is explicable but not
expected to the original goal. Given these facts, we conclude that the
straightforward formulation is not effective.

#### Problem 2:

How to successfully use SDM Training for NAT models on distilled data, while
the distilled data brings bias comparing with the original data?

##  4 Approach

###  4.1 Self-Distillation Mixup with Pre-Rerank and Fine-Tune Training

We add two stages to the classic SDM Training process. One is Pre-Rerank on
self-distilled data, the other is Fine-Tune on filtered teacher-distilled
data. The overall training process is seen in Figure 3. We define our enhanced
training strategy as Self-Distillation Mixup with Pre-Rerank and Fine-Tune
(SDMRT) Training.

Let $\mathcal{D}$ denote the original parallel data:
$\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$, where $\mathcal{X}$ is the
original source-side corpus and $\mathcal{Y}$ is the original target-side
corpus. Let $\mathcal{\hat{D}}$ denotes the distilled parallel data:
$\mathcal{\hat{D}}=\\{\mathcal{X};\mathcal{\hat{Y}}\\}$, where
$\mathcal{\hat{Y}}$ is the distilled target-side corpus. For a clearer
presentation, we summarise the SDMRT training process concretely in Appendix.

When translation model $M$ generates $k$-candidates
$\mathbbm{Y}=\\{y_{0}^{{}^{\prime}},...,y_{k-1}^{{}^{\prime}}\\}$ on the input
sentence $x$, we can use another model $M_{r}$ to re-score the candidates
instead of scoring them by the model $M$ itself, which is known as Rerank.
Reranking is an effective method to change the data distribution. We define
the stage of re-score the self-distilled data in our SDMRT as Pre-Rerank. We
demonstrate two rerank strategies. One is using Language Model trained on the
original target-side corpus $M_{r}=LM(\mathcal{Y})$ to evaluate the Perplexity
of the conditional candidates. The other is using Alignment Model such as an
AT teacher model $M_{r}=AT(\mathcal{X},\mathcal{Y})$ to score the Conditional
Entropy between the source input and the prediction output.

Fine-Tune on filtered teacher-distilled data is an effective method to reduce
the bias gap mentioned earlier. The key of the Fine-Tune stage in our SDMRT is
to filter incorrect train set out. We calculate the score TER$(\hat{y},y)$ of
the target-side of the teacher-distilled data $\mathcal{\hat{D}}$, where
$(x,\hat{y})\in\hat{D}$; $(x,{y})\in{D}$, and only keep the sentence pairs
$(x,\hat{y})$ with TER score lower than $\tau$.

###  4.2 Theoretical Analysis

(Zhou et al., 2020) has used Bayesian Decision Theory to analyse distilled
data with access to the true distribution. Furthermore, we use Bayesian
Decision Theory to explain the reasons for the effectiveness and invalidity of
SDM in §3.1 and §3.2. We also give the reasons why our SDMRT is effective.

#### Background: Bayesian Decision Theory

In the problem of translations, let $x$ denote the input sequence and $y$
denote the output label sequence, where $x\in\mathcal{X}$ and
$y\in\mathcal{Y}$. Let $\mathcal{H}$ denote all the possible hypothesis
functions from the input to the output space:
$\mathcal{H}=\\{h:\mathcal{X}\rightarrow\mathcal{Y}\\}$. Let $r(y|x)$ denote
the conditional risk on the input $x$, which is the expected loss of
predicting y based on the posterior probabilities:

$\displaystyle r(y|x)=\mathbbm{E}_{P(y^{\prime}|x)}[L(y,y^{\prime})],$ (4)

where $L(y,y^{\prime})$ is the loss function that penalizes predicting the
true target $y^{\prime}$ as $y$. The classification task aims to find a
hypothesis function $h$ that minimizes the overall risk $R$ given by

$\displaystyle R(h)=\mathbbm{E}_{P(x)}[r(h(x)|x)]$ (5)

This is known as the Bayes risk. To minimize the overall risk, obviously we
need to minimize the conditional risk for each input $x$. The Bayesian
decision rule states that the global minimum of $R(h)$ is achieved when the
classifier makes predictions that minimize each conditional risk given $x$ and
this gives the Bayes optimal classifier:

$\displaystyle h^{*}(x)=\mathop{\arg\min}\limits_{y\in\mathcal{Y}}r(y|x)$ (6)

We define $h^{*}$ as the best function in theory. So the parallel data can be
considered as $\mathcal{D}=\\{\mathcal{X};h^{*}(\mathcal{X})\\}$. So:

  * •

For baseline training, we train models on
$\\{\mathcal{X};h^{*}(\mathcal{X})\\}$ and get $h_{1}$ close to $h^{*}$.

  * •

For SD training, we train models on $\\{\mathcal{X};h_{1}(\mathcal{X})\\}$ and
get another $h_{1}$.

  * •

For SDM training, we train models on
$\\{\mathcal{X};h^{*}(\mathcal{X})\\}\cup\\{\mathcal{X};h_{1}(\mathcal{X})\\}$
and get a new $h_{2}$ between $h^{*}$ and $h_{1}$, which means $h_{2}$ is
closer to $h^{*}$.

#### Modeling Diversity

Let $\mathcal{M}$ denote a trainable network model and $\Theta$ denote the
parameters of the model $\mathcal{M}$. Then the $\mathcal{D}$ can be defined
as:

$\displaystyle\mathcal{D}=\\{\mathcal{X};h(\mathcal{X},M^{*},\Theta^{*})\\}$
(7)

, where the theoretically best function $h^{*}$ can be replaced by the
theoretically best model $M^{*}$ and the theoretically best parameters
$\Theta^{*}$.

Let $M_{s}$ and $\Theta_{s}$ denote the student model and parameters. Let
$M_{t}$ and $\Theta_{t}$ denote the teacher model and parameters. Then SDM
Training can be considered to find the best $(M_{s},\Theta_{s})$ for:

$\displaystyle\\{\mathcal{X};h(\mathcal{X},M^{*},\Theta^{*})\\}\cup\\{\mathcal{X};h(\mathcal{X},M_{t}^{*},\Theta_{t}^{*})\\}$
(8)

, where $(M_{s},\Theta_{s})\subseteq(M^{*},\Theta^{*})$ and
$(M_{t},\Theta_{t})\subseteq(M^{*},\Theta^{*})$. If the model is the same one,
the training objective is simplified to only find the best $\Theta_{s}$. So
SDM Training is effective when AT or NAT models trained on original data. But
when training NAT models on distilled data using SDM, $M_{t}$ and $M_{s}$ is
different which makes the goal complex.

Modeling Diversity between different models can also be regarded as Domain
Specificity (Sennrich et al., 2016a) between the predictions of these models.
We can use Language Model (LM) to evaluate the Perplexity (PPL) of the
predictions. We use original target-side of train set to train the LM and
evaluate the average Perplexity of test set on IWSLT14 EN$\rightarrow$DE task.

Metrics | Original | Model Predictions
---|---|---
|  | AT | NAT-CRF | CMLM
Perplexity | 23.59 | 28.78 | 127.59 | 37.22

Table 5: Average PPL of original Test Set and multiple model predictions on
IWSLT14 EN$\rightarrow$DE.

As shown in Table 5, it is clear that different models occur diversity. Rerank
is an effective method to change the domain distribution of the distilled data
generated by $M_{s}$ to $M_{t}$. We further analyse the effect of Rerank in
§5.3.

#### Confirmation Bias

Network predictions are, of course, sometimes incorrect. In self-training,
this situation is reinforced when incorrect predictions are used as labels for
unlabeled samples, as it is the case in pseudo-labeling. Overfitting to
incorrect pseudo-labels predicted by the network is known as confirmation bias
(Arazo et al., 2020).

Multi-modality is typically characterized by token repetitions. NAT models may
consider many possible translations at the same time and have more likely to
generate inconsistent outputs than AT models. We use the frequency of repeated
tokens as a proxy for measuring NAT multi-modality (Ghazvininejad et al.,
2020). We evaluate the token repetitions of test set under classic SDM
training on IWSLT14 EN$\rightarrow$DE task. Multi-modality can be regard as
confirmation bias in NMT.

Metrics | Model Predictions
---|---
| AT | NAT-CRF | CMLM
Frequency of | 0.072% | 0.937% | 0.461%
repeated tokens | - | 13.01$\times$ | 6.42$\times$

Table 6: The percentage of repeated tokens of different data on IWSLT14
EN$\rightarrow$DE task.

Model | Iter | WMT14 | WMT16 | IWSLT14
---|---|---|---|---
|  | En $\rightarrow$ De | De $\rightarrow$ En | En $\rightarrow$ Ro | Ro $\rightarrow$ En | En $\rightarrow$ De
AT models
Transformer-base(Vaswani et al., 2017) | - | 27.3 | - | - | - | -
Transformer-base(Our Implementation) | - | 27.67 | 31.17 | 34.37 | 33.98 | 28.34
Transformer-big(Our Implementation) | - | 28.83 | - | - | - | -
NAT models
V-NAT(Gu et al., 2018) | - | 19.17 | 23.20 | 29.79 | 31.44 | -
NAT-CRF(Sun et al., 2019) | - | 23.32 | 25.75 | - | - | 26.39
iNAT(Lee et al., 2018) | 10 | 21.61 | 25.48 | 29.32 | 30.19 | -
CMLM*(Ghazvininejad et al., 2019) | 10 | 27.03 | 30.53 | 33.08 | 33.31 | -
Our Implementation
V-NAT | - | 16.32 | 20.01 | 24.83 | 25.57 | 15.79
NAT-CRF | - | 23.76 | 25.56 | 29.31 | 29.17 | 22.41
iNAT | 10 | 21.19 | 25.37 | 29.17 | 30.09 | 22.75
CMLM | 10 | 26.03 | 29.61 | 32.94 | 33.07 | 26.78
CMLM* | 10 | 27.06 | 30.94 | 33.07 | 33.36 | -
SDMRT
V-NAT | - | 17.31 (+0.99) | 21.07 (+1.06) | 25.87 (+1.04) | 26.64 (+1.07) | 16.9 (+1.11)
NAT-CRF | - | 24.68 (+0.92) | 26.47 (+0.91) | 30.35 (+1.04) | 30.26 (+1.09) | 23.51 (+1.10)
iNAT | 10 | 21.97 (+0.78) | 25.94 (+0.57) | 29.94 (+0.77) | 30.9 (+0.81) | 23.64 (+0.89)
CMLM | 10 |  26.77 (+0.74) |  30.45 (+0.84) |  33.57 (+0.63) |  33.69 (+0.62) |  27.49 (+0.71)
CMLM* | 10 |  27.72 (+0.66) |  31.65 (+0.71) |  33.72 (+0.65) |  33.94 (+0.58) | -

Table 7: Performance of BLEU (Papineni et al., 2002) score of multiple NAT
models on WMT14 En$\leftrightarrow$De, WMT16 En$\leftrightarrow$Ro and IWSLT14
En$\rightarrow$De. Note that we train CMLM baseline using Transformer-base as
teacher model, but their (Ghazvininejad et al., 2019) teacher model is
Transformer-big. The results of CMLM* are under Transformer-big teacher model.

As shown in Table 6, predictions of NAT models have more repeated tokens than
those of AT models, which proves the confirmation bias occurring between the
self-distilled data and the teacher-distilled data. To reduce the confirmation
bias, Fine-Tune the model on the correct data is necessary. We further analyse
the effect of Fine-Tune in §5.3.

##  5 Experiments & Results

We evaluate our SDMRT Training strategy on five standard NMT benchmarks
including WMT14 En$\leftrightarrow$De and WMT16 En$\leftrightarrow$Ro in both
directions, and IWSLT14 En$\rightarrow$De in single direction.

###  5.1 Experiments Setting

#### Datasets

The sizes of the datasets are (train=4.5M / valid=3k / test=3k / dict=42k),
(train=610k / valid=2k / test=2k / dict=40k) and (train=160k / valid=7k /
test=6k / dict=10k) for WMT14 En$\leftrightarrow$De, WMT14
En$\leftrightarrow$Ro and IWSLT14 En$\rightarrow$De respectively. We use
Transformer-base (Vaswani et al., 2017) as teacher model to create the AT-
distilled data, and initialize NAT model with the same model size. We use BPE
(Sennrich et al., 2016b) to tokenize the sentences and create the vocabulary.

#### Model Configurations

For WMT tasks, we use the following hyperparameters: $n_{layers}$ = 12,
$n_{heads}$ = 8, $d_{hidden}$ = 512, $d_{FFN}$ = 2048, $lr$ = 5e-4,
$n_{warmsup}$ = 10k, $bz$ = 8192, $n_{gpu}$ = 8. For IWSLT tasks, smaller
hyperparameters are set: $n_{layers}$ = 12, $n_{heads}$ = 4, $d_{hidden}$ =
512, $d_{FFN}$ = 1024, $lr$ = 5e-4, $n_{warmsup}$ = 6k, $bz$ = 10240,
$n_{gpu}$ = 1. Our models are trained on Tesla V100 GPUs. Adam (Kingma and Ba,
2015) is used as the optimizer and the inversed-sqrt scheduler is used for lr
decaying. All models are trained for 300k steps, and last 4 checkpoints are
averaged. For the Pre-Rerank stage, we use AT teacher model to re-score the
predictions instead of a language model. For the Fine-Tune stage, we filter
the distilled data with TER score less than 0.8, and pre-train models for 240k
steps and further-train for 60k steps. We implement models in the experiment
with fairseq (Ott et al., 2019).

###  5.2 Main Results

Table 7 shows the translation performance of our SDMRT method. For all tasks,
SDMRT forms the baselines by a large margin (0.6$\sim$1.2 BLEU) on distilled
data.

###  5.3 Study

#### SDM v.s. SDMRT on different NAT models

As shown in Table 8, the performance of our SDMRT outperforms that of classic
SDM. Surprisingly, NAT-CRF with SDM can achieve positive results on smaller
datasets, while CMLM achieves positive results on larger datasets.

Model | WMT14 EN$\rightarrow$DE | IWSLT14 EN$\rightarrow$DE
---|---|---
Baseline
NAT-CRF | 23.76 | 22.41
CMLM | 26.03 | 26.78
SDM
NAT-CRF | 23.15(-0.61) | 22.95(+0.54)
CMLM | 26.23(+0.20) | 26.67(-0.11)
SDMRT
NAT-CRF | 24.68(+0.92) | 23.51(+1.10)
CMLM | 26.77(+0.74) | 27.49(+0.71)

Table 8: SDM v.s. SDMRT on IWSLT14 EN$\rightarrow$DE task and WMT14
EN$\rightarrow$DE task.

STEP | Baseline | SDMRT |  $\Delta$BLEU
---|---|---|---
1 | 19.54 | 21.58 | 2.04
5 | 25.57 | 26.24* | 0.67
10 | 26.03 | 26.77* | 0.74

Table 9: BLEU scores of every step with maxiter=10 for CMLM model on WMT14
EN$\rightarrow$DE task. SDMRT can overstep baseline earlier in step of 5.
Figure 4: Results of different training strategies for CMLM model on IWSLT14
EN$\rightarrow$DE task.

#### SDM v.s. SDMRT on different steps of IR NAT models

It is clear in Figure 4, SDM only has positive influence at the first several
steps, but our SDMRT can improve translation quality at any step.

#### Acceleration for Iterative Refinement NAT

For Iterative Refinement (IR) NAT such as CMLM, it usually requires more than
4 steps to make the refinement converge, which causes negative influence on
the speedup. Table 9 shows, on AT-distilled data, our SDMRT training strategy
can overstep baseline within 5 steps which means 2$\times$ acceleration
compared to baseline.

Metrics | Rerank | Original | NAT-CRF | CMLM
---|---|---|---|---
Perplexity | w/o | 23.59 | 127.59 | 37.22
| w/ | - | 76.55$\downarrow$ | 30.47$\downarrow$

Table 10: Re-visit average PPL of multiple models predictions on IWSLT14
EN$\rightarrow$DE.

Model | SDM | SDMRT
---|---|---
AT | 0.072% | - | - | -
NAT-CRF | 0.937% | 13.01$\times$ | 0.533%$\downarrow$ | 7.40$\times$
CMLM | 0.461% | 6.42$\times$ | 0.195%$\downarrow$ | 2.71$\times$

Table 11: Re-visit the percentage of repeated tokens of different data on
IWSLT14 EN$\rightarrow$DE task.

Model | SDMRT
---|---
| w/ LM-Rerank | w/ AT-Rerank
NAT-CRF | 24.72 | 24.68
CMLM | 26.33 | 26.77

Table 12: Different Rerank strategies on WMT14 EN$\rightarrow$DE task.

Model | SDMRT
---|---
| Data-Filtering | No Data-Filtering
NAT-CRF | 24.68 | 23.83 (-0.85)
CMLM | 26.77 | 26.13 (-0.64)

Table 13: Different Data-Filtering strategies on WMT14 EN$\rightarrow$DE task.

#### Re-visit Modeling Specificity

Following the preliminary experiments before, we re-visit the average
Perplexity of test set on IWSLT14 EN$\rightarrow$DE. The results in Table 10
prove that the distribution of NAT models’ predictions under Rerank is more
closer to that of original data.

#### Re-visit Confirmation Bias

As mentioned earlier, confirmation bias in NMT is Multi-modality, and Multi-
modality is typically characterized by token repetitions. We re-measure the
frequency of repeated tokens. As shown in Table 11, the percentage of repeated
tokens of with rerank stage sharply decreases, which proves that our SDMRT is
an effective method.

#### AT-Rerank v.s. LM-Rerank

We compare the effect of different rerank strategies. As shown in Table 12,
AT-Rerank performs overall better than LM-Rerank. Specially, LM-Rerank can
achieve similar performance as AT-Rerank for NAT-CRF model.

#### Effect of Data-Filtering on Fine-Tune Stage

To measure the effect of filtered teacher-distilled data, we change our SDMRT
training process and fine-tune the models on teacher-distilled data. The
results in Table 13 prove that Data-Filtering is important in our SDMRT
training process.

##  6 Conclusion

In this work, we empirically find directly adopting the classic Self-
Distillation Mixup (SDM) Training strategy to NAT models is not as effective
as in AT models. Through careful analysis, we observe Modeling Diversity and
Confirmation Bias between the AT teacher model and the NAT student models,
causing SDM theoretically invalid for NAT models when trained on AT-distilled
data. Furthermore, we propose an enhanced training strategy named SDMRT and
achieve new SOTA performance for multiple NAT models compared to the
baselines. In addition, it’s demonstrated to be an effective method to reduce
the iteration number for Iterative Refinement NAT.

## References

  * Arazo et al. (2020) Eric Arazo, Diego Ortego, Paul Albert, Noel E. O’Connor, and Kevin McGuinness. 2020\.  Pseudo-labeling and confirmation bias in deep semi-supervised learning.  In _2020 International Joint Conference on Neural Networks, IJCNN 2020, Glasgow, United Kingdom, July 19-24, 2020_ , pages 1–8. IEEE.
  * Bahdanau et al. (2015) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015.  Neural machine translation by jointly learning to align and translate.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Du et al. (2021) Cunxiao Du, Zhaopeng Tu, and Jing Jiang. 2021.  Order-agnostic cross entropy for non-autoregressive machine translation.  In _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 2849–2859. PMLR.
  * Freitag et al. (2017) Markus Freitag, Yaser Al-Onaizan, and Baskaran Sankaran. 2017.  Ensemble distillation for neural machine translation.  _CoRR_ , abs/1702.01802.
  * Ghazvininejad et al. (2020) Marjan Ghazvininejad, Vladimir Karpukhin, Luke Zettlemoyer, and Omer Levy. 2020\.  Aligned cross entropy for non-autoregressive machine translation.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 3515–3523. PMLR.
  * Ghazvininejad et al. (2019) Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. 2019.  Mask-predict: Parallel decoding of conditional masked language models.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019_ , pages 6111–6120. Association for Computational Linguistics.
  * Gu et al. (2018) Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. K. Li, and Richard Socher. 2018\.  Non-autoregressive neural machine translation.  In _6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings_. OpenReview.net.
  * Gu and Kong (2021) Jiatao Gu and Xiang Kong. 2021.  Fully non-autoregressive neural machine translation: Tricks of the trade.  In _Findings of the Association for Computational Linguistics: ACL/IJCNLP 2021, Online Event, August 1-6, 2021_ , volume ACL/IJCNLP 2021 of _Findings of ACL_ , pages 120–133. Association for Computational Linguistics.
  * Hinton et al. (2015) Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015.  Distilling the knowledge in a neural network.  _CoRR_ , abs/1503.02531.
  * Kim and Rush (2016) Yoon Kim and Alexander M. Rush. 2016.  Sequence-level knowledge distillation.  In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016_ , pages 1317–1327. The Association for Computational Linguistics.
  * Kingma and Ba (2015) Diederik P. Kingma and Jimmy Ba. 2015.  Adam: A method for stochastic optimization.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Lee et al. (2021) Ann Lee, Michael Auli, and Marc’Aurelio Ranzato. 2021.  Discriminative reranking for neural machine translation.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021_ , pages 7250–7264. Association for Computational Linguistics.
  * Lee et al. (2018) Jason Lee, Elman Mansimov, and Kyunghyun Cho. 2018.  Deterministic non-autoregressive neural sequence modeling by iterative refinement.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 - November 4, 2018_ , pages 1173–1182. Association for Computational Linguistics.
  * Ott et al. (2019) Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019.  fairseq: A fast, extensible toolkit for sequence modeling.  In _Proceedings of NAACL-HLT 2019: Demonstrations_.
  * Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002.  Bleu: a method for automatic evaluation of machine translation.  In _Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, July 6-12, 2002, Philadelphia, PA, USA_ , pages 311–318. ACL.
  * Saharia et al. (2020) Chitwan Saharia, William Chan, Saurabh Saxena, and Mohammad Norouzi. 2020.  Non-autoregressive machine translation with latent alignments.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , pages 1098–1108. Association for Computational Linguistics.
  * Sennrich et al. (2016a) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016a.  Controlling politeness in neural machine translation via side constraints.  In _NAACL HLT 2016, The 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, San Diego California, USA, June 12-17, 2016_ , pages 35–40. The Association for Computational Linguistics.
  * Sennrich et al. (2016b) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016b.  Neural machine translation of rare words with subword units.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume 1: Long Papers_. The Association for Computer Linguistics.
  * Sun et al. (2019) Zhiqing Sun, Zhuohan Li, Haoqing Wang, Di He, Zi Lin, and Zhi-Hong Deng. 2019\.  Fast structured decoding for sequence models.  In _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pages 3011–3020.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008.
  * Zhou et al. (2020) Chunting Zhou, Jiatao Gu, and Graham Neubig. 2020.  Understanding knowledge distillation in non-autoregressive machine translation.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.

##  Appendix A Appendix

1: procedure TRAIN_RERANK($\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$):

2: If ( $M_{r}$ is Language Model )

3: Train randomly initialized $M_{r}$ on $\mathcal{Y}$

4: Elseif ( $M_{r}$ is Alignment Model )

5: Train randomly initialized $M_{r}$ on $\mathcal{D}$

6: Return $M_{r}$

7:

1: procedure FILTER($\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$,
$\mathcal{\hat{D}}=\\{\mathcal{X};\mathcal{\hat{Y}}\\}$):

2: Filter $\mathcal{D}^{*}=$ {($x$, $\hat{y}$) $|$ TER($\hat{y}$, $y$)
$<\tau$; $(x,\hat{y})\in\hat{D}$; $(x,{y})\in{D}$ }

3: Return $\mathcal{D}^{*}$

4:

1: procedure TRAIN($\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$):

2: Train randomly initialized $M$ on $\mathcal{D}$

3: Return $M$

4:

1: procedure RERANK_DIS($\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$, $M$,
$M_{r}$):

2: For ($x$, $y$) in $\mathcal{D}$:

3: Generate $k$-candidates
$\mathbbm{Y}=\\{y_{0}^{{}^{\prime}},...,y_{k-1}^{{}^{\prime}}\\}$ by $M(x)$

4: If ( $M_{r}$ is Language Model )

5: Select the best $y^{{}^{\prime}}$ order by sorted
$\\{M_{r}(y^{{}^{\prime}})|y^{{}^{\prime}}\in\mathbbm{Y}\\}$

6: Elseif ( $M_{r}$ is Alignment Model )

7: Select the best $y^{{}^{\prime}}$ order by sorted
$\\{M_{r}(x,y^{{}^{\prime}})|y^{{}^{\prime}}\in\mathbbm{Y}\\}$

8: Return $\mathcal{D}^{{}^{\prime}}=\\{(x,y^{{}^{\prime}})\\}$

9:

1: procedure FINE_TUNE($M$, $\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$):

2: Further train $M$ on $\mathcal{D}$

3: Return $M$

4:

1: procedure SDMRT($\mathcal{D}=\\{\mathcal{X};\mathcal{Y}\\}$,
$\mathcal{\hat{D}}=\\{\mathcal{X};\mathcal{\hat{Y}}\\}$):

2: $M_{r}\leftarrow$TRAIN_RERANK($\mathcal{D}$)

3: $\mathcal{D}^{*}\leftarrow$ FILTER($\mathcal{D}$, $\mathcal{\hat{D}}$)

4: $M\leftarrow$TRAIN($\mathcal{\hat{D}}$)

5: $\hat{\mathcal{D}}_{N}\leftarrow$RERANK_DIS($\mathcal{D}$, $M$, $M_{r}$)

6: $M\leftarrow$TRAIN($\mathcal{\hat{D}}\cup\hat{\mathcal{D}_{N}}$)

7: $M\leftarrow$FINE_TUNE($M,\hat{\mathcal{D}^{*}}$)

8: Return $M$

Algorithm 1 SDMRT Training