Multiple-Masks Error Correction Code Transformer for Short Block Codes

With the broadening applications of deep learning, neural decoders have emerged as a key research focus, specifically aimed at improving the decoding performance of conventional decoding algorithms. In particular, error correction code transformer (ECCT), which utilizes the transformer architecture, has achieved state-of-the-art performance among neural network-based decoders. We present three technical contributions to significantly enhance the performance of ECCT. First, we propose a novel transformer architecture of ECCT, termed the multiple-masks ECCT (MM ECCT). We employ multiple masked self-attention blocks with different mask matrices in a parallel manner to learn diverse relationships among the codeword bits. Second, we discover that constructing mask matrices based on systematic parity check matrices (PCMs) can make the attention maps sparse, which not only enhances the decoding performance but also reduces computational complexity. Finally, we propose using complementary mask matrices derived from cyclic permutations of the systematic PCM. These complementary mask matrices are specifically designed to enhance the decoding of cyclic codes. Our extensive simulation results show that the proposed MM ECCT architecture with carefully designed mask matrices outperforms the original ECCT by a large margin, achieving state-of-the-art decoding performance among neural decoders. The source code is available at https://github.com/iil-postech/mm-ecct.