#  Neural Network-based Two-Dimensional Filtering for OTFS Symbol Detection

Jiarui Xu, Karim Said, Lizhong Zheng, and Lingjia Liu J. Xu, K. Said, and L.
Liu are with Wireless@Virginia Tech, the Bradley Dept. of ECE at Virginia
Tech. L. Zheng is with the EECS Department at the Massachusetts Institute of
Technology. The work of J. Xu, K. Said, L. Zheng, and L. Liu was supported by
the National Science Foundation under Grant NSF/CCF-2003059.

###### Abstract

Orthogonal time frequency space (OTFS) is a promising modulation scheme for
wireless communication in high-mobility scenarios. Recently, a reservoir
computing (RC) based approach has been introduced for online subframe-based
symbol detection in the OTFS system, where only the limited over-the-air (OTA)
pilot symbols are utilized for training. However, the previous RC-based
approach does not design the RC architecture based on the properties of the
OTFS system to fully unlock the potential of RC. This paper introduces a novel
two-dimensional RC (2D-RC) approach for online symbol detection on a subframe
basis in the OTFS system. The 2D-RC is designed to have a two-dimensional (2D)
filtering structure to equalize the 2D circular channel effect in the delay-
Doppler (DD) domain of the OTFS system. With the introduced architecture, the
2D-RC can operate in the DD domain with only a single neural network, unlike
our previous work which requires multiple RCs to track channel variations in
the time domain. Experimental results demonstrate the advantages of the 2D-RC
approach over the previous RC-based approach and the compared model-based
methods across different modulation orders.

###### Index Terms:

2D-RC, OTFS, online learning, deep learning, signal detection, channel
equalization, two-dimensional

##  I Introduction

Next-generation wireless communication systems are expected to support
reliable communication quality in high-mobility scenarios, such as high-speed
railways, unmanned aerial vehicles, and low earth orbit [1]. Orthogonal time
frequency space (OTFS) modulation [2] has emerged as a potential technology to
combat the high Doppler effect in such scenarios. Different from orthogonal
frequency division multiplexing (OFDM) which multiplexes information symbols
in the time-frequency (TF) domain, OTFS is a two-dimensional (2D) modulation
scheme that transmits information symbols in the delay-Doppler (DD) domain. In
the DD domain, the channel has a sparse and compact representation [2], making
it intrinsically suitable for reliable data transmission in time-varying
channels.

The benefits of adopting OTFS modulation in high-mobility scenarios have
attracted substantial interest in investigating equalization techniques for
the OTFS system. Existing approaches can be mainly categorized into two
classes: model-based methods and learning-based approaches. Model-based
approaches are designed based on analyzing the input-output relationship and
the structure of the equivalent channel matrix in the OTFS system.
Specifically, linear minimum mean square error (LMMSE) based linear
equalizers, such as the work in [3], are introduced to conduct low-complexity
detection by taking advantage of the channel structure. To achieve better
performances, non-linear detectors, such as the message passing algorithm
(MPA) based approaches [4, 5, 6], are developed to approximate the maximum a
_a posteriori_ (MAP) performance with low complexity by leveraging the channel
sparsity. While model-based approaches are explainable and easy to analyze,
the assumptions of simplified system models and perfect channel knowledge may
degrade the performance when such assumptions do not hold.

Learning-based detection approaches offer a promising solution for OTFS symbol
detection in high-mobility scenarios by leveraging the power of neural
networks (NNs). For example, in [7], a two-dimensional convolutional neural
network (2D CNN) method is developed which adopts MPA for data augmentation
and then processes the received signal as a 2D image with the CNN. However,
such approaches require extensive training samples along with a long training
time. While recent advances, such as GAMP-NET [8], have attempted to unfold
existing iterative model-based algorithms with NNs to reduce the number of
trainable parameters, they still have a high demand for the amount of training
data. Moreover, deep unfolding approaches also rely on perfect knowledge of
channel state information (CSI), limiting their usage in practical scenarios.

Recently, reservoir computing (RC) has risen as an appealing approach due to
its efficient training procedure [9]. RC is a particular type of recurrent
neural network (RNN) that only contains a few numbers of trainable parameters
and can be learned with a limited amount of training data. Such properties
enable its application in the online symbol detection task for both the OFDM
system [10, 11, 12, 13] and the OTFS system [14], where over-the-air (OTA)
training data is extremely scarce. Particularly, an RC-based approach for the
OTFS system is first introduced in [14] to perform the online subframe-based
detection by adopting multiple RCs in the time domain. This approach, however,
directly applies the existing RC structure designed for the OFDM system in
[10] and does not incorporate the domain knowledge of the OTFS system to
reveal the full potential of RC.

In this work, we introduce a novel 2D-RC structure for the online subframe-
based symbol detection task in the OTFS system. The introduced 2D-RC retains
the advantages of RC that can be learned with only a limited number of
training pilots and has a low training complexity, which differentiates it
from other NNs that require a large amount of training data and a long
training time. Furthermore, 2D-RC is novelly designed to have a 2D filtering
structure tailored to equalizing the DD-domain 2D circular channel effect. The
2D structure allows it to operate in the DD domain with a single NN, which is
different from the previous RC-based approach in [14] that requires multiple
RCs in the time domain to track channel changes within one subframe. As the
number of RCs does not need to be configured, the 2D-RC can be more easily
adopted in practice, making it a more general approach. Simulation results
show that 2D-RC can achieve substantial performance improvement over the
previous multiple-RC approach under different modulation orders. The results
also demonstrate the advantages of 2D-RC over compared model-based schemes.

Notations: Non-bold letter, bold lowercase letter, bold uppercase letter, and
bold Euler script letter, i.e., $x$, $\bm{x}$, $\bm{X}$, and $\mathbcal{X}$,
denote scalar, vector, matrix, and tensor, respectively. $\bm{F}_{M}$ is the
normalized $M$-point discrete Fourier transform (DFT). $(\cdot)^{\dagger}$
denotes the Moore–Penrose inverse. $\langle\cdot\rangle_{M}$ is the modulo
operator of divider $M$. $\mathrm{vec}(\cdot)$ denotes the operation of
vectoring the matrix by stacking along the columns, and
$\mathrm{vec}^{-1}(\cdot)$ denotes the reverse operation. $\odot$ denotes the
matrix Hadamard product operation. The $n$-mode Hadamard product between the
matrix $\bm{U}\in\mathbb{C}^{I_{n}\times I_{n+1}}$ and the $N$-dimensional
tensor $\mathbcal{X}\in\mathbb{C}^{I_{1}\times I_{2}\times\dots\times I_{N}}$
is defined as

$(\bm{U}\odot_{n}\mathbcal{X})[i_{1},\dots,i_{n},i_{n+1},\dots,i_{N}]\\\
=U[i_{n},i_{n+1}]\cdot X[i_{1},\dots,i_{n},i_{n+1},\dots,i_{N}],$

where $U[i_{n},i_{n+1}]$ is the $(i_{n},i_{n+1})$-th element in $\bm{U}$, and
$X[i_{1},\dots,i_{n},i_{n+1},\dots,i_{N}]$ is the
$(i_{1},\dots,i_{n},i_{n+1},\dots,i_{N})$-th element in $\mathbcal{X}$. The
concatenation of two tensors $\mathbcal{X}_{1}$ and $\mathbcal{X}_{2}$ along
the $n$-th dimension is represented by
$\text{cat}_{n}(\mathbcal{X}_{1},\mathbcal{X}_{2})$.

##  II System Model

The transmitter and receiver structures in the OTFS system are shown in Fig.
1. $Q$-th order quadrature amplitude modulation ($Q$-QAM) symbols are
modulated in the DD domain, which forms the transmitted signal $\bm{X}$ of
size $M\times N$. $M$ and $N$ denote the number of delay bins and Doppler
bins. In this work, we consider the CP-OTFS system with practical rectangular
transmit and received pulse shaping waveforms. The CP-OTFS system can be
implemented as an overlay of the OFDM system, where a cyclic prefix (CP) is
added for each OFDM symbol in the subframe, i.e., $N$ CPs for one OTFS
subframe.

###  II-A OTFS Transmitter and Receiver

The transmitted signal $\bm{X}$ in the DD domain is converted to the TF domain
through inverse symplectic finite Fourier transform (ISFFT) operation, which
can be written as

$\displaystyle\bm{X}_{tf}=\mathrm{ISFFT}(\bm{X})=\bm{F}_{M}\bm{X}\bm{F}_{N}^{H},$
(1)

where $\bm{X}_{tf}$ represent the TF domain signal. The TF domain signal is
then transformed to the time domain signal $\bm{S}\in\mathbb{C}^{M\times N}$
for transmission by the Heisenberg transform and appended with CP. The
transmitted time-domain signal after adding CP can be expressed as

$\displaystyle\bm{S}=\bm{A}_{cp}\bm{F}_{M}^{H}\bm{X}_{tf}=\bm{A}_{cp}\bm{X}\bm{F}_{N}^{H},$
(2)

where
$\bm{A}_{cp}=[\bm{G}_{cp}^{T},\bm{I}_{M}^{T}]^{T}\in\mathbb{R}^{(N_{cp}+M)\times
M}$ is the CP addition matrix, $\bm{G}_{cp}$ with a size of $N_{cp}\times M$
consists of the last $N_{cp}$ rows of the identity matrix $\bm{I}_{M}$, and
$N_{cp}$ is the CP length. The vector form can be written as
$\bm{s}=\mathrm{vec}(\bm{S})\in\mathbb{C}^{MN\times 1}$.

Figure 1: OTFS system diagram.

Let $\bm{R}=\mathrm{vec}^{-1}(\bm{r})\in\mathbb{C}^{M\times N}$ be the
unfolded time domain received signal. After removing CP, the received time
domain signal $\bm{R}$ is converted back to the TF domain $\bm{Y}_{tf}$
through the Winger transform, which can be formulated by

$\displaystyle\bm{Y}_{tf}=\bm{F}_{M}\bm{E}_{cp}\bm{R},$ (3)

where $\bm{E}_{cp}=[\bm{0}_{M\times
N_{cp}},\bm{I}_{M}]\in\mathbb{R}^{M\times(N_{cp}+M)}$ is the CP removal
matrix. The DD domain received signal $\bm{Y}$ is obtained by applying the
SFFT to the $\bm{Y}_{tf}$, which is expressed as

$\displaystyle\bm{Y}=\mathrm{SFFT}(\bm{Y}_{tf})=\bm{F}_{M}^{H}\bm{Y}_{tf}\bm{F}_{N}.$
(4)

###  II-B Channel

The channel response of the time-varying channel in the DD domain can be
represented by

$\displaystyle
h(\tau,\nu)=\sum_{i=0}^{P-1}h_{i}\delta(\tau-\tau_{i})\delta(\nu-\nu_{i}),$

where $h_{i}$, $\tau_{i}$, and $\nu_{i}$ represent the complex path gain,
delay, and Doppler shift of the $i$-th path; $P$ is the number of propagation
paths. The normalized delay shift $\mathcal{l}_{i}$ and Doppler shift
$\kappa_{i}$ are given by $\tau_{i}=\frac{\mathcal{l}_{i}}{M\Delta f}$ and
$\nu_{i}=\frac{\kappa_{i}}{T_{f}}$, where $\mathcal{l}_{i}$ and $\kappa_{i}$
are not necessarily integers, $\Delta f$ is the subcarrier spacing, and
$T_{f}$ is the total subframe duration. In the time domain, the received
signal can be expressed as [2]

$\displaystyle r(t)=\int\int h(\tau,\nu)s(t-\tau)e^{j2\pi\nu(t-\tau)}d\tau
d\nu+w(t),$

where $s(t)$ denotes the transmitted signal, and $w(t)$ is the additive
Gaussian noise.

###  II-C Input-Output Relationship

The DD-domain input-output relationship in the CP-OTFS system can be expressed
as [15]

$\begin{split}Y[l,k]&=\sum_{i=0}^{P-1}\sum_{l^{\prime}=0}^{M-1}\sum_{k^{\prime}=0}^{N-1}h_{i}{z}^{k_{p}(N_{cp}+l-\mathcal{l}_{i})}\psi_{M}(l^{\prime}-\mathcal{l}_{i})\\\
&\times\psi_{N}(\kappa_{i}-k^{\prime})X[\langle
l-l^{\prime}\rangle_{N},\langle k-k^{\prime}\rangle_{M}],\end{split}$ (5)

where $Y[l,k]$ represents the ($l,k$)-th element of $\bm{Y}$, ${z}\triangleq
e^{j\frac{2\pi}{(M+N_{cp})N}}$, and
$\psi_{M}(x)\triangleq\frac{1}{M}e^{j\pi\frac{M-1}{M}x}\frac{\sin\pi
x}{\sin\pi x/M}$. For simplicity, we omit the noise term. When
$\mathcal{l}_{i}$ and $\kappa_{i}$ are integers, the (5) simplifies to

$\displaystyle{Y}[l,k]=\sum_{i=0}^{P-1}h_{i}{z}^{k_{i}(N_{cp}+l-l_{i})}{X}[\langle
l-l_{i}\rangle_{M},\langle k-k_{i}\rangle_{N}],$ (6)

where $l_{i}$ and $k_{i}$ are the integer delay and integer Doppler.

###  II-D Problem formulation

(a) Blockwise pilot pattern (b) Spike pilot pattern

Figure 2: Pilot patterns. The green grids are filled with known pilot symbols.
The green grid with a square marker denotes the spike pilot. The cross markers
are guard symbols. The blank region represents data symbol positions.

The OTFS symbol detection task is to recover the transmitted DD-domain symbol
$\bm{X}$ in one OTFS subframe from the received signal. To aid the detection
of the unknown data symbols, the pilot symbols, which are known at both the
transmitter and receiver sides, are inserted in each subframe. In this paper,
we consider two pilot structures: the blockwise pilot pattern and the spike
pilot pattern, which are shown in Fig. 2. Specifically, the pilot symbols are
placed in a block of the subframe in the blockwise pilot structure. In the
spike pilot structure, a spike pilot is transmitted along with guard symbols
surrounding it. For learning-based approaches, the blockwise pilot pattern is
adopted to learn the interference between the data symbols and pilot symbols.
For model-based schemes that require knowledge of the CSI, the spike pilot
pattern is utilized for channel estimation [16]. More details about the choice
of pilot patterns for learning-based and model-based approaches are provided
in [14].

Denote $\bm{\Omega}$ as the pilot position indication matrix with $1$
indicating the pilot positions and $0$ otherwise. For the introduced learning-
based approach, the input to the NN is the received DD-domain signal $\bm{Y}$.
The training target is composed of the pilot symbols modulated in the DD
domain, which can be denoted as
$\bm{X}_{\mathrm{train}}\triangleq\bm{\Omega}\odot\bm{X}$.

##  III Introduced Approach – 2D-RC

In this section, we introduce the 2D-RC approach for online subframe-based
symbol detection in the OTFS system. RC is a class of RNNs for processing
temporal or sequential data. It consists of an input layer, an RNN-based
reservoir, and an output layer. The characteristic feature of RC is that the
input layer and the reservoir weights are fixed after random initialization
and only the output layer is updated through a simple linear regression. The
introduced 2D-RC retains the same fast and simple training process as RC,
enabling it to perform online symbol detection on a subframe basis. Moreover,
it is uniquely designed to facilitate online symbol detection tailored towards
the OTFS system. Specifically, the DD-domain channel works as a 2D circular
operation over the transmitted signal in the OTFS system as shown in (5). To
equalize this 2D circular channel effect, 2D-RC is designed to have a circular
padding procedure and a 2D filtering process. As 2D-RC works in the DD domain,
it only requires a single NN for detection, as opposed to previous work [14]
that exploits multiple single-dimensional RCs (1D-RCs) to track the channel
variations in the time domain. For the remainder of this paper, we refer to
the RC structure in [14] as “1D-RC”.

Figure 3: 2D-RC Structure. For simplicity, the nonlinear function and the
extended state are ignored here.  Figure 4: The 2D windowing process in 2D-RC.

###  III-A Pre-processing

The input to 2D-RC is the DD-domain received signal
$\bm{Y}\in\mathbb{C}^{M\times N}$. Two pre-processing procedures, including
windowing and padding, are employed before the processing of 2D-RC.

####  III-A1 2D Windowing

As discussed in [10], adopting a sliding window to the input can increase the
short-term memory of RC. Following the same philosophy, we adopt the 2D
windowing process, where a 2D sliding window with size $M_{w}\times N_{w}$ is
employed to process the input. The $M_{w}$ is the window size along the delay
dimension and $N_{w}$ is the window size along the Doppler dimension. For each
${Y}_{c}[l,k]$, the windowing region is obtained by
$\bm{Y}_{w}[l,k]=\bm{Y}_{c}[l-M_{w}+1:l,k-N_{w}+1:k]\in\mathbb{C}^{M_{w}\times
N_{w}}$. When the $l<M_{w}-1$ or $k<N_{w}-1$, zeros are filled in the
windowing region to maintain the window size of $M_{w}\times N_{w}$. The
windowed input is formed by
$\bm{y}_{w}[l,k]=\mathrm{vec}(\mathrm{rev}(\bm{Y}_{w}[l,k]^{T}))\in\mathbb{C}^{N_{i}}$,
where $\mathrm{rev}(\cdot)$ stands for reserving the values in the matrix
along both dimensions, $\mathrm{vec}(\cdot)$ represents vectoring the matrix
by stacking along the columns, and $N_{i}=M_{w}N_{w}$. By collecting all the
$\bm{y}_{w}[l,k]$, we obtain an input tensor
$\mathbcal{Y}_{w}\in\mathbb{C}^{N_{i}\times M\times N}$. Fig. 4 visualizes the
2D windowing process.

####  III-A2 2D Circular Padding

RC requires a degree of forgetfulness to remove the impact from the random
initialization of the internal state [17]. Therefore, 2D-RC also needs to
learn the optimal forget length to eliminate the impact of the initial state.
The 2D circular padding process is designed to facilitate the learning process
of the optimal forget length. Let $M_{f}$ and $N_{f}$ be the maximum forget
length along the delay and Doppler dimension, respectively. The 2D padded
input
$\tilde{\mathbcal{Y}}\in\mathbb{C}^{N_{i}\times(M+M_{f})\times(N+N_{f})}$ is
obtained by concatenating the $\mathbcal{Y}_{w}$ along the second and third
dimensions as the following:

$\displaystyle\tilde{\mathbcal{Y}}=$
$\displaystyle\text{cat}_{2}(\text{cat}_{3}(\mathbcal{Y}_{w},\;\;\mathbcal{Y}_{w}[:,:,0:N_{f}-1]),$
$\displaystyle\text{cat}_{3}(\mathbcal{Y}_{w}[:,0:M_{f}-1,:],\mathbcal{Y}_{w}[:,0:M_{f}-1,0:N_{f}-1])).$

Specifically, the 2D circular padding exploits the values at the start to pad
at the end of the corresponding dimension. The circular channel operation
inspires the utilization of circular padding to initialize the states of
2D-RC.

###  III-B Structure of 2D-RC

Denote $\tilde{\bm{y}}[m,n]\in\mathbb{C}^{N_{i}}$ as the ($m,n$)-th element
along the second and third dimensions of the pre-processed input
$\tilde{\mathbcal{Y}}$. We design the state transition equation for 2D-RC as

$\displaystyle\bm{u}[m,n]$
$\displaystyle=f(\bm{W}_{i}\;\tilde{\bm{y}}[m,n]+\bm{W}_{r}\;\bm{u}[m-1,n]$
$\displaystyle+\bm{W}_{d}\;\bm{u}[m-1,n-1]+\bm{W}_{c}\;\bm{u}[m,n-1]),$ (7)

where $\bm{u}[m,n]\in\mathbb{C}^{N_{n}}$ represent state vector for the
$(m,n)$-th input; $N_{n}$ stands for the number of neurons;
$\bm{W}_{i}\in\mathbb{C}^{N_{n}\times N_{i}}$ is the input weight matrix;
$N_{i}$ denote the input dimension; $\bm{W}_{r}\in\mathbb{C}^{N_{n}\times
N_{n}}$, $\bm{W}_{c}\in\mathbb{C}^{N_{n}\times N_{n}}$, and
$\bm{W}_{d}\in\mathbb{C}^{N_{n}\times N_{n}}$ denote the reservoir weights
along the row, column, and diagonal directions, respectively; $f(\cdot)$ is
the hyperbolic tangent nonlinear activation function. The input weights and
reservoir weights are all randomly initialized by sampling from a uniform
distribution. In line with the 1D-RC approach, all reservoir weights are
configured to be sparse with spectral radii less than $1$. The initial states
$\bm{u}[-1,n]$, $\bm{u}[m,-1]$, and $\bm{u}[-1,-1]$ are all initialized as
zero vectors. The estimated output is obtained through the following output
equation:

$\displaystyle\hat{O}[m,n]=\bm{W}_{o}\;\tilde{\bm{u}}[m,n],$ (8)

where
$\tilde{\bm{u}}[m,n]=[\;\tilde{\bm{y}}[m,n]^{T},\bm{u}[m,n]^{T}]^{T}\in\mathbb{C}^{N_{n}+N_{i}}$
is the extended state formed by concatenating the input and the state,
$\bm{W}_{o}\in\mathbb{C}^{1\times(N_{n}+N_{i})}$ stands for the output
weights. By collecting all the state vectors ${\bm{u}}[m,n]$, the extended
state vectors $\tilde{\bm{u}}[m,n]$, and the estimated output $\hat{O}[m,n]$,
we can obtain the state tensor
${\mathbcal{U}}\in\mathbb{C}^{N_{n}\times(M+M_{f})\times(N+N_{f})}$, the
extended state tensor
$\tilde{\mathbcal{U}}\in\mathbb{C}^{(N_{n}+N_{i})\times(M+M_{f})\times(N+N_{f})}$
and the estimated output matrix
$\hat{\bm{O}}\in\mathbb{C}^{(M+M_{f})\times(N+N_{f})}$. The structure is shown
in Fig. 3.

###  III-C Learning Algorithm

Like other variants of RC, the input and reservoir weights of 2D-RC are all
fixed after random initialization. Only the output weights are updated during
training. The training objective for 2D-RC can be written as

$\displaystyle\min_{m_{f}\in\mathcal{L}_{m},n_{f}\in\mathcal{L}_{n},\bm{W}_{o}}||\mathrm{vec}(\bm{\Omega}\odot\hat{\bm{O}}_{m_{f},n_{f}})-\mathrm{vec}(\bm{X}_{\mathrm{train}})||_{2}^{2},$
(9)

where
$\hat{\bm{O}}_{m_{f},n_{f}}=\hat{\bm{O}}[m_{f}:m_{f}+M-1,n_{f}:n_{f}+N-1]\in\mathbb{C}^{M\times
N}$ represents the truncated output, $m_{f}$ is a forget length in the delay
forget length set $\mathcal{L}_{m}$ with $M_{f}$ as the maximum delay forget
length, and $n_{f}$ is a forget length in the Doppler forget length set
$\mathcal{L}_{n}$ with $N_{f}$ as the maximum Doppler forget length. Let
$\tilde{\mathbcal{U}}_{m_{f},n_{f}}=\tilde{\mathbcal{U}}[:,m_{f}:m_{f}+M-1,n_{f}:n_{f}+N-1]\in\mathbb{C}^{(N_{n}+N_{i})\times
M\times N}$ be the truncated extended state. The masked truncated extended
state tensor is denoted as
$\bar{\mathbcal{U}}_{m_{f},n_{f}}=\bm{\Omega}\odot_{2}\tilde{\mathbcal{U}}_{m_{f},n_{f}}$,
where $\odot_{2}$ represents conducting the Hadamard product along the second
and third dimensions. The masked truncated extended state matrix
$\bar{\bm{U}}_{m_{f},n_{f}}=\mathrm{vec}_{2}(\bar{\mathbcal{U}}_{m_{f},n_{f}})\in\mathbb{C}^{(N_{n}+N_{i})\times
MN}$ is formed by vectoring the last two dimensions of
$\bar{\mathbcal{U}}_{m_{f},n_{f}}$ with $\mathrm{vec}_{2}(\cdot)$ denoting
vectoring along the second and third dimensions. Then by substituting (8) into
(9), the learning objective function becomes

$\displaystyle\min_{m_{f}\in\mathcal{L}_{m},n_{f}\in\mathcal{L}_{n},\bm{W}_{o}}||\bm{W}_{o}\bar{\bm{U}}_{m_{f},n_{f}}-(\mathrm{vec}(\bm{X}_{\mathrm{train}}))^{T}||_{2}^{2}.$
(10)

The joint optimization problem is solved by alternatively minimizing the
objective with the forget length and the output weights. Specifically, we
first fix the forget length $m_{f}$ and $n_{f}$ and obtain the trained output
weights by the LS solution

$\displaystyle\hat{\bm{W}}_{o}^{(m_{f},n_{f})}=(\mathrm{vec}(\bm{X}_{\mathrm{train}}))^{T}\bar{\bm{U}}_{m_{f},n_{f}}^{\dagger}.$
(11)

Then the optimal forget lengths along the delay dimension and Doppler
dimension are learned by finding the length that minimizes the loss with

$\displaystyle\hat{m}_{f},\hat{n}_{f}\\!=\\!\\!\\!\\!\operatorname*{argmin}_{m_{f}\in\mathcal{L}_{m},n_{f}\in\mathcal{L}_{n}}\\!\\!||\hat{\bm{W}}_{o}^{(m_{f},n_{f})}\bar{\bm{U}}_{m_{f},n_{f}}\\!-\\!(\mathrm{vec}(\bm{X}_{\mathrm{train}}))^{T}||_{2}^{2}.$
(12)

Instead of searching through all the possible delay and Doppler forget length
pairs, we first find the optimal Doppler forget length and then find the
optimal delay forget length to reduce the training complexity.

###  III-D Testing with 2D-RC

At the testing stage, the transmitted symbols
$\hat{\bm{x}}\in\mathbb{C}^{1\times MN}$ are estimated by

$\displaystyle\hat{\bm{x}}=\mathcal{Q}(\hat{\bm{W}}_{o}^{(\hat{m}_{f},\hat{n}_{f})}\;\tilde{\bm{U}}_{\hat{m}_{f},\hat{n}_{f}}),$
(13)

where $\hat{\bm{W}}_{o}^{(\hat{m}_{f},\hat{n}_{f})}$ is the trained output
matrix when utilizing the forget length $\hat{m}_{f}$ and $\hat{n}_{f}$,
$\tilde{\bm{U}}_{\hat{m}_{f},\hat{n}_{f}}=\mathrm{vec}_{2}(\tilde{\mathbcal{U}}_{\hat{m}_{f},\hat{n}_{f}})\in\mathbb{C}^{(N_{n}+N_{i})\times
MN}$ is obtained by vectoring the truncated extended state tensor
$\tilde{\mathbcal{U}}_{\hat{m}_{f},\hat{n}_{f}}$ with forget length
$\hat{m}_{f}$ and $\hat{n}_{f}$, and $\mathcal{Q}(\cdot)$ is the quantization
operation that maps the output to the constellation points. The transmitted
data symbols are extracted with
$\hat{\bm{X}}_{\mathrm{data}}=\bar{\bm{\Omega}}\odot\hat{\bm{X}}$, where
$\hat{\bm{X}}=\mathrm{vec}^{-1}(\hat{\bm{x}})\in\mathbb{C}^{M\times N}$ and
$\bar{\bm{\Omega}}$ is the complement of $\bm{\Omega}$.

##  IV Numerical Experiments

In this section, we evaluate the performance of 2D-RC for symbol detection in
the OTFS system. We set $N=14$ and $M=1024$ following the 3GPP 5G NR standard
[18, 19]. The carrier frequency is $4$ GHz and subcarrier spacing is $15$ kHz.
The 3GPP 5G NR clustered delay line (CDL) channel with delay profile “CDL-C”
[20] is considered with a delay spread of $10$ ns and user velocity of $150$
km/h. The pilot overhead is $4.69\%$, which is set to satisfy the pilot
overhead requirement specified in [18, 19]. All the compared approaches adopt
the same training overhead for a fair comparison. The parameters of 2D-RC are
set as $N_{n}=6$, $M_{w}=4$, $N_{w}=14$, and $l_{c}=7$. The delay forget
length and Doppler forget length are searched in the range of $7$ to $8$ and
the range of $13$ to $14$, respectively. The spectral radii of all the
reservoir weights are configured as $0.9$ and the sparsities are set as $0.6$.
The quantization operation is set as the minimum distance mapping.

Figure 5: BER comparison under QPSK.

We compare the performance of 2D-RC with the 1D-RC [14], LMMSE detectors,
least squares minimum residual (LSMR) based method [21], and MPA [4]. For the
1D-RC approach, the number of neurons is $12$, and $7$ 1D-RCs are adopted for
detection. The LMMSE detector in the OTFS system is implemented in the time
domain with the block-wise channel inverse to reduce the computational
complexity [22]. The LMMSE method in the OFDM system performs detection in the
TF domain with the LMMSE estimated channel in the TF domain [23]. For the
iterative LSMR-based scheme, the number of iterations for interference
cancellation is set as $5$ and $10$ for QPSK and $16$ QAM, respectively. The
number of iterations for LSMR is $15$ for both modulation schemes. For the MPA
method, the number of iterations is $30$ and the damping factor is set as
$0.6$. All the estimated CSI in the OTFS system is obtained by the approach in
[16] with the spike pilot pattern. The power of the spike pilot is set to
ensure that the OTFS subframe with the spike pilot pattern has approximately
the same PAPR as utilizing the blockwise pilot pattern. The reason is that a
high peak-to-average power ratio (PAPR) may compel the power amplifier (PA) to
operate in the non-linear region, resulting in signal distortion and spectral
spreading [24]. Depending on the tested SNR and modulation order, the received
pilot signal-to-noise ratio (SNR) ranges from around $20$ dB to $47$ dB, which
covers the commonly considered pilot SNRs in existing works, e.g., [14, 5].

In Fig. 5 and Fig. 6, we compare the bit error rate (BER) performance of
different approaches in the CP-OTFS system under the QPSK and 16 QAM
modulations, respectively. Compared with the existing learning-based 1D-RC
method, 2D-RC is demonstrated to have better performance under both the QPSK
and 16 QAM modulations, especially in the high SNR regime. Note that $7$ RCs
are utilized in the 1D-RC approach, while only a single NN is exploited for
2D-RC. The reason is that the 1D-RC method directly adopts the existing RC
architecture in the time domain and does not leverage domain knowledge of the
OTFS system for its design. When operating in the time domain, multiple RCs
are required to track the changes in the time-varying channel. Instead, 2D-RC
incorporates the 2D circular structure in the DD-domain input-output
relationship into its design. By incorporating structural knowledge, even with
a single NN, 2D-RC is more effective than the 1D-RC method that adopts
multiple RCs. The 2D-RC also outperforms compared model-based approaches,
i.e., LMMSE, MPA, and the LSMR-based approach, when employing the estimated
channel. Different from the model-based approaches that rely on the knowledge
of CSI, the introduced learning-based 2D-RC approach does not require channel
knowledge. Therefore, the performance of 2D-RC is not affected by the accuracy
of channel estimates and can be more easily adopted in practical scenarios
when it is hard to obtain an accurate CSI. Furthermore, all the considered
OTFS-based detectors in the CP-OTFS system are shown to perform better than
the LMMSE approach in the OFDM system in mid to low SNR regimes.

Figure 6: BER comparison under 16 QAM.

##  V Conclusion

In this paper, we introduce a learning-based 2D-RC approach for the symbol
detection task in the OTFS system. As with other variants of RC, the
introduced 2D-RC approach can conduct online subframe-based symbol detection
with a limited amount of training data. Extending beyond our previous RC-based
approach, the 2D-RC is based on a novel structure tailored to equalizing the
2D circular channel effect in the DD domain of the OTFS system. By employing
the 2D structure, the 2D-RC approach is able to perform online subframe-based
detection in the DD domain with a single NN. Simulation results demonstrate
that 2D-RC significantly outperforms the previous 1D-RC approach that requires
multiple RCs to track the channel changes across time. Furthermore, compared
with the model-based approaches, 2D-RC does not require any channel knowledge
or system assumptions. It is shown to have better performance than the LMMSE
and MPA with the estimated CSI across different modulation orders.

## References

  * [1] M. Series, “IMT Vision–Framework and overall objectives of the future development of IMT for 2020 and beyond,” _Recommendation ITU_ , vol. 2083, no. 0, 2015.
  * [2] R. Hadani, S. Rakib, M. Tsatsanis, A. Monk, A. J. Goldsmith, A. F. Molisch, and R. Calderbank, “Orthogonal time frequency space modulation,” in _2017 IEEE Wireless Commun. Netw. Conf._ , 2017, pp. 1–6.
  * [3] T. Zou, W. Xu, H. Gao, Z. Bie, Z. Feng, and Z. Ding, “Low-complexity linear equalization for OTFS systems with rectangular waveforms,” in _2021 IEEE Intl. Conf. on Commun. (ICC)_. IEEE, 2021, pp. 1–6.
  * [4] P. Raviteja, K. T. Phan, Y. Hong, and E. Viterbo, “Interference cancellation and iterative detection for orthogonal time frequency space modulation,” _IEEE Trans. Wireless Commun._ , vol. 17, no. 10, pp. 6501–6515, 2018.
  * [5] F. Liu, Z. Yuan, Q. Guo, Z. Wang, and P. Sun, “Message passing-based structured sparse signal recovery for estimation of OTFS channels with fractional Doppler shifts,” _IEEE Trans. Wireless Commun._ , vol. 20, no. 12, pp. 7773–7785, 2021.
  * [6] H. Zhang and T. Zhang, “A low-complexity message passing detector for OTFS modulation with probability clipping,” _IEEE Wireless Commun. Lett._ , vol. 10, no. 6, pp. 1271–1275, 2021.
  * [7] Y. K. Enku, B. Bai, F. Wan, C. U. Guyo, I. N. Tiba, C. Zhang, and S. Li, “Two-dimensional convolutional neural network-based signal detection for OTFS systems,” _IEEE Wireless Commun. Lett._ , vol. 10, no. 11, pp. 2514–2518, 2021.
  * [8] X. Zhang, L. Xiao, S. Li, Q. Yuan, L. Xiang, and T. Jiang, “Gaussian AMP aided model-driven learning for OTFS system,” _IEEE Commun. Lett._ , vol. 26, no. 12, pp. 2949–2953, 2022.
  * [9] G. Tanaka, T. Yamane, J. B. Héroux, R. Nakane, N. Kanazawa, S. Takeda, H. Numata, D. Nakano, and A. Hirose, “Recent advances in physical reservoir computing: A review,” _Neural Netw._ , vol. 115, pp. 100–123, 2019.
  * [10] Z. Zhou, L. Liu, and H.-H. Chang, “Learning for detection: MIMO-OFDM symbol detection through downlink pilots,” _IEEE Trans. Wireless Commun._ , vol. 19, no. 6, pp. 3712–3726, 2020.
  * [11] J. Xu, Z. Zhou, L. Li, L. Zheng, and L. Liu, “RC-Struct: A structure-based neural network approach for MIMO-OFDM detection,” _IEEE Trans. Wireless Commun._ , vol. 21, no. 9, pp. 7181–7193, 2022.
  * [12] J. Xu, L. Li, L. Zheng, and L. Liu, “Detect to learn: Structure learning with attention and decision feedback for MIMO-OFDM receive processing,” _IEEE Trans. Commun._ , pp. 1–1, 2023.
  * [13] L. Li, J. Xu, L. Zheng, and L. Liu, “Real-time machine learning for multi-user massive MIMO: Symbol detection using Multi-Mode StructNet,” _IEEE Trans. Wireless Commun._ , pp. 1–1, 2023.
  * [14] Z. Zhou, L. Liu, J. Xu, and R. Calderbank, “Learning to equalize OTFS,” _IEEE Trans. Wireless Commun._ , vol. 21, no. 9, pp. 7723–7736, 2022.
  * [15] S. S. Das, V. Rangamgari, S. Tiwari, and S. C. Mondal, “Time domain channel estimation and equalization of CP-OTFS under multiple fractional dopplers and residual synchronization errors,” _IEEE Access_ , vol. 9, pp. 10 561–10 576, 2020.
  * [16] P. Raviteja, K. T. Phan, and Y. Hong, “Embedded pilot-aided channel estimation for OTFS in delay–Doppler channels,” _IEEE Trans. Veh. Technol._ , vol. 68, no. 5, pp. 4906–4917, 2019.
  * [17] M. Lukoševičius and H. Jaeger, “Reservoir computing approaches to recurrent neural network training,” _Comput. Sci. Review_ , vol. 3, no. 3, pp. 127–149, 2009.
  * [18] _5G; NR; Physical channels and modulation_ , 3GPP Std. TS 36.211, Rev. 16.2.0, 2020.
  * [19] _5G; NR; Multiplexing and channel coding_ , 3GPP Std. TS 36.212, Rev. 16.2.0, 2020.
  * [20] _5G; Study on channel model for frequencies from 0.5 to 100 GHz_ , 3GPP Std. TR 38.901, Rev. 16.1.0, 2020.
  * [21] H. Qu, G. Liu, L. Zhang, S. Wen, and M. A. Imran, “Low-complexity symbol detection and interference cancellation for OTFS system,” _IEEE Trans. Commun._ , vol. 69, no. 3, pp. 1524–1537, 2021.
  * [22] Y. Hong, T. Thaj, and E. Viterbo, _Delay-Doppler Communications: Principles and Applications_. Academic Press, 2022.
  * [23] P. Hoeher, S. Kaiser, and P. Robertson, “Two-dimensional pilot-symbol-aided channel estimation by wiener filtering,” in _1997 IEEE int. conf. on acoust., speech, and signal process._ , vol. 3, pp. 1845–1848.
  * [24] Y. Rahmatallah and S. Mohan, “Peak-to-average power ratio reduction in OFDM systems: A survey and taxonomy,” _IEEE Commun. Surveys Tuts._ , vol. 15, no. 4, pp. 1567–1592, 2013.


# Effective color charge in high temperature
gluon plasma at $A_{0}$ background

V. Skalozub
Oles Honchar Dnipro National University, 49010 Dnipro, Ukraine e-mail:
<EMAIL_ADDRESS>
###### Abstract

In high temperature $SU(2)$ gluodynamics, in a relativistic background
$R^{ext.}_{\xi}$ gauge, the effective color charge $\bar{g}^{2}(A_{0})$ is
calculated in the presence of the $A_{0}$ condensate which is related to
Polyakov’s loop and also serves as the order parameter of the deconfinement
phase transition at high temperature. The gauge-invariant value of the
condensate $(A_{0}^{cl})_{0}$ is obtained from the two-loop effective
potential of order parameter $W_{L}(A_{0}^{cl})$ which is independent of the
gauge-fixing parameter $\xi$ and has a nontrivial minimum position. Due to the
smallness of the expansion parameter $\sim\frac{g^{2}}{8\pi^{2}}$, this
potential is applicable at high temperatures beginning from the deconfinement
temperature $T_{D}$. The effective charge accounts for a correlation between
the one- and the two-loop terms of the $W_{L}(A_{0}^{cl})$. It results in
decreasing of the effective gluon interactions in the plasma, compared to the
zero condensate case. Temperature dependence of
$\bar{g}^{2}((A_{0}^{cl})_{0})$ is investigated. The importance of the found
dependence for different applications is noted.

##  1 Effective potential of order parameter

In quantum chromodynamics at high temperature, a new matter state - quark-
gluon plasma - is created [1]. Standard information on the deconfinement phase
transition is adduced, in particular, in [2]. The order parameter of the phase
transition is the Polyakov loop. In $SU(2)$, it is defined as

$\left<L\right>=\frac{1}{2}TrP\exp\bigl{(}ig\int_{0}^{\beta}dx_{4}~{}A_{0}(x_{4},\vec{x})\bigr{)}.$
(1)

Here, $g$ is coupling constant,
$A_{0}(x_{4},\vec{x})=A^{a}_{0}(x_{4},\vec{x})\frac{\tau^{a}}{2}$ is the zero
component of the gauge field potential, $\tau^{a}$ is Pauli matrix,
$\beta=\frac{1}{T}$ is inverse temperature and integration is going along the
fourth direction in the Euclidean space-time.

The background of the gluon plasma is formed out of so-called $A_{0}$
condensate which is a constant constituent of the gluon field potential
$Q_{\mu}=[Q_{0}=A_{0}+(Q_{0})_{rad.},(Q_{i})_{rad.}]$ and $Q_{rad.}$ are
radiation fields. This condensate is correlated to the Polyakov loop and also
could serve as the order parameter for the phase transition. The $A_{0}$ order
parameter is a special type solution of field equations whereas the latter one
can be calculated for an arbitrary field configuration. At low temperatures
$T<T_{d},$ $A_{0}=0$. For $T>T_{d},A_{0}\not=0,$ where $T_{d}$ is a phase
transition temperature. In this phase color is opened and colored states are
realized.

The generation of $A_{0}$ has been derived by different methods in various
models from a two-loop effective potential
$W(A_{0})=W^{(1)}(A_{0})+W^{(2)}(A_{0})$. Review paper on this topic is [3]
where detailed information is presented. The results of calculation for SU(2)
gluodynamics in the relativistic background $R_{\xi}(A_{0})$ gauge read [4],
[5], [6],

$\displaystyle W(x)$ $\displaystyle=$ $\displaystyle W^{(1)}(x)+W^{(2)}(x),$
(2) $\displaystyle\beta^{4}W^{(1)}(x)$ $\displaystyle=$
$\displaystyle\frac{2}{3}\pi^{2}[B_{4}(0)+2B_{4}(\frac{x}{2})],$
$\displaystyle\beta^{4}W^{(2)}(x)$ $\displaystyle=$
$\displaystyle\frac{1}{2}g^{2}[B_{2}^{2}(\frac{x}{2})+2B_{2}(0))B_{2}(\frac{x}{2})]+\frac{2}{3}g^{2}(1-\xi)B_{3}(\frac{x}{2})B_{1}(\frac{x}{2}),$

where $B_{i}(x)$ are Bernoulli’s polynomials, $x=\frac{gA_{0}\beta}{\pi}$. The
Bernoulli’s polynomials defined $modulo$ 1 are

$\displaystyle B_{1}(x)$ $\displaystyle=$ $\displaystyle
x-\frac{x}{2|x|},~{}~{}B_{2}(x)=x^{2}-|x|+\frac{1}{6},$ (3) $\displaystyle
B_{3}(x)$ $\displaystyle=$ $\displaystyle
x^{3}-\frac{3}{2}\frac{x^{3}}{|x|}+\frac{1}{2}x,$ $\displaystyle B_{4}(x)$
$\displaystyle=$ $\displaystyle x^{4}-2|x|^{3}+x^{2}-\frac{1}{30}.$

At $x=0$ the $B_{1}(x)$ is defined to be 0. Because of the symmetry of
polynomials, a periodic structure in internal space is realized (for details
see [3]). Bellow, we consider the main interval of $x,~{}0\leq x\leq 2$.

To find minima of $W(x)$ we apply an expansion in powers of $g^{2}$ and get

$\displaystyle\beta^{4}W_{min}$ $\displaystyle=$
$\displaystyle\beta^{4}W(0)-\frac{1}{192\pi^{2}}(3-\xi)^{2}g^{4},$ (4)
$\displaystyle x_{0}=g^{2}\frac{(3-\xi)}{8\pi^{2}},$

where the first term is the value at $x=0$. As we see, both the minimum
position and the minimum energy value are gauge-fixing dependent. Hence the
gauge invariance of the $A_{0}$ condensation phenomenon is questionable.

In [6], to obtain a gauge-fixing-independent effective potential related to
Polyakov’s loop, the following procedure has been developed. The expression
(2) was calculated on a special orbit in the $(x,\xi)$-plain where the
$\xi$-dependence is cancelled in the total in the order $g^{2}$, due to the
variation of $A_{0}$ which also ensures expressing $W(x)$ through Polyakov’s
loop. This relation looks as (eq.(17) in [6])

$x=x_{cl}+\frac{g^{2}}{4\pi^{2}}B_{1}(\frac{x_{cl}}{2})(\xi+1),$ (5)

where $x$ is ”nonphysical nonrenormalized” value and $x_{cl}$ is observable
”physical” one. This formula follows from the one-loop correction to the
Polyakov loop. To obtain $\xi$-independent expression we have to substitute
(5) in (2) and expand in powers of $g^{2}$. As a result, $\xi$-dependent parts
are mutually cancelled and we find the effective potential of the order
parameter $W_{L}(x_{cl})$.

Formula (5) identically coincides with the characteristic line in the
$(\xi,x)$-plane coming through the point $\xi=-1,x_{cl}$. This curve is
obtained by the Nielsen identity method for the effective potential [6]. In
such a way these two gauge independent approaches are related. This important
correspondence is also discussed in details for full QCD with finite chemical
potential in [7].

A posteriori described procedure can be realized if we set in (2) $x\to
x_{cl},~{}\xi\to-1$. The same, of course, refer to the minimum values (4).
They are the gauge invariant characteristics of the background state in the
plasma. In particular, Polyakov’s loop in the minimum equals to
$<L>=\cos(\frac{g^{2}}{4\pi})$.

##  2 Effective color charge

Now we turn to investigation of the influence of $A_{0}^{cl}$ on the coupling
parameter. To find such dependence, we introduce the definition

$\frac{1}{\tilde{g}^{2}}=\frac{1}{g^{2}}\frac{W}{W^{(1)}}=\frac{1}{g^{2}}\Bigl{(}1+\frac{W^{(2)}(x)}{W^{(1)}(x)}\Bigr{)},$
(6)

where $\tilde{g}^{2}$ is an effective color charge squared, which is generated
in the $A_{0}$ presence due to the correlation of the two-loop, $W^{(2)}(x)$,
and the one-loop, $W^{(1)}(x)$, contributions in (2). Note again that here
$x=(x_{cl})_{0}$ (4), and for
$\xi=-1,~{}~{}(x_{cl})_{0}=g^{2}\frac{1}{2\pi^{2}}$. In one loop order the
condensate is absent, and $\tilde{g}^{2}=g^{2}$.

Formula (6) corresponds to the definition of the effective electric charge
$\tilde{e}^{2}$ in a magnetic field in QED, or gluon field coupling
$\bar{g}^{2}$ in color magnetic field, at finite temperature [8]. The main
difference is that in both latter cases the influence of vacuum polarization
follows from the correlations of the one-loop and tree-level terms in
corresponding effective potentials. That is, we have to replace in (6)
$W^{(1)}(x)\to\frac{H^{2}}{2}$ and $W^{(2)}(x)\to W^{(1)}(H).$

At first glance, this is a simple mathematical difference which does not
depend on the way of field generations. In the case of magnetic field, one may
assume that it is produced by some external source which maintains the field
of arbitrary strengths. Contrary, in the case of $A_{0}$ any external source
is present basically. This field is a solution to Yang-Mills field equations
subjected to the periodicity conditions in the imaginary time. It has no
sources. So, the only mechanism for generating this field is a spontaneous
condensation happening when higher order contributions are taken into
consideration. The definition (6) accounts for this possibility.

Other important point is that the coupling $g^{2}$ could also depend on
different external parameters - temperature, fields, density, etc. Relation
(6) accounts for the presence of them through the value of $(x_{cl})_{0}$. As
a result, we obtain the effect of the condensate in the plasma.

Moreover, because the ratio of the one- and two-loop quantum corrections is
proportional to $\frac{g^{2}}{8\pi^{2}}$, we can verify that the same value of
the condensate follows if we take into consideration in Bernoulli’s
polynomials the terms of the order $g^{2}$ in (2). That is
$x^{2}-\frac{1}{30}$ in $B_{4}(x)$ and the linear in $x$ terms in the
$W^{(2)}$ which is of the order $g^{2}$. This effective potential is

$W^{as.}=\frac{2}{3}\pi^{2}\Bigl{(}\frac{x^{2}}{2}-\frac{1}{30}\Bigr{)}-\frac{g^{2}}{12}\Bigl{(}(3-\xi)x-\frac{1}{6}\Bigr{)}.$
(7)

To investigate the behavior of $\tilde{g}^{2}$ as function of $x_{cl}$, we
insert $(x_{cl})_{0}=g^{2}\frac{1}{2\pi^{2}},\xi=-1$. Since the charge $g^{2}$
is temperature dependent and (due to asymptotic freedom) it decreases with
temperature increase, we calculate the effective charge for series of
decreasing values. The results are adduced in Table 1

Table 1: Effective charge squared $\tilde{g}^{2}$ as function of $g^{2}$ $g^{2}$ | 10 | 7 | 6 | 5 | 3 | 1 | 0.1
---|---|---|---|---|---|---|---
$\tilde{g}^{2}$ | 4.3637 | 4.6670 | 4.5697 | 4.2933 | 3.0517 | 1.0461 | 0.1006

In first line we give a selected series of the charge squared. In the second
one we present the corresponding values $\tilde{g}^{2}$ calculated according
(6), (7). As we see, the effective interactions in the plasma are considerably
changed. At high temperatures, the presence of the condensate is not
important. But at temperatures close to the deconfinement $T_{d}$, effective
coupling $\tilde{g}^{2}$ significantly decreases due to the condensate. Such
unexpected behavior should be taken into account when the deconfinement phase
transition and the high temperature plasma are investigated. In particular,
the values of $\tilde{g}^{2}$ have to be used in all perturbation calculations
of radiation corrections. It is important that the expansion parameter of the
theory remains small, $\frac{g^{2}}{8\pi^{2}}\sim\frac{1}{10}$. So, the
effective potential (2) is the reliable approximation beginning the
deconfinement phase transition temperatures.

Similar behavior of effective couplings takes also place in other models, in
particular, in full QCD. The actual values of $\bar{g}^{2}$ have to be
computed additionally. They must be used in treating experiment data obtained
at high energy collisions of heavy ions, in problems of the early Universe
where the high temperature phases with either different $A_{0}$ condensates or
temperature dependent magnetic fields determined by relevant gauge groups had
existed. These questions we left for the future.

## Acknowledgments

Author indebted Alexander Pankov and Oleg Teryaev for useful discussions and
suggestions.

## References

  * [1] Yasumichi Aoki, Szabolcs Borsanyi, Stephan Duerr, Zoltan Fodor, Sandor D. Katz, Stefan Krieg, and Kalman Szabo, JHEP 6 , JUN, (2009 ).
  * [2] M. Le Bellack Termal field theory. Cambridge Univ. Press 1996, Cambridge.
  * [3] O. A. Borisenko, J. Boháčik and V. V. Skalozub, Fortch. Phys. 43, 301 (1995).
  * [4] V. M. Belyaev, Physics Letters B 254 (1), 153 (1991).
  * [5] V. Skalozub. Mod. Phys. Lett. A 7 (31), 2895 (1992).
  * [6] V. V. Skalozub, Phys. Part. Nucl. Letts 18 (7), 738 (2021). [ArxiV: 2006.05737 hep-ph]
  * [7] M. Bordag and V. Skalozub, Eur. Phys. J. C81 , 998 (2021). [ArxiV: 2009.11734 v4 hep-ph]
  * [8] V. V. Skalozub, Int. J. Mod. Phys. A 11 (32), 5643 (1996).


# Are Straight-Through gradients and Soft-Thresholding
all you need for Sparse Training?

Antoine Vanderschueren1
UCLouvain, Belgium
<EMAIL_ADDRESS>Christophe De Vleeschouwer1
UCLouvain, Belgium
<EMAIL_ADDRESS>
###### Abstract

Turning the weights to zero when training a neural network helps in reducing
the computational complexity at inference. To progressively increase the
sparsity ratio in the network without causing sharp weight discontinuities
during training, our work combines soft-thresholding and straight-through
gradient estimation to update the raw, i.e. non-thresholded, version of zeroed
weights. Our method, named ST-3 for straight-through/soft-thresholding/sparse-
training222Source code and weights available at
https://github.com/vanderschuea/stthree , obtains SoA results, both in terms
of accuracy/sparsity and accuracy/FLOPS trade-offs, when progressively
increasing the sparsity ratio in a single training cycle. In particular,
despite its simplicity, ST-3 favorably compares to the most recent methods,
adopting differentiable formulations [42] or bio-inspired neuroregeneration
principles [25]. This suggests that the key ingredients for effective
sparsification primarily lie in the ability to give the weights the freedom to
evolve smoothly across the zero state while progressively increasing the
sparsity ratio.

††footnotetext: 1Part of this work has been funded by the Walloon Region
project SmartGate N°1910087, and by the Fonds de la Recherche Scientifique –
FNRS

##  1 Introduction

State-of-the-art neural networks are composed of a few million parameters,
leading to a few billion computations per inference. To limit these
computations, sparse networks have been thoroughly investigated in the past
few years [43, 4, 27, 41, 21, 38, 3, 6, 32, 22, 42, 25], and significant
efforts have been devoted to their efficient hardware implementation [8, 29].
Sparse networks reduce the _inference complexity_ by setting a majority of
their weight parameters to zero. Straightforward approaches to train sparse
networks are based on pruning [43, 9, 4], meaning links are removed from the
network (which is equivalent to setting their weight to zero) with no chance
of being re-activated later along the training. Since the pruning decision is
generally based on weight magnitudes, this approach penalizes the weights that
have to _change signs_ during training, since those weights end up being
pruned when they cross the zero border [33]. In contrast, modern
sparsification methods favor the emergence of zero weights, without definitely
canceling the corresponding links, allowing weights to freely switch between
an active and inactive state. The wide variety of approaches that have been
proposed for this purpose are surveyed in Section 2. Our work proposes a
conceptually and computationally simple solution that (surprisingly) quite
favorably compares to this rich array of prior art.

Our proposal mitigates computational complexity at training by restricting the
learning to a single round of gradient descent iterations. This is in contrast
with the iterative post-training pruning solutions [6, 32], which achieve SoA
accuracy/sparsity trade-off, but suffer from an extremely large computation
cost due to the multiple rounds of training required to progressively increase
the pruning ratio.

As many previous works building on a single round of training, our method
adopts a dynamic weight thresholding procedure to progressively increase the
ratio of weights that are set to zero along the training, and to balance
zeroed weights across layers. Our work is however specific in that (i) it
continuously updates zeroed weights all along the training, and (ii) it sets
the weights to zero based on soft-thresholding, instead of hard-thresholding,
popular in previous works. Those specificities are detailed as follows.

To continuously update the zeroed weights, our method takes inspiration from
the training of quantized neural networks (QNN), where the raw (unquantized)
weights are updated with a Straight-through-estimator (STE) to avoid a
gradient that goes to zero due to the step-function of the quantizer [14, 23].
We apply this same principle to the weights that are ’quantized’ to zero
during network sparsification. This means that raw and non-thresholded weights
are maintained and updated in the backward path, even if their thresholded
version are zeroed in the forward path. STE allows gradients and momentum to
switch positive weights towards negative values (and vice-versa), without
getting stuck in an inactive state when being close to zero.

To turn low magnitude weights to zero, thereby sparsifying the network, soft-
thresholding is adopted. This prevents the abrupt changes of forward-path
weight values induced by hard-thresholding operators, which may turn a small
gradient update into a sharp weight discontinuity. Such abrupt discontinuities
hamper the network accuracy, especially at high sparsity rates, where they
lead to premature layer collapse [36], as illustrated in Figure 1.

Despite its simplicity, our approach surpasses the accuracy/sparsity trade-
offs achieved by the most recent (and generally more sophisticated) single
training cycle alternatives. Thereby, it reveals that the key ingredients
required to train effective sparse networks primary consist in (1) giving the
weights the capability to evolve freely and smoothly between active and
inactive states during training, especially in the early epochs; and (2)
progressively increasing the sparsity ratio without causing sharp
discontinuities in weight values along the training.

Another convincing experimental argument in favor of our method lies in the
fact that, when combined with iterative and thus quite complex (due to
multiple rounds of training cycles) post-training pruning solutions proposed
in [6, 32], our ST-3 defines a novel SoA accuracy/sparsity trade-off. Overall,
due to its effectiveness but also computational and conceptual simplicity, our
ST-3 provides a novel and valuable baseline to assess future sparsification
methods. Despite the key ingredients of ST-3 having been independently
proposed in earlier works (generally to implement multi-cycle pruning or
quantization), our work is novel and original in the way it integrates them in
a method specifically suited to single-cycle training of sparse networks. Our
work is also the first one to demonstrate that those conceptually simple
ingredients are sufficient to reach SoA performance, when properly combined.

##  2 Related Works

While zeroing the weights in a structured manner might lead to higher
acceleration when considering specific hardware implementations, the accuracy
obtained with unstructured sparsity is generally higher at a given sparsity
level[24, 12, 37], and significant speed up can be achieved with unstructured
sparsity on off-the-shelf hardware [15, 31, 7, 29]. Hence, this paper
concentrates on _unstructured_ sparsity. Initial approaches to learn sparse
networks simply removed links from the network at the end of an entire
training cycle, based on their respective L1 magnitude [19, 9]. This however
comes at the cost of a substantial loss in prediction accuracy compared to a
dense network, including when the remaining _active_ weights undergo a
_finetuning_ step (i.e. a shorter training cycle, with small learning rate).
Modern approaches do not consider the definitive and abrupt removal of links.
They instead implement partial pruning at the end of multiple training cycles
or design solutions for progressive weight zeroing.

Reference approaches to progressively increase the zeroed weights ratio have
been studied in [43] and [17]. Global Magnitude Pruning (GMP) [43]
progressively prunes the links along the training. A related method, named
STR, sets the weights that lie below a certain threshold to zero, using soft-
thresholding. In contrast to our ST-3, only the non-zeroed weights are updated
by gradients.

A few previous methods have revealed the benefit of a non-permanent zeroing of
weights. RIGL [3] builds on gradient momentum to resurrect some of the zeroed
weights. GraNet [25] boosts the pruning plasticity through bio-inspired
neuroregeneration principles. ProbMask [42] learns a probabilistic pruning
mask that is transformed into a binary one by sampling (multiple times, each
requiring gradient computations) a Gumbel distribution, thus making it
trainable via gradient descent. Our ST-3 is conceptually and computationally
simpler since it relies on straight-through gradient estimation to
continuously update zeroed weights. Our experiments also reveal that when
combined with soft-thresholding, it reaches better accuracy-sparsity trade-
offs than earlier works.

Beyond its specificities regarding weight zeroing (with soft-thresholding) and
zeroed weight updates (based on straight-through estimation, as done when
training quantized networks [14, 23]), our ST-3 adopts conventional solutions
regarding three other important design questions :

Weight significance definition. Most methods use gradient-magnitude [3],
weight-magnitude [9, 43, 17], or a mix of both [36] to select the weights to
set to zero after some preliminary training. As an alternative, DNW [38] uses
neural architecture search to discover neural wirings, i.e. independent
critical connections between channels, thereby finding a sparse subnetwork of
the complete graph. Our work adopts the main trend, zeroing weights based on
their magnitude.

Distributing sparsity across layers. It has been widely documented [30, 3, 20]
that the network prediction accuracy is more impacted when pruning occurs
close to the input. In a CNN, it is however more beneficial from a
computational point of view to prune the first layers since they correspond to
high-resolution channels. Two heuristics have been recently proposed to
address this trade-off. First, Erdos-Rényi-Kernel (ERK) [30, 3] proposes to
scale the global sparsity ratio (= ratio of zeros to the total parameter
count) with a layer-wise factor. so as to induce a higher (smaller) sparsity
for layers with more (less) parameters, e.g. the first convolution layers of a
ResNet-50 remain dense, thereby preserving accuracy. Second, Layer-Aware-
Pruning (LAMP) [20] selects the weights to prune globally, but relies on a
score that is assigned to each weight rather than on its magnitude. This score
is computed on the sorted and flattened array of weights
($\mathbf{W^{l}}_{\text{sorted}}$) of each layer ($l$).

Both ERK and LAMP target higher accuracy at the cost of an increased
complexity (as measured in FLOPS). However, our experiments did not reveal
significant accuracy gains when using one of those heuristics. Hence, unless
said otherwise, weights are considered independently of their layer in our
experiment. As detailed in Section 3, to handle use cases where the gain in
number of operations (FLOPS) is more important than the accuracy, we propose
to normalize the weights by the size of their kernel, which approximates the
standard deviation of the weights in the layer kernel. Consequently, this
normalization reduces the risk to set a whole kernel to zero, thereby
preserving accuracy at high sparsity ratio.

Controlling the global level of sparsity along the training. This third
question is the one that most impacts the training computational cost. It is
therefore central to our paper, which targets simplicity without penalizing
accuracy. Most authors [6, 5, 32, 18] agree on the fact that: (i) the pruning
should not start at the beginning of training, due to the _heavy changes_ a
model undergoes in the early stage of training; and (ii) the increase of
sparsity should be progressive, to limit the interference between weight
update and weight zeroing. Two strategies co-exist to progressively increase
the sparsity ratio.

The first one recommends decoupling training iterations and weight zeroing, to
promote training consistency and thus convergence quality. Therefore, it
builds on multiple training cycles, and is thus intrinsically more complex. It
increments the amount of zero weights by a constant fraction of non-zeroed
weights at the end of each full training cycle. The learning rate - and
optionally the weights’ values - are rewinded (the acronym LRR is used to
denote this rewinding process), i.e. reset, to an anterior value, before the
training resumes. Doing so, SoA accuracy are obtained, but the training cost
largely grows with the sparsity ratio [6, 32]. When adopting the same post-
training sparsity increase than LRR, and thus iterating over multiple training
cycles, our method outperforms the original LRR method, and defines a novel
accuracy/sparsity SoA.

The second strategy primarily targets low computational load, and therefore
increases the sparsity ratio along a single training cycle. GMP [43] proposes
a cubic increase of the sparsity ratio during training, and has been adopted
by ProbMask [42] and GraNet [25]. This increase starts after a set amount of
epochs, and finishes before the loss, and thus the gradients, become too
small. Among the methods running in a single training cycle, ProbMask [42] and
GraNet [25] achieve highest accuracies. As mentioned above, [42], transforms
the discrete optimization problem associated to the pruning mask definition
into a constrained expected loss minimization problem over a continuous
probability space, which is solved using the Gumbel-Softmax trick. This
however comes at the cost of multiple computations of gradients in each
training iteration, and our experiments demonstrate that the accuracy in [42]
fall behind our approach. GraNet [25], and other similar approaches like [3,
26], update the pruning mask every few thousands iterations, by regenerating
connections when pruning others. This plasticity can be seen as a discrete
approximation of the regeneration of weights that naturally occurs in our
method when adopting soft-thresholding on top of the gradient accumulation
inherent to our recommended straight-through estimator.

##  3 Our Method

This section introduces ST-3 as a new baseline for sparse training. It
provides a relevant baseline candidate, because it is easy to implement (see
Supplementary File), has very little overhead compared to regular dense
training, and performs well with the default dense training parameters. ST-3
adopts weight magnitude soft-thresholding with threshold growing along
training to progressively increase the sparsity ratio in the forward path. To
preserve the opportunity for a zeroed weight to become significant again,
similarly to what is done for quantized network training, it considers a
straight-through gradient estimator to update a non-thresholded version of
each weight maintained in the backward path. Improved training stability is
achieved by scaling the weights in the forward path to compensate for the loss
of magnitude inherent to soft-thresholding. Despite ST-3 builds on a set of
existing or straightforward solutions (progressive increase of sparsity ratio
with magnitude loss compensation, straight-through estimator, and soft-
thresholding), it is novel in the sense that no previous work has combined all
the ST-3 ingredients in a sparse training context. Our experimental study
demonstrates the importance of including all elements to preserve ST-3
performance. They are detailed below, and all contribute to a stable and
consistent training, aiming at minimizing the impact of weight zeroing on the
gradient descent updates.

##### Straight-Through-estimation

To control the sparsity ratio, our method sets to zero the weights that lie
below some adaptive threshold. As mentioned before, a main requirement in
sparse training is that the weights that are zeroed at some point during the
training, because they are below the threshold, get the opportunity to become
significant again in subsequent training iterations. Therefore, our method
maintains a dense and continuously updated version of the weights in parallel
to their sparse thresholded version, to be considered in the forward path but
also to back-propagate the gradients. Following [1, 39], a straight-through
estimator is considered to define the gradient with respect to the dense
weight as a copy of the gradient with respect to the output of the
thresholding operator, namely with respect to the sparse weight. This
differentiation of raw (to update) and thresholded (to use in forward and
backward paths) weights is standard in the training of quantized neural
networks [14, 28], but is not common in the pruning community.

Figure 1: Accuracy of a sparse ResNet-20 on Cifar-10. Conventional straight-
through-estimation (STE), with hard-thresholding, fails at 99% sparsity, as
already pointed in [36]. Replacing hard- by soft-thresholding reduces the
gradient mismatch and aleviates the issue (Dense accuracy of 91.2% orange
dashed line)

##### Soft Thresholding.

Figure 2: Hard-thresholding (left) induces a discrepancy between the forward
weight update and the gradient received in the backward path. Soft-
thresholding (right) does not suffer from such inconsistency, since it
preserves a smooth relationship between update and gradient value. This
prevents sharp weight discontinuities along training.

Most of the persistent pruning literature makes use of hard thresholding to
increase the pruning ratio [9, 4, 3]. However, in sparse training, since the
weights that are zeroed (because they are below the threshold) at some point
during the training might become significant again in subsequent training
iterations, they are also likely to cross the pruning threshold, and be
subject to the hard thresholding function discontinuity. This might lead to
abrupt changes in the forward path, which might be inconsistent with the
gradient received in the backward path. Figure 2 depicts a case where the
pruning threshold is large compared to the gradient magnitude. In this case,
hard thresholding induces severe discontinuities in the forward weight update,
while soft thresholding preserves a smooth and consistent evolution of this
forward weight. The results shown in Figure 1 confirm that soft thresholding
leads to a more stable training, and better accuracy of trained models.

Figure 3: (best viewed in color) (top) Accuracy/sparsity curves obtained, on
Cifar10 w/ ResNet-20, (bottom) Accuracy/sparsity curves obtained on Cifar10 w/
VGG-11, and on Cifar100 w/ WideResNet-34x2. The methods starting w/ LRR are
recursively trained, i.e. each point on the curve needs every point on its
left to be computed first. All points represent the average of 3 runs (with
different seeds common to all methods) and standard deviation is shown. The
x-axis is depicted in log-scale.

##### Weight rescaling.

The weights that contribute to the same output neuron are referred to as being
part of the same _filter_. To preserve the mean absolute magnitude of a neuron
in presence of weight zeroing, we propose to rescale the active weights
filter-wise, based on the ratio of lost weight-magnitude introduced by the
zeroed weights. Formally, letting $W^{l}\in R^{Nout\times Nin}$ denote the
dense weight matrix (before soft-thresholding) of layer $l$ and $th$ be the
sparsity threshold, then the factor multiplying the soft-thresholded vector
associated to the $j^{th}$ output is defined as follows:

$\text{scale}^{l}_{j}=\displaystyle\frac{\sum\limits_{i}^{Nin}~{}\|W^{l}_{ji}\|}{\sum\limits^{Nin}_{i}~{}\begin{cases}\|W^{l}_{ji}\|&\quad\text{if
}\|W^{l}_{ji}\|>th\\\ 0&\quad\text{otherwise}\end{cases}}$

This is similar to the rescaling used by dropout [35].

##### Increase and layer-wise allocation of sparsity.

To mitigate the instabilities induced by the thresholding, and leave the
network the time to adapt to the deactivation of some weights in the forward
path, we increase the sparsity ratio progressively along training. To control
the increase of global sparsity, the cubic pattern introduced in [43] is
adopted. Two variants of our method have been studied in our experiments to
turn the global sparsity ratio into a layerwise sparsity.
The first variant, denoted ST-3, is adopted by default. It simply considers
the network weights globally. Hence, no specific attention is drawn to the
balance of sparsity across layers. It has the advantage of being
straightforward, and our experiments reveal that it results in better
accuracy/sparsity tradeoffs as ERK or LAMP.
The second variant, denoted ST-3σ, has been designed to bias the training
towards increased FLOPS gains. It arises from the observation that to save
FLOPS, it is better to prune weights in layers where the feature maps are
larger, since large maps require more operations. However, large feature maps
generally manipulate small kernels (i.e. due to the adoption of a small number
of high-resolution channels in common architectures), whose initialization
results in large weight magnitudes when adopting the popular Kaiming He
initialization [11]. As a consequence, as shown experimentally [3], after
global thresholding, those kernels with larger weights end-up in being
relatively less sparse than the kernels associated to the smaller resolutions
channels of wider layers. This aspect is generally overlooked in previous work
[6, 4, 32, 20], where only the model compression is considered, disregarding
the number of operations affected by the zeroed weights. Hence, most works
[17, 42] treat inference acceleration as a by-product, while still optimizing
only for a specified sparsity ratio. To favor the zeroing of weights in layers
manipulating (few) high resolution channels, and improve the computational
gain associated to sparsity, we propose, before soft-thresholding, to multiply
each weight by the square root of the number of weights in its corresponding
kernel. This can be interpreted as normalizing the weights by their standard
deviation in individual kernel since the variance of the weights following the
Kaiming He initialization distribution is known to be inversely proportional
to the size of the kernel [11].

##  4 Experiments

###  4.1 Methods used as baselines

Our approach is compared to a set of previous methods that are representative
of the SoA and of the recent trends in the field. We now introduce those
methods, and explain how they have been configured to provide a fair
comparison, and help in quantifying the actual benefit brought by our proposed
ST-3.

As a recognized SoA upper bound in terms of accuracy/sparsity trade-off, the
Learning-Rate-Rewind (LRR) recursive pruning method is considered [32]. This
method however has the obvious downside of requiring significantly more
training cycles than every other method considered in this paper. Typically 10
consecutive training cycles are needed to reach 90% sparsity, which isn’t
practically scalable to larger datasets. In our experiments, LRR is also
combined with our ST-3 as follows: the cubic increase of sparsity ratio
adopted in ST-3 is replaced by the increase of sparsity ratio implemented by
LRR between training cycles, with learning rate rewind. This hybrid approach
is denoted LRR+ST-3. It is as expensive to train than LRR, but leads by far to
the best accuracy/sparsity trade-offs, which confirms the benefit of straight-
through estimation combined to soft-thresholding with magnitude loss
compensation.

Among the methods adopting a single training cycle, the GMP [43], STR [17],
and ProbMask [42] methods, presented in Section 2, are probably the most
relevant to compare with: GMP and STR because they are conceptually simple
threshold-based pruning and sparsification methods, respectively; ProbMask
because it corresponds to the current SoA in one-training cycle methods (even
if its training time is twice as long as the ST-3 one, due to gradient
accumulation required by the Gumbel-Softmax trick at each training step).

For both GMP and STR methods, the _sparsity ratio_ increases following the
same cubic polynomial function as the one adopted by our method (see Sec. 3),
and global thresholding is adopted to achieve the required sparsity ratio, as
for ST-3 333Homogenizing the methods configuration does not significantly
impact the method performance (as demonstrated by our results on ImageNet in
Sec. 4.3). For ProbMask, the sparsity increase ratio and batch-sizes are kept
as in the original paper, as changing them led to significant degradation in
accuracy.

For completeness, comparison is also provided with DNW [38], RIGL [3] and
GraNet [25] (see Section 2 for the presentation of those methods), implemented
as recommended by their respective authors.

Figure 4: Accuracy on ImageNet w/ ResNet-50 (top) and MobileNetv1 (bottom) as
a function of the global sparsity (left), and of the FLOPS (right). All graphs
use a log-scale for the x-axis.

###  4.2 Cifar10 and Cifar100

The methods introduced in Section 4.1 have been run on Cifar-10 and Cifar-100
[16]. Three different networks are considered: ResNet-20 [10], VGG-11 [34])
(as simplified by [27]), and WideResNet-34x2 [40]. Similar to what is done in
related papers, one training cycle consists in 160 epochs and, in the single
training cycle scenario, the progressive increase of the sparsity ratio along
a training cycle starts at the 5th epoch and ends at the 80th epoch.

##### One training cycle.

In Figure 3, our ST-3 achieves the best accuracy/sparsity tradeoffs among
methods trained with a single training cycle, including compared to
ProbMask[42]444Note that, in contrast to what is implemented in the code
released with the original paper, our results have been generated by adapting
the network’s batch-normalization moving-mean and moving-variance on the
training – and not the testing - set. Adapting on the testing set indeed
raises a methodological issue. This difference explains the slight discrepancy
between our plots and the ones in [42]. and GraNet[25]. The soft-thresholding
mechanism adopted in STR [17] makes it possible for some weights to move
between inactive and active states, which results in a significant boost in
performance compared to GMP [43]. However, STR falls short compared to ST-3.
This is because STR does not account for gradient information in any way,
except for momentum-like terms stored in the optimizer. ST-3 solves this
issue, by allowing the accumulation of the gradient information conveyed by
the straight-through-estimator.

##### Multiple training cycles.

The computational cost inherent to the multiple training cycles imposed by LRR
makes it impractical. However, LRR remains a de-facto upper bound to compare
with when considering the training of sparse network. Without surprise, we
observe in Figure 3 that LRR performs better than ST-3. More interestingly,
LRR+ST-3 method is even better, defining a new SoA. The superiority of
LRR+ST-3 over LRR provides an additional clue that the use of a straight-
through estimator to continuously and smoothly update the raw (non-
thresholded) version of zeroed weights leads to a significant increase in
performance when training sparse networks. A deeper comparison between ST-3
and LRR reveals that ST-3 gets closer to LRR on Cifar10 than with Cifar100. We
attribute this phenomenon to the fact that the increased data augmentation
associated to the longer training implemented by LRR is more beneficial for
Cifar100 (only 500 training images per class) than for Cifar10 (5000 training
images per class). To challenge this hypothesis, we have trained ST-3 for 4x
more epochs than a reference cycle at a sparsity threshold of 67.23% on
ResNet-20. This longer training has resulted in a test accuracy of 92.38%
which is better than the 92.22% obtained by LRR with 6x the number of epochs
of a reference training cycle. The role of data augmentation in the
performance improvement observed with longer training is also explored on
ResNet-20 in Figure 3 (top right). LRR is applied as described earlier, but
the augmentation seed is reset between each training cycle, ensuring that LRR
manipulates the same augmented samples than one-cycle methods. With this
additional constraint, we observe that the LRR performance degrades up to the
one of our ST-3.

###  4.3 ImageNet

ResNet-50 [10] and MobileNetv1 [13] have been trained on ImageNet [2] using
the same standard hyper-parameters, number of epochs (=100), and data-
augmentation as previous related works [3, 17, 42, 25], leading to a _dense_
baseline accuracy of 77.1%. To train sparse models, only one-cycle methods
have been considered, due to the excessive training time required by the LRR
iterations. The sparsity progressively increases between the 5th and the 50th
epochs. Values for curves corresponding to related works with ResNet-50 on
ImageNet are directly copied from their respective papers 555In particular,
this means that the ProbMask accuracy is the one obtained in [42], i.e. when
using test samples to define the batch-normalization parameters. Hence, they
correspond to an upper bound of the results that would be obtained when
defining batch-normalization parameters with training samples only. .

##### Accuracy/Sparsity

In Figure 4, ST-3 reaches the best accuracy/sparsity tradeoffs. The closest
contenders are the most recent results provided by GraNet and ProbMask, but a
consistent gap remains that widens up to 5% at 95% sparsity. Our methods
unequivocally pull ahead when optimising sparsity only. It is interesting to
note that the gap between the two variants of our method, namely ST-3 and
ST-3σ, is moderate and stays more or less constant across all sparsity ratios.
The superiority of ST-3 over ST-3σ confirms previous observations indicating
that pruning earlier layers penalizes accuracy (slightly) more than zeroing
later ones [3].

##### Accuracy/FLOPS

FLOPS is a widespread metric to assess the computational gain corresponding to
an ideal scenario where every multiplication with zero can be avoided. In
practice, inference speedup is measured as the FLOPS reduction factor. As
depicted in Figure 4, our ST-3σ method, a ST-3 variant proposed to bias the
distribution of sparsity across layers in favor of layers with higher
resolution, significantly improves sparsity/speedup tradeoffs, making our
approach competitive in that respect as well.

###  4.4 Ablation studies

Figure 5: A study dissecting the impact of straight-through estimation, soft-
thresholding and weight-rescaling during sparse training. A log-scale is used
for the x-axis.

##### Isolating ST-3 components

Figure 5 evaluates the importance of the three components of our method,
namely straight-through-estimation (STE), soft-thresholding, and weight
rescaling for mean magnitude loss compensation. Straight-through-estimation
and soft-thresholding appear to be the main contributors to the success of our
method. The rescaling factor helps in some respect but only when soft-
thresholding is used. When used with pure STE, meaning with hard-thresholding,
it tends to amplify the number of sharp discontinuities in the evolution of
weights, thereby causing a severe penalty.

##### Weight switching between active/inactive states

Figure 6: (best viewed in color) Number of switches between active and
inactive state for the weights of a ResNet-20 trained with ST-3 for 160 epochs
on Cifar-10 at 90% sparsity. The global sparsity ratio is fixed (left) or
cubicly increasing (right). (Weights that switch once or never are not
included)

Figure 6 investigates how the soft-thresholded weights in ST-3 _switch_
between zero and non-zero values during training. It presents how the counts
of switches per weight are distributed in four sequences of epochs, obtained
by partitioning the total number of epochs in groups of 40 epochs. For
clarity, weights that switch less than two times are ignored in the plot. When
plotting these distributions, a color code is used to distinguish the weights
that are active (i.e. non-zero) at the end of training from the ones that are
zeroed.

On the left, the sparsity ratio is set to its final value from the beginning
of training, while it progressively increases along training on the right. We
observe that progressively increasing the sparsity ratio fundamentally changes
the weight switching pattern. Progressive increase stabilize the training, as
attested by the reduced amount of switches between states, especially in the
second group of epochs. It also makes the state switching pattern closer to
the typical sign switching pattern observed in dense training (the signs of
weights that remain significant at the end of training change more often in
the first epochs). A progressive increase of the sparsity is thus desired
since it complements the stability offered by straight-through estimation and
soft-thresholding.

##  5 Conclusion

Throughout this paper the relevance of our straightforward and computationally
simple approach, ST-3(σ), has been widely demonstrated. Its efficiency in
training very sparse networks has been assessed on both cifar-10(0) and
ImageNet with a variety of architectures, attaining single-cycle SoA accuracy.
The fact that our simple and relatively straightforward approach achieves the
highest accuracies among approaches with comparable complexity suggests that
the key towards effective sparsification does not hold in a sophisticated
formulation of the optimization problem but rather in the ability to give the
weights the freedom to evolve across the zero state while progressively
increasing the sparsity ratio.

## References

  * [1] Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville.  Estimating or propagating gradients through stochastic neurons for conditional computation.  CoRR, abs/1308.3432, 2013.
  * [2] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [3] Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen.  Rigging the lottery: Making all tickets winners.  37th International Conference on Machine Learning, ICML 2020, PartF16814:2923–2933, 2020.
  * [4] Jonathan Frankle and Michael Carbin.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  7th International Conference on Learning Representations, ICLR 2019, pages 1–42, 2019.
  * [5] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin.  Pruning neural networks at initialization: Why are we missing the mark?  ICLR, 2021.
  * [6] Jonathan Frankle, David J Schwab, and Ari S Morcos.  The Early Phase of Neural Network Training.  In International Conference on Learning Representations, pages 1–20, 2020.
  * [7] Trevor Gale, Erich Elsen, and Sara Hooker.  The state of sparsity in deep neural networks.  CoRR, abs/1902.09574, 2019.
  * [8] Scott Gray, Alec Radford, and Diederik P Kingma.  GPU Kernels for Block-Sparse Weights.  OpenAI ’17, (1):12.
  * [9] Song Han, Jeff Pool, John Tran, and William J. Dally.  Learning both weights and connections for efficient neural networks.  Advances in Neural Information Processing Systems, 2015-Janua:1135–1143, 2015.
  * [10] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep Residual Learning for Image Recognition.  37(50):1951–1954, dec 2015.
  * [11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  CoRR, abs/1502.01852, 2015.
  * [12] Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi Yang.  Filter pruning via geometric median for deep convolutional neural networks acceleration.  Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2019-June:4335–4344, 2019.
  * [13] Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  arXiv preprint arXiv:1704.04861, 2017.
  * [14] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio.  Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations.  2016\.
  * [15] Nal Kalchbrenner, Erich Elsen, Karen Simonyan, Seb Noury, Norman Casagrande, Edward Lockhart, Florian Stimber, Aäron Van Den Oord, Sander Dieleman, and Koray Kavukcuoglu.  Efficient neural audio synthesis.  35th International Conference on Machine Learning, ICML 2018, 6:3775–3784, 2018.
  * [16] Alex Krizhevsky.  (CIFAR10) Learning Multiple Layers of Features from Tiny Images.  … Science Department, University of Toronto, Tech. …, pages 1–60.
  * [17] Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi.  Soft threshold weight reparameterization for learnable sparsity.  pages 5544–5555, 2020.
  * [18] Guillaume Leclerc and Aleksander Madry.  The Two Regimes of Deep Network Training.  2020\.
  * [19] Yann LeCun, John S Denker, and Sara A. Solla.  Optimal Brain Damage (Pruning).  Advances in neural information processing systems, pages 598–605, 1990.
  * [20] Jaeho Lee, Sejun Park, Sangwoo Mo, Sungsoo Ahn, and Jinwoo Shin.  Layer sparsity for the magnitude-based pruning.  CoRR, abs/2010.07611, 2020.
  * [21] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip H.S. Torr.  SnIP: Single-shot network pruning based on connection sensitivity.  In 7th International Conference on Learning Representations, ICLR 2019, 2019.
  * [22] Bailin Li, Bowen Wu, Jiang Su, and Guangrun Wang.  EagleEye: Fast Sub-net Evaluation for Efficient Neural Network Pruning.  Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics), 12347 LNCS:639–654, 2020.
  * [23] Fengfu Li, Bo Zhang, and Bin Liu.  Ternary Weight Networks.  (Nips), 2016.
  * [24] Hao Li, Hanan Samet, Asim Kadav, Igor Durdanovic, and Hans Peter Graf.  Pruning filters for efficient convnets.  5th International Conference on Learning Representations, ICLR 2017 - Conference Track Proceedings, (2016):1–13, 2017.
  * [25] Shiwei Liu, Tianlong Chen, Xiaohan Chen, Zahra Atashgahi, Lu Yin, Huanyu Kou, Li Shen, Mykola Pechenizkiy, Zhangyang Wang, and Decebal Constantin Mocanu.  Sparse Training via Boosting Pruning Plasticity with Neuroregeneration.  (NeurIPS), 2021.
  * [26] Shiwei Liu, Lu Yin, Decebal Constantin Mocanu, and Mykola Pechenizkiy.  Do we actually need dense over-parameterization? in-time over-parameterization in sparse training.  In International Conference on Machine Learning, pages 6989–7000. PMLR, 2021.
  * [27] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell.  Rethinking the value of network pruning.  7th International Conference on Learning Representations, ICLR 2019, pages 1–21, 2019.
  * [28] Naveen Mellempudi, Abhisek Kundu, Dheevatsa Mudigere, Dipankar Das, Bharat Kaul, and Pradeep Dubey.  Ternary Neural Networks with Fine-Grained Quantization.  2017\.
  * [29] Asit K. Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius.  Accelerating sparse deep neural networks.  CoRR, abs/2104.08378, 2021.
  * [30] Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H. Nguyen, Madeleine Gibescu, and Antonio Liotta.  Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science.  Nature Communications, 9(1):1–12, 2018.
  * [31] Jongsoo Park, Hai Li, Sheng Li, Wei Wen, Yiran Chen, Ping Tak Peter Tang, and Pradeep Dubey.  Faster cnns with direct sparse convolutions and guided pruning.  5th International Conference on Learning Representations, ICLR 2017 - Conference Track Proceedings, 2017:1–12, 2017.
  * [32] Alex Renda, Jonathan Frankle, and Michael Carbin.  Comparing Rewinding and Fine-tuning in Neural Network Pruning.  (2019):1–31, 2020.
  * [33] Victor Sanh, Thomas Wolf, and Alexander M. Rush.  Movement pruning: Adaptive sparsity by fine-tuning.  Advances in Neural Information Processing Systems, 2020-Decem(NeurIPS):1–12, 2020.
  * [34] K. Simonyan and A. Zisserman.  Very Deep Convolutional Networks for Large-Scale Image Recognition.  CoRR, 2014.
  * [35] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  Journal of Machine Learning Research, 15(56):1929–1958, 2014.
  * [36] Hidenori Tanaka, Daniel Kunin, Daniel L.K. Yamins, and Surya Ganguli.  Pruning neural networks without any data by iteratively conserving synaptic flow.  Advances in Neural Information Processing Systems, 2020-Decem(NeurIPS), 2020.
  * [37] Huan Wang, Can Qin, Yulun Zhang, and Yun Fu.  Neural Pruning via Growing Regularization.  pages 1–15, 2020.
  * [38] Mitchell Wortsman, Ali Farhadi, and Mohammad Rastegari.  Discovering neural wirings.  Advances in Neural Information Processing Systems, 32(NeurIPS), 2019\.
  * [39] Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley J. Osher, Yingyong Qi, and Jack Xin.  Understanding straight-through estimator in training activation quantized neural nets.  ICLR, 2019.
  * [40] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  CoRR, abs/1605.07146, 2016.
  * [41] Hattie Zhou, Janice Lan, Rosanne Liu, and Jason Yosinski.  Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask.  (NeurIPS), 2019.
  * [42] Xiao Zhou, Weizhong Zhang, Hang Xu, and Tong Zhang.  Effective Sparsification of Neural Networks with Global Sparsity Constraint.  Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pages 3598–3607, 2021.
  * [43] Michael H. Zhu and Suyog Gupta.  To prune, or not to prune: Exploring the efficacy of pruning for model compression.  6th International Conference on Learning Representations, ICLR 2018 - Workshop Track Proceedings, 2018.

## Appendix

Figure 7: (left) A study dissecting the impact of straight-through estimation,
soft-thresholding and weight-rescaling during sparse training with LAMP-based
sparse-weight-selection. (right) FLOPS/Sparsity trade-off comparing l1 and
LAMP sparse training. A log-scale is used for the x-axes.

This appendix contains the following additional information:

  1. 1.

An additional ablation study on the use of LAMP combined with ST-3

  2. 2.

The full hyper-parameter list for our training

  3. 3.

Our ST-3(σ)algorithm in pseudo-code format

  4. 4.

The numerical values for the ImageNet results

###  A Comparison to LAMP pruning

Figure 7 repeats the same experiences as Figure 5 in the main text, but with
the addition of LAMP instead of l1-based pruning. The conclusions about the
usefulness of our added rescaling and soft thresholding remain the same:
rescaling is only useful when combined with soft-thresholding. LAMP doesn’t
show any additional performance increase aver traditional l1-based pruning
while being more expensive FLOPS-wise. Although not shown on this graph, LAMP,
combined with STE alone, has the advantage of not suffering from layer
collapse due to its formal definition.

###  B Training parameters

All training configuration files and precise network structures are uploaded
made open-source at https://github.com/vanderschuea/stthree. A quick rundown
is provided in Table 1 for easy comparison.

####  B.1 Hyper-parameters

The hyper-parameters used for our training of the networks are presented in
Table 1. All results on Cifar-10(0) are reported from a train/val/test split
of sizes 45k/5k/10k. The results on ImageNet are reported on the validation
set.

Model | ResNet-20/VGG-11/WideResNet-34x2 | ResNet-50/MobileNetv1
---|---|---
epochs | 160 | 100
optimizer | SGD
lr | 0.1 | 0.2
momentum | 0.9
batch-size | 128 | 256
weight-decay | 1e-4 | see below
lr-decay | Step-lr@[80,120] | cosine-lr (+ warmup)
grad. clip. | 3.0
Table 1:  Hyper-parameters used for training. (left: networks on Cifar-10,
right: networks on ImageNet)

Due to the double regularization-effect of pruning and weight-decay, the
weight-decay parameter is slowly decreased as follows at different sparsity
ratios:

  * •

ResNet50: 1e-4 (80-95), 1e-5 (96.5-98), 0.0 (99)

  * •

MobileNetv1: 1e-4 (80), 1e-5 (89-95)

We do not apply decaying on the weight-decay parameter for the cifar-10(0)
models as they don’t benefit much from such optimizations.

####  B.2 Augmentations

##### Cifar-10(0)

The standard 4x4 padding is applied before a 32x32 crop is taken. The data is
normalized and also randomly horizontally flipped.

##### ImageNet

Images are randomly cropped and resized to 224x224. The data is normalized and
also randomly horizontally flipped. (Test images are resized to 256x256 and a
centercrop of size 224x224 is taken for evaluation)

###  C Pseudo-code

Algorithm 1 PyTorch-like Pseudo-code for ST-3σ without and with Sparsity
Distribution Suggestion

1:for $i=1\to\text{n\\_epochs}$ do

2: for $batch=1\to\text{len(train\\_loader)}$ do

3: sp_ratio = get_sp_ratio(i, batch) $\triangleright$ update sparsity ratio

4: weights = torch.tensor([]) $\triangleright$ aggregate all weights

5: for module in model.modules do

6: w = module.weight

7: weights.append(w.flatten() * $\sqrt{\text{w.shape.sum()}}$)

8: th = weights.abs().quantile(sp_ratio)

9: W${}_{\text{sparse}}$ = {} $\triangleright$ aggregate sparsified weights

10: for name, module in model.named_modules do

11: mth = th / $\sqrt{\text{w.shape.sum()}}$

12: with STE(): $\triangleright$ straight-through gradients

13: w = module.weight $\triangleright$ dense weights

14: w${}_{\text{sparse}}$ = sign(w) * ($\|\text{w}\|$-mth).clip(0)
$\triangleright$ soft-th

15: w${}_{\text{sparse}}$ *= get_scale(w, mth) $\triangleright$ apply
rescaling

16: W${}_{\text{sparse}}[\text{name}]$ = w${}_{\text{sparse}}$

17: inputs = trainloader[batch]

18: loss = model(W${}_{\text{sparse}}$, inputs)

19: loss.backward() $\triangleright$ \+ update all dense model weights

20:

###  D Table of results

Table 2 and 3 contain the exact numerical values used for the graphs
presenting the results on ImageNet. The values for the other methods, are
taken from their respective papers which use the same training times and data-
augmentations. No ’best’ values are put in bold as it would depend on whether
comparison is made on Accuracy or GFLOPS as the common denominator. For a
better understanding, we refer to the figures presented in the paper’s main
body.

Method |  Accuracy $\left[\%\right]$ |  Sparsity $\left[\%\right]$ | GFLOPS
---|---|---|---
GMP | 67.7 | 74.11 | 163
STR | 68.35 | 75.28 | 101
STR | 66.52 | 79.07 | 81
RigL+ERK | 68.39 | 75 | 296
ST-3σ (Ours) | 68.67 | 80 | 135
ST-3 (Ours) | 69.1 | 80 | 181
STR | 64.83 | 85.8 | 55
GMP | 61.8 | 89.03 | 82
STR | 62.1 | 89.01 | 42
STR | 61.51 | 89.62 | 40
RigL+ERK | 62.33 | 90 | 154
ProbMask | 65.19 | 89 | 63
ST-3σ (Ours) | 66.15 | 89 | 86
ST-3 (Ours) | 66.67 | 89 | 119
RigL+ERK | 56.27 | 95 | 103
ProbMask | 60.1 | 94.1 | 34
ST-3σ (Ours) | 60.94 | 94.1 | 52
ST-3 (Ours) | 61.19 | 94.1 | 74
Table 2: Numerical Values for the figures with results on ImageNet w/ MobileNetv1 Method |  Accuracy $\left[\%\right]$ |  Sparsity $\left[\%\right]$ | GFLOPS
---|---|---|---
RigL | 74.6 | 80 | 940
RigL+ERK | 75.1 | 80 | 1717
DNW | 76 | 80 | 818
GMP | 75.6 | 80 | 818
STR | 76.19 | 79.55 | 766
STR | 76.12 | 81.27 | 705
GraNet | 76 | 80 | 1431
ST-3σ (Ours) | 76.44 | 80 | 739
ST-3 (Ours) | 76.95 | 80 | 1215
RigL | 72.0 | 90 | 531
RigL+ERK | 73.0 | 90 | 981
DNW | 74 | 90 | 409
GMP | 73.91 | 90 | 409
STR | 74.73 | 87.7 | 402
STR | 74.31 | 90.23 | 343
STR | 74.01 | 90.55 | 341
GraNet | 74.5 | 90 | 818
ProbMask | 74.68 | 90 | 381
ST-3σ (Ours) | 75.28 | 90 | 397
ST-3 (Ours) | 76.03 | 90 | 764
RigL | 67.5 | 95 | 327
RigL+ERK | 70 | 95 | 491
DNW | 68.3 | 95 | 204
GMP | 70.59 | 95 | 204
STR | 70.97 | 94.8 | 182
STR | 70.4 | 95.03 | 159
STR | 70.21 | 95.15 | 162
ProbMask | 71.5 | 95 | 205
ST-3σ (Ours) | 73.69 | 95 | 219
ST-3 (Ours) | 74.46 | 95 | 436
RigL | 64.5 | 96.5 | 164
RigL+ERK | 67.2 | 96.5 | 286
STR | 67.78 | 96.11 | 127
STR | 67.22 | 96.53 | 117
ST-3σ (Ours) | 72.62 | 96.5 | 167
ST-3 (Ours) | 73.31 | 96.5 | 351
DNW | 58.2 | 98 | 82
GMP | 57.9 | 98 | 82
STR | 62.84 | 97.78 | 80
STR | 61.46 | 98.05 | 73
STR | 59.76 | 98.22 | 64
ProbMask | 66.83 | 98 | 105
ST-3σ (Ours) | 69.75 | 98 | 116
ST-3 (Ours) | 70.46 | 98 | 220
GMP | 44.78 | 99 | 41
STR | 54.79 | 98.79 | 54
STR | 51.82 | 98.98 | 47
STR | 50.35 | 99.1 | 44
ProbMask | 61.07 | 99 | 55
ST-3σ (Ours) | 63.25 | 99 | 69
ST-3 (Ours) | 63.88 | 99 | 120
Table 3: Numerical Values for the figures with results on ImageNet w/
ResNet-50