Accelerating Error Correction Code Transformers

## Abstract

Abstract Error correction codes (ECC) are crucial for ensuring reliable information transmission in communication systems. Choukroun & Wolf ( 2022b ) recently introduced the Error Correction Code Transformer (ECCT), which has demonstrated promising performance across various transmission channels and families of codes.
However, its high computational and memory demands limit its practical applications compared to traditional decoding algorithms. Achieving effective quantization of the ECCT presents significant challenges due to its inherently small architecture, since existing, very low-precision quantization techniques often lead to performance degradation in compact neural networks. In this paper, we introduce a novel acceleration method for transformer-based decoders.
We first propose a ternary weight quantization method specifically designed for the ECCT, inducing a decoder with multiplication-free linear layers.
We present an optimized self-attention mechanism to reduce computational complexity via code-aware multi-heads processing.
Finally, we provide positional encoding via the Tanner graph eigendecomposition, enabling a richer representation of the graph connectivity.
The approach not only matches or surpasses ECCT’s performance but also significantly reduces energy consumption, memory footprint, and computational complexity.
Our method brings transformer-based error correction closer to practical implementation in resource-constrained environments, achieving a 90% compression ratio and reducing arithmetic operation energy consumption by at least 224 times on modern hardware.

## 1 Introduction

Reliable digital communication systems rely heavily on ECC to ensure accurate decoding in the presence of noise. Developing efficient decoding techniques for these codes remains a complex challenge in communications research. In recent years, the application of machine learning to communications has driven the development of advanced decoding methods, leveraging deep learning architectures (Nachmani et al., ; ; Gruber et al., ; Kim et al., ; Nachmani & Wolf, ; Buchberger et al., ; Choukroun & Wolf, ; ). Notably, the work of Choukroun & Wolf () introduced a Transformer-based decoder (Vaswani et al., ) adapted to the ECC setting, demonstrating significant improvements over traditional methods across multiple code families.

Despite these advancements, the ECCT and similar neural decoders face significant challenges due to their high memory requirements, energy consumption, and computational complexity. These resource-intensive solutions pose substantial barriers to deployment in many physical communication systems, where efficiency and practicality are paramount, thus constraining the broader adoption and further refinement of these advanced decoding techniques.

Neural network (NN) quantization offers a promising approach to addressing these challenges. Recent research has shown that constraining NN weights to 1-bit and ternary representations can be effective (Ma et al., ; Wang et al., ), particularly when combined with 8-bit activations. This approach replaces multiplication operations with integer addition, significantly reducing energy consumption and memory footprint. However, applying extreme quantization techniques to smaller models presents considerable challenges. Wang et al. () demonstrated that while the performance gap between BitNet and FP16 Transformers narrows as model size increases, this gap is particularly pronounced in smaller models. For instance, a BitNet model with 100M parameters (considered small) showed a 10% higher loss than its full-precision counterpart. This disparity would be even more severe for ECCT, whose largest version contains only 2 million parameters. While Ma et al. () improved upon this method, significant performance gaps remain in smaller models, partly due to the use of absolute mean quantization, which lacks flexibility in dynamically adjusting weight sparsity during training.

Recent research on self-attention mechanisms has focused on reducing complexity and memory usage, particularly in large language models. Two main approaches have emerged: sparse attention methods (Beltagy et al., ; Zaheer et al., ; Child et al., ) and attention approximations (Choromanski et al., ).
However, these techniques were not designed to optimize smaller models, such as ECCT, which are also more sensitive to the information loss that occurs when applying sparse attention or self-attention approximations than larger ones, due to the limited number of layers. In addition to the capacity-related challenges, ECCT’s unique architecture poses additional ones. ECCT’s inherently sparse code-aware mask is incompatible with sparse attention methods, since it cannot be reduced further without modifying the information brought by the code. Similarly, attention approximation methods are incompatible because they bypass the step where attention masks are applied, making them mask incompatible.

To address these challenges, we propose a novel approach aimed at significantly reducing the memory footprint, computational complexity, and energy consumption of ECCT, thereby enhancing its viability for real-world applications. Our method introduces three key innovations: (i) Weight quantization to the ternary domain through *Adaptive* Absolute Percentile (AAP) quantization. (ii) Head Partitioning Self Attention (HPSA), an efficient multi-head self-attention mechanism tailored for bipartite graph message passing (MP), designed to reduce computational complexity and runtime. (iii) Spectral positional encoding (SPE) of the Tanner graph by processing its Laplacian eigenspace. The Tanner graph Laplacian eigenspace forms a meaningful local coordinate system, providing structural information that is lost with ECCT’s binary masking, without affecting inference runtime.

Our experimental results, conducted across a diverse range of codes, demonstrate that this approach not only matches, and in some cases exceeds, the performance of ECCT, but also offers computational complexity comparable to that of Belief Propagation (BP) Pearl (). These findings represent a significant step towards making transformer-based error correction practical for communication systems with limited computational resources, potentially bridging the gap between advanced neural decoding techniques and traditional efficient algorithms, such as BP.

## 2 Related Work

Neural decoders for ECC have evolved from model-based methods, which implement parameterized versions of classical BP (Nachmani et al., ; ; Nachmani & Wolf, ; Caciularu et al., ), to model-free approaches utilizing general NN architectures (Kim et al., ; Gruber et al., ; Bennatan et al., ; Cammerer et al., ; Choukroun & Wolf, ).
A significant advancement in this field is the ECCT (Choukroun & Wolf, ; ; ), which, along with its extension using a denoising diffusion process (Choukroun & Wolf, ), has achieved SOTA performance across various codes. These neural decoders primarily target short to moderate-length codes, addressing scenarios where classical decoders may not achieve optimal performance. Subsequently, Park et al. (; ) demonstrated improved performance, but at the expense of increased computational cost.

Transformers, while being powerful architectures, are resource-intensive. In response to the need to optimize large language models (LLMs), numerous quantization methods have been developed (Gholami et al., ; Wan et al., ; Zhu et al., ). These techniques fall into two categories: post-training quantization (PTQ) (Choukroun et al., ; Frantar et al., ; Chee et al., ) and quantization-aware training (QAT). Due to the resource-intensive nature of LLMs, recent studies have focused mainly on PTQ because of its low computational requirement and training overhead. However, PTQ often utilizes high-precision parameters, making it difficult to fully exploit the efficiency of quantization. In contrast, QAT has higher potential for accuracy but generally requires more resources and time, leaving research on QAT of LLMs in its preliminary stages (Jeon et al., ). Despite these challenges, notable work has emerged in QAT for LLMs. Wang et al. () demonstrated effective quantization of weights to {-1, 1} values and activations to 8-bit integers. An enhanced approach by Ma et al. () introduced an additional zero weight value and utilized Abs-mean quantization, highlighting a correlation between model size and performance degradation post-quantization.

Recent efforts to address the computational limitations of self-attention mechanisms in transformers have focused on acceleration techniques. One approach approximates the self-attention function, reducing its computational cost from quadratic to linear time complexity (Choromanski et al., ). Other methods, such as those proposed by Child et al. () and Beltagy et al. (), combine local and global attention to improve efficiency. Zaheer et al. () further refines these methods by incorporating random global connections. The Reformer (Kitaev et al., ) explores Locality-Sensitive Hashing attention, while Shazeer () introduces multi-query attention with shared keys and values across attention heads. Building on this, Ainslie et al. () presents grouped-query attention, which uses fewer key-value heads to achieve results comparable to multi-head attention, but with faster computation. Additionally, Pope et al. () introduces an optimized key-value cache mechanism to accelerate inference time.

Transformers have also been applied to graph-structured data, introducing graph structure as a soft inductive bias to address limitations of Graph neural networks (GNNs), such as over-squashing (Alon & Yahav, ; Topping et al., ). Dwivedi & Bresson () proposed using Laplacian eigenvectors as PEs, while Kreuzer et al. () incorporated Laplacian eigenvalues and used a dedicated Transformer for structural encoding. Building on these approaches, Rampášek et al. () further improved performance by integrating innovations such as Signet (Lim et al., ), which addresses the sign ambiguity of eigenvectors, random-walk PE (Dwivedi et al., ), and PE based on the gradients of eigenvectors (Beaini et al., ).

## 3 Setting and Background

#### Problem Settings

We assume a standard transmission protocol that uses a linear code $C\subset\{0,1\}^{n}$. The code is defined by a binary generator matrix $G\in\{0,1\}^{k\times n}$ and a binary parity check matrix $H\in\{0,1\}^{(n-k)\times n}$, satisfying $GH^{T}=0$ over $GF(2)$. The parity check matrix bipartite graph representation is referred to as the Tanner graph, which consists of $(n-k)$ check nodes and $n$ variable nodes.
The transmission process begins with a $k$-bit input message $m\in\{0,1\}^{k}$, transformed into an $n$-bit codeword $x\in C$ via $G$, satisfying $Hx=0$. This codeword is transmitted via a Binary-Input Symmetric-Output channel, resulting in a channel output $y=x_{s}+\epsilon$, where $x_{s}$ represents the Binary Phase Shift Keying modulation of $x$, and $\epsilon$ denotes random noise.
The decoding function $f:\mathbb{R}^{n}\to\mathbb{R}^{n}$ aims to provide a soft approximation $\hat{x}=f(y)$ of the original codeword.
Following Bennatan et al. (); Choukroun & Wolf (), a preprocessing step is applied to ensure codeword invariance and prevent overfitting present in model-free solutions. This yields a $(2n-k)$-dimensional vector $\tilde{y}=h(y)=[|y|,s(y)]$, where $|y|$ denotes $y$’s magnitude, and $s(y)\in\{0,1\}^{(n-k)}$ is the binary syndrome, computed as $s(y)=Hy_{b}:=H\text{bin}(y):=H(0.5(1-\text{sign}(y)))$.
The codeword soft prediction takes the form $\hat{x}=y\odot\hat{\tilde{\epsilon}}$, where $\hat{\tilde{\epsilon}}$ denotes the prediction of *multiplicative* noise $\tilde{\epsilon}$ defined such that $y=x_{s}\odot\tilde{\epsilon}=1+\epsilon\odot x_{s}$.
In our framework, the parameterized model is explicitly defined as $\tilde{x}_{s}=y\odot f_{\theta}(h(y))$, where $f_{\theta}$ represents our parameterized decoder.

#### Error Correction Code Transformer

The ECCT (Choukroun & Wolf, ) is a neural error decoder based on the Transformer encoder architecture (Vaswani et al., ). Its input $h(y)$ is defined as $h(y)=[|y|,1-2s(y)]\in\mathbb{R}^{2n-k}$, where $|y|$ represents the magnitude and $s(y)$ the syndrome. Each element is embedded into a high-dimensional space, resulting in an embedding matrix $\Phi\in\mathbb{R}^{(2n-k)\times d}$.
The embedding vectors are processed by $N$ Transformer encoder blocks using Code-Aware Self-Attention (CASA), defined by $A_{H}(Q,K,V)=\text{Sotfmax}(d^{-\frac{1}{2}}(QK^{T}+g(H)))V$.
Here, $g(H):\{0,1\}^{(n-k)\times n}\to\{-\infty,0\}^{(2n-k)\times(2n-k)}$ is a fixed binary mask that eliminates connections between bits that are more than two steps apart in the Tanner graph induced by the parity check matrix $H$. The Transformer encoder’s final block output undergoes two consecutive projections: $\hat{\tilde{\epsilon}}=W_{o}^{T}(W_{d\to 1}\Phi)$, where $W_{d\to 1}\in\mathbb{R}^{d\times 1}$ and $W_{o}\in\mathbb{R}^{(2n-k)\times n}$, yielding the noise prediction $\hat{\tilde{\epsilon}}$.

## 4 Method

Our proposed method enhances ECCT through several key modifications designed to improve both performance and efficiency. The primary enhancements are as follows:

- 1.
We replace all linear layers within the Transformer blocks with our novel *Adaptive* Absolute Percentile (AAP) Linear layers. This modification introduces an adaptive quantization approach, achieving ternary weight representation and thereby improving the model’s efficiency.
- 2.
We introduce a novel self-attention mechanism, HPSA, which supersedes the CASA used in ECCT (Choukroun & Wolf, ).
HPSA significantly reduces memory footprint, computational complexity, and runtime, thus enhancing the overall efficiency of the model.
To the best of our knowledge, our approach is the first to map the structure of the graph into patterns, with each group of heads within the multihead self-attention mechanism applying a specific pattern.
- 3.
We incorporate the SPE derived from the Tanner graph’s Laplacian eigenspace. This approach is inspired by Kreuzer et al. ()’s method of injecting a soft inductive bias of the graph’s structure into the model, enabling the integration of a fine-grained connectivity absent in ECCT’s binary mask.
- 4.
To further optimize the model’s efficiency, we replace Gaussian Error Linear Units (GeLUs) (Hendrycks & Gimpel, ) with Rectified Linear Units (ReLUs).
- 5.
We introduce a two-phased training process to enhance the model’s performance.

This change simplifies the activation function to a thresholding operator which further contributes to complexity reduction.

Figure: Figure 1: AAP Linear Layer: (a) QAT: Training with quantization noise; (b) Inference: Matrix multiplication using only integer additions with fixed ternary weights and fixed weight scale.
Refer to caption: extracted/5910077/images/aap_training_2.png

### 4.1 Adaptive Absolute Percentile Quantization

Ternary quantization of a single precision tensor involves the element-wise assignment to one of three bins: {-1, 0, +1}. This results in $3^{n}$ possible arrangements for each weight tensor, where $n$ is the tensor’s number of elements. In NNs with numerous weights, finding the optimal arrangement becomes infeasible due to this highly exponential number of options. Existing approaches, such as abs-mean quantization (Ma et al., ) often struggle to achieve the right sparsity for precise management of feature retention and elimination, making certain desirable weight distributions extremely difficult to attain during training.

To address this challenge, we propose a novel method that provides maximum flexibility to the model.
Our *Adaptive* Absolute Percentile (AAP) quantization method aims to identify the appropriate percentile of absolute values to use as a scaling factor. This percentile is optimized during training, thereby defining the desired sparsity and structure at the finest granularity.
For each weight tensor (excluding biases) during each training forward pass, we calculate the $p$-th percentile, for a predefined $p$, of the absolute values of the weights, denoting this value as $\gamma$. The value of $\gamma$ depends solely on the current weight distribution and changes with each training iteration. The scale is then computed as $\gamma\delta$, where $\delta$ is a learnable parameter initialized to one. This approach allows $\delta$ to adjust the percentile dynamically throughout training, helping the model effectively balance sparsity and information retention for each weight matrix.

In contrast to existing methods, which either rely on a weight distribution-based scale (e.g., Ma et al. ()) or use a learnable scale that may be initialized with a calibration set (e.g., Jeon et al. ()), we combine both approaches. The computed scale $\gamma\delta$ is then used to scale the entire weight matrix. Finally, each scaled weight is rounded to the nearest integer among {-1, 0, +1}.

$$ $\displaystyle\text{Ternary}(W)$ $\displaystyle=\text{RoundClip}\bigg{(}\frac{W}{\gamma\cdot\delta+\varepsilon}, -1,1\bigg{)}$ (1) $\displaystyle\text{RoundClip}(x,a,b)$ $\displaystyle=\text{max}(a,\text{min}(b,\text{round}(x)))$ $\displaystyle\gamma$ $\displaystyle=\text{Percentile}(\text{Abs}(W),0.5)$ $$

where $\text{Percentile}(x,p)$ returns the $p$-th percentile value of $x$, and $\text{Abs}(x)$ computes the element-wise absolute values of $x$.
The activations undergo Absmax quantization to INT8 as follows:

$$ $\text{Quant}(x)=\text{RoundClip}\bigg{(}x\times\frac{Q_{b}}{\alpha},-Q_{b},Q_{ b}\bigg{)}$ (2) $$

where $\alpha=\lVert x\rVert_{\infty}$ and $Q_{b}$ is the maximum value for the INT8 quantization range. Similarly to Wang et al. (); Ma et al. (), $\alpha$ is not fixed during inference.
The complete quantization scheme, incorporating both weight and input quantization, operates as follows:

$$ $\text{AAPLinear}(x,W,b)=\text{Quant}(x)\cdot\text{Ternary}(W)\cdot\frac{\gamma \delta\alpha}{Q_{b}}+b$ (3) $$

Figure: Figure 2: Head Partitioning Self-Attention: (a) First-ring and (b) second-ring head attention mechanisms. $Q$, $K$, $V$ denote query, key, and value tensors for variable (v) or check (c) nodes. $v=$ and $c=$ indicate new representations for variable and check nodes, respectively. $M_{cv}$, $M_{vc}$, $M_{cc}$, $M_{vv}$ are HPSA masks (see Fig. 3). $\sigma$ denotes the Softmax function.
Refer to caption: extracted/5910077/images/first_ring_hpsa_figure_3.png

Here, $x\in\mathbb{R}^{m\times n}$ is the FP32 layer’s input, $W\in\mathbb{R}^{n\times p}$ is the FP32 weight, and $b\in\mathbb{R}^{p}$ is the FP32 bias.
The product of the quantized weights and input is dequantized before bias addition.
All scaling factors, $\delta$, $\gamma$, and $\alpha$, are scalars, which enhances computational efficiency. Figure 1 illustrates the AAP mechanism during both the training and inference phases.
The method avoids floating-point matrix multiplication, relying primarily on integer addition and subtraction operations, significantly reducing computational complexity.

Figure: Figure 3: Code-aware masks of Hamming(4,7). AECCT masks utilize two distinct patterns, with each head applying only one: either first-ring or second-ring MP. First-ring MP uses c-to-v and v-to-c masks, while second-ring MP employs v-to-v and c-to-c masks. In contrast, the ECCT mask (on the left) applies both first and second rings for all heads. AECCT masks exhibit greater sparsity compared to ECCT, leading to reduced computational complexity.
Refer to caption: extracted/5910077/images/aecct_masks.png

### 4.2 Head Partitioning Self Attention

While the CASA mechanism of ECCT has demonstrated effective performance in decoding, we aim to further optimize its computational efficiency since we seek to develop neural decoders with complexity comparable to their classical counterparts such as BP.
To this end, we introduce Head Partitioning Self Attention (HPSA), which maintains the effectiveness of CASA while significantly reducing computational complexity. HPSA strategically divides ECCT’s masking via the attention heads into two groups: first-ring and second-ring MP heads. This division not only enhances efficiency but also introduces a graph-structure inductive bias by distinguishing between neighbors and second-ring connections, in contrast to the Code-Aware mask in ECCT. An illustration of HPSA is provided in Figure 2.

#### First Group: First-Ring Message Passing

This group of heads performs attention between nearest neighbors in the Tanner graph. This process, which we term first-ring MP, facilitates communication between variable nodes and check nodes. The corresponding attention masks are the $c\to v$ and $v\to c$ in Figure 3, demonstrating the increased sparsity of HPSA compared to the Code-Aware mask from ECCT.

#### Second Group: Second-Ring Message Passing

The second group focuses on what we call second-ring connections. These heads apply attention only between nodes at a distance of two in the Tanner graph. This allows for MP between variable nodes and other variable nodes, as well as between check nodes and other check nodes. The corresponding attention masks are the $c\to c$ and $v\to v$ in Figure 3, further illustrating the sparsity enhancement of HPSA.

By structuring the attention mechanism, HPSA achieves results comparable to CASA while drastically reducing complexity. This approach brings the computational efficiency of our method closer to that of the BP algorithm, moving us significantly closer to practical implementation in resource-constrained environments.

### 4.3 Positional Encoding of the Tanner Graph

Figure: Figure 4: Tanner PE. (a) The SPE matrix is concatenated to the initial nodes’ embedding matrix. (b) Creation of the SPE vector for individual node $j$, which is then concatenated with the node’s embedding. $\lambda_{i}$ denotes the i-th smallest eigenvalue of the Tanner graph. $\phi_{i}$ denotes the eigenvector corresponding to the i-th smallest eigenvalue and $\phi_{i,j}$ is its j-th element.
Refer to caption: extracted/5910077/images/spe_3.png

Although the two-rings connectivity code-aware mask has proven effective in ECCT, it provides the model with limited information about the Tanner graph’s structure. By design, it does not distinguish between first-ring and second-ring connections (Choukroun & Wolf, ). To enhance the decoder’s performance beyond this limitation, we propose incorporating a soft inductive bias through SPE induced by the Tanner graph. This approach, inspired by Kreuzer et al. (), injects information from the Laplacian eigenspace, which serves as a meaningful local coordinate system, thereby enriching the model’s understanding of the graph’s topology.
The following procedure is applied for each node $j$ in the Tanner graph, as illustrated in Figure 4:

$$ $\text{SPE}_{j}=W_{(2n-k)\to 1}\text{MHSA}(Q_{j},K_{j},V_{j})$ (4) $$

$$ $Q_{j}=K_{j}=V_{j}=W_{2\to d_{\text{PE}}}\Phi_{j}$ (5) $$

where $\Phi_{j}\in\mathbb{R}^{(2n-k)\times 2}$ is constructed by concatenating the graph’s eigenvalues with the $j$-th node’s corresponding values in the eigenvectors, $d_{\text{SPE}}$ is a hyperparameter, $W_{2\to d_{\text{SPE}}}$ is a learnable tensor, $W_{(2n-k)\to 1}$ is a reduction operator (e.g., linear projection, max/average pooling), and MHSA denotes Multi-Head Self-Attention.
The resulting vector $\text{SPE}_{j}$ serves as the PE for node $j$ and is concatenated with the node’s embedding. This process is repeated for all nodes in the graph. At inference, the learned SPE vectors remain fixed, removing the extra runtime computation present during training.

## 5 Analysis

#### Compression Rate

The linear layers in the ECCT model constitute over 95% of the total weight count, including the channel’s output embedding. By employing ternary values, which theoretically require only 1.58 bits for representation, we achieve significant compression. Replacing FP32 values with ternary values results in a 95% reduction in the memory footprint of these layers. Consequently, the AECCT’s overall memory footprint is reduced to approximately 10% of the original ECCT, achieving a compression rate of around 90%.

#### Energy Consumption

Energy consumption is a critical factor, especially when deploying the AECCT on edge devices or in data centers, as it directly impacts battery life and operational costs. We base our analysis on energy consumption models for addition and multiplication operations on 7nm and 45nm chips for FP32 and INT8, as outlined by Horowitz (); Zhang et al. (); Wang et al. (). Our findings indicate that the AECCT achieves substantial energy savings. Specifically, it reduces the energy consumption of arithmetic operations in linear layers by at least 224 times on 7nm chips and 139 times on 45nm chips, compared to the original ECCT.

Figure: Figure 5: Comparison of attention sparsity levels for HPSA with $h_{f}=h_{s}=4$. Sparsity level represents the proportion of query-key dot products avoided relative to a full pairwise attention mechanism.
Refer to caption: extracted/5910077/images/casa_vs_hpsa_sparsity.png

#### Complexity

Dedicated hardware optimized for this approach avoids attention calculations that are subsequently masked out by the code-aware mask. Assuming such hardware, for a Tanner graph $T=(V,E)$, the ECCT CASA’s complexity in a single Transformer encoder block is $\mathcal{O}(d(\sum_{x_{i}\in V}d_{x_{i}}+\beta_{x_{i}}))$, where $d$ is the embedding vector size, $d_{x_{i}}$ is the degree of vertex $x_{i}$, and $\beta_{x_{i}}$ denotes the number of vertices with a distance of two from $x_{i}$, derived from applying both first- and second-ring MP in each head. In contrast, our HPSA approach reduces this complexity to $\mathcal{O}(d_{h}h_{f}(\sum_{x_{i}\in V}d_{x_{i}})+d_{h}h_{s}(\sum_{x_{i}\in V
}\beta_{x_{i}}))$, where $h_{f}$ and $h_{s}$ correspond to the number of first-ring and second-ring heads, respectively, and $d_{h}$ is the head dimension. This reduction stems from the partitioning of attention heads, where each head is dedicated exclusively to either first-ring or second-ring MP. Figure 6 illustrates this complexity reduction by comparing the number of query-key dot products avoided in CASA and HPSA for various codes. The sparsity level is defined as the percentage of dot products avoided relative to quadratic pairwise attention. As shown, HPSA achieves sparsity levels of at least 78% across most codes, while the CASA’s sparsity ranges from 30% to 78%. This visual representation corroborates our theoretical analysis, demonstrating the significant computational efficiency gained through HPSA. By strategically partitioning attention heads and dedicating them to specific ring levels, HPSA dramatically reduces the number of necessary dot product calculations, resulting in a more efficient and optimized attention mechanism.

The AECCT model’s complexity is governed by parameters $N$, $d$, $h_{f}$, and $h_{s}$, offering exceptional flexibility in balancing accuracy and computational efficiency. As demonstrated by Choukroun & Wolf (), even the most modest ECCT architectures (e.g., $N=2$, $d=32$) consistently outperform BP across several codes. This performance advantage extends to AECCT, which not only maintains this superior decoding capability but does so with complexity comparable to BP. As illustrated in Figure 7, the shallowest AECCT architecture, with complexity comparable to BP, outperforms BP with 50 iterations (L=50). This showcases AECCT’s ability to offer superior performance even at its most basic configuration, achieving a balance between computational efficiency and decoding capability that ECCT could not attain. The performance gap widens as we increase $N$ and $d$ in AECCT, since increasing the number of BP iterations beyond 50 yields only marginal improvements. Further analysis of AECCT’s complexity is provided in Appendix A.

#### AAP Sparsity

We analyzed the sparsity level of the Adaptive Absolute Percentile (AAP) linear layers in the AECCT for a model trained on BCH(63,51) code. Figure 6 illustrates our findings, revealing that the percentage of zero-valued weights ranges from approximately 40% to 50%.
Importantly, this sparsity effectively reduces the dimension of the embedding vectors to around $0.45d$, further amplifying the efficiency gains discussed in our complexity analysis.

Figure: Figure 7: Accuracy vs. Complexity analysis for (a) BCH(63,51), (b) LDPC(49,24), at $E_{b}/N_{0}=5$. We compare AECCT with BP and ECCT in terms of decoding quality versus complexity. AECCT($N$, $d$) denotes an architecture with $N$ Transformer encoder blocks with embedding size $d$, with a similar notation used for ECCT. BP($L$) denotes BP with $L$ iterations.
Refer to caption: extracted/5910077/images/bch_5db.png

## 6 Experiments * * * Code is available at https://github.com/mlaetvayn/AECCT

#### Training & Inference

We utilize a Post-Layer Normalization (Post-LN) Transformer architecture, consistent with the original Transformer design in Vaswani et al. (), distinguishing it from the ECCT approach, as we empirically found it to be more effective.
The training process is divided into two phases.

In the first phase, we train the ECCT from scratch, incorporating several modifications: ReLU activations replace GeLUs, HPSA is used instead of CASA, and learnable Tanner graph PE is integrated. The model undergoes training for 1000 epochs (1500 for $N=10$), with each epoch consisting of 1000 batches. We employ the Adam optimizer with a batch size of 128.

In the second phase, the linear layers within the encoder blocks are replaced with AAP-linear layers, initialized using the weights obtained from the first phase. QAT is then applied using the same configuration as in the first phase. Upon completion, the weights of the AAP-linear layers are fixed as ternary values, and their corresponding scales are also fixed.

Throughout both phases, following the approach of Choukroun & Wolf (), we use a zero codeword with a Gaussian channel sampled from a normalized SNR $(E_{b}/N_{0})$ range of 3 to 7. The learning rate is initialized at $10^{-4}$ and decays to $5\times 10^{-7}$ following a cosine schedule. The cross-entropy loss is used to guide the model in learning the multiplicative noise (Bennatan et al., ).

**Table 1: We present the performance of our proposed method against established baselines, measured using -log(BER) across three normalized SNR levels. The negative logarithm transformation of BER is employed for clearer visualization, with larger values representing superior error correction capabilities. We compare our AECCT to the ECCT as our baseline. Additionally, we compare it to BP Pearl () with $L=5$ (first row) iterations and $L=50$ (second row) iterations. We separate the comparison between ECCT and AECCT according to the number of encoder blocks, $N$. For each $N\in\{6,10\}$, bold text indicates the best results between ECCT and AECCT for that specific $N$ value. Notably, for BCH codes, $N=10$ for AECCT was unnecessary, as AECCT with $N=6$ outperforms ECCT with $N=10$.**
| Method | BP | ECCT $N=6$ | AECCT $N=6$ | ECCT $N=10$ | AECCT $N=10$ |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | 4 | 5 | 6 | 4 | 5 | 6 | 4 | 5 | 6 | 4 | 5 | 6 | 4 | 5 | 6 |
| Polar(64,48) | 3.52<br>4.26 | 4.04<br>5.38 | 4.48<br>6.50 | 6.36 | 8.46 | 11.09 | 6.43 | 8.54 | 11.12 | 6.43 | 8.40 | 11.00 | 6.54 | 8.51 | 11.12 |
| Polar(128,86) | 3.80<br>4.49 | 4.19<br>5.65 | 4.62<br>6.97 | 6.31 | 9.01 | 12.45 | 6.04 | 8.56 | 11.81 | 7.26 | 10.60 | 14.80 | 7.28 | 10.60 | 14.59 |
| Polar(128,96) | 3.99<br>4.61 | 4.41<br>5.79 | 4.78<br>7.08 | 6.31 | 9.12 | 12.47 | 6.11 | 8.81 | 12.15 | 6.85 | 9.78 | 12.90 | 6.79 | 9.68 | 12.93 |
| LDPC(49,24) | 5.30<br>6.23 | 7.28<br>8.19 | 9.88<br>11.72 | 5.79 | 8.13 | 11.40 | 6.10 | 8.65 | 12.34 | 6.35 | 9.01 | 12.43 | 6.67 | 9.35 | 13.56 |
| LDPC(121,60) | 4.82<br>- | 7.21<br>- | 10.87<br>- | 5.01 | 7.99 | 12.78 | 5.17 | 8.32 | 13.40 | 5.51 | 8.89 | 14.51 | 5.71 | 9.31 | 14.90 |
| LDPC(121,70) | 5.88<br>- | 8.76<br>- | 13.04<br>- | 6.19 | 9.89 | 15.58 | 6.38 | 10.1 | 16.01 | 6.86 | 11.02 | 16.85 | 7.05 | 11.40 | 17.30 |
| LDPC(121,80) | 6.66<br>- | 9.82<br>- | 13.98<br>- | 7.07 | 10.96 | 16.25 | 7.27 | 11.50 | 16.90 | 7.76 | 12.30 | 17.82 | 7.98 | 12.60 | 18.10 |
| BCH(31,16) | 4.63<br>- | 5.88<br>- | 7.60<br>- | 6.39 | 8.29 | 10.66 | 7.01 | 9.33 | 12.27 | 6.41 | 8.30 | 10.77 | 7.21 | 9.47 | 12.45 |
| BCH(63,36) | 3.72<br>4.03 | 4.65<br>5.42 | 5.66<br>7.26 | 4.68 | 6.65 | 9.10 | 5.19 | 6.95 | 9.33 | 5.09 | 6.96 | 9.43 | 4.90 | 6.64 | 9.19 |
| BCH(63,45) | 4.08<br>4.36 | 4.96<br>5.55 | 6.07<br>7.26 | 5.60 | 7.79 | 10.93 | 5.90 | 8.24 | 11.46 | 5.72 | 7.99 | 11.21 | 5.83 | 8.15 | 11.52 |
| BCH(63,51) | 4.34<br>4.50 | 5.29<br>5.82 | 6.35<br>7.42 | 5.66 | 7.89 | 11.01 | 5.72 | 8.01 | 11.24 | 5.38 | 7.40 | 10.50 | 5.68 | 7.88 | 11.04 |

#### Results

We evaluate our proposed method on three types of linear block codes: Low-Density Parity Check (LDPC) codes (Gallager, ), Polar codes (Arikan, ), and Bose–Chaudhuri–Hocquenghem (BCH) codes (Bose & Ray-Chaudhuri, ), using parity check matrices from Helmling et al. (). The architecture is defined by two key parameters: the number of encoder layers ($N$) and the embedding dimension ($d$). Performance is assessed by measuring bit error rates (BER) across a range of $E_{b}/N_{0}$ values, followed by Choukroun & Wolf (). Table 1 presents our results, showing the negative natural logarithm of the BER. We compare the performance of our AECCT to the ECCT and BP (Pearl, ) across two different architectures: $N=6$ and $N=10$, both with an embedding dimension of $d=128$. The results indicate that the AECCT performs on par with the ECCT, and in some cases, even exceeds it for certain codes, while remaining much more efficient.

**Table 2: Evaluation of AECCT components against (Post-Ln) ECCT. For added generality, we used ($N=6$, $d=64$) for LDPC(49,24) while maintaining ($6$, $128$) for other codes as in Tab. 1.**
| Model | POLAR(64,48) | BCH(31,16) | LDPC(49,24) |
| --- | --- | --- | --- |
| ECCT | 6.40 8.50 11.10 | 6.95 9.21 12.04 | 5.97 8.44 12.01 |
| ECCT + HPSA | 6.40 8.52 11.17 | 7.00 9.24 12.12 | 5.98 8.48 12.12 |
| ECCT + SPE | 6.43 8.53 11.10 | 7.01 9.21 12.31 | 5.97 8.46 12.09 |
| ECCT + HPSA + SPE | 6.50 8.61 11.15 | 7.00 9.25 12.07 | 6.01 8.48 12.01 |
| ECCT + AP | 6.39 8.49 11.14 | 7.05 9.27 12.33 | 5.90 8.34 11.70 |
| ECCT + AAP | 6.41 8.51 11.13 | 7.06 9.37 12.37 | 5.91 8.35 11.74 |
| AECCT: HPSA + SPE + AAP | 6.43 8.54 11.12 | 7.01 9.33 12.27 | 5.89 8.33 11.67 |

#### Ablation Study

Our comprehensive ablation study evaluates the key components of our proposed model, with results detailed in Table 2. We use an ECCT model with Post-Ln architecture as our baseline, then separately incorporate each AECCT component to assess its individual impact. First, we examine the impact of HPSA. The results demonstrate that HPSA maintains or improves performance relative to the CASA-based baseline, while simultaneously reducing computational complexity. This dual benefit of preserved or enhanced effectiveness coupled with increased efficiency underscores HPSA’s value as a key component of our model. Next, we investigate the influence of the SPE. We find that integrating positional and structural information from the Tanner graph’s Laplacian through SPE significantly boosts overall model performance. To ensure a fair comparison, we maintain consistent total embedding dimensions by reducing the size of the channel’s output embedding vectors before concatenating the SPE vectors. Finally, we assess the impact of AAP quantization. Our analysis shows that AAP quantization outperforms absolute percentile (AP) quantization. The adaptive approach introduces a learnable parameter $\delta$, enabling dynamic adjustment of weight sparsity and effectively controlling feature filtration.

Appendices B and C present additional ablation studies. The former evaluates AECCT with varying numbers of first and second ring heads, revealing their similar importance with optimal performance when $h_{f}=h_{s}$. The latter compares a binary weighted version of AECCT to (our) ternary weighted AECCT, both using AAP quantization. Results demonstrate the superiority of ternary representation, achieving substantial performance gains with minimal bit usage increase (1.58 vs 1), justifying our choice of ternary quantization.
We analyze in Appendix D the necessity of $\delta$ in the AAP method, demonstrating its importance for dynamic thresholding across different AAP layers.

## 7 Conclusions

We introduced the AECCT, an enhanced version of the ECCT initially proposed by Choukroun & Wolf (). The AECCT integrates several novel techniques: Adaptive Absolute Percentile Quantization, which compresses the linear layer weights in the Transformer encoder blocks to ternary values; Head Partitioning Self-Attention, which replaces the code-aware self-attention module, significantly reducing complexity; and Tanner Graph Positional Encoding, which improves the model’s overall effectiveness. The AECCT achieves a complexity level comparable to BP while reducing memory usage, energy consumption, and computational complexity, all while delivering performance on par with the ECCT. Altogether, these enhancements bring transformer-based error correction decoders closer to practical deployment in real-world communication systems, offering notable improvements in the reliability of physical layer communications. As future work, we wish to explore learned Tanner-graph-based positioning techniques and apply pattern-based head partitioning to other structured learning problems.

In a broader context, our work addresses a gap in the literature regarding the acceleration of small Transformers, particularly those where attention patterns are dictated by the problem domain. The novel quantization method we propose enables exact localized adaptations, and the head partitioning method we propose addresses any hierarchical or structured data.

## Appendix A Complexity Analysis

In this section, we provide a detailed breakdown of the complexity for various components of our AECCT model, focusing on the AAP linear layer, the Head Partitioning Self-Attention (HPSA) mechanism, and the second-ring degree $\beta$.

#### AAP Linear Complexity

We analyze the complexity of the AAP linear layer by separating it into multiplication and addition components. The complexity for FP32 multiplications, which arises from the quantization of the input activation matrix and the dequantization of the output activation matrix, is given by

$$ $2(2n-k)d=2\lvert V\rvert d=\mathcal{O}(\lvert V\rvert d),$ (6) $$

where $T=(V,E)$ is the Tanner graph, $d$ is the embedding vector size, $k$ denotes the input message size, and $n$ is the output vector size of the channel. Matrix multiplication, which involves only additions and subtractions, results in an INT8 addition complexity of

$$ $(2n-k)d^{2}=\mathcal{O}(\lvert V\rvert d^{2}).$ (7) $$

The bias addition, performed in FP32, is $\mathcal{O}(\lvert V\rvert d)$.

#### HPSA Complexity

Similarly, we decompose the complexity of HPSA into multiplications and additions. Assuming an equal number of first- and second-ring heads, the total number of FP32 multiplications for all first-ring heads in a single Transformer encoder block is

$$ $\left(\sum_{i=1}^{n}d_{i}+\sum_{i=1}^{n-k}\tilde{d}_{i}\right)\frac{d}{2}=2 \lvert E\rvert\frac{d}{2}=\mathcal{O}(\lvert E\rvert d),$ (8) $$

where $d_{i}$ denotes the degree of the $i$-th variable node and $\tilde{d}_{i}$ denotes the degree of the $i$-th parity check node. The number of FP32 additions is similar.

Figure: Figure 8: The expected values of $\beta_{x_{i}}$, representing the number of vertices two edges away from $x_{i}$, are compared to the expected values of $d_{x_{i}}$, the degree of $x_{i}$, for $d=32$ and $d_{h}=4$. The size of $\beta_{x_{i}}$ significantly affects the complexity of second-ring heads in the HPSA.
Refer to caption: extracted/5910077/images/degree_comparison.png

The total number of FP32 multiplications required for all second-ring heads in a single Transformer encoder block is

$$ $\frac{d}{2}\sum_{x_{i}\in V}\beta_{x_{i}},$ (9) $$

where $\beta_{x_{i}}$ represents the number of vertices at a distance of two edges from $x_{i}$. Again, the number of FP32 additions is similar.

#### AECCT Complexity

Having examined the complexities of individual components, we now combine these to determine the total complexity of AECCT. The results of this combined analysis are presented in Table 3.

**Table 3: Complexity comparison between AECCT and (sum-product) BP, assuming for simplicity that the number of first-ring heads equals the number of second-ring heads in the HPSA. We analyze the magnitude of $\beta$ for several codes in A.1. The AECCT involves FP32 multiplications, FP32 additions, and INT8 additions. Therefore, we analyze each of these operations separately.**
| Operation | AECCT | BP |
| --- | --- | --- |
| FP32 MUL | $\mathcal{O}(Nd(\lvert V\rvert+\lvert E\rvert+\sum_{x_{i}\in V}\frac{\beta_{x_{<br>i}}}{2}))$ | $\mathcal{O}(L\lvert E\rvert)$ |
| FP32 ADD | $\mathcal{O}(Nd(\lvert E\rvert+\sum_{x_{i}\in V}\frac{\beta_{x_{i}}}{2}))$ | $\mathcal{O}(L\lvert E\rvert)$ |
| INT8 ADD | $\mathcal{O}(Nd^{2}\lvert V\rvert)$ | - |

Figure: Figure 9: The impact of $h_{f}$ and $h_{s}$ on HPSA’s complexity is analyzed for three codes: BCH(63,45), LDPC(49,24), and POLAR(128,86). We calculate the number of multiplications required for the CASA and compare it to HPSA with all possible combinations of $h_{f}$ and $h_{s}$, where $h_{f}$ represents the number of first-ring heads and $h_{s}$ the number of second-ring heads.
Refer to caption: x1.png

### A.1 Impact of β 𝛽 \beta italic_β

We analyze the expected value of $\beta_{x_{i}}$ to evaluate its influence on computational complexity, as illustrated in Figure 8. Our analysis indicates that $\mathbb{E}[\beta_{x_{i}}]$ is approximately $\frac{\mathbb{E}[d_{x_{i}}]^{2}}{2}$. Given that second-ring heads exhibit higher complexity, it is feasible to employ more first-ring heads. Figure 9 demonstrates the complexity of the HPSA for various combinations of $h_{f}$ and $h_{s}$, compared to the CASA mechanism. The results clearly show that HPSA significantly reduces complexity compared to CASA.

Figure: Figure 10: -ln(BER) vs. Complexity for LDPC(49,24) at $E_{b}/N_{0}$ values of (a) 4 and (b) 6, and BCH(63,51) at $E_{b}/N_{0}$ values of (c) 4 dB and (d) 6 dB. We compare AECCT, BP, and ECCT in terms of performance and complexity. AECCT($N$, $d$) denotes an architecture with $N$ Transformer encoder blocks and an embedding size of $d$, with a similar notation used for ECCT. BP($L$) refers to BP with $L$ iterations.
Refer to caption: extracted/5910077/images/ldpc_4db.png

### A.2 Complexity vs BER

We analyze the trade-off between complexity and performance for AECCT, ECCT, and BP in Figure 10. The results are presented for $E_{b}/N_{0}$ of 4 and 6. Our findings indicate that AECCT is comparable to BP in both complexity and performance, while demonstrating better scalability.
Moreover, AECCT architectures consistently outperform ECCT architectures with significantly lower complexity. For example, AECCT with $N=6$ and $d=64$ consistently surpasses ECCT with $N=6$ and $d=32$, while also being more computationally efficient.

**Table 4: Ablation of $h_{f}$ and $h_{s}$**
| Code | $N$, $d$ | 1st ring, 2nd ring | Neg ln(BER) |
| --- | --- | --- | --- |
| BCH(31,16) | 6, 128 | 2, 6 | 6.91 9.24 12.1 |
| BCH(31,16) | 6, 128 | 4, 4 | 7.00 9.25 12.1 |
| BCH(31,16) | 6, 128 | 6, 2 | 6.95 9.25 12.1 |

## Appendix B First & Second Ring Heads Balance

We analyze the optimal configuration of first-ring and second-ring heads in the HPSA mechanism. We evaluate AECCT without the AAP contribution, varying the number of first-ring heads $h_{f}\in\{2,4,6\}$ and setting $h_{s}=8-h_{f}$, where $h_{s}$ represents the number of second-ring heads. Table 4 presents these findings, indicating that both types of heads contribute similarly, with $h_{f}=h_{s}=4$ yielding the best results.

**Table 5: AECCT binary vs AECCT ternary**
| Code | $N$, $d$ | AECCT ternary | AECCT binary |
| --- | --- | --- | --- |
| LDPC(49,24) | 6, 128 | 6.10 8.65 12.3 | 5.96 8.42 11.9 |
| BCH(31,16) | 6, 128 | 7.01 9.33 12.1 | 6.52 8.55 11.0 |
| POLAR(64,48) | 6, 128 | 6.37 8.52 11.1 | 6.12 8.20 10.6 |

## Appendix C Ternary vs Binary Precision Choice

We evaluate our AECCT method using binary precision with AAP quantization as an alternative to the ternary precision AAP quantization. The results are listed in Table 5. Evidently, the ternary quantization significantly outperforms the binary one in terms of precision. This substantial performance improvement, achieved with only a minimal increase in bit usage (1.58 vs 1), strongly supports our decision to use ternary over binary quantization.

Figure: Figure 11: Analysis of $\delta$. The average $\delta$ value of each type of linear layer across the AECCT trained on (a) BCH(63,51); (b) LDPC(121,70).
Refer to caption: extracted/5910077/images/delta_bch_63_51.png

## Appendix D AAP Dynamic Feature Control

We present the post-training values of $\delta$ for two AECCT models in Figure 11. Notably, the $\delta$ values are higher for the self-attention projections, leading to a greater elimination of features, whereas in the feed-forward layers, $\delta$ retains more information. This behavior may be explained by the fact that self-attention (through the query, key, and value projections) focuses on a specific subset of features for each attention head, while the feed-forward layers primarily reduce redundancy between blocks without drastically limiting the feature set.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Covered as part of weaknesses
soundness: 3
strengths: 1. Tackles an important problem of reducing the complexity of neural decoders. 

2. The performance improvements in terms of memeory and compute are non-trivial and very impressive. 

3. All the three main ideas proposed are novel and interesting. Specifically, splitting the masking based on the node structure is a clever example of leveraging domain structure. 

4. Set of experiements are exhaustive and includes sufficient number of results as well as ablation studies to show the merits of proposed ideas.
summary: This paper introduces simplifications to the Error Correction Code Transformer (ECCT) neural decoder model to reduce the memory and compute requirements. The authors claim that conventional quantization methods result in a significant drop in performance and instead propose alternative approaches to reduce the complexity. 

First is Adaptive Absolute Percentile (AAP) quantization, to optimize the percentile of scaling factor (percentage of absolute values). 

Second it head partitioning self-attention to limit the attention to either between similar nodes (check/variable) or between different nodes (check<->variable), dividing the attention heads into two groups. This significantly reduces the number of active elements in the attention mask.

Finally, uses spectral positional encoding (SPE) to introduce an inductive bias about the tarnner graph into the model.
weaknesses: 1. While a lot of interesting ideas were discussed in the paper to improve the efficiency of ECCT, I am not fully convinced about the practical relevance of a "universal decoder" architecture that simply utilizes a parity check matrix for all codes. In coding theory, each well-known family of codes has highly specific representations and special propoerties that cannot be fully captured by a simple parity check matrix. For instance, in Polar codes, the reliability sequence and the sequential decoding of information bits plays a crucial role in deciding the performance of the code. Claiming better performance for all family of codes w.r.t BP decoder is a bit misleading. Again, I appreciate the line of ECCT and related works from an academic curiousity point of view, but in my opinion, the merits and drawbacks should be highlighted more clearly without any overpromising. 

2. One common troubling trend in comparing the performance of neural decoders is the choice of weak classical baselines. While I understand the difficulty of a common decoder architecture outperforming the highly specialised decoders for each of the class of codes, the comparison should neverthless be done for completeness. For instance, it is known that BP is not optimal for many codes. For instance, when comparing with BCH codes, Berlekamp-Massey decoding algorithm should be used as the classical non-learning baseline and similarly for Polar codes, successive cancellation list decoding should be used. I strongly suggest authors to include the best classical baselines for each code to clearly show the gap with respect to the SOTA neural decoder, whether it is positive or negative.

3. I beleive authors chose the format of the negative log BER for presenting the results based on the original ECCT paper. But this format is not very informative and very uncommon in information and coding theory literature. While it is easy to identify which scheme is doing better, the gap between the performances is hard to interpret as the scale is not very intuitive. Rather, the standard format of SNR gap in dB for a target BER/BLER is more informative.

4. "Our analysis shows that AAP outperforms AP quantization" --> I see from ablation studied in Table 2 that the performace difference between ECCT + AP vs ECCT + APP is very small (0.01 - 0.1). Again, this scale of -log(BER) is very non-intuitive but in my opinion, this is negligible difference in performance and brings into question the value of doing AAP over AP. 

5. While the head partitioning idea is interesting, is it also similar to the idea proposed in the paper "CrossMPT: Cross-attention Message-Passing Transformer for Error Correcting Codes", specifically the second-ring MP. 

6. The details about SPE are unclear and I do not understand the "spectral" positioning and the "Laplacian eigenspace" ideas discussed. Rewriting the seubsection to simplify the discussion of these ideas would be helpful.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Discuss any potential challenges or modifications needed to implement HPSA and SPE on current hardware platforms used for error correction.
2. Compare the implementation complexity of HPSA and SPE to that of the original ECCT approach.
3. If possible, provide estimates or examples of how these techniques might impact hardware design or utilization in practical ECC systems.
soundness: 3
strengths: 1. The paper is well written and easy to follow, and the contributions are relevant.
2. The techniques introduced demonstrate significant compression without performance degradation compared to ECCTs.
3. The technique is tested on 3 different ECC types - Polar, Low-Density Parity Check (LDPC), and Bose–Chaudhuri–Hocquenghem (BCH) codes - demonstrating robustness and applicability across various error correction scenarios.
4. The authors conduct ablation studies evaluating the impact of each of the three components.
summary: This paper introduces Accelerated Error Correction Code Transformer (AECCT), which reduces the memory footprint and computational complexity of Error Correction Code Transformers (ECCT). The proposed technique introduces three distinct novel components:
1. Adaptive Absolute Percentile (AAP) Quantization: This method compresses model weights to ternary values, reducing memory and energy use while maintaining model accuracy. Unlike traditional quantization, AAP dynamically adjusts sparsity to retain essential features.
2. Head Partitioning Self-Attention (HPSA): This self-attention mechanism reduces computational complexity by dividing attention heads to focus separately on first-ring (neighbor) and second-ring (more distant) connections within ECC’s Tanner graph. HPSA achieves higher sparsity, thus lowering computational costs compared to ECCT’s Code-Aware Self-Attention (CASA).
3. Spectral Positional Encoding (SPE): By embedding structural information from the Tanner graph’s Laplacian eigenspace, SPE enriches the model’s positional representation, enhancing the decoding ability without increasing runtime.

Using these components, the authors demonstrate a 90% compression ratio and reduce energy consumption by over 224 times compared to ECCT, approaching the efficiency of traditional algorithms like Belief Propagation (BP), making practical deployment possible.
weaknesses: 1. Two out of three proposed techniques (HPSA and SPE) may have limited applicability outside of ECC.
2. It is not clear how easy it will be to implement these techniques in existing hardware. Ternary quantization has been explored by previous work. Beyond quantization, does HPSA and SPE introduce additional implementation complexities on existing hardware? 
3. The energy efficiency numbers are estimates, the paper does not provide any real measured numbers on actual hardware.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1.In table2, why are the results of ECCT+AAP better than AECCT for datasets BCH and LDPC? Does it mean HPSA and SPE bring degradation for quantized ECCT? 

2.Are the energy consumption saving rate(224 times on 7nm chips and 139 times on 45nm chips) tested on real chips, or just an estimation?
soundness: 2
strengths: This paper presents contributions to enhancing ECCT's practical applicability through a comprehensive optimization approach. The key strengths lie in its multi-faceted architectural improvements and efficient implementation strategies.

1. The paper introduces three well-designed technical innovations that directly address ECCT's structural limitations. The Head Partitioning Self Attention (HPSA) mechanism optimizes the self-attention computation specifically for bipartite graph message passing, while the Spectral Positional Encoding (SPE) leverages Tanner graph eigendecomposition to provide richer structural information without runtime overhead. These architectural improvements demonstrate a deep understanding of ECCT's underlying structure and successfully enhance its performance.

2. Except for two structural improvements of ECCT, the paper makes a step further to explore the possibility of low-bit quantization of ECCT. The paper uses a ternary weight quantization method (AAP) specifically designed for ECCT's compact architecture, achieving multiplication-free linear layers while maintaining model accuracy. 

The comprehensive experimental validation demonstrates that these improvements maintain ECCT's decoding performance while reducing its computational complexity to a level comparable with traditional Belief Propagation methods. This work contributes to making transformer-based error correction more practical in resource-constrained environments and helps narrow the efficiency gap between neural decoding techniques and traditional algorithms.
summary: This paper addresses the challenges of high computational complexity and memory requirements faced by Error Correction Code Transformers (ECCT) in practical applications by introducing an innovative acceleration method. The approach centers on three technical innovations: 
- a specially designed ternary weight quantization method (AAP) that enables multiplication-free linear layers.
- an optimized self-attention mechanism (HPSA) based on code-aware multi-head processing that significantly reduces computational complexity.
- a positional encoding scheme (SPE) through Tanner graph eigendecomposition that provides richer graph connectivity representation without affecting inference runtime. 

Experimental results demonstrate that this method not only matches or surpasses the original ECCT's performance but also achieves a 90% compression ratio while reducing arithmetic operation energy consumption by at least 224 times on modern hardware, making transformer-based error correction more practical in resource-constrained environments.
weaknesses: The weaknesses of this work are manifested in the following three aspects:

1. Limited Scope and Generalization.
The method's exclusive focus on ECCT architecture raises concerns about its broader impact. As ECCT has not gained significant attention in academia nor found applications in industry, optimizations specific to this architecture may have limited practical value.

2. Lack of Technical Novelty in Quantization.
The proposed AAP quantization method, while showing good performance, primarily combines existing techniques - weight distribution analysis and learnable quantization steps. The approach does not present substantial innovation in quantization methodology, making it more of an implementation enhancement than a theoretical advancement.

3. Insufficient Validation of Acceleration Claims.
Despite positioning itself as an acceleration method, the paper's experimental validation concentrates on compression ratios, sparsity metrics, and energy consumption reduction. While these metrics demonstrate improved efficiency, the absence of direct inference speed measurements and real-world acceleration benchmarks leaves the actual acceleration effects unverified. This gap between claimed acceleration and demonstrated results weakens the paper's central premise.

Suggestions:

1. Enhanced Validation of Quantization Method

The paper would benefit from more comprehensive experiments comparing the proposed AAP method with existing quantization approaches. Specifically, comparative studies with methods like GPTQ and OmniQuant that focus solely on weight distribution or learnable scales would better demonstrate the advantages of combining these techniques. Such experiments would strengthen the paper's claims about AAP's effectiveness and innovation in quantization methodology.

2. Demonstration of Actual Acceleration Effects

The paper needs to provide concrete evidence of inference speedup to support its acceleration claims. We suggest adding experiments that measure and analyze the actual inference acceleration achieved after implementing all proposed optimizations. If practical acceleration is currently limited by implementation constraints, this limitation should be explicitly acknowledged, and the path toward achieving actual speedup in future work should be discussed. This transparency would better align the paper's claims with its demonstrated results.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: Merely out of curiosity: have the authors considered non transformer based architectures for this task? Presumably denoising diffusion models might have a role to play given the nature of the problem?
soundness: 3
strengths: Through it's experimental results, the paper convincingly demonstrates the benefit of ternary weight quantization, Spectral Position Embedding and tailoring the attention mask for bipartite graph message passing over the previous state of the art transformer based method - the error correction code transformer (ECCT).
summary: This paper introduces a novel transformer based neural error correction architecture. It builds on the idea of the Error Correction Code Transformer (ECCT) by performing ternary quantization of the parameters in linear layers, by modifying the masking mechanism in the attention operation and also by introducing positional embeddings. The architecture is empirically validated against the original ECCT and the more traditional Belief Propagation algorithm and is shown to outperform both in terms of accuracy.
weaknesses: As an outsider to the field of error correction, I found that the description of the problem domain was a bit too brief. The paper would benefit from expanding the problem setting section. It would also benefit the paragraph describing the ECCT architecture if the equations weren't inlined and that the sequence of transformations defining the architecture were serialized.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 4
questions: 1) See point 5 in Weaknesses - is the sparsity structured, and if so, does your assumed hardware architecture exploit that?

2) See (4) in Weaknesses - Where is this assumption coming from and what edge platform architecture was used? Is there a citation or a cycle-accurate simulation of the architecture? Analytical complexity numbers that assume optimal hardware are unrealistic.

3) What are the authors' thoughts on the contribution of each component of this framework to the overall acceleration? Is there an ablation study that can be presented in the main doc? If not, a small table and rough numbers would be good so that we can see how much general approaches like quantization contribute and how much the domain-specific tweaks such as Tanner graph masking contribute.

See discussion - the authors have conditionally addressed much of the questions above and paper weaknesses. I am editing the review to reflect this.
soundness: 3
strengths: 1) Solid presentation, all major thrusts are clearly presented and detailed. The paper also provides diagrams and examples to ease understanding and reading in what can be presented as a very dense, mathematical topic that a reviewer would have trouble grasping or reading through.
2) Analysis and experimentation include results on hardware for energy consumption reduction, a key metric in model compression. It also includes ablation studies for different encoding methods and acceleration due to model compression. All in all it analyzes most of the ablations relevant to the topic. I particularly commend the hardware results for different transistor nodes.
3) Contributions while building clearly off of a prior design and prior work, are cleverly tuned to the problem domain and a very clear niche. In particular the use of the Tanner graph and selective masking shows adaptation of the transformer compression problem to the given domain.
summary: The paper has three main thrusts that I can see, and it targets these towards specifically, acceleration and compression of small transformer model architectures used for error-correction code decoding in communications systems. This is an incremental but significant upgrade to earlier work on ECC transformer decoders. The main thrusts are:
1) Quantization: The paper uses an adaptive layer-based quantization scheme that combines maximum-based quantization with distribution-based quantization schemes to provide a ternary quantization for the network, greatly accelerating its computation (claimed) while also compressing model size.
2) Head Partitioning: Masking attention heads based on the Tanner graph node connections at varying levels of relation (first- and second-ring masks) allows greater sparsity in attention computation while also preserving performance.
3) Positional Encoding: A soft positional encoding is also introduced as an inductive bias in the tokenizer/learnable encoder (from what I can see) that further preserves signal positional information and compensates for information loss in the attention head masking.
weaknesses: 1) Concluding statements mention that the paper shows a general approach for the compression problem in small transformers - I would say that this is a bit of a stretch - the Tanner subgraph masking and the adjustments to positional encoding are not applicable to all problem domains, leaving the model quantization formulation the only really general cross-domain contribution. That formulation has not been benchmarking in the experimental section against prior art, and its acceleration alone (as opposed to the entire framework) has not been tested in the main document. I would say that ablation results for each component of the framework should be added in order to make this general claim.

2) A key contribution is replacing GeLU with ReLU in the bullet points of the front matter. This is absolutely not novel, as has been examined in [1] just to name a single source. I would strongly advise omitting this claim or qualifying it to the given problem domain.
[1] https://arxiv.org/abs/2310.04564 [Addressed]

3) While this work builds on the prior ECCT work, I would want a presentation of the model architecture in the main draft and clearly marked so that the scale of the 'small transformer' is visible, rather than having to dig through the ECCT paper and assume that this is the paper used. [Addressed]

4) The work assumes dedicated hardware for the transformer decoder in Section 5 when calculating complexity. This may not be a realistic assumption in practice due to the specialized and expensive design process for accelerator ASICs. I would want to see complexity calculations on a common edge accelerator platform such as an NVIDIA Hopper architecture or a Coral edge TPU, or a citation for the architecture-level simulator presumably used to generate these numbers (e.g. ScaleSim). Numbers for hardware without that backing are unreliable.

5)  The sparsity induced by masking is structured, and therefore may be applicable to standard accelerators. Why is there no mention of this? I do not see an analysis of the benefits of the structured sparsity from the masks in Section 4.2 compared to unstructured sparsity, which is a major plus point that has not been exploited.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. As mentioned above, the authors should report BLER curves.

2. The authors provide results for a number of LDPC codes. However, there are no details regarding these codes. Are them the best-available LDPC codes in the literature, or some designed by the authors? 

3. Related to Comment 2 and my comments in "Weaknesses," the authors should rework their results section and compare the performance of the proposed approach with that of state-of-the-art codes/decoders
soundness: 3
strengths: The paper proposes interesting methods like HPSA and adaptive absolute percentile quantization, which could be valuable outside of coding applications.
summary: This paper presents a method to enhance the efficiency and practicality of transformer-based error correction code (ECC) decoders, specifically targeting the error correction code transformer (ECCT) model introduced in prior work. The authors propose several modifications to ECCTs, including ternary weight quantization, a novel attention mechanism called head partitioning self-attention (HPSA), and a spectral positional encoding inspired by the Tanner graph's eigenspace. 
Jointly, these improvements lead to reduce memory, energy, and computational requirements while maintaining or improving the performance of the original ECCTs.
weaknesses: While the proposed approach introduces certain novel elements, I believe the paper falls short in critical areas, particularly regarding contribution scope and experimental results. Therefore, I recommend rejection for the following reasons:

1. Limited contribution: The paper proposes interesting methods like HPSA and adaptive absolute percentile quantization, which could be valuable outside of coding applications. However, since this paper is fundamentally a paper in coding theory, these contributions remain peripheral to the main focus (furthermore, per se they are not enough to warrant publication at ICLR).
And within the realm of coding theory, the contribution of the paper is limited. 

While relevant for the coding community, the contribution lacks the depth and broader impact expected at a major machine learning conference or in a top coding journal. A more appropriate fit would be a coding conference such as ISIT or ITW.

Furthermore, although ML-based decoding of conventional codes is a growing research area, with already many published papers, this paper does not provide a method that competes with conventional state-of-the-art codes/decoders. 


2. Weak comparisons and benchmarking: The experimental results are insufficient to substantiate the relevance of the proposed decoder.

Specifically, the authors compare the performance of their transformer-based decoder to that of belief propagation decoding for various codes, including polar codes and BCH codes. BP decoding of BCH codes and polar codes is known to be highly suboptimal. There are well-known, more efficient decoding algorithms for these codes (note that the considered BCH codes are very short). Hence, the current results provide no hint on whether the proposed decoders are good or not. 

To establish relevance, the paper would need to demonstrate comparable or superior performance to state-of-the-art decoding techniques in terms of accuracy or complexity. However, even if such comparisons were provided, in my opinion the contribution is still insufficient for a high-profile venue.

3. Use of non-standard metrics: The paper presents performance using the negative natural logarithm of the BER. Is there any reason for not adhering to standard metrics such as bit error rate curves or block error rate (BLER) curves (the latter being preferable)?  
The use of this unconventional (to day the least) metric complicates unnecessarily the interpretation of the results. I do not think it is wise to come up with an awkward way of showing results, when there is a established way to present results for error correcting codes.

Furthermore, reporting performance for only a few specific values of $E_b/N_0$ provides an incomplete view; it's unclear whether phenomena like error floors might occur at higher values.

Including conventional BLER curves over a range of $E_b/N_0$ values would make the results easier to interpret and align with community standards. Without such curves, it's challenging to assess the effectiveness of the proposed decoder and, in my opinion, the paper should be accepted in any venue.

Meta-review: This paper presents a method to enhance the efficiency and practicality of transformer-based error correction code decoders, focusing on the ECCT model introduced in prior work. Specifically, the authors propose: (i) an adaptive layer-based quantization scheme, (ii) masking attention heads based on the Tanner graph, and (iii) a soft positional encoding.

The reviewers have appreciated the detailed hardware and acceleration studies and noticed that some ideas (eg HPSA) may find general applicability. Nevertheless, important issues were raised during the review process. Specifically, reviewer mhUK criticized the limited scope, the lack of technical novelty in quantization and the insufficient validation of acceleration claims. While the first point was adequately addressed during the rebuttal, the remaining two (specifically the last one) were not. Additionally, reviewer mcns noted how the proposed method is still far from being competitive w.r.t. decoders that are pretty standard in the coding theory literature (SC decoding with a list size of 1). Now, I agree with the authors that the field of neural decoders is rather novel and, as such, it would be unfair to expect performance comparable with the state of the art. However, one would still expect a performance improvement from a paper on ECCT, since that is a critical issue hindering applicability. 

The weaknesses discussed above lead me to a 'reject' decision.

In the course of the discussion between reviewers and AC, reviewer 3y8a suggested to re-frame the contribution in the context of hardware acceleration for edge transformer architectures (adding revised test cases), which is a valid suggestion. In any case, given the strengths of the submitted manuscript, I would encourage the authors to improve their manuscript and re-submit to a future venue.

Decision: Reject