Realizing Neural Decoder at the Edge with Ensembled BNN

The future wireless communication system 6G will not only be equipped with multi-band high-speed transmission but also energy-efficient communication, low latency, and high security. In a digital communication system, different physical layer encryption algorithms like LDPC, Polar, Turbo codes [9, 3] are used as channel coding methods[4] to prevent the data from getting corrupted by channel noise. When the channel deviates from the Gaussian setting in a practical scenario, to exploit the power of the encoder, Neural Networks (NN) have been used to design the decoder while the encoder is fixed as a near-optimal code [5]. Deploying decoders for these codes takes up huge computation which is only possible because of recent advancements in signal processing methods. With a surge in the number of devices in the network, the interactions among themselves may result in excessive signal processing at the user end that gives rise to huge power consumption. Therefore economic energy usage to have elongated battery life in mobile devices has been a research direction of utmost importance[17, 10, 11]. In a noisy channel, encoding has been challenging even though the decoders have good performance[16]; so the authors in [13, 12, 18, 14] proposed neural code where the encoder and decoder are jointly trained. To overcome the problem of convergence to a local minimum in joint optimization, [7] proposed TurboAE that uses Convolutional Neural Network (CNN) based over-complete Auto Encoder (AE) model incorporating interleavers and de-interleavers to achieve the performance of State Of The Art (SOTA) channel codes under the AWGN scenario. All the existing neural AEs have real-valued network parameters and perform floating-point operations during deployment. For instance, the TurboAE architecture has nearly $26e5$ parameters that take up a memory of $20.84$ MB considering a $64$ bit floating-point representation. Out of total $26e5$ parameters, the encoder has nearly $1.5e5$ parameters whereas the decoder has nearly $25e5$ parameters. Because of the huge number of parameters in the AEs, deploying it in a resource-limited Internet Of Thing (IoT) setup is a challenging task. Furthermore, with the advent of edge computing in IoT scenarios, the computation is decentralized to edge devices where the data is processed locally. Realizing a Neural Decoder such as TurboAE[7] at a user end that has limited memory and computing power is not practically feasible.

We denote a real valued NN $g_{\bm{\phi}}(.)$ where $\bm{\phi}$ represents the real valued network parameters. The output from the NN is given by $\mathbf{y}=g_{\bm{\phi}}(\mathbf{x})$ where $\mathbf{x}$ is the input features to the NN and can be real valued. The neural network $g_{\bm{\phi}}(.)$ can be of any type: a fully connected, a CNN or a Recurrent Neural Networks (RNN). As TurboAE uses a CNN for the Neural Decoder, we now focus on CNNs. For $g_{\mathbf{\phi}}(.)$ a CNN of $L$ layers, the parameters are the filters of the CNN and are given by $\mathbf{\phi}=\{\mathbf{W}_{1},\dots,\mathbf{W}_{L}\}$ where $\mathbf{W}_{l}\in\mathbb{R}^{c_{o}\times c_{i}\times k}$ for $l^{th}$ layer of one dimensional CNN. Here $c_{i}$ and $c_{o}$ represents number of input and output channels and $k$ is the dimension of the filter. For a one dimensional CNN as used in TurboAE, if the input to $l^{th}$ layer of CNN has spatial features of dimension $h_{in}$, then input to $l^{th}$ layer is $\mathbf{a}_{l}\in\mathbb{R}^{c_{i}\times h_{in}}$. The output of $l^{th}$ layer is $\mathbf{a}_{l+1}\in\mathbb{R}^{c_{o}\times h_{out}}$ where $h_{out}$ is the dimension of the output. For a Binary Neural Network (BNN), the weights and activations ($\mathbf{W}$ and $\mathbf{a}$) are binarized using the $sign$ function before taking convolution.

The binarized parameter $\mathbf{W}^{b}_{l}$ and $\mathbf{a}^{b}_{l}$ is given by:

The real-valued convolution is approximated with binary weights and activations as $\mathbf{W}_{l}\ast\mathbf{a}_{l}\approx\mathbf{W}^{b}_{l}\circledast\mathbf{a}^{b}_{l}$ where $\circledast$ is convolution performed with bitwise operations. Even though the binarized weights are used for the forward pass, only the real-valued latent weights are updated with the real-valued gradients during backpropagation. However, during inference, these latent weights can be dropped and a binary network with the binary weights and activations can be deployed. The $sign$ function is non-differentiable and has gradients as zero almost everywhere; thus it is not appropriate for the backpropagation during the training. Therefore a straight-through estimator [2] was proposed that binarizes in the forward pass but during backpropagation, it just passes the gradients as it is to the previous layers. For instance, if $b=sign(r)$, then $grad_{r}=grad_{b}\mathbf{1}_{|r|\leq 1}$ where $grad_{r}=\frac{\partial C}{\partial r}$, $grad_{b}=\frac{\partial C}{\partial b}$ and $C$ is the cost function of the NN. To have a stable update during the training, the updated real-valued weights are clipped between $[-1,1]$.

If a real-valued network $g_{\bm{\phi}}(.)$ is deployed in a $64$ bit system, then its binary version will occupy $64$ times lesser memory and all the floating-point operations can be converted to just xnor and popcount operations. However, because of this extreme compactification, the performance generally degrades significantly. So [8] proposed to use Ternary Neural Network (TNN) where $3$ bits $\{-1,0,1\}$ are used. Therefore the ternarized parameter $t$ is given by:

where $\Delta\approx 0.7E(|\mathbf{r}|)$ in our architecture where $\mathbf{r}$ the set parameters of the real network. The introduction of zero as another bit along with $\{+1,-1\}$ gives a better representation power and therefore better performance than BNN. But the zero weights need not to be saved during deployment. So the memory requirement of TNN is same as that of the BNN. Note that the activation is still binary and thus the computational complexity is also same as the BNN. Therefore with TNN, an improvement in performance over BNN is possible without any degradation in memory requirement or computation.

The method of channel coding in TurboAE can be divided into three sub-problems: an encoder $f_{\theta}(.)$ at the transmitter, a channel $c(.)$ and a decoder $g_{\phi}(.)$ at the receiver. In a communication system, the encoder $x=f_{\theta}(u)$ encodes the binary bits $\mathbf{u}=(u_{1},\dots,u_{K})\in\{+1,-1\}^{K}$ of block length $K$ to get the codeword $\mathbf{x}=(x_{1},\dots,x_{N})$ of length $N$ such that the codeword satisfies the power constraints. The code rate is $R=\frac{K}{N}$, where $N>K$. The i.i.d. AWGN channel corrupts the encoded bits and generates $z_{i}=x_{i}+w_{i}$ such that $w_{i}\sim\mathcal{N}(0,\sigma^{2})$ for $i=1,\dots,K$. The noise in the AWGN channel is represented by the signal to noise ratio $\text{SNR}=-10\log_{10}\sigma^{2}$. After transmission through the channel, the decoder $g_{\phi}(z)$ receives the real valued noisy encoded bits $z$ and map them to an estimate of the actual message sequence $\hat{\mathbf{u}}=(\hat{u}_{1},\dots,\hat{u}_{K})\in\{+1,-1\}^{K}$ using a decoding algorithm. Channel coding aims to minimize the Bit Error Rate (BER) or the BLock Error Rate (BLER) of the recovered message signal $\hat{u}$ given by $BER=\frac{1}{K}\sum_{1}^{K}Pr(\hat{u}_{i}\neq u_{i})$ and $BLER=Pr(\hat{\mathbf{u}}\neq\mathbf{u})$. Naively applying deep learning models by replacing encoder and decoder with general purpose neural networks does not perform well. So in [7], authors have proposed a TurboAE with interleaved encoding and iterative decoding using 1D convolutional neural networks. To make the Neural Decoder utilizable at the edge, we first propose to binarize and ternarize the iterative decoder of TurboAE and inspect its performance. We briefly describe the TurboAE architecture before discussing the proposed compressing techniques.

Turbo code is one of the first capacity approaching codes based on recursive systematic convolutional (RSC) code that has an optimal decoding algorithm namely the Bahl-Cocke-Jelinek-Raviv (BCJR)[1]. To add long-range memory to the code, interleaving is used: out of two copies of input bits, the first one passes through the RSC code and the second goes through the interleaver before passing through the same RSC code as shown in Fig. 1(left). After the transmission through the channel, this code is then decoded by repeating (i) and (ii) alternatively: (i) soft decoding based on the signal received from the first copy (ii) using the de-interleaved version as a prior for decoding the second copy as shown in Fig. 1(right). This iterative decoding method keeps re-estimating the posterior distribution on the transmitted bits. Both the interleaved encoder and the iterative decoder are learnable as proposed in TurboAE [7]. The interleaver $\mathbf{x}^{\pi}=\pi(\mathbf{x})$ and the de-interleaver $\mathbf{x}=\pi^{-1}(\mathbf{x}^{\pi})$ shuffles and un-shuffles the input sequence with a random interleaving array known to both encoder and decoder respectively. A code rate of $1/3$ is considered for the interleaved encoder $f_{\theta}$ that has three learnable blocks $f_{1,\theta},f_{2,\theta}$ and $f_{3,\theta}$. The first two takes the original message bit $\mathbf{u}$ to produce $\mathbf{x}_{1}$ and $\mathbf{x}_{2}$ whereas the third block takes the interleaved message $\pi(\mathbf{u})$ to return $\mathbf{x}_{3}$ as shown in Fig. 1. The encoded messages are transmitted through the channel and the received noisy messages are $\mathbf{z}_{1}$,$\mathbf{z}_{2}$ and $\mathbf{z}_{3}$. Our focus is mostly on the compression of the iterative decoder part so that it can be deployed at the edge devices. Thus we do not discuss much on the encoder part in this work. Interested readers may refer to [7] for more details on the encoder.

To validate the usefulness of the proposed compression techniques, we consider the setting used in [7] to train the encoder and decoder networks. A large batch size, preferably greater than or equal to $500$, is used to average the channel noise effects. We train the encoder and decoder separately to avoid any possible local optima. BinTurboAE and TernTurboAE need a smaller learning rate than the real-valued TurboAE. Hence we reduced the learning rate by 10 times whenever the validation loss gets saturated for higher training epochs. The hyper-parameters used in our experiment are shown in Table I.

In summary, we propose BinTurboAE and TernTurboAE intending to deploy the end-to-end channel coding in the targeted low-power edge devices by reducing the memory requirement and the computations by nearly $64$ times at the cost of acceptable performance degradation. We then propose BinTurboAE-bag and TernTurboAE-bag to improve the performance offered by a single BinTurboAE or single TernTurboAE respectively and achieve the performance close to the real network. The ensembled technique implemented with four such weak learners is shown to consume $16$ times less memory and computing power than the real-valued TurboAE with nearly similar performance.