GNN-based Auto-Encoder for Short Linear Block Codes: A DRL Approach

TL;DR

This work addresses the challenge of designing effective short-length channel codes for URLLC by proposing a GNN-DRL auto-encoder that jointly optimizes encoder and decoder. A DRL-based code designer generates parity-check matrices via a lattice-graph MDP, while an EW-GNN decoder provides BP-like decoding with edge-weighted message passing, enabling scalable, high-performance decoding. Through iterative training, the auto-encoder produces codes that outperform traditional LDPC/BCH schemes at short lengths and achieves substantial coding gains under MLD, with the EW-GNN decoder offering robust performance across different code lengths. The approach offers offline, adaptable code design and a decoder that generalizes to longer codes without retraining, yielding practical impact for URLLC systems.

Abstract

This paper presents a novel auto-encoder based end-to-end channel encoding and decoding. It integrates deep reinforcement learning (DRL) and graph neural networks (GNN) in code design by modeling the generation of code parity-check matrices as a Markov Decision Process (MDP), to optimize key coding performance metrics such as error-rates and code algebraic properties. An edge-weighted GNN (EW-GNN) decoder is proposed, which operates on the Tanner graph with an iterative message-passing structure. Once trained on a single linear block code, the EW-GNN decoder can be directly used to decode other linear block codes of different code lengths and code rates. An iterative joint training of the DRL-based code designer and the EW-GNN decoder is performed to optimize the end-end encoding and decoding process. Simulation results show the proposed auto-encoder significantly surpasses several traditional coding schemes at short block lengths, including low-density parity-check (LDPC) codes with the belief propagation (BP) decoding and the maximum-likelihood decoding (MLD), and BCH with BP decoding, offering superior error-correction capabilities while maintaining low decoding complexity.

Figures (13)

• Figure 1: A parity-check matrix in systematic form and its corresponding Tanner graph. The path marked by blue lines represents a 4-length cycle.
• Figure 2: System model.
• Figure 3: An example of the full graph-based MDP generating a parity-check matrix. The flipping threshold $\alpha_f$ is set as $0.5$.
• Figure 4: The block diagram of the DRL-based code designer.
• Figure 5: Illustration of updating the node embedding in the GNN-based decoder. The variable node $v_i$ receives messages from $\mathcal{M}(v_i)=\{u_a, u_b, u_c, u_d\}$. $g(\cdot\mid\bm{\theta}_g)$ is a trainable neural network.