NECST: Neural Joint Source-Channel Coding

For reliable transmission across a noisy communication channel , classical results from information theory show that it is asymptotically optimal to separate out the source and channel coding processes . However , this decomposition can fall short in the finite bit-length regime , as it requires non-trivial tuning of hand-crafted codes and assumes infinite computational power for decoding . In this work , we propose Neural Error Correcting and Source Trimming ( NECST ) codes to jointly learn the encoding and decoding processes in an end-to-end fashion . By adding noise into the latent codes to simulate the channel during training , we learn to both compress and error-correct given a fixed bit-length and computational budget . We obtain codes that are not only competitive against several capacity-approaching channel codes , but also learn useful robust representations of the data for downstream tasks such as classification . Finally , we learn an extremely fast neural decoder , yielding almost an order of magnitude in speedup compared to standard decoding methods based on iterative belief propagation . 1 INTRODUCTION . We consider the problem of encoding images as bit-strings so that they can be reliably transmitted across a noisy communication channel . Classical results from ( Shannon ( 1948 ) ) show that for a memoryless communication channel , as the image size goes to infinity , it is optimal to separately : 1 ) compress the images as much as possible to remove redundant information ( source coding ) and 2 ) use an error-correcting code to re-introduce redundancy , which allows for reconstruction in the presence of noise ( channel coding ) . This separation theorem has been studied in a wide variety of contexts and has also enjoyed great success in the development of new coding algorithms . However , this elegant decomposition suffers from two critical limitations in the finite bit-length regime . First , without an infinite number of bits for transmission , the overall distortion ( reconstruction quality ) is a function of both source and channel coding errors . Thus optimizing for both : ( 1 ) the number of bits to allocate to each process and ( 2 ) the code designs themselves is an extremely difficult task . Second , maximum-likelihood decoding is in general NP-hard ( Berlekamp et al . ( 1978 ) ) . Thus obtaining the accuracy that is possible in theory is contingent on the availability of infinite computational power for decoding , which is impractical for real-world systems . Though several lines of work have explored various relaxations of this problem to achieve tractability ( Koetter & Vontobel ( 2003 ) , Feldman et al . ( 2005 ) , Vontobel & Koetter ( 2007 ) ) , many decoding systems rely on heuristics and early stopping of iterative approaches for scalability , which can be highly suboptimal . To address these challenges we propose Neural Error Correcting and Source Trimming ( NECST ) codes , a deep learning framework for jointly learning to compress and error-correct an input image given a fixed bit-length budget . Three key steps are required . First , we use neural networks to encode each image into a suitable bit-string representation , sidestepping the need to rely on hand-designed coding schemes that require additional tuning for good performance . Second , we simulate a discrete channel within the model and inject noise directly into the latent codes to enforce robustness . Third , we amortize the decoding process such that after training , we can scale our model to massively large datasets . These components pose significant optimization challenges due to the inherent non-differentiability of discrete latent random variables . We overcome this issue by leveraging recent advances in unbiased low-variance gradient estimation for variational learning of discrete latent variable models , and train the model using a variational lower bound on the mutual information between the images and their binary representations to obtain good , robust codes . At its core , NECST can also be seen as an implicit deep generative model of the data derived from the perspective of the joint-source channel coding problem . In experiments , we test NECST on several grayscale and RGB image datasets , obtaining improvements over industry-standard compression ( e.g , JPEG ) and error-correcting codes ( e.g. , low density parity check codes ) . We also learn an extremely fast neural decoder , yielding almost an order of magnitude ( two orders for magnitude for GPU ) in speedup compared to standard decoding methods based on iterative belief propagation . Additionally , we show that in the process we learn discrete representations of the input images that are useful for downstream tasks such as classification . 2 SOURCE AND CHANNEL CODING . 2.1 PRELIMINARIES . We consider the problem of reliably communicating data across a noisy channel using a system that can detect and correct errors . Let X be a space of possible inputs ( e.g. , images ) , and pdata ( x ) be a source distribution defined over x ∈ X . The goal of the communication system is to encode samples x ∼ pdata ( x ) as messages in a codeword space Ŷ = { 0 , 1 } m. The messages are transmitted over a noisy communication channel , where they are corrupted to become noisy codewords in Y . Finally , a decoder produces a reconstruction x̂ ∈ X of the original input x from a received noisy code y . The goal is to minimize the overall distortion ( reconstruction error ) ‖x− x̂‖ in ` 1 or ` 2 norm while keeping the message length m as small as possible ( rate ) . In the absence of channel noise , low reconstruction errors can be achieved using a compression scheme to encode inputs x ∼ pdata ( x ) as succinctly as possible ( e.g. , JPEG ) . In the presence of channel noise , longer messages are typically needed to redundantly encode the information and recover from errors ( e.g. , parity check bits ) . 2.2 THE JOINT SOURCE-CHANNEL CODING PROBLEM . Based on fundamental results due to Shannon ( Shannon , 1948 ) , existing systems adhere to the following pipeline . A source encoder compresses the source image into a bit-string with the minimum number of bits possible . A channel encoder re-introduces redundancies , part of which were lost during source coding to prepare the codeword ŷ for transmission . The decoder then leverages the channel coding scheme to infer the original signal ŷ , producing an approximate reconstruction x̂ . In the separation theorem , Shannon proved that the above approach is optimal in the limit of infinitely long messages . That is , we can minimize distortion by optimizing the source and channel coding processes independently . However , given a finite bit-length budget , the relationship between rate and distortion becomes more complex . The more bits that we reserve for compression , the fewer bits we have remaining to construct the best error-correcting code and vice versa . Balancing the two sources of distortion through the optimal bit allocation , in addition to designing the best source and channel codes , makes this joint source-channel coding ( JSCC ) problem challenging . In addition to these design challenges , real-world communication systems face computational and memory constraints that hinder the straightforward application of Shannon ’ s results . Specifically , practical decoding algorithms rely on approximations that yield suboptimal reconstruction accuracies . The information theory community has studied this problem extensively , and proposed different bounds for finite bit-length JSCC in a wide variety of contexts ( Pilc ( 1967 ) , Csiszar ( 1982 ) , Kostina & Verdú ( 2013 ) ) . Rather than relying on hand-crafted coding schemes that may require additional tuning in practice , we propose to learn the appropriate bit allocations and coding mechanisms using a flexible deep learning framework . 3 NEURAL SOURCE AND CHANNEL CODES . Given the above considerations , we explore a learning-based approach to the JSCC problem . To do so , we consider a flexible class of codes parameterized by a neural network and jointly optimize for the encoding and decoding procedure . This approach is inspired by recent successes in training ( discrete ) latent variable generative models , such as VAEs ( Neal ( 1992 ) , Rolfe ( 2016 ) , van den Oord et al . ( 2017 ) ) . We explore the connection to different types of autoencoders in section 4.1 . 3.1 CODING PROCESS . Let X , Ŷ , Y , X̂ be random variables denoting the inputs , codewords , noisy codewords , and reconstructions respectively . We model their joint distribution p ( x , ŷ , y , x̂ ) using the following graphical model X → Ŷ → Y → X̂ as : p ( x , ŷ , y , x̂ ) = pdata ( x ) qencoder ( ŷ|x ; φ ) pchannel ( y|ŷ ; ) pdecoder ( x̂|y ; θ ) ( 1 ) In equation 1 , pdata ( x ) denotes the distribution over inputs X . It is not known explicitly in practice , and only accessible through samples . pchannel ( y|ŷ ; ) is the channel model , where we specifically focus on the binary symmetric channel ( BSC ) case . The BSC independently flips each bit in the codeword with probability ( e.g. , 0→ 1 ) . Therefore , Ŷ takes values in Ŷ = Y = { 0 , 1 } m and pchannel ( y|ŷ ; ) = m∏ i=1 yi⊕ŷi ( 1− ) yi⊕ŷi⊕1 where ⊕ denotes addition modulo 2 ( i.e. , an eXclusive OR ) . A stochastic encoder qencoder ( ŷ|x ; φ ) generates a codeword ŷ given an input x . Specifically , we model each bit ŷi in the code with an independent Bernoulli random vector . We model the parameters of this Bernoulli with a neural network fφ ( · ) ( an MLP or a CNN ) parameterized by φ : qencoder ( ŷ|x , φ ) = m∏ i=1 σ ( fφ ( xi ) ) ŷi ( 1− σ ( fφ ( xi ) ) ) ( 1−ŷi ) where σ ( z ) denotes the sigmoid function . Similarly , we posit a probabilistic decoder pdecoder ( x̂|y ; θ ) parameterized by an MLP/CNN that , given y , generates a decoded image x̂ . We model each pixel as a factorized Gaussian with a fixed , isotropic covariance to yield : x̂|y ∼ N ( fθ ( y ) , σ2I ) , where fθ ( · ) denotes the decoder network . 4 END-TO-END VARIATIONAL LEARNING OF NEURAL CODES . In order to learn an effective coding scheme , we maximize the mutual information between the input X and the corresponding noisy codeword Y ( Barber & Agakov ( 2006 ) ) . That is , the code ŷ should be robust to partial corruption ; even its noisy instantiation y should preserve as much information about the original input x as possible ( MacKay ( 2003 ) ) . First , we note that we can analytically compute the encoding distribution qnoisy enc ( y|x ; , φ ) after it has been perturbed by the channel by marginalizing over ŷ : qnoisy enc ( y|x ; , φ ) = ∑ ŷ∈Ŷ qencoder ( ŷ|x ; φ ) pchannel ( y|ŷ ; ) which yields : qnoisy enc ( y|x ; φ , ) = m∏ i=1 ( σ ( fφ ( xi ) ) − 2σ ( fφ ( xi ) ) + ) yi ( 1− σ ( fφ ( xi ) ) + 2σ ( fφ ( xi ) ) − ) ( 1−yi ) It is relatively straightforward to handle other types of channel noise ; we provide additional examples of the binary erasure channel ( BEC ) and deletion channel in the Appendix . We get the following optimization problem : max φ I ( X , Y ; φ , ) = H ( X ) −H ( X|Y ; φ , ) = Ex∼pdata ( x ) Ey∼qnoisy enc ( y|x ; , φ ) [ log p ( x|y ; , φ ) ] + const . ( 2 ) ≥ Ex∼pdata ( x ) Ey∼qnoisy enc ( y|x ; , φ ) [ log pdecoder ( x|y ; θ ) ] + const . ( 3 ) where p ( x|y ; , φ ) is the true ( and intractable ) posterior from Eq . 1 and pdecoder ( x|y ; θ ) is an amortized variational approximation . The true posterior p ( x|y ; , φ ) — the posterior probability over possible inputs x given the received noisy codeword y — is the best possible decoder . However , it is also often intractable to evaluate and optimize . We therefore use a tractable variational approximation pdecoder ( x̂|y ; θ ) . Crucially , this variational approximation is amortized ( Kingma & Welling ( 2013 ) ) , and is the inference distribution that will actually be used for decoding at test time . Because of amortization , decoding is guaranteed to be efficient , in contrast with existing error correcting codes , which typically involve NP-hard MPE inference queries . Given any encoder ( φ ) , we can find the best amortized variational approximation by maximizing the lower bound ( 3 ) as a function of θ . Therefore , the NECST training objective is given by : max θ , φ Ex∼pdata ( x ) Ey∼qnoisy enc ( y|x ; , φ ) [ log pdecoder ( x|y ; θ ) ] In practice , we approximate the expectation of the data distribution pdata ( x ) with a finite dataset D : max θ , φ ∑ x∈D Ey∼qnoisy enc ( y|x ; , φ ) [ log pdecoder ( x|y ; θ ) ] ≡ L ( φ , θ ; x , ) ( 4 ) Thus we can jointly learn the encoding and decoding scheme by optimizing the parameters φ and θ . In this way , the encoder ( φ ) is “ aware ” of the computational limitations of the decoder ( θ ) , and is optimized accordingly . However , a main learning challenge is that we are not able to backpropagate directly through the discrete latent variable y ; we elaborate upon this point further in Section 5.1 . The solution will depend on the number of available bitsm , the noise level , and the structure of the input data x . For intuition , consider the case m = 0 . In this case , no information can be transmitted , and the best decoder will fit a single Gaussian to the data . When m = 1 and = 0 ( noiseless case ) , the model will learn a mixture of 2m = 2 Gaussians . However , if m = 1 and = 0.5 , again no information can be transmitted , and the best decoder will fit a single Gaussian to the data . Adding noise forces the model to decide how it should effectively partition the data such that ( 1 ) similar items are grouped together in the same cluster ; and ( 2 ) the clusters are ” well-separated ” ( even when a few bits are flipped , they do not ” cross over ” to a different cluster that will be decoded incorrectly ) . Thus , from the perspective of unsupervised learning , NECST attempts to learn robust binary representations of the data .

Summary: of paper: For the finite-bit case of the noisy communication channel model, it is suboptimal to optimize source coding (compression of input) and error correction (fault tolerance for inherent noise in the channel) separately. The authors propose a neural network model (NECST) that is very similar to the standard VAE, except using binary latents with corruption (e.g.,  random bit flipping in the style of a binary symmetric channel). They use VIMCO to optimize through the discrete units. In their experiments, they show that they can outperform a JPEG+ideal channel code model, but perform similarly to a VAE+LDPC (LDPC is a classic error correcting code) setup.