基于稀疏率缩减的“白盒”纠错码Transformer:一种可解释的神经网络解码器
作者与机构: 本文由郑子彦(清华大学深圳国际研究生院)、刘振华(Chin Wa (Ken) Lau,华为技术有限公司、香港中文大学)、郭念(华为技术有限公司)、石翔(华为技术有限公司、清华大学)和黄绍伦(Shao-lun Huang,清华大学深圳国际研究生院)共同完成。本文发表于第二届简约与学习会议(CPAL 2025)。通讯作者为黄绍伦。
学术背景: 本研究属于通信与信息理论领域,具体聚焦于信道编码中的纠错码(Error Correcting Codes, ECCs)解码技术。在现代通信系统中,为确保数据在噪声信道中可靠传输,发送端会使用纠错码为信息添加冗余,接收端则需通过解码器从受噪声干扰的信号中恢复原始信息。最大似然(Maximum Likelihood, ML)解码虽是最优的,但其计算复杂度随码长呈指数增长,在实际中难以实现。传统的主流解码算法,如基于置信传播(Belief Propagation, BP)的消息传递算法,虽然高效,但其性能受限于Tanner图中的短环问题,且存在收敛不确定性。
近年来,基于深度学习的方法为解码带来了新思路。其中,基于Transformer架构的神经网络解码器(如ECCT, CrossMPT)展现出了优异的性能,甚至超越了传统BP算法。然而,这些“黑盒”模型内部工作机制不透明,其强大的解码能力缺乏理论解释,这阻碍了研究人员对其性能的进一步理解和改进。与此同时,Yu等人提出的Coding-Rate Transformer(Crate)为理解通用Transformer提供了一个“白盒”视角,将其解释为优化一个“稀疏率缩减”(Sparse Rate Reduction)目标。本研究旨在将这一可解释的“白盒”理论框架引入到纠错码解码这一特定任务中,构建一个既具有强大性能又具备理论可解释性的新型解码器。
研究目的: 本研究旨在提出并验证一个名为“白盒纠错码Transformer”(White-box Error Correction Code Transformer, WeCCT)的新型解码器框架。其核心目标是:1)从稀疏率缩减的理论视角重新形式化解码问题;2)设计一个新颖的、融合了码字结构的多头Tanner子空间自注意力(Multi-head Tanner-Subspaces Self Attention, MTSA)机制;3)构建一个参数高效且理论驱动的解码架构,在保持与先进解码器(如ECCT, CrossMPT)竞争性性能的同时,提供对其内部运作机制的理论洞察。
详细工作流程: WeCCT的整体架构是一个深度神经网络,其设计完全由稀疏率缩减的优化目标推导而来。工作流程主要分为以下几个阶段:
1. 问题形式化与理论推导: 研究首先将AWGN信道下的最大似然解码问题,通过Tweedie公式和基于能量的模型视角,近似转化为一个在比特和校验子(Syndrome)联合表示空间上的优化问题。具体目标是最大化一个“稀疏率缩减”目标函数:max [R(Z) - R_c(Z | U_[k]) - λ||Z||_1]。其中,R(Z)度量整个令牌集的整体信息率(编码率),R_c(Z | U_[k])度量令牌在一组低维子空间U_[k]下的编码率,λ||Z||_1项促进表示的稀疏性。这个目标可以分解为两个子问题:通过注意力机制压缩特征(最小化R_c)和通过前馈网络促进稀疏性(最小化λ||Z||_1 - R(Z))。这为后续的神经网络模块设计提供了严格的理论依据。
2. 输入嵌入(Input Embedding): 与先前工作不同,WeCCT对输入进行了专门设计。对于长度为n的接收信号y,模型生成n个比特令牌和n-k个校验子令牌(k为信息位长度)。 * 比特令牌: 对接收信号y进行对数似然比(LLR)缩放,即y_llr = 2y/σ^2,其中σ^2是噪声方差。然后,每个缩放后的值乘以一个可学习的、与比特位置相关的嵌入向量w_emb,i,得到初始比特表示z_i^0。 * 校验子令牌: 计算“软校验子”。对于每个校验方程,使用其涉及的所有比特中幅度最小的那个作为该方程的可靠性权重,乘以传统的硬判决校验子向量。这个加权后的软校验子值再乘以另一个可学习的、与校验子位置相关的嵌入向量w_emb,n+j,得到初始校验子表示z_{n+j}^0。
3. 解码器层(Decoder Layers)的核心操作: 初始化的比特和校验子表示将通过L个相同的解码器层进行迭代优化。每一层包含两个核心模块,分别对应理论推导中的两个优化子问题,并采用交替更新的策略: * A. 多头Tanner子空间自注意力(MTSA): 此模块用于实现特征压缩(最小化R_c)。其关键创新在于将码字的Tanner图结构显式地融入注意力机制。 * Tanner子空间定义: 基于奇偶校验矩阵H,定义一个扩展的连接矩阵M(H),用于表示Tanner图中所有节点(包括比特节点和校验节点)之间的连接关系。每个节点的表示向量z_i所属的“Tanner子空间”T_i,被定义为与其在Tanner图中直接相连的所有邻居节点表示所张成的空间。这为信息流动提供了几何解释和结构约束。 * 注意力计算: MTSA机制计算注意力时,引入了基于M(H)的掩码函数φ(M(H))。该掩码将不相连的节点对之间的注意力权重设置为负无穷(经Softmax后为0),从而强制信息仅在Tanner图中相连的比特与校验子节点之间传递。这与传统BP算法中消息沿边传递的精神一致,但在高维表示空间中通过自注意力实现了更丰富、更全局的交互。理论推导表明,MTSA操作近似于对子空间编码率R_c的梯度下降步骤。 * 交替更新: 在每一层中,首先固定校验子表示,利用比特和校验子的当前表示,通过MTSA更新比特表示;然后,利用更新后的比特表示和旧的校验子表示,通过另一个MTSA更新校验子表示。比特和校验子使用各自独立的、可学习的子空间基U_b,[k]和U_s,[k],这允许两个域学习适应其不同角色的特征。 * B. 迭代收缩阈值算法(ISTA)前馈网络: 此模块用于实现结构化稀疏化(最小化λ||Z||_1 - R(Z))。在MTSA更新得到中间表示Z^{l+1/2}后,ISTA模块通过解决一个LASSO类型的稀疏编码问题来进一步精炼表示:Z^{l+1} ≈ argmin { λ||Z||_1 + 1/2 ||Z^{l+1/2} - D^l Z||_F^2 }。其中D^l是一个可学习的(完备的)非相干或正交字典。这个优化问题通过ISTA迭代算法求解,其更新步骤可以表示为带有ReLU(实验中用GELU)激活函数的线性变换,从而可以无缝集成到神经网络的前馈层中。这个过程促进了表示的稀疏性,有助于提取最本质的特征并抑制噪声。
每个MTSA和ISTA操作前后都包含层归一化(LayerNorm)以稳定训练。整个流程在算法1中进行了详细描述。
4. 输出预测: 经过L层处理后的最终比特表示{z_i^L},不再像ECCT或CrossMPT那样与校验子表示拼接后通过全连接层预测乘性噪声,而是直接通过一组独立的、与比特位置相关的线性投影w_out,i和偏置θ_i,映射为每个比特的判决概率p_i。最终通过阈值比较(如0.5)得到硬判决的译码码字x̂。
5. 实验设置: * 训练目标: 使用二元交叉熵损失函数,直接最小化预测比特与真实传输码字之间的误差。 * 模型配置: 评估了两种配置:WeCCT(L=6)和WeCCT(L=12)。嵌入维度d=128,注意力头数h=8。使用Adam优化器,学习率采用余弦衰减。 * 对比基线: 包括传统BP解码器(50次迭代)、自回归BP(AR BP)以及先进的Transformer解码器ECCT和CrossMPT。 * 测试数据集: 在多种码族上进行了广泛测试,包括BCH码、Polar码和LDPC码(含MacKay和CCSDS标准中的码)。测试时,在每个信噪比(SNR)下收集至少500个帧错误以计算误比特率(BER),重点关注4dB, 5dB, 6dB等实用SNR范围。性能度量指标为-ln(BER),值越大表示性能越好。
主要结果: 1. 性能对比: 如表1和图2所示,WeCCT在绝大多数测试的码字上取得了极具竞争力的性能。尽管参数量大幅减少(见下文),但WeCCT(L=6)的性能普遍优于ECCT。当将层数增加到12层(WeCCT-L12)以匹配先前模型的计算预算时,WeCCT的性能达到或超过了当前最先进的CrossMPT解码器,同时在多数情况下仍保持更低的参数量。这证明了其架构和理论框架的有效性。 2. 参数与计算效率: WeCCT展现了显著的参数效率。如表2所示,对于BCH(63,45)码,WeCCT-6的参数量比ECCT和CrossMPT减少了约64%。即使层数翻倍的WeCCT-12,其参数量仍比CrossMPT少31%。计算复杂度(FLOPs)也显著降低,WeCCT-6对于BCH(63,45)码的FLOPs比ECCT少40%,比CrossMPT少29%;对于更长的LDPC(121,70)码,减少比例更高。效率提升主要源于:a) 注意力模块中键(Key)和值(Value)的投影矩阵共享;b) 用高效的ISTA操作替代了标准Transformer中参数量庞大的前馈网络。 3. 消融实验验证: 如表3所示,对LDPC(121,60)码的消融研究表明,移除MTSA中的Tanner子空间机制(即取消基于Tanner图的注意力掩码)会导致性能急剧下降。这有力地证明了将码字结构显式编码到注意力机制中对实现优异解码性能至关重要。 4. 理论目标验证: 如图3和表4所示,实验监测了训练过程中编码率和稀疏度的变化。结果表明,随着网络层数的加深,子空间编码率R_c显著下降,同时表示的稀疏度(非零元素比例)增加。这直接验证了WeCCT确实在优化其理论推导的核心目标——稀疏率缩减。更重要的是,表4显示,在训练过程中,更深层的相对率缩减程度与最终的解码性能(-ln(BER))同步提升,建立了理论优化目标与实际解码效能之间的强相关性。
结论: 本研究成功提出了首个用于纠错码解码的“白盒”Transformer架构——WeCCT。通过从稀疏率缩减的角度形式化解码问题,并设计了新颖的多头Tanner子空间自注意力机制,WeCCT为Transformer在解码任务中的成功提供了一个清晰的理论解释。实验证明,该框架不仅实现了参数和计算上的高效率,而且在多种码族上取得了与最先进方法相媲美甚至更优的性能。其“白盒”特性使得模型的内部运作(如特征压缩和稀疏化过程)变得可分析、可解释,弥合了深度学习模型与经典编码理论之间的隔阂。
研究亮点与价值: 1. 理论创新性: 首次将“白盒”Transformer理论(稀疏率缩减)系统地应用于纠错码解码领域,为神经网络解码器提供了坚实的数学基础。 2. 结构创新: 提出了“多头Tanner子空间自注意力”机制,创造性地将码字的图结构(Tanner图)以子空间约束的形式嵌入到注意力计算中,实现了结构感知的消息传递。 3. 模型高效性: 在保持高性能的同时,大幅降低了模型的参数量和计算复杂度,这对于解码器在资源受限的实际通信设备(如终端、物联网设备)中的部署具有重要意义。 4. 可解释性: 整个模型的设计,从注意力到前馈网络,都有明确的优化目标(压缩R_c、促进稀疏性)对应。训练过程中可观测的率缩减和稀疏化趋势,为模型为何有效提供了直观证据。 5. 应用价值: 为开发高效、可靠且可解释的下一代信道解码器提供了新范式。其框架易于扩展至其他信道模型、更长的码字(通过分层Tanner子空间)或探索更针对性的稀疏化策略。
其他有价值内容: 文章附录提供了关键理论推导(近似1和近似2)的完整证明,详细阐述了从最大似然解码到稀疏率缩减目标的转化过程,以及MTSA机制作为梯度下降近似的推导。此外,附录还包含了更详细的复杂度分析、消融实验以及编码率与稀疏度随训练演化的可视化分析,进一步支撑了文章的主论断。