基于低复杂度Mamba架构的短码长LDPC码交叉选择性扫描消息传递解码器

《Electronics》：CrossMPM: Cross-Selective Scan Message-Passing Mamba for Short-Length LDPC Codes

　　

纠错码是实现噪声信道上可靠通信的基础技术。基于Transformer的神经解码器取得了优异的纠错码解码性能，但其注意力机制导致计算复杂度随码长显著增长。为解决这一问题，本研究提出了一种低复杂度的基于Mamba的神经解码器——交叉选择性扫描消息传递Mamba网络。CrossMPM用坦纳图投影和交叉选择性扫描状态空间模型替代了交叉注意力消息传递Transformer中基于交叉注意力的幅度向量与伴随式向量之间的信息交换。坦纳图投影根据奇偶校验矩阵聚合相连节点的表示，从而融入码结构信息；而交叉选择性扫描SSM则从聚合后的表示中动态生成选择参数，以控制状态更新和输出。因此，CrossMPM能够在无需注意力操作的情况下建模局部奇偶约束和长程依赖关系，并降低了计算复杂度。仿真结果表明，与CrossMPT和混合Mamba-Transformer解码器相比，CrossMPM在实现具有竞争力的解码性能的同时拥有更低的计算复杂度，且在更长码长下仍能保持这一性能。这些结果表明，CrossMPM在纠错性能与计算效率之间实现了良好的权衡。

研究背景与问题

纠错码是确保噪声信道中数据传输可靠性的基础技术。传统的迭代解码算法如置信传播和最小和算法已被广泛使用，而近年来基于深度学习的神经解码器也受到关注。其中，基于Transformer的神经解码器虽然解码性能优异，但其注意力操作的计算复杂度随码长增加而急剧上升，限制了其在较长码字上的可扩展性。例如，交叉注意力消息传递Transformer将幅度和伴随式作为分离表示并通过掩码交叉注意力迭代更新，但计算开销高昂。混合Mamba-Transformer解码器虽引入Mamba架构，但未保持两者分离表示，反而将其合并为单一序列，导致序列长度增加，未能有效降低整体计算成本。因此，亟需一种既能有效利用幅度与伴随式互补信息、又能避免高计算复杂度的Mamba基解码器架构。

研究内容与结论

为此，研究人员提出了交叉选择性扫描消息传递Mamba网络，一种基于Mamba的低复杂度神经解码器。CrossMPM保持幅度和伴随式为分离表示，并用坦纳图投影和交叉选择性扫描状态空间模型替代CrossMPT中的交叉注意力进行信息交换。实验表明，CrossMPM在多种短码长LDPC码上实现了优于CrossMPT的解码性能，并在大多数条件下优于ECCM。对于更长码长的无线区域网LDPC码和LDPC码，CrossMPM同样保持了性能优势。在计算效率方面，对于LDPC码，CrossMPM的浮点运算次数和峰值训练GPU内存消耗分别比ECCM降低了65.91%和15.22%；对于更长的LDPC码，相比CrossMPT分别降低了48.90%和54.94%。这些结果证明CrossMPM在纠错性能与计算效率之间提供了有利的权衡。

关键技术方法

1. 1.
   初始嵌入层：对接收信号的幅度和伴随式的每个元素，通过节点特定的可训练嵌入向量和偏置，将其投影为d维表示。
2. 2.
   坦纳图投影：根据奇偶校验矩阵定义的连接关系，对相连节点的表示取均值进行聚合，并将它们投影到目标节点的表示域中。在幅度更新时，聚合与每个变量节点相连的校验节点的伴随式表示；在伴随式更新时，聚合与每个校验节点相连的变量节点的已更新幅度表示，从而将局部奇偶约束的结构信息纳入后续更新过程。
3. 3.
   交叉Mamba块：包含SSM路径和门控路径。在SSM路径中，对当前节点表示进行归一化、线性投影和一维卷积及激活函数处理后得到SSM输入；随后应用交叉选择性扫描SSM，其选择参数由坦纳图投影得到的另一表示的聚合结果动态生成，用于控制状态更新和输出。同时，门控路径生成门控表示，与SSM输出进行逐元素相乘后再经输出投影得到Mamba输出。该块在正向和反向两个方向分别处理，最终输出为两方向输出的和。
4. 4.
   门控残差连接与前馈网络：通过门控残差连接选择性融合交叉Mamba块的输出与原有表示，再经过前馈网络和残差连接生成最终的幅度和伴随式表示。
5. 5.
   输出预测与训练目标：将所有解码器块的幅度和伴随式表示沿节点维度拼接，经归一化和两个全连接层后得到乘性噪声预测。训练时对所有解码器块施加逐层二元交叉熵损失，各层损失求和构成最终训练目标。

研究结果

• •
  解码性能评估：在多种LDPC码上，CrossMPM在所有条件下均优于CrossMPT，并在多数条件下优于ECCM。对于更长码长的WRAN LDPC码和LDPC码，CrossMPM也持续优于CrossMPT，表明性能提升不限于短码。
• •
  计算与内存效率分析：对于短码，CrossMPM的FLOPs低于CrossMPT和ECCM，例如在LDPC码上分别降低28.96%和65.91%；峰值训练GPU内存消耗低于ECCM但略高于CrossMPT。对于更长码长，CrossMPM相比CrossMPT在FLOPs和GPU内存上均有大幅降低，例如在LDPC码上分别降低48.90%和54.94%，表明其效率优势在更长码长下更为显著。
• •
  消融研究：移除坦纳图投影或将交叉选择性扫描SSM替换为常规选择性扫描SSM会严重损害解码性能，在6dB时性能下降超过一半。移除反向扫描或逐层损失也会导致性能下降，而将门控残差连接替换为标准残差连接影响较小。这表明坦纳图投影和交叉选择性扫描SSM是CrossMPM的核心贡献组件。

总结与讨论

研究人员提出了一种名为CrossMPM的低复杂度基于Mamba的神经解码器。该模型通过坦纳图投影和交叉选择性扫描SSM在幅度与伴随式之间高效交换信息，避免了现有Transformer解码器中注意力机制的高计算负担。实验表明，CrossMPM在多种LDPC码上实现了优于CrossMPT和ECCM的解码性能，同时在计算和内存效率上具有显著优势，特别是在更长码长下。未来工作将评估CrossMPM在更长LDPC码和更宽码率范围内的性能，并探索其在其他线性分组码如BCH码和极化码以及非AWGN信道如瑞利衰落信道上的适用性。此外，还计划在多用户无源随机接入场景中研究将主动传输检测和信道估计后的信号作为CrossMPM的解码输入。