Latent-Attention-Based Transformer for Near ML Polar Decoding in Short-Code Regime

Transformer architectures have emerged as promising deep learning (DL) tools for modeling complex sequence-to-sequence interactions in channel decoding. However, existing transformer-based decoders for error correction codes (ECCs) demonstrate inferior decoding performance and generalization capabilities compared to conventional algebraic decoders, especially in short-code regimes. In this work, we propose a novel latent-attention based transformer (LAT) decoder for polar codes that addresses these limitations through three pivotal innovations. First, we develop a latent-attention mechanism that supersedes the conventional self-attention mechanism. This architectural modification enables independent learning of the Query and Key matrices for code-aware attention computation, decoupling them from the Value matrix. As a result, the LAT decoder emphasizes position-wise decoding interactions while mitigating interference from context correlations that are misaligned with the inherent code structure. Second, we devise an advanced training framework incorporating three synergistic components: entropy-aware importance sampling that emphasizes low-probability regions in the signal constellation space, experience reflow that introduces empirical labels to improve characterization of decoding boundaries, and dynamic label smoothing for likelihood-based regularization. Third, we propose a code-aware masking scheme which allows dynamic adaptation to varying code configurations without modifying the network architecture. Numerical evaluations demonstrate that the proposed LAT decoder achieves near maximum-likelihood (ML) performance in terms of both bit error rate (BER) and block error rate (BLER) for short-length polar codes. Moreover, the architecture exhibits robust generalization capabilities across a wide range of code rates and code lengths.