Error Correction Code Transformer
Authors: Yoni Choukroun
Lior Wolf
Venue: NeurIPS 2022 Accept (2022)
Keywords: ECC
Deep Learning
Transformers

Abstract: Error correction code is a major part of the physical communication layer, ensuring the reliable transfer of data over noisy channels.
Recently, neural decoders were shown to outperform classical decoding techniques.
However, the existing neural approaches present strong overfitting, due to the exponential training complexity, or a restrictive inductive bias, due to reliance on Belief Propagation.
Recently, Transformers have become methods of choice in many applications, thanks to their ability to represent complex interactions between elements.
In this work, we propose to extend for the first time the Transformer architecture to the soft decoding of linear codes at arbitrary block lengths.
We encode each channel's output dimension to a high dimension for a better representation of the bits' information to be processed separately.
The element-wise processing allows the analysis of channel output reliability, while the algebraic code and the interaction between the bits are inserted into the model via an adapted masked self-attention module.
The proposed approach demonstrates the power and flexibility of Transformers and outperforms existing state-of-the-art neural decoders by large margins, at a fraction of their time complexity.

Error Correction Code Transformer
Yoni Choukroun Tel Aviv University choukroun. yoni @gwail...mom com
Lior Wolf Tel Aviv University lioroolf@gmall. com
Abstract
Error correction code is major part of the physical communication layer, ensuring the reliable transfer of data over noisy channels. Recently, neural decoders were shown outperform classical decoding techniques. However, the existing neural approaches present strong overfitting, due to the exponential training complexity, restrictive inductive bias, due to reliance on Belief Propagation. Recently, Transformers have become methods of choice in many applications, thanks to their ability to represent complex interactions between elements. In this work, we propose to extend for the first time the Transformer architecture to the soft decod- ing of linear codes at arbitrary block lengths. We encode each channel's output dimension to a high dimension for better representation the bits' information to be processed separately. The element-wise processing allows the analysis of channel output reliability, while the algebraic code and the interaction between the bits inserted into the model via adapted masked self-attention module. The proposed approach demonstrates the power and flexibility of Transformers and outperforms existing of their time complexity.
neural decoders by large margins, at fraction
Introduction
Reliable digital communication is of major importance in the modern information age and involves the design of codes to be robustly decoded under noisy transmission channels. While optimal decoding defined by the NP-hard maximum likelihood rule, the efficient decoding algebraic block codes is an open problem. Recently, powerful cearning-based techniques have been introduced to this field.
At present, deep learning models that implement parameterized versions of the legacy Belief Propaga- tion (BP) decoders are dominant 19. 20 18, 21 4 This approach is appealing, since the exponential size of the input space is reduced by ensuring that certain symmetry conditions
met. In this case, is sufficient to train the network on noisy version the zero codeword. However, this approach remains extremely restrictive, due to the strong inductive bias induced by the Tanner graph.
employ variants of generic neural networks and may potentially benefit from the application of powerful architectures that have emerged in recent years in various fields. However, such decoders suffer from the curse of dimensionality, since the training grows exponentially with the number of information bits [31). A major challenge of model-free decoders, therefore, the difficulty the network in lcarning the code and identifying its most reliable components, especially for the basic multi-layer perceptron architectures suggested in the literature.
Model-free decoders 23 10
In this work, we propose model-free decoder built upon the Transformer architecture 29]. As far can ascertain, this is the first adaptation of Transformers to eтror correction codes. Our Transformer employs a high-dimensional scaled element-wise embedding of the input which can be seen as positional and reliability encoding. Our decoder employs an adapted-mask self-attention mechanism, in which the interaction between bits follows the code's parity check matrix, enabling the incorporation of domain knowledge into the model.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Applied to wide variety of codes, our method outperforms the state-of-the-ar learning- based solutions by very large margins, even with very shallow architectures, enabling potential deployment. This is the first time decoder designed de-novo outperforms neural BP-based decoders.
2 Related Works
Over the past few years, the emergence of deep learning has demonstrated the advantages of Neural Networks in many communication applications, such as channel equalization, modulation, detection, quantization, compression, and decoding [12]. In [23) the entire coding-dccoding channel trans- mission pipeline was abstracted as a fully parameterized autoencoder. In many scenarios, neural network-based methods outperform existing decoders, as well as popular codes (20) [O].
Two main classes of Neural decoders can be established in the current literature 2.5 Model-free
decoders employ general types of neural network architectures, as in obtained in [0 for polar codes (n 16) are similar to the maximum a posteriori (MAP) decoding results. However, the exponential number of possible codewords make the decoding of larger codes unfeasible. In (2), preprocessing of the channel output allows the decoder to remain provably invariant to the codeword and enables the deployment of any kind of model without any overfitting cost. The model-free approaches generally make use of stacked fully connected networks or recurrent neural networks in order to simulate the iterative process existing in many legacy decoders. However, these architectures have difficulties in learning the code and analyzing the reliability of the output, and generally require prohibitive parameterization or expensive graph permutation preprocessing (2].
1 2 The results
The second class is that of model-based decoders [9 20 implementing augmented parameterized versions of classical BP decoders, where the Tanner graph is unfolded into a neural network which weights are assigned to each variable edge, resulting in improvement in comparison to the baseline BP method. This approach remains limited since is built on top of a restrictive model, Recently, (25) presented data-driven framework for codeword permutation selection pre-decoding stage. A self-attention mechanism is applied over several permutation-speeifii embeddings to estimate the best permutation (among given set) to be decoded. A combination of the BP model with hyper-graph network was shown to enhance performance [21]. Further improvements over the latter were obtained by introducing an autoregressive signal based on previous state conditioning and SNR estimation (22].
Transformer neural networks were originally introduced for machine translation dominate the performance tables for most applications in the field of Natural Language Processing (NLP), e.g.. [29] [7). Transformer encoders primarily rely on the self-attention operation in conjunction with feed-forward layers, allowing manipulation of variable-size sequences and learning of long-range dependencies. Many works extended the Transformer architecture from NLP to various applications, such as computer vision [8], speech analysis [17 and molecular properties prediction [6].
29 and now
3 Background
We provide the necessary background error correction coding and the Transformer architecture.
3.1 Coding
We assume a standard transmission that uses linear code C. The code is defined by the binary generator matrix G of size k x n and the binary parity check matrix H of size (n- k) x n defined such that G over the two elements' Galois field.
The input message m E {0, 13 is encoded by G to a codeword x E C' C. f0, 1]" satisfying Hrr = 0 and transmitted via a Binary-Input Symmetric-Outpul channel (c.g. AWGN channel). Let y denote the channel output represented as y F2 + where v, denotes the Binary Phase Shift Keying (BPSK) modulation of (i.e., over {=1)), and is random noise independent of the transmitted z.
The main goal of the decoder f RD - RD to provide a soft approximation = f of the codeword. An illustration of the coding framework is presented in Figure
We wish to employ the Transformer architecture as part of a model-free approach. We, therefore, follow the preprocessing and post-processing of [2] to remain provably invariant to the transmitted

---

Encoder
Decoder
x E [0, 1)"
X2 E [11)a
y IR
(0, 1
BPSK
X, +
fo())
X R
Figure I: Illustration the communication system. Our work focuses on the design and training of the parameterized decoder fg.
codeword and to avoid overfitting. This step involves a transformation of the channel output which entails loss of information without implying any intrinsic performance penalty in decoding, i.c. the preprocessing coupled with an appropriately designed decoder can achieve MMSE decoding.
The preprocessing replaces y with vector of dimensionality 2n k defined
(I) where, [. denotes vector concatenation, lyl denotes the absolute value (magnitude) of y and s(y) Hys {0, 1]"=* denotes the binary code syndrome obtained via the multiplication of the parity check matrix by the binary mapping of y such that = 0.5(1-
j = hly) |1,l, s
- sign(y)).
The post-prose pest-processig predicts the transmitted codeword by multiplying the elements of y with the predicted multiplicative noise defined such that y = I, Namely, the prediction takes the form - y f(h(y))
In our setting, the decoder is parameterized by 0 such that i = fo(lul. Hyb).
3.2 Transformers
The Transformer was introduced novel, attention- based building block for machine trans- lation [29]. The input sequence is first embedded into a high-dimensional space, coupled with positional embedding for each element. The embeddings are then propagated through multiple normalized self-attention and feed-forward blocks.
The self-attention mechanism introduced by Transformers is based on a trainable associative memory with (key, value) vector pairs, where query vector IRd is matched against a set of k key vectors using scaled inner products, follows
QKT Vd
A(Q. K, V) = Softmax
V.
where E Nxd. K E Rexd and V Rexd represent the packed N queries, & keys and values tensors respectively. Keys, queries and values are obtained using linear transformations of the sequence's elements. A multi-head self-attention layer is defined by extending the self-attention mechanism using h attention heads, i.e. h self-attention functions applied to the input, reprojected to values via dh D linear layer.
4 Error Correction Code Transformer
We present the elements of the proposed Transformer decoder, the complete architecture, and the training procedure.
4.1
Positional Reliability Encoding
We first consider each dimension of (0 separately and project each one to a high dimensional embedding 2n 2 such that
W.
if i otherwise
Φi (1 2(s(())

---

where [W, E R P= denotes the one-hot encoding defined according to the bit position. The embedding modulated by the magnitude and syndrome values, such that less reliable elements (i.e. low magnitude) collapse to the origin. This property becomes especially appealing when applied to the standard dot product of the self-attention module: for the first layer and two distinct information bits embedding @i 2 we have
[wll,,, [u:1(1 2(s(()) (W,, Wj), isncj.
ij
( 03)
(5)
where an unreliable information bit vanishes and a non-zero syndrome entails a negative scaling, potentially reducing the impact on the softmax aggregation, illustrated in Sec. 6
In contrast to (2). which requires permutation the code order to artificially provide to the fully connected network indications about the most reliable chamnel outputs, in our construction, the channel output reliability is directly obtained and maintained in the network via the scaled bit-wise embedding. The proposed scaled encoding can be thought of positional encoding according to the input reliability since the bit positions are fixed. We note that representing the codeword by set of one-hot vectors is very different from any of the existing methods.
4.2 Code-Aware Self-Attention
In order to detect and correct errors, a decoder must analyze and compare the received word bits via the parity check equations, such that non-zero syndrome indicates that channel errors have occurred. our Transformer-baaed architecture, bit comparisons become natural via the self-attention interaction mechanism.
However, comparing every pair of elements, is generally performed in Transformer architectures, is sub-optimal, since not every bit is necessarily related to all the others. We propose to make use of the self-atte afffattention mechanism in order to incorporate fundamental domain knowledge about the relevant code. Specifically, we use to indicate that the syndrome values should only be dependent on the corresponding parity check bits.
Thus, for given algebraic code defined by the matrix H we propose to define a function g(lt): : {0, ()(m=))xA + 0-20 defining the mask to be applied to the self-attention mechanism, such that
QKK" + V. Vd
An K,V) Softmax
(6)
We propose to build the symmetric mask such that contains information about every pairwise bit relation, follows. The mask is first initialized the identity matrix. For each row of the binary parity check matrix H, we unmask at the locations of every pair ones the row, since those bits connected and may impact each other well as the syndrome in the decoding procedure. We also unmask the location of pairs of ones with the corresponding syndrome bit at + i, since they define the parity check equations. Illustrations of typical self-attention maps given in Appendix A.
Algorithm I: Mask construction Pseudo Code
function (H) 1,n shape
k n-1 mask (2ekkn)
for i in range(0) n-k) do
idx where (=1]===== for j in idx do mask (nti,,,,,1] for k in idx do mask [j k] mask [k,j]
mask[j,nt
11
return (- mask)

---

Mask
PC Matrix
X
X
X
x
1-
6-
2-
%0
G
1
2
4
5
(a)
(6)
(c)
Figure 2: We present the corresponding parity-check matrix (a), the induced Tanner graph (c) and the proposed masking (b) on the Hamming(7,4) code.
We summarize the construction of the mask in Algorithm 1 This construction enables greater freedom in decoding than the relations enabled by the Tanner graph, since related bits may impact each other beyond the parity check equations depicted in Figure While regular Transformers can be assimilated to a neural network applied on complete graph, the proposed mask can be seen as the adjacency matrix of the Tanner graph extended two rings connectivity. In contrast to BP, which collapses the information via interleaved variable and check layers, the masked self-attention allows the simultancous aross-analysis of related elements.
Most importantly, since the mask is fixed and computed once, the self-attention quadratic complexity bottleneck (nndd is now reduced to the density of the code ofE,, (H)) This property is especially appealing for low-density codes (e.g. [9], while in all experiments the original complexity of the self-attention is reduced by 84% on average, as presented in Sec.
4.3
Architecture and Training
The initial encoding is defindd as d dimensional one hot encoding of the 2n k input elements. The decoder is defined as concatenation of N decoding layers composed of self-attention and feed-forward layers interleaved by normalization layers. The output module is defined by two fully connected layers. The first layer reduces the element-wise embedding to a one-dimensional 2n vector and the second to a dimensional vector representing the soft decoded noise. An illustration of the model given Figure
The dimension of the feed-forward network is four times that of the embedding (29 and is composed of GEGLU layers (27], with layer normalization set the pre-layer norm setting, in [S]] 32 We use an eight-head self-attention module in all experiments. colorblue The impact of the number of heads on the accuracy further explored in Appendix C.
The training objective the cross-entropy function where the goal is to learn to predict the mud- riplicative noise (2). Denoting the soft multiplicative noise 2, such that y x,2a, obtain E. = 2.0² y. Thus, the binary multiplicative noise to be predicted is defined by ž = bin (y w.). such that the loss computed for single received word y
L
Σ & log(fo(y)) 1 log(1 fo()))
(7)
The estimated hard-decoded codeword is straightforwardly obtained as d. = bin(sign(Jo()) v)).
The Adam optimizer [4]] is used with 128 samples per minibatch, for 1000 epochs, with 1000 minibatches per epoch. For N = 10 architectures we trained the models for 1500 epochs. We note that using more epochs can improve performance. However, the current setting is already more than enough to reach new SOTA performance. Due to the construction of the model and its input preprocessing, the zero codeword enough for training. The additive Gaussian noise sampled randomly per batch in the (3..... 7} normalized SNR (i.e. Eg/No range. Additional experiments with the non-Gaussian Rayleigh fading channel are provided Appendix E. We initialized the learning rate to 10-1 coupled with cosine decay scheduler down to 5 10-7 at the end of the training. No warmup was employed [32]. Training time range from 12 24 hours depending on the code length, and no optimization of the self-attention mechanism have been employed. Training and
5

---

biul (sigu())(()
y
Initial Embedding
Decoding
Output
Nx
Feed Fwd
FC
Norm
2n k
FC
2H bin(sign())
MH-SA
d
Embedding
Norm
H
Norm
Figure 3: Illustration of the proposed Transformer architecture. The main differences from other Transformers are the initial scaled bit-embedding, the construction of the code-aware masked self- attention, and the output module.
experiments have been performed on 12GB Titan V GPU. The training time was in the range of 19-40, 40-102, and 56-171 seconds per epoch for the 10 architectures, respectively.
5
Experiments
!
To evaluate our method, we train the proposed architecture with three classes of linear block codes: Low-Density Parity Check (LDPC) codes [9], Polar codes [ and Bose-Chaudee Houudhuri-Hocquenghenn (BCH) codes 3. All parity check matrices are taken from [I]] The proposed architecture is defined solely by the number of encoder layers N and the dimension of the embedding d. We compare our method to the BP algorithm, to the augmented hypernetwork BP algorithm of (2 (hyp BP), to the RNN architecture of [2 and to the very recent SOTA performance of [22] (AR BP). All the results were obtained from the corresponding papers. comparison to the SCL Polar decoder [28] given in Appendix D. The results reported as bit error rates (BER) for different normalized SNR values (dB). We follow the testing benchmark 21 22]. During testing, our decoder decodes at least 1 random codewords, to obtain at least 500 frames with errors at each SNR value.
The results are reported in Tab. [ where we present the negative natural logarithm of the BER. For each code we present the results of the BP-based competing methods for 5 and 50 iterations (first and second rows), corresponding to a neural network with and 100 layers, respectively. We present our framework performance for six different architectures with N (2,6) d {32, 64, 128) respectively (first third rows). Since LDPC codes are optimized for BP decoding (26), we also present for these codes the performance of N 10, d 128 architecture. As can be seen, approach outperforms current SOTA results by very large margins on several codes, at fraction of the number of iterations, and is able outperform legacy methods (e.g. BP) at convergence, even with extremely shallow architectures (e.g. N 2, 32). We remark that the proposed method seems to perform better for high rate (i.e. 4 codes.
We provide plots of more BER values for some of the codes in Figure (ab) for Polar(64,32) and BCH(633,51) codes, respectively. Our method is able outperform the converged AR-BP by up to two orders of magnitude for high SNR values. Performance on larger codes is studied in Appendix H.
In Fig. 4 4c), We compare our method with the best model of 2 the only code provided BCH(127,64) 64). Their network is built as five-layer stacked GRU model, with five iterations with an embedding of 5(2n k) 950 trained for 5000 epochs, with batch size of 1000. As can be seen, our method outperforms this model by up to 21 %. with only half the number of layers, one-tenth as many parameters, and fraction of the training requirements.
'Code available htttt:::/ Bttpthhhhwb.cccoooco///CCC/

---

Table A comparison of the negative natural logarithm of Bit Error Rate (BER) for three normalized SNR values of our method with literature baselines. Higher is better.
Concurrent results are obtained after and at convergence results in second row obtained after I. = 50 BP iterations (i.e. 100 layers neural network). Best results in bold, second best underlined, and the minimal Transformer architecture to outperform every other competing method is in italic.
BP iterations in first row (i.e. 10 layers neural network)
Our performance is presented for seven different architectures: for N = {2, 6). we present results for d = (32, 64, 128) (first to third rows), and for N = 10 we run only the d 128 configuration.
Method
BP 24
Hyp 2P
AR BP 22
Ours N=2
Ours N=6
Ours N=10
4
5
6
5
4
5
6
4
5
5
4.27
6.95
5.71
7.63
9.94
Polar(64.32)
3.52 4.04
4.48
4.25 5.49
7.02
4.77
6.30
8.19
4.57
5.86 7.50
6.48
8.60
11.43
4.26 5.38 6.50
6.10
5.57
7.43
9.82
4.87 6.2
7.93
6.99
9.44
12.32
4,92
5,82
10.24
Polar(64,48)
4.15 4.15 4.68 5.31
4.91
6.48
8.41
5.25
6.96
9.00
5.14 6.78 8.9
6.15
8.20
10.86
5.94 7:42
4,92
6.44
8.39
5,41
7.19
9.30
5,36 7.12
9.39
6.36 8.46
11.09
3.51 4.52 5.93
4.47
6.34
8.89
Polar(128,644
3.38 3.80 4.15
3.89 5,18
6.94
4.02
5.48
7.55
3.83 5.16 7.04
5.12
2.36
10.48
4,10 6.15
4.52 6.12
8.25
4.84
6.78
9.30
5.92 8.64
12.18
Pol.r(128.86)
3.80 4.19 4.62
4.57
6.18
8.27
4.81
6.57
9.04
4,30 5.58 7.34
4.49 5,90 7.75
5.36
5.75
7.45
8.16
11:29
10.22
4.49 6.97
4.95 6.84
9.28
5.39
7.37
10.13
4.75 8.29
6.31
9.01
12:45
4.56 5.98 7.93
5.39 7.62
10:45
Polar(128,96)
3.99 4.41 4.78
4.73
6.39
8.57
4.92
6.73
9.30
4.69 6.20 8.30
5.88
8.33
11.49
4.61 5.79 7.08
4.94 6.76
9.09
5.27
7.44
10.2
4.88 6.58
8.93
6.31
9.12
12.47
5.30 7.28 9.88
5.76
7,90
11.17
6.05
8.13
11.68
4,51 6.07 8.11
5.74 8.13
11:30
LDPC(49.24)
6.23 8.19 11.72
6,23 8.54
11.95
6.58
9.39
12.39
4.58 6.18 8.46
5.91 8.42
11.90
6.68 9.53 13.30
4.71 6.38 8.73
6.13
8.71
12.10
3.88 5.51 8.06
4.98 7.91 12.70
4.82 7,21 10.87
5.22 8.29 13.00
5.22
8.31
13.07
LDPC(121.60)
5.02 7.94 12.72
5.73 9.35 15.01
3,93 5.66 8.51
5.17 8.31
13:30
LDPC(121.70)
5.88 8.76 13:04
6.39 9.81 14.04
6.45
10.01 14.77
4.63 6.68 9.73
6.28
6.11 9.62 15.10
15.57
7.13 11.50 17.92
4.67 6.79 9.98
6.40 10.21
16.11
5.27 7.59 10.08
6.92
10.74 15.10
LDPC(12180)
6.66 9.82 13.98
6,95
10.68
15.80
7.22
11.03 15.90
5.29 7.63 10.90
7.17 11.21 16.37
8.12 12.90 17.82
5,30 7.65 11.03
7.41 11.51 16.44
4.95 6.67 8.94
6.88 9.86 13.40
MacKay(96,48)
6.84 9.40 12.57
7.19
10.02
13.16
744 10.65 14.65
5.04 6.80 9.23
7.10 10.12 14.21
8.39 12.24 16.41
5.17 7.07 9.64
7.38 10.72 14.83
8.30
6.34
14.40
6.55 9.65 13.78
6.99
10.57
15.27
7.25
10.99 16.36
CCSDS(128.64)
4,41 6.09 8.49
6.65
10.40 15.46
8.02 12.60 17.75
4.59 6.42 9.02
6.88 10.90 15.90
5.74
9.59
1CH(31,16)
4.63 5.88 7.60
5.05
6.64
8.80
5.48
7.37 9.61
4.7% 6.15 7.98
5.85
7.52 10.08
5,18 6.82 8.91
6.39
8.29
10.66
3.62
5.66
3.96
5.35
7.20
4.33
5.94
8.21
3,79 4.87 6.35
4.42
5:91
8.01
BCH(63,36)
4.05 5.28 7.01
4.62 6.24
8.44
4.03 5.42 7.26
4.29
5,91
8.01
4.57
6.39
8.92
4.21 7.25
4,86
9.10
BCH(63445)
4.36 5.55 7.26
4.08 4,96
6.07
4.64 6.27
6.07
8.45
8.51
4.97
4.80
6.43
6.90
8.69
9.41
4.66 6.16 8.17
4.47 5.88 7.81
8.49
5.60
5.41
5.16
7.79
7.49
7.02
10.93
10:25
9.75
BCH(63,51)
4.34 5.29
6.35
4.64
6.08
8.16
4.95
6.69
9.18
4,60 6.05 8.05
4.78 6.34 8.49
5,20
5.46
7.08
7.57
10.51
9.65
4.5 5.82 7.42
4.80 6.44
8.58
5.17
7.16
9.53
5.01 6.72 9.03
5.66
7.89
11.01
Analysis
We study the impact of the proposed embedding and the masking procedure; furthermore, we analyze and compare the complexity of our method with other methods.
7

---

MN) BX
125
10
10
10"h
1844
10a3
10"
10"
3.5
4.0
4.5
5.0
3.5
BANAKO
Eshuod81
(a)
(b)
(c)
Figure 4: BER for various values of SNR for (a) Polar(64,32) and (b) BCH(63,51) codes. Comparison with the best model of Bennatan et al. 2 on BCH((27,64)) code (c).
Layer
Layer
Thansmpped CN
UNoomppheldd dVl
Gemplell DV
6:3-
0:15
C119
6:00
600
090
Trpaal
Topal
Layer
Layer
6.40
Creminidd the
6:10
Crematiill TM
Compard CA
Sultmas
4.15
6:6
605
Tedlex
(a)
(b)
Figure 5: Analysis of the scaled embedding values of the first column of the masked self-attention map with corrupted and uncorrupted codeword (CW) for the standardized Hamming code, using the (a) (N, d) = (4,32) 32) and (b) (N, d) = (2, 32) architecture. The impact of the noisy embedding on the self-attention's first pariky-check element (index 7) is visible.
6.1 Impact of the Reliability Embedding
We depict in Figure the impact of the scaled embedding on the self-attention mechanism. We chose the the popular Hamming(7, 4) code, corrupting the zero codeword with additive noise at the first bit (zero bit index), involving non-zero syndrome at the first parity check equation. We present the masked softmax values of the self-attention map at index zero, i.e. we visualize the influence of the first (unreliable) bit on the self-attention aggregation (first column of the map). As can be seen, the first bit embedding is impactless when it corrupted and then detectable, while its impact on the syndrome embedding is considerably increased (index 7). Once the NN corrects the bit (last layer(s)), the values return to normal. A similar analysis for larger codes can be found in in Appendix B.
6.2 Ablation Study: Masking
We present in Figure the impact of the proposed mask on the convergence. We show an unmasked architecture, where we leave the network to learn the code by itself (i.e. g(H) = 0), and compare it with our proposed masking approach. The benefit obtained with the mask is clear and understandable in light of the difficulties the Transformer may have in learning the code by itself. Therefore, the connectivity provided by the proposed masking framework reduces the loss by 76%, 69% and 66% for the BCH(633, 36), POLAR 32) and LDPC(49, 49, 24) codes, respectively.
8

---

Convergence
Masked UNisolted
Selt Attentian Complexity Sparsity
Brbetracn
Spanilyy
Socglesitt
Cunn Hexil (16)
Sporsity (%)
8
POLARI
201J8/
I2P01 10010.200
ED 83)
EOPOIN 70
0
200
400
000
800
1000
Epach
Dace
Figure 6: The effect of the proposed masking procedure on the convergence for BCH(63, 36),
Figure 7: Sparsity and complexity ratio of the
POLAR(64, 32) and LDPC(49, 24) using
proposed self-attention map for several codes.
(N, d) (6, 64).
6.3 Complexity Analysis
We first present, in Figure 7 the sparsity ratio of the masked self-attention map, well as the induced complexity ratio of the symmetric map for several codes presented the Results section. The ratio is computed with respect to the Onn²) legacy self-attenties map size. We can that sparsity reach up to more than 80% while the computation ratios range from 5 to 35% only. This improvement can lead to sizable reduction in power and computation time dedicated hardware.
The complexity of the network is defined by (df - k) + hd), where h < n² denotes the fixed number of computations of the self-attention module. The real-world complexity for different codes explored in Appendix G. As with most neural decoding approaches, while being extremely effective, the proposed method may remain inefficient compared to non-learning solutions, such as [24], in terms of memory requirement, power consumption, and computational resources, limiting their deployment potential. Acceleration of the proposed Transformer via architectural modification is left for future work.
or NN acceleration (e.g. pruming, quantization) [3 [6]
In comparison, the existing AR-BP SOTA method of (22 requires O(2N dnd n dry + (n k)² + dy)) more operations than Hyper BP(21). where du, dy, dg are the number of variable nodes, the capacity of the hypernetoork and primary network, respectively. Typical networks f and g are a 128 dimensional FC network with 4 layers and 16 dimensional FC network with 2 layers respectively. With these hyperparameters, the d part of the complexity is approximately 1282 Moreover, the method of [22) cannot be parallelized, due the hypernetwork structure, i.e. each sample defines different set of parameters. Thus, our network far surpasses the existing state- of-the-art error rates, with much less time complexity. Hyper-BP (21) is slightly less computationally intensive, since adds complexity term of O (2.N (dNnddddd) ) on top of BP, but this approach as well the original model-based network of 19 are far less competitive terms of accuracy.
Conclusions
We present a novel Transformer architecture for decoding algebraic block codes. The proposed model allows effective representation of interactions, based on the nigh-dimensional embedding of the channel output and dedeedeeendent masking scheme of the self-attention module. The proposed framework is more efficient than standard Transformers, due to the sparse design of existing block codes, making its training and deployment extremely affordable. Even with a limited number of computational blocks, it outperforms popular (neural) decoders on broad range of code families. We believe this new architecture will allow the development and adoption new families of performing Transformer-bast based codes capable of setting new standards in this field.

---

Acknowledgments
This project has received funding from the European Research Council (ERC) under the Euro- pean Union's 's Horizon 2020 research and innovation programme (grant ERC CoG 725974). The contribution of the first author is part of a PhD thesis research conducted at Tel Aviv University.
References
[1] Erdal Arikan. Channel polarization: A method for constructing capacity-achieving codes, In 2008 IEEE International Symposium on Information Theory, pages 1173-1177. IEEE, 2008.
[2] Amir Bennatan, Yoni Choukroun, and Pavel Kisilev. Deep learning for decoding of linear codes-a syndrome-basee approach. 2018 IEEE International Symposium on Information Theory (ISIT), pages 1595-1599. IEEE, 2018.
[3] Raj Chandra Bose and Dwijendra K hay-Chaudhuri. On class of error correcting binary group codes. Information and control, 3(1):68-79, 1960.
[4] Andreas Buchberger, Christian Häger, Henry D Pfister, Laurent Schmalen, et al. Learned decimation for neural belief propagation decoders. arXiv preprint arXiv:2011.02166 2020.
[5] Sebastian Cammerer, Tobias Gruber, Jakob Hoydis, and Stephan ten Brink. Scaling deep learning-based decoding of polar codes via partitioning. In GLOBECOM 2017-2017 IEEE Global Communications Conference, pages 1-6. IEEE, 2017.
[6] Yoni Choukroun and Lior Wolf. Geometric transformer for end-to-end molecule properties prediction. arXiv preprint arXiv:2110.13721 2021.
[7] Jacob Devlin, Ming-Wös Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 181000088 2018.
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint
arXiv:2010.119299 2020.
[9] Robert Gallager. Low-density parity-check codes. IRE Transactions on information theory, 8(1):21-28, 1962.
[10] Tobias Gruber, Sebastian Cammerer, Jakob Hoydis, and Stephan ten Brink. On deep learning- based channel decoding. In 20/7 51st Annual Conference on Information Sciences and Systems (CISS), pages 1-6. IEEF, 2017.
[11] Michael Helmling, Stefan Scholl, Florian Gensheimer, Tobias Dietz, Kira Kraft, Stefan Ruzika, and Norbert Wehn. Database of Channel Codes and ML Simulation Results. www. uni-kn kl. de/ channel-codes! 2019.
[12] Mohamed Ibnkahla. Applications of neural networks to digital communications-a survey. Signal processing, 185-::115
[13] Hyeji Kim, Yihan Jiang, Ranvir Rana, Sreeram Kannan, Sewoong Oh, and Pramod Viswanath. Communication algorithms via deep learning. Sixth International Conference on Learning Representutions (ICLR), 2018.
[14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:1412.6980,, 2014.
[15] Guillaume Klein, Yoon Kim, Yuntian Deng, Jean Senellart, and Alexander M. Rush. Opennmt: Open-source toolkit for neural machine translation. In Proc. ACL, 2017.
[16] Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. survey of transformers. arXiv preprint arXiv:2106.045: 04554, 2021.
10

---

[17] Andy T Liu, Shang- Wen Li, and Hung-yi Lee. Tera: Self-supervised learning of transformer encoder representation for speech. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 29:2351-2366, 2021.
[18] Loren Lugosch and Warren Gross. Neural offset min-sum decoding. In 2017 IEEE Interna- tional Symposium on Information Theory (ISIT), pages 1361-1365. IEEE, 2017.
[19] Eliya Nachmani, Yair Be' Be'ery, and David Burshtein. Learning to decode linear codes using deep learning. In 2016 54th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pages 341-346. IEEE, 2016.
[20] Eliya Nachmani, Elad Marciano, Loren Lugosch, Warren Gross, David Burshtein, and Yair Be' ery. Deep learning methods for improved decoding of linear codes. IEEE Journal of Selected Topics in Signal Processing, 12(1):119-331, 2018.
[21] Eliya Nachmani and Lior Wolf. Hyper-graph-networt decoders for block codes. In Advances in Neural Information Processing Systems, pages 2326-2336, 2019.
[22] Eliya Nachmani and Lior Wolf. Autoregressive belief propagation for decoding block codes. arXiv preprint arXiv:2103.11788 2021.
[23] Timothy J O'Shea and Jakob Hoydis. An introduction to machine learning communications systems. arXiv preprint arXiv: 2017.
[24] Judea Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan kaufmann, 1988.
[25] Nir Raviv, Avi Caciularu, Tomer Raviv, Jacob Goldberger, and Yair Be' eeyery. perm2vec: Graph permutation selection for decoding of error correction codes using self-attention. arXiv preprint arXiv:2002.02311 02.2000...22.,,.... 2020.
[26] Thomas Richardson, Mohammad Amin Shokrollahi, and Rüdiger L Urbanke. Design of aapaaity-approcchiin irregular low-density parity-check codes. IEEE transactions informa- tion theory, 47(2))619-637, 2001.
[27] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 2020.
[28] Ido Tal and Alexander Vardy. List decoding of polar codes. IEEE Transactions on Information Theory, 61(5):2213-222) 2015.
[29] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
[30] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04 047688 2020.
[31] Xiao-An Wang and Stephen Wicker. An artificial neural net viterbi decoder. IEEE Transac- lions on communications, 44(2): 166-11,, 1996.
[32] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu. On layer normalization the transformer architecture. arXiv preprint arXiv:2002.04745, 2020.
11


Reviews:
[{"review_id": "yL2lOIGjFh", "reviewer": "Reviewer_MQDw", "summary": "This paper studies neural decoding of (linearly coded) transmissions in the physical communication layer. It proposes a novel architecture that has three main features: it is based on transformers, contains a scaled element-wise embedding of the input, and has an adapted-mask self-attention mechanism. The proposed architecture is compared empirically to existing neural decoders and is shown to outperform them with higher decoding success and less computational complexity.", "questions": "Question 1: in related work, a dichotomy between model-based on model-free was presented and the authors of the present work place theirs’ in the model-free category. Since the present work uses expert knowledge in how the adapted-mask is generated, one could argue that it belongs somewhere in the middle of these categories. Could the authors elaborate on this?\n\nQuestion 2: Would there be a way to compare neural decoders to other non-learned SOTA decoders? Would it be possible to do more quantitive comparison as well?", "limitations": "See above.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n----------------\n\nStrength 1: The presented method is the first transformer-based decoder and it also gives an efficient architecture.\n\nStrength 2: The proposed method has better decoding performance than other neural network based approaches. \n\nStrength 3: The proposed method allows to input expert knowledge to the system and has lower complexity than other neural network based methods.\n\n\nWeaknesses:\n---------------------\n\nMinor Weakness 1: Only AWGN channel is considered in this work. This may be in line with other neural decoder works, but it would be important to consider also wireless channels with fast fading and e.g., with interference, which change the noise statistics. There is a possibility that the learned methods might provide even further gains in these scenarios by learning the noise correlation statistics.\n\nMinor Weakness 2: Especially with wireless fast fading channels, other modulations than BPSK should be considered, at least 16QAM.\n\nAfter the author feedback (rebuttal):\n------------------------------------------\nI have reviewed the author feedback and they have done good job in clarifying most of the questions the reviewers have posed. I have changed the score accordingly. \n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "i2n0exlmbhT", "reviewer": "Reviewer_AayN", "summary": "- The paper presents a novel Transformer based generic decoding procedure for typical linear error correction code families (LDPC, Polar etc..). \n- The method does achieve SOTA performance on mid block length codes (~100-200)\n- The key novel contribution of the work was the introduction of positional reliability embeddings and the attention mask, both of which (particularly the mask) uses some clever information theoretic domain understanding. These ideas improve the convergence and hence the performance of the model\n- One key selling point of the paper is that the method is general enough to be applied to any linear code without any modification. ", "questions": "- The use of mask is quite nice, and the authors claim that it reduces the complexity by 84%.. does this reduction in complexity lead to reduction in FLOPS/run-time.. or just a symbolic reduction which leads to lower complexity and better convergence? Please clarify\n\n- The intuition behind positional reliability encoding is intuitively clear to me, at least for the elements corresponding to |y_i|.. can the authos explain the intuition behind the encoding for the syndrome bits? \n\n- Comparison with List-SC decoding for polar codes might be interesting. ", "limitations": "No societal limitations", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- The paper presents a novel way of using transformers and self attention to achieve strong results in terms of the BER\n- The architecture isn't just using a ML-model and blindly using it for the decoding problem. The training uses some clever ideas from the code construction to improve the convergence.\n\nWeaknesses:\n- The explanation was overall good, but in some places a bit lacking. More detailed comments at the end\n- The authors briefly mention that the runtime/power numbers might be sub-par as compared with non-ML methods.. still it might be great to see the comparison to know where the field is at. \n\nSpecific suggestions:\n- A couple of typos.. line 117 propriety -> property\n- $y_b$ -> although clear, it is not defined in the manuscript\n- The postprocessing an pre-processing is not clear.. please describe more clearly. For example: it is not clear to me (without reading the cited reference, what you mean by \"The post-processing step plugs back the vector elements of y....\"\n- The use of mask is quite nice, and the authors claim that it reduces the complexity by 84%.. does this reduction in complexity lead to reduction in FLOPS/run-time.. or just a symbolic reduction which leads to lower complexity and better convergence? Please clarify", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "h7TXKmjTNll", "reviewer": "Reviewer_FPkL", "summary": "This paper proposes a Transformer based that employs relaxed inductive bias compared tgo Tanner graph-based methods while utilizing domain knowledge compared to standard model-free decoders based on fully-connected graph.\n", "questions": "- Figure 3: Is the multi-head self-attention block still necessary? Given that we’re looking for the localized bits that meet each row of the parity check (PC) matrix, will single-head self-attention suffice? In other words, if you permute the rows of the PC matrix, you will get the same result, indicating that the problem has some sort of equivariance property that the authors may utilize in the Transformer design\n\n- Figure 2: This is a good illustration; it would be fascinating to see how the attention maps visualized after the training and how they compare to the restrictive Tanner graph; it would also be interesting to see if the mask-based decoders genuinely help attention to more targeted interactions or not \n\n- Ln 147: Meant Algorithm 1?\n", "limitations": "- Ln 154-157: Is the complexity reduction in 4.2 from code-aware self-attention adequate enough to apply the method to longer LDPC codes, such as greater than 512, say 1K or 4K lengths? If it’s not, the authors should have stated that clearly as a limitation or future work of the study", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper's main strength is the well-defined scope of the problem. \n\nThe paper's main weakness is that the proposal’s benefits were not studied thoroughly. For examples, I have a few questions below\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YecrtciLAeq", "reviewer": "Reviewer_Nt2o", "summary": "This paper considers a very important and interesting topic: the decoding of error-correcting codes. \nOver the past few years, there has been a lot of work on applying learning to the decoding problem. Various neural architectures have been considered as well, starting from feedforward networks to recurrent and convolutional neural networks. There also have been a long line of work which introduced learnable parameters to the existing decoder architecture, some of which are considered as a baseline in this paper. \n\nThere has not been a transformer-based channel decoder yet; training a transformer that can successfully decode error-correcting codes is empirically quite challenging. This paper addresses the challenge by utilizing the knowledge of codes in a way; using the appropriate masking and modulating the embedding by syndrome values. This approach is novel and interesting - so I would like to advocate the acceptance of this paper in that regard. \n", "questions": "I'd be curious to see the comparison between the reliability of neural-augmented BP and the transformer  (vanilla, no weight adaptation network) -based decoder when their complexities are similar. \n\nI'd be also curious to see the performance of these decoders for non-Gaussian channels. Is there a difference in terms of robustness? ", "limitations": "The authors mention that the complexity could be potentially improved by several techniques. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n(Strength) \nAs mentioned in detail above, the paper successfully demonstrates a transformer-based channel decoder for the first time. \n\n(Weakness) \nDespite the idea being novel and the results being promising, I have two major concerns. \n\nThe first is whether the comparison between the BP families and the transformer families is fair. Given that the transformer-based decoder is not a purely data-driven decoder, we would expect to see gains compared to the existing decoder. For example, L-layered augmented BP decoders typically have complexity similar to the L-layered BP decoders. Hence, the complexity of the augmented BP decoder is similar to the complexity of the traditional BP decoder but is superior in terms of reliability. [Hyper-BP might be computationally expensive due to the 'hyper network' which selects the weights. However, there are other neural augmented BP algorithms that do not include the weight selection module but are pretty good.]\n\nSecond, the presentation and ablation studies can be further improved. For example, section 6.1 and the description of Figure 5 lack details. Also, I wonder if any similar ablation study can be done for the polar or LDPC or BCH codes. Section 6.2 is interesting and insightful. In Section 3, providing the number of operations would be more informative (and potentially mention vanilla augmented-BP algorithms as well). \n\nFor these reasons, my current score is 4, but I am really at the exact borderline. I would like to listen to the authors; I will go through the rebuttal and fix my score afterwards. \n\n--- \nAfter the rebuttal: The new experiment results (e.g., Appendix G) are exactly are very informative and satisfactory. The reliability gap between the transformer-based decoders and neural BP decoders does not seem huge. (It is a bit hard to interpret the gain from the negative logarithm of the BER, and the SNR gain might be more indicative.) Nevertheless, I still really like the idea of utilizing the transformer architecture for decoding, and the updated manuscript is comprehensive. I'd updated my score accordingly. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept