Neural Joint Source-Channel Coding via Bernoulli Latent Straight-Through Estimator

I. INTRODUCTION

BASED on Shannon’s source-channel separation theorem [1], modern communication systems are composed of multiple modules in which source coding and channel coding are treated independently as different modules. The practical communication systems for data transmission across noisy channel consist of source coding part that compresses target data to eliminate the inherent redundancy of the data, and channel coding part that re-introduces redundancy for error detection and correction purpose. Recently, these coding blocks have been block-wisely developed to further improve their performance. For instance, maximum likelihood (ML) decoding process is regarded as an NP-hard problem, which means that there is a limitation to computational power [2]. Although prior studies on decoding process such as linear programming (LP) of LDPC code are studied, these iterative methods still have limitations by assuming an optimal system block [3].

In particular, for image data transmission, source coding with JPEG [4] or WebP along with channel coding via low density parity check (LDPC) [5], Turbo [6], or Polar [7] coding have demonstrated the performance that reaches the Shannon capacity. However, since Shannon’s source-channel separation theorem only holds for long enough codes, to achieve Shannon’s capacity also for short codes, joint sourcechannel coding (JSCC) model which combines source coding and channel coding process has been proposed [8].

In the meantime, deep learning (DL) has been utilized as a powerful tool for various fields including computer vision (CV) [9], natural language processing [10], and wireless communication systems [11], [12], [13], [14], [15] including JSCC that considered wireless image transmission scheme over the additive white Gaussian noise (AWGN) channel [16] and Rayleigh fading channel [17] [18] based on the autoencoder structure of machine learning.

While the works in [16], [17], [18] considered continuous channels, reference [19] proposed JSCC via neural network for discrete channels, coined neural error correction and source trimming (NECST). Specifically, the authors considered binary symmetric channel (BSC) and binary erasure channel (BEC), along with a variational autoencoder (VAE) model to account for stochastic encoding. By maximizing mutual information between image source data and noisy latent codeword, the authors showed that their model outperforms conventional LDPC encoder/decoder.

Learning the discrete-distributed domain has been considered one of the most challenging tasks in machine learning framework. While [19] presented a novel framework for discrete representation learning, it suffers from noisy gradient estimation that generally has large variance due to the nature of the score function estimator [20].

In this paper, we aim at further improving JSCC over discrete channels by proposing low-variance gradient estimator for NECST by replacing score function estimator to pathwise gradient estimator [20]. Key idea is to adopt a straightthrough (ST) estimator [21] that allows gradient flow over non-differentiable functions, e.g., discrete sampling of the stochastic message, akin to the way that ST estimator has been utilized in Gumbel-softmax [22] and concrete distributions [23] for modeling VAE.

The rest of the paper is organized as follows. In Sec. II, we describe the considered system model and summarize the objective function of neural JSCC. Then, we introduce a naïve, vanilla score function estimator to train the neural JSCC, followed by an improved version that uses multi-samples to reduce the variance of the gradient estimator, which is essentially the NECST training framework. Then, in Sec. III, the proposed training scheme for neural JSCC with path-wise gradient computation with the aid of ST estimator is organized, and the corresponding experimental results are presented in Sec. IV. Lastly, Sec. V concludes the paper.

II. SYSTEM MODEL

In this section, we review NECST [19]. As mentioned, since discrete latent variables are not directly differentiable, the authors [19] used score function based stochastic gradient estimation techniques (also known as REINFORCE [24]) with variational inference for Monte Carlo objectives (VIMCO) estimator [25]. In this section, we briefly summarize 1) the overall neural NJSCC system model and coding scheme; and 2) the VIMCO-based mutual information optimization in NECST [19].

A. Coding Process over Discrete Channel

Coding process of JSCC is compromised with stochastic encoder [TeX:] $$\left(q_\phi\right)$$ and decoder [TeX:] $$\left(p_\theta\right)$$ that are parameterized by vectors [TeX:] $$\phi \text { and } \theta \text {, }$$ As illustrated in Fig. 1, the encoder network [TeX:] $$q_\phi(\mathbf{x})$$ maps an input image data x to an 𝑚-bit codeword y as

with parameterized function [TeX:] $$f_\phi(\cdot)$$ defined from multi-layer neural network with [TeX:] $$L^e$$ layers as [26]

where [TeX:] $$f_{\phi^{(l)}}(x)=g\left(W^{(l)} x+b^{(l)}\right)$$ represents the non-linear activation function of the 𝑙th layer with parameter [TeX:] $$\phi^{(l)}=\left\{W^{(l)}, b^{(l)}\right\}$$ with weight matrix [TeX:] $$W^{(l)}$$ and bias vector [TeX:] $$b^{(l)}$$ given the vector of neural network parameters [TeX:] $$\phi=\left\{\phi^{(l)}\right\}_{l=1, \cdots, L^e-1},$$ and [TeX:] $$[\cdot]_i$$ stands for 𝑖th element of the corresponding vector. The non-linear function [TeX:] $$g(\cdot)$$ can be, e.g., a Rectified Linear Unit (ReLU) or a hyperbolic tangent function. While we design the last layer of parameterized function [TeX:] $$f_\phi(\cdot)$$ of (2) to set the size of the output [TeX:] $$f_\phi(\mathbf{x}) \in \mathbb{R}^m,$$ so that it can represent the 𝑚-bit codeword, still, it is in real domain which cannot be directly expressed in 𝑚-bit codeword y. Accordingly, we transform the length-𝑚 real vector [TeX:] $$f_\phi(\mathbf{x})$$ into probability vector of length 𝑚 with sigmoid activation function

Note that this function is applied element-wise (see (1) and maps negative infinity into probability value of zero, and positive infinity to probability of one. Now, treating [TeX:] $$\sigma\left(\left[f_\phi(\mathbf{x})\right]_i\right)$$ as the probability of 𝑖th bit (among 𝑚 bits) being 1, we have (1), the stochastic encoder that computes probability of generating random bit y given input x. Note that all-one vector y would have probability of [TeX:] $$q_\phi(\mathbf{y} \mid \mathbf{x})=\prod_{i=1}^m \sigma\left(\left[f_\phi(\mathbf{x})\right]_i\right)^{\mathbf{y}_i=1},$$ while all-zero vector y has the probability [TeX:] $$\prod_{i=1}^m\left(1-\sigma\left(\left[f_\phi(\mathbf{x})\right]_i\right)\right)^{\left(1-\mathbf{y}_i=0\right)}.$$

From (1), the 𝑚-bit codeword [TeX:] $$\mathbf{y} \sim q_\phi(\cdot \mid \mathbf{x})$$ will be transmitted over discrete channel model, e.g., binary symmetric channel (BSC). We focus on BSC channel in this paper since BSC is commonly treated as a complex channel compared to binary erasure channel (BEC) [27]. In case of BSC, each element of the encoded codeword [TeX:] $$\mathbf{y} \in\{0,1\}^m$$ is independently flipped with probability [TeX:] $$\epsilon \text { (e.g., } 0 \rightarrow 1 \text { with probability } \epsilon \text {; }0 \rightarrow 0 \text { with probability } 1-\epsilon)$$ to form a noisy codeword [TeX:] $$\hat{\mathbf{y}}$$ as shown in Fig. 2.

Mathematically, the distribution of the noisy codeword [TeX:] $$\hat{\mathbf{y}}$$ with BSC channel [TeX:] $$q_{\text {channel }}$$ can be formulated as

with channel noise level 𝜖. Given received noisy codeword [TeX:] $$\hat{\mathbf{y}}$$ with length 𝑚, the decoder network [TeX:] $$p_\theta(\cdot \mid \hat{\mathbf{y}})$$ then maps [TeX:] $$\hat{\mathbf{y}}$$ to reconstructed source data [TeX:] $$\hat{\mathbf{x}}$$. The decoder network is also defined by parameterized multi-layer neural network analogously to the encoder network as

with [TeX:] $$L^d$$ layers assuming Gaussian distribution [TeX:] $$\mathcal{N}(\cdot, \cdot).$$ In practice, mean of the probabilistic decoder [TeX:] $$p_\theta(\mathbf{x} \mid \mathbf{y})$$ in (6) is used as the estimated input [TeX:] $$\hat{\mathbf{x}}$$.

B. NECST optimization with score function estimator

Estimating and optimizing the mutual information between unknown random variables (e.g., dataset, noisy codeword) is intractable which hinders precise optimization of neural JSCC over discrete channel. To this end, in [19], NECST approximately optimized the variational bound based on amortized inference [28], [29] by maximizing the mutual information between input data x and noisy codeword [TeX:] $$\hat{\mathbf{y}}$$ as

where [TeX:] $$H(\cdot)$$ denotes differential entropy. To logically denote the encoding distribution with the noisy channel in (7), we define [TeX:] $$q_{\text {channel }}(\mathbf{x}, \hat{\mathbf{y}} ; \phi, \epsilon)$$ as

where 𝑝(x) is the ground-truth of input data distribution. Since the expectation of (7) cannot be computed explicitly, (7) is approximated to form an objective function for NECST as

with cardinality of the training dataset [TeX:] $$|\mathcal{D}|.$$ As optimization of (9) requires gradient estimation due to discrete sampling, for which [19] adopted score function estimator which suffers from large-variance issue, NECST considers multi-sample objective named VIMCO [25]. Precisely, VIMCO is a score function-based gradient estimator that modifies the discrete latent vector y into the average of the 𝐾-sampled variable [TeX:] $$\frac{1}{K} \sum_{i=1}^K \mathbf{y}^i$$ for learning.

1) Vanilla Score Function Approach: First, we introduce the vanilla score function estimator that solves (9) which suffers from large variance, yet being an an unbiased gradient estimator for (9). As computing the gradient with respect to 𝜙 for the NECST loss [TeX:] $$\mathcal{L}_{\mathrm{NECST}}(\theta, \phi ; x, \epsilon)$$ is the tricky part, we first compute the gradient with respect to the decoder parameter vector 𝜃 which is easier to estimate.

Above (24) can be easily estimated with Monte Carlo sampling as

In detail, we sample [TeX:] $$\mathbf{y}^n \text {, }$$ hence the [TeX:] $$\hat{\mathbf{y}}^n$$ at forward pass from the conditional stochastic encoder distribution [TeX:] $$q_\phi(\mathbf{y} \mid \mathbf{x})$$ given the input x, then the function [TeX:] $$p_\theta(\mathbf{x} \mid \mathbf{y}) \text { given } \hat{\mathbf{y}}^n$$ can be evaluated deterministically. Now we estimate the gradient [TeX:] $$\nabla_\phi \mathcal{L}_{\mathrm{NECST}}(\theta, \phi ; x, \epsilon)$$ with respect to the parameters 𝜙, which is more tricky part when training the NECST model. The difficulty comes from dependency of parameter 𝜙 to the distribution of the generated samples y through [TeX:] $$q_\phi(\cdot \mid \mathbf{x}),$$ while the loss [TeX:] $$\mathcal{L}_{\text {NECST }}$$ is computed using the samples y. Mathematically, the corresponding gradient can be computed as

Similar to (10), we can also use Monte Carlo sampling to estimate the gradient with respect to the encoder parameter as

2) Multi-Sample Score Function Approach: However, the variance of score function estimator scales with the sample vector dimension, it generally suffers from high gradient variance issues [22] and [23]. Accordingly, [19] apply a Monte Carlo variance reduction technique, VIMCO, leverages this large variance via multi-samples in (14). Key idea is to sample 𝐾 latent variables y instead of single draw in (14) to reduce the variance of the gradient estimate in (14). The NECST optimizing objective with VIMCO can be written as

From this modified objective function, the gradient with respect to the decoder parameter vector 𝜃 can be written as

and hence estimate with Monte Carlo sampling as

Similarly, the gradient with respect to the encoder parameter vector 𝜙 can be computed with lower variance thanks to multi sampling as

Finally, this gradient can be estimated with Monte Carlo sampling as

III. LOW-VARIANCE GRADIENT ESTIMATOR FOR NEURAL JSCC

Even though NECST utilized VIMCO to reduce the variance of the estimated gradient as compared to vanila score function estimator, it still suffers from large variance due to limitations of score function estimator while it also requires increase computational complexity due to multi sampling. Furthermore, it requires different objective function for decoder (15) and encoder (17) that requires careful design of the loss functions. Recalling that the reason for such multi samples and different objective function stemming from score function estimation comes from the non-differentiable nature of sample y which is a realization of sample draw from parameterized distribution [TeX:] $$q_\phi(\mathbf{y} \mid \mathbf{x}),$$ we simplify the problem by making the codeword y a differentiable function, and hence reduce the variance of the gradient estimation. Key idea is to use path-wise gradient estimation instead of score-function estimation by replacing the sampled codeword y with soft codeword y′ that can be directly differentiated with respect to the encoder parameter vector 𝜙.

1) Path-wise Gradient Approach: As mentioned, to optimize the encoder [TeX:] $$\left(q_\phi\right) \text { and decoder }\left(p_\theta\right)$$ at once, we introduce an end-to-end optimized neural JSCC model without using VIMCO that requires multi-sampling Bernoulli-distributed binary codeword during training step. Accordingly, we directly minimize the loss function [TeX:] $$\mathcal{L}_{\text {NECST }}(\theta, \phi ; x, \epsilon).$$ However, the problem lies in the discrete nature of the latent codeword y.

The Straight-Through (ST) estimator [21] is a novel method to simply avoid its limitation. This technique is widely used in recent discrete VAE researches [30] [31] to allow gradient flow over non-differentiable backward pass. To this end, we propose to utilize lower variance gradient estimator via pathwise gradient estimation to solve (14) via ST estimator. Precisely, the ST estimator replaces non-differentiable 𝑚-bit latent codeword y into differentiable codeword y′ as

where the function stop_gradient() removes dependency with respect to the learnable parameters, i.e., 𝜃 and 𝜙 here. Note that the value itself is same for y′ and y. Thus, the gradient of Bernoulli-sampled noisy codeword [TeX:] $$\nabla_\phi \hat{\mathbf{y}}$$ can be replaced by [TeX:] $$\nabla_\phi \hat{\mathbf{y}}^{\prime}.$$

ST gradient estimator in (20) lead to ignore the Bernoulli distributed sampling in encoding part (1). Finally, we can redefine end-to-end optimizing objective (9) for our neural JSCC model with Bernoulli ST (BernST) method as

where ˆy′ is the output of the channel using the differentiable codeword y′ generated from (18).

Given the loss function, the gradient with respect to the decoder parameter 𝜃 can be computed as

which can be estimated using Monte Carlo sampling as

Similarly, the gradient for the encoder parameter vector 𝜙 can be obtained as

where the expectation E𝛿 is taken with respect to inherent randomness that does not depend on the learnable parameter 𝜙 and 𝜃, e.g., channel noise in 𝑞channel (·), and we particularly emphasize the direct dependency of ˆy′ with respect to parameter 𝜙 and the inherent random variable 𝛿 by ˆy′ (𝜙, 𝛿). Furthermore, the gradient can be estimated using Monte Carlo sampling as

It is worth noting that, the objective of NECST (15) and BernST (21) are essentially same except for that the proposed objective (21) has lower variance with less computational complexity which leads to an efficient and effective learning of entire neural JSCC system.

IV. EXPERIMENTAL RESULTS

In this section, to validate our method, we demonstrate the compression and error correction performance of BernST and compare them to NECST. Our comparison is on the MNIST [32] and CIFAR-10 [33] datasets using 50000/10000/10000 and 45000/5000/10000 split into the training, validation and test sets, respectively. For fair comparison, the architecture and hyperparameters of the encoder and decoder are set as in [19]. All of our experiments are implemented using Pytorch, which allows for automatic differentiation through the gradient updates.

To begin with, Table I and Table II compare the performance of compression and error correction with fixed number of bits for MNIST of 28×28 gray scale images of handwritten digits and CIFAR-10 of 32×32 RGB images dataset, an RGB image dataset of 32x32 size, respectively. The results reported the L2 distance between the original and reconstructed image in test set per single image that showed the best performance on the validation set during training. It can be easily checked that the proposed method achieves better reconstructed image compared to original NECST thanks to the proposed variance reduction technique.

To investigate the effect of number of bits for transmission, Fig. 3 shows reconstruction performance over BSC channel as a function of the number of bits with fixed noise level 𝜖 = 0.1. In Fig. 3, we can see that our proposed method outperforms NECST over all regions while the gap becomes enlarged in the sufficient bits regime due to the ability of neural joint source coding that handles compression and noise correction simultaneously.

Figs. 4 and 5 shows the reconstruction error during the training process for BernST and NECST according to the channel noise level. There is not much difference in the early stages of learning, but this is because the encoder is not well trained in the early stages and there is not enough information in the encoded codewords. Fig. 4 directly shows that reconstruction performance of BernST is much higher than NECST during entire training process. As shown in Fig. 5, the advantages of BernST methods are not clearly evident as the channel noise level is increased (𝜖 > 0.3), it is meaningful that it consistently product good performance at overall channel noise levels.

Lastly, as mentioned, while the proposed BernST encodes source data with single-sampled Bernoulli codewords, NECST with VIMCO requires multi samples. To examine the effect of multi samples, we plot the reconstruction error with respect to training iterations for different multiple sample numbers. In Fig. 6, increase in multi sample indeed achieves better reconstruction performance thanks to variance reduction also as reported in [19], our method outperforms NECST with single sample thanks to the low-variance nature of pathwise gradient estimator.

V. CONCLUSION

In this paper, neural JSCC for discrete channels has been proposed with reduced variance for the gradient estimator that leads to improved performance as compared to the state-ofthe- art neural JSCC, NECST. By utilizing path-wise gradient computation with straight-through estimator, we were able to show that the proposed BernST reconstructs the original image better than NECST for almost all noise levels and number of bits. Also, the proposed approach shows faster convergence with lower computational complexity thanks to the accurate estimation of the training gradient. Experimental results on both MNIST and CIFAR-10 dataset demonstrate that the proposed neural JSCC shows better compression and error correction performance as compared to the previous neural JSCC system for various channel conditions. As a final remark, one of the limitations of the existing (including ours) neural JSCC system is that the model must be trained separately on various code lengths and dataset. Future work may consider meta-learning [26], [34] of neural JSCC system for a variablelength code to adapt to the new environment quickly.