# Adaptive Learned Belief Propagation for Decoding Error-Correcting Codes
## 5 Adaptive Learned Message Passing Algorithms
### Two-Stage Decoder  
In a parallel decoder, the weights are restricted in a discrete set. The number of parallel decoders depends exponentially on the size of this set. The two-stage decoder predicts arbitrary non-negative weights, without the exponential complexity of the parallel decoders. Further, since the weights are arbitrary, the two-stage decoder can improve upon the performance of the parallel decoders, when the output LLRs are sensitive to the weights. _Architecture_: Recall that we wish to find a function \(g(\mathbf{y})\) that minimizes the BER in (13). In a two-stage decoder, this function is expressed by an NN \(\boldsymbol{\hat{\gamma}}=g_{\theta}(\mathbf{y})\) parameterized by vector \(\boldsymbol{\theta}\). Thus, the two-stage decoder is a combination of an NN and a WBP. First, the NN takes as input either the LLRs at the channel output \(\mathrm{L}(\mathbf{y})\) or \((\mathrm{L}(\mathbf{y}),\mathbf{y})\) and outputsthe vector of weights \(\mathbf{\hat{\gamma}}\). Then, the WBP decoder takes the channel LLRs \(\mathrm{L}(\mathbf{y})\) and weights \(\mathbf{\hat{\gamma}}\) and outputs the decoded word \(\mathbf{\hat{y}}\). The parameters \(\mathbf{\theta}\) are found using a dataset of examples \(\big{\{}\left(\mathbf{y}^{(i)},\mathbf{\gamma}^{(i)}\right)\big{\}}_{i}\), where \(\mathbf{\gamma}^{(i)}\) is the target weight. This dataset can be obtained through a simulation, i.e., transmitting a codeword \(\mathbf{c}^{(i)}\), receiving \(\mathbf{y}^{(i)}\), and using, e.g., an offline parallel decoder to determine the target weight \(\mathbf{\gamma}^{(i)}\). In this manner, \(g\mathbf{(y)}\) is expressed in a functional form instead of being determined by real-time search, which may be more expensive. In this paper, the NN is a CNN consisting of a cascade of two one-dimensional convolutional layers \(\mathbf{Conv}_{1}\) and \(\mathbf{Conv}_{2}\), followed by a dense layer \(\mathbf{Dens}\). \(\mathbf{Conv}_{i}\) applies \(F_{i}\) filters of size \(S_{i}\) and stride 1, and the rectified linear unit (ReLU) activation, \(i=1,2\). The output of \(\mathbf{Conv}_{2}\) is flattened and passed to \(\mathbf{Dense}\), which produces the vector of weights \(\mathbf{\hat{\gamma}}\) of length \(T\). This final layer is a linear transformation with ReLU activation to produce non-negative weights. The model is trained by minimizing the quantile loss function  
\[\ell_{\xi}(\mathbf{\gamma},\mathbf{\hat{\gamma}})=\text{mean}\Big{(}\text{max}(\xi( \mathbf{\gamma}-\mathbf{\hat{\gamma}}),(\xi-1)(\mathbf{\gamma}-\mathbf{\hat{\gamma}}))\Big{)},\]  
where \(\xi\in(0,1)\) is the quantile parameter and max is applied per entry and mean over vector entries. The choice of loss is obtained by cross-validating the validation error over a number of candidate functions. This is an asymmetric absolute-like loss, which, if \(\xi\geq 1/2\) as in Section 6, encourages entries of \(\mathbf{\hat{\gamma}}\) to be close to entries of \(\mathbf{\gamma}\) from above rather than below. Complexity of the two-stage decoderThe computational complexity of the two-stage decoder in the inference mode is the sum of the complexity of the CNN and WMS decoder  
\[\text{RM}=T(m+n)+\text{RM}(\text{CNN}),\]  
where the computational complexity of the CNN is  
\[\text{RM}(\text{CNN}) =\text{RM}(\mathbf{Conv}_{1})+\text{RM}(\mathbf{Conv}_{2})+\text{ RM}(\mathbf{Dense})\] \[=F_{1}(n-S_{1}+1)S_{1}+F_{1}F_{2}(n-S_{1}-S_{2}+2)S_{2}+F_{2}(n-S _{1}-S_{2}+2). \tag{17}\]  
The complexity can be significantly reduced by pruning the weights, for example, by setting to zero the weights below a threshold \(\tau_{\text{prun}}\). **Remark 4**: _Neural decoders are sensitive to distribution shifts and often require retraining when the input distribution or channel conditions change. To lower the training complexity, [18] proposed a decoder that learns a mapping from the input SNR to the weights in WBP, enabling the decoder to operate across a range of SNRs.