基于互信息量估计的几何与概率联合整形技术

梁家熙，牛泽坤，胡卫生，义理林

(上海交通大学 区域光纤通信网与新型光通信系统国家重点实验室，上海 200240)

0 引 言

在加性高斯白噪声(Additive White Gaussian Noise，AWGN)信道中，使用高阶的正交振幅调制(Quadrature Amplitude Modulation，QAM)在高信噪比(Signal Noise Ratio，SNR)的情况下，互信息量(Mutual Information，MI)与香农极限相差1.53 dB[1]，QAM的方形星座图是造成容量差距的原因。星座整形是一种在QAM基础上提高MI的方式，分为几何整形和概率整形，其中几何整形指调整星座点的几何分布，概率整形指调整星座点在发送信号中出现的概率。几何整形和概率整形通过一些传统的方式实现，例如几何整形可以通过一定的规则调整星座图点与点的位置，充分利用空间，但是无法证明这种方法可以得到使系统的信道容量最大的编码方式[2];概率整形可以通过麦克斯韦-玻尔兹曼(Maxwell- Boltzmann，MB)分布调整星座点的概率分布，在AWGN信道下，证明MB分布可以达到性能最佳，但是对于其他信道而言，MB分布无法取得最佳性能[3]。

近年来，深度学习被引入物理层通信中，通信系统可以被视为一个自动编码器，此时发端与收端都可以由神经网络的全连接层担任，这就是端到端的深度学习(End-to-End Deep Learning，E2EDL)通信系统[4]。E2EDL相较于传统的通信算法优势在于从一个全局的角度出发考虑优化，而不是将局部最优组合到一起，得到一种非全局最优的方案[5]。通过E2EDL，可以实现星座整形，训练得到几何整形[4]和概率整形[6]的编码器，性能接近或超过传统算法的最佳性能。通过E2EDL在几何整形星座图的基础上做概率整形，可以取得比单独几何整形和概率整形更高的增益[6]。但是，目前的方法都是在固定的几何整形上做概率整形，没有进行联合优化，且这些方法无法证明当前整形结果的信息量是最优的。

本文提出了一种基于互信息量估计(Mutual Information Neural Estimation，MINE)的几何与概率联合整形方法。MINE可以估计两种不同分布的MI[7]，本文提出的系统利用MINE估计当前通信系统的MI，并且训练发端编码器以满足最大化MI。基于MINE的系统无需训练解码器，可以避免解码器误差对编码器训练带来的损失。本文提出了几何整形与概率整形的联合优化方式，通过训练迭代优化几何整形与概率整形的编码器，使系统接近到几何整形与概率整形所能达到的MI的极限。本文将系统在不同SNR的AWGN信道中进行仿真，比较了几何与概率联合整形的系统与其他方案之间的MI。仿真结果表明，几何与概率联合整形系统的MI高于单独进行几何或概率整形的系统，在SNR为10 dB的AWGN信道中，几何与概率联合整形的系统MI与单独做概率整形的系统相比，有0.027 9 bit/symbol的增益；与单独做几何整形的系统相比，有0.041 7 bit/symbol的增益。

1 MINE原理

根据香农信息论的内容，信道输入和输出之间的MI的极大值就是信道容量，可以通过最大化MI的方式找到可以逼近信道容量的编码方式。但从历史上看，MI往往是难以计算的[8]，精确的计算仅仅适用于离散变量(可以直接求出总和)，或是用于已知概率分布的有限问题，对于更加广泛的应用场景，MI的精确计算通常是不可能的。

2018年，来自蒙特利尔大学的Belghazi等人提出了一个基于 Kullback-Leibler(KL)散度的对偶表示的通用参数神经网络互信息估计器，MINE[7]。MINE是一种数据驱动估计MI下界的方式，可以估计任意两种不同分布的MI的下界，并且可以将估计值作为损失函数，将梯度沿着神经网络估计器进行反向传播。

MINE的原理如下：两个变量X和Z的MI可以通过下式表示：

式中：X和Z为两种不同的分布；I(X;Z)为两种分布间的MI；PXZ为X和Z的联合分布；PX和PZ分别为X和Z的边际分布；⊗为两种分布间的张量积。另外引入KL散度的概念，KL散度又叫相对熵，是两个不同分布之间的差异的非对称性度量，在信息论中，KL散度是两个不同分布的信息熵的差值，表示如下：

式中：P和Q为两种不同的分布；DKL(P||Q)为P和Q的KL散度；EP为变量在P分布下的数学期望值。P和Q都是严格连续的分布。根据式(1)和(2)可以推导出，MI和KL散度的关系为

接下来，MINE的重点在于KL散度的对偶表示，这份工作由Donsker和Varadhan于1983年完成，因此KL散度的对偶表示又称Donsker-Varadhan表示[9]，具体如下

式中：e为自然对数底；T为所有函数集合R中的一个函数，T:Ω→R表示在实数域函数中存在函数T；sup为函数的上界。式(4)的意义在于，在一个由所有实函数组成的域中，一定能找到一个函数对应关系T，按照函数对应关系将P和Q的分布转化为标量，而这个标量使得等号右侧函数的值最大，而这个最大值就是P和Q的KL散度。函数T可以由神经网络函数来表示，假设Θ为所有神经网络所有参数的集合，因此TΘ可以看作所有神经网络函数的集合，而TΘ是R的一个子集，因此神经网络估计的MI与真实的MI间有如下关系：

式中，IΘ(X,Z)为通过神经网络计算的MI的估计值，对于参数为θ∈Θ的神经网络，根据式(3)、(4)和(5)可以推导得到：

通过式(6)，在两个随机分布里随机选取联合样本和边界样本，再通过训练找到使得样本的MI估计值最大的神经网络参数，这个MINE值就是两个分布MI的下界。

因此，计算两个不同分布的MI，可以转变为求一个神经网络在特定损失函数下的最大值。定义一个神经网络，用于将两种不同分布的变量转化为一个标量，即Tθ(X,Z)，通过将Tθ代入式(6)可以计算得到在当前神经网络参数θ下的MI估计值，然后将式(6)等号右边的相反数设置为损失函数，试图通过最小化损失函数，即最大化MI的估计值，即可取得MI的下界。

在本文中，系统使用MINE作为估计系统MI的方法，取得信道两段的符号，通过训练的方式找到使得估计值最大的参数θ，再通过式(6)计算得到通信系统MI的估计值。

2 基于MINE的几何与概率联合整形系统设计

2.1 几何整形系统

几何整形是通过设计星座点的位置，使得星座点之间的欧氏距离增大，从而降低解码端分辨星座点的难度，降低误码率。从信息论角度考虑，几何整形的目的是寻找一种使得信道容量最大的调制方法，获得更高的通信速率。在基于MINE的几何整形系统中，训练MINE神经网络计算信道两端符号的MI，然后通过最大化MI的方式训练编码器，从而得到使得信道容量最大的调制方式。

基于MINE的几何整形系统如图1所示，其中编码器和MINE是两个全连接神经网络，分别表示调制使用的神经网络和估计MI使用的神经网络。编码器输入原始比特，输出调制后的符号；MINE输入信道两端的符号，输出用于估计MI的标量。m为原始比特，输入编码器之后进行功率归一化，得到输入信道的符号x，x经过信道之后得到有损伤的符号y，将x和y输入MINE神经网络，由式(6)计算MI。本文的神经网络，编码器拥有3个隐藏层，每个隐藏层有256个神经元，使用Leaky ReLU激活函数；MINE拥有3个隐藏层，每个隐藏层有128个神经元，使用Leaky ReLU激活函数。

图1 基于MINE的几何整形系统设计

训练过程如图2所示，共训练L个迭代周期，首先构建编码器和MINE神经网络。每个训练周期中，先生成原始比特，然后根据图1进行数据正向传播并且计算MI。训练编码器和MINE的目的都是通过式(6)最大化MINE，因此损失函数设置为MINE的相反数。每个迭代周期训练1次MINE神经网络，每10个迭代周期训练1次编码器神经网络。

图2 基于MINE的几何整形系统训练流程

2.2 概率整形系统

概率整形是在星座图的基础上，通过改变星座符号的出现概率，使得功率低的星座点出现概率大，功率高的星座点出现概率小，从而降低发送信号的整体功率，使得噪声的功率也随之下降，起到提高信道容量的作用。通常概率整形在标准方形QAM下进行，不同星座点的概率分布符合MB分布。

图3 基于MINE的概率整形系统设计

基于MINE的概率整形系统如图3所示，其中概率生成器和MINE为两个神经网络，MINE与2.1节相同，概率生成器输入固定的数，输出每个星座点的出现权重，本系统中固定输入设置为1。将权重经过Gumbel-Softmax激活层[10]得到含有随机性的不同星座点的输出概率，Gumbel-Softmax表达式为

式中：α为星座点的权重向量；G为Gumbel随机数向量；h为通过softmax(柔性最大值)计算得到的不同的星座点出现概率的向量；temperature为温度常数，当温度常数接近0时，h接近独热编码分布；当温度常数接近正无穷时，h接近均匀分布。

每个星座点的概率与星座点所对应的符号相乘，再将所有星座点的结果相加，即可得到概率整形编码器的输出符号，然后对符号进行功率归一化，得到输入信道的符号x，x经过信道得到有损伤的符号y，将x和y输入MINE神经网络，经过式(6)计算MI。概率生成器有两个隐藏层，每个隐藏层有64个神经元，使用Leaky ReLU激活函数。训练过程如图4所示，共训练M个迭代周期，首先构建概率生成器和MINE神经网络。每个训练周期中，首先通过概率生成器生成概率权重，经过Gumbel-Softmax函数激活后与星座点集相乘得到符号，再根据图3进行数据正向传播并计算MI。训练概率生成器的损失函数设置为MINE的相反数。每个迭代周期训练1次MINE神经网络，每20个迭代周期训练1次概率生成器神经网络。

图4 基于MINE的概率整形系统训练流程

2.3 几何与概率联合整形系统

基于MINE的几何与概率联合整形系统的具体训练流程如图5所示。首先进行几何整形和概率整形的预训练，训练的方法分别参照2.1与2.2节，其中训练概率整形所使用的星座图为编码器产生的星座图。训练结束后计算该状态下的MI，记为maxMI。

接下来进行迭代的几何整形和概率整形训练，即分别训练系统的编码器和概率生成器。首先进行几何整形，先将现有的编码器模型保存，然后根据不同星座点的概率随机生成不同数目的几何整形星座点符号，将这些符号通过信道进行最大化MI的训练，一次迭代训练编码器若干次，在训练过程中不断计算新的几何整形星座图的MI，如果新的MI比maxMI高，则停止训练，将新的MI记为nowMI，然后保存新的编码器模型。如果训练超过一定次数之后MI一直小于maxMI，则回退到保存的编码器模型，重新进行训练。

概率整形的训练部分，先保存概率发生器模型，训练概率发生器并且计算MI，超过maxMI则保存MI和概率发生器模型并且停止训练，训练超过一定次数则回退概率发生器模型。就这样训练，迭代一定次数，保证每次迭代的几何整形和概率整形的模型的MI都大于上一次的MI。

将训练周期的总数设置为K，每个训练周期中，编码器和概率发生器最多训练L次迭代，可以保证每次编码器和概率发生器训练后，所得新整形系统的MI高于上一个周期的MI。在训练过程中，系统的MI逐步增加，逼近几何与概率整形联合优化的性能极限。

图5 基于MINE的几何与概率整形联合系统设计

几何整形预训练设置10 000个迭代周期，概率整形预训练设置100 000个迭代周期。，暂停训练并回退模型的阈值设置为2 000，即训练2 000个迭代周期内，如果编码器或概率发生器的性能没有使得系统的信道容量有所提升，就回退并重新进行训练。在2 000个迭代周期中，前500迭代周期用于预训练MINE，剩下的1 500个迭代周期中，每100个周期采集一次MI，如果发现新的MI超过了过去的maxMI，则将maxMI更新为nowMI，并且立刻停止训练。

3 系统仿真与结果展示

3.1 系统仿真设置

为了验证基于MINE的几何与概率联合整形系统的性能，在AWGN信道下进行了仿真验证，在不同SNR下对QAM、几何整形、概率整形以及几何与概率联合整形的MI进行对比，仿真系统中每个符号对应4位比特。几何整形系统按照2.1节设置，L设置为20 000。概率整形按照2.2节设置，M设置为100 000。几何与概率联合整形按照2.3节设置，K设置为60。计算MI的方法采用蒙特卡洛估计的方式[11]，估计采样的样本量足够多的情况下，这种方式得到的MI是接近精确的。

3.2 仿真结果

基于MINE的几何与概率联合整形在不同SNR的AWGN下的整形方案热力图如图6所示。图中横坐标和纵坐标分别为信号在I路和Q路的归一化能量大小，右侧坐标轴表示不同星座点的出现概率权重，权重值经过softmax函数计算可以得到星座输出点的概率。

在热力图中，星座点的亮度越高，代表该星座点的权重越高，该星座点发送的可能性越大。在SNR较低，噪声较大的情况下，中心星座点的概率较高，周围的概率较低；随着SNR的升高，周围星座点的权重逐渐升高，逐渐接近中心星座点的概率。

图6 基于MINE的几何与概率联合整形在不同SNR下的AWGN信道中整形方案的热力图

图7所示为QAM、几何整形、概率整形以及几何与概率联合整形在不同SNR下的AWGN信道中的MI变化曲线。由图可知，在SNR较低的情况下，所有调制格式和香农极限的MI较接近，由于曲线比较接近，因此选取了SNR=10 dB处的曲线进行放大，由图可知，几何与概率联合整形的系统性能最优，其次是概率整形，再次是几何整形，16QAM调制相较于星座整形而言性能是最差的。由图可知，几何与概率联合整形的MI相较于概率整形有0.027 9 bit/symbol的增益，相较于几何整形有0.041 7 bit/symbol的增益，相较于16QAM调制有0.112 6 bit/symbol的增益。与香农极限相比，几何与概率联合整形在SNR=10 dB处仍有0.183 4 bit/symbol的差距，但是与16QAM与香农极限间的0.296 0 bit/symbol差距相比，有了较大的性能提升，如果需要继续提升性能，需要考虑在更高阶的调制格式上进行几何与概率整形。

因此，在AWGN信道中，基于MINE的几何与概率联合整形系统可以训练得到一种性能优于几何整形与概率整形中任意一种整形方案的编码方案。在几何整形与概率整形的基础上，通过这种方法，可以让信道两端的MI更加接近香农极限。

图7 4种不同编码方式在不同SNR下的AWGN信道中MI的变化趋势

4 结束语

为解决QAM在高SNR的情况下，MI与香农极限相差1.53 dB的问题，本文提出了一种基于MI估计的几何与概率联合整形的方法，在QAM的基础上通过星座整形提升系统的MI。本文分别介绍了基于MI的几何整形、概率整形以及将两种整形方法联合训练的步骤，并且将这套方法在不同SNR下的AWGN信道中进行了仿真验证。仿真结果表明，几何与概率联合整形方法的性能优于单独进行几何整形和概率整形中性能。在SNR为10 dB时，系统的MI与单独做几何整形相比有0.041 7 bit/symbol的增益，与单独做概率整形相比有0.027 9 bit/symbol的增益。