Accelerating Error Correction Code Transformers

Thank you for your thoughtful feedback.

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.

The paper is not a coding theory paper and should not be judged as such. It belongs to a substantial and growing body of work on neural decoders, which is published in the ML community. This paper addresses a critical goal: bringing neural decoders closer to real-world applicability. The required tools for this task are developed within the deep learning community. The proposed methods, including HPSA and adaptive absolute percentile quantization, are novel and generic machine learning tools with broad applicability beyond coding theory.

While the proposed approach introduces certain novel elements, I believe the paper falls short in critical areas, particularly regarding contribution scope and experimental results.

To address concerns about experimental results, we have added BER and BLER curves for a more comprehensive evaluation. Our contributions encompass both acceleration and error correction, introducing spectral positional encoding in this context for the first time and demonstrating that our quantization method outperforms state-of-the-art ternary quantization techniques. The results for the ablation study models are presented below:

We also show that AAP can generalize to other domains. We took the TinyBERT model https://arxiv.org/abs/1909.10351, a 14M-parameter model well-suited for evaluating small-model quantization, and used the SST-2 dataset as in its original paper. The results are summarized below:

Given that the FP32 accuracy of TinyBERT in our experiments was 88.7%, these results clearly demonstrate that AAP outperforms Abs-Mean, the SOTA method for ternary quantization, in terms of accuracy and applicability to small transformers.

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

Thank you for raising this concern. First and foremost, our focus is on reaching and surpassing ECCT's performance while ensuring the model is significantly faster and more lightweight. We successfully achieved this, demonstrating that our approach matches or exceeds ECCT's effectiveness while drastically reducing runtime and memory requirements.

We appreciate your feedback and have addressed it by adding additional results comparing our method to SCL for polar codes, which we believe provide stronger evidence of the decoder's relevance. While SCL decoding is specifically optimized for Polar codes, AECCT demonstrates comparable performance with low L values. This highlights the versatility of AECCT, as it can be trained to decode any code, yet still achieves results on par with an algorithm tailored specifically for Polar codes.

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.

ECCT is regarded as the state-of-the-art among neural decoders for many families of codes. Additionally, we have included a comparison to SCL in the paper, presented through BER and BLER curves. A detailed complexity analysis is provided in Appendix A to address this concern.

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)?

Thank you for your comment. We have added BER and BLER curves to the paper as requested. The use of the negative natural logarithm of the BER aligns with a line of previous work in the field.

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?

Thank you for pointing this out. All the codes used were sourced from https://rptu.de/channel-codes and represent various families of sparse codes.

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.

Answered above.

See point 5 in Weaknesses - is the sparsity structured, and if so, does your assumed hardware architecture exploit that?

We appreciate the reviewer’s insightful comment regarding the structure of sparsity. The sparsity in both AECCT and ECCT is not structured. While we attempted to permute the sequence to create clusters of unmasked entries https://arxiv.org/abs/2001.04451, this proved to be a highly non-trivial task within the context of Tanner graphs.

Can the authors include a clearly marked discussion of the ECCT architecture used in Experiments? I would prefer to have that marked off and visible.

Thank you for the suggestion. We have added a clearly marked discussion on the ECCT architecture used in the experiments, as recommended.

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.

Thank you for your comment. A complexity analysis detailing the contributions of different components, assuming dedicated hardware, is provided in Appendix A.

Thank you for your thoughtful feedback and for raising these points.

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.

As highlighted in the paper, HPSA is a versatile method that can be applied across many domains to achieve both acceleration and compression of self-attention mechanisms. This makes it broadly applicable to data-structured problems in machine learning. For example, in multi-modal problems, a subgroup of self-attention heads can be allocated to each modality, showcasing its adaptability to various challenges.

Furthermore, we emphasize that our Adaptive Absolute Percentile (AAP) quantization method represents a significant contribution to the compression of small Transformers. To address your concern about benchmarking, we conducted additional experiments comparing AAP with Abs-Mean quantization by Ma et al. (2024) https://arxiv.org/abs/2402.17764, recognized as the state-of-the-art method for ternary quantization. These results demonstrate AAP's advantages over Abs-Mean, establishing its effectiveness as a general-purpose quantization method for small Transformers.

The experiments included the three models from the ablation study (POLAR(64,48), BCH(31,16), and LDPC(49,24)) and the TinyBERT model https://arxiv.org/abs/1909.10351, a 14M-parameter model well-suited for evaluating small-model quantization. For TinyBERT, we used the SST-2 dataset as in its original paper. The results are summarized below:

Results for Ablation Study Models:

Results for TinyBERT:

Given that the FP32 accuracy of TinyBERT in our experiments was 88.7%, these results clearly demonstrate that AAP outperforms Abs-Mean, the SOTA method for ternary quantization, in terms of accuracy and applicability to small Transformers.

I would strongly advise omitting the ReLU to GeLU claim, or qualify it against prior art.

We have followed the reviewer's advice and included a citation to https://arxiv.org/abs/2310.04564.

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.

We revised the ECCT section to improve clarity, following the reviewer’s request. To illustrate how small ECCT is, consider the following: (1) it uses small embedding sizes (32 to 128), (2) it employs a limited number of encoder blocks (2 to 10), (3) the input sequences, representing codewords, are relatively short, with a maximum length of ~160—much smaller than those in LLMs, and (4) the final two linear projections are minimal, containing only ~20k parameters. Altogether, these factors result in a model with at most ~2M parameters.

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.

Regarding complexity calculations on common edge accelerator platforms, it is important to note that ECC is not designed to run efficiently on GPUs or TPUs, as this would be extremely wasteful. Obtaining accurate complexity measurements for AECCT requires implementing dedicated hardware to fully leverage the AAP linear layers.

To utilize the ternary-weight linear layer, we attempted to use the BitBLAS library https://github.com/microsoft/BitBLAS, which implements a 2-bit kernel, as demonstrated by Ma et al. (2024) https://arxiv.org/abs/2402.17764. Unfortunately, no improvement was observed due to the small sequence size. This issue is documented https://github.com/microsoft/BitBLAS/issues/118 , as the overhead of activation dequantization and quantization becomes significant when both the input and output features are small.

We sincerely thank the reviewer for their thoughtful and constructive comments, which have helped us improve the clarity and presentation of our paper.

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.

We have expanded the Problem Setting section to provide a more accessible description of the problem domain. Additionally, we have revised the paragraph describing the Error Correction Code Transformer architecture by serializing the sequence of transformations and presenting the equations in a non-inline format, as suggested.

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?

We appreciate your interest in exploring alternative architectures for this task. Indeed, we are actively considering applying the AECCT method to a new approach called DDECC https://arxiv.org/abs/2209.13533, which builds on the ECCT by incorporating denoising diffusion models. We believe this expansion has the potential to further enhance the decoding process and address challenges posed by noise in communication channels.

Thank you again for your valuable input, and we hope the revisions align with your expectations.

In Table 2, 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?

First and foremost, we can observe from the ablation study that SPE improves the model's performance, demonstrating its effectiveness as a standalone enhancement. HPSA is a method that significantly reduces the complexity of the self-attention mechanism. Combining it with ternary quantization, which also reduces complexity, can impact the model's performance—a trade-off for achieving significant acceleration and compression. While the model can adapt to each method individually, it may struggle when both are applied together. This is supported by the performance of ECCT+HPSA, which achieves strong results.

Are the energy consumption saving rates (224 times on 7nm chips and 139 times on 45nm chips) tested on real chips, or just an estimation?

The energy consumption saving rate is an estimation: We follow a line of prior work, including Wang et al. (2023) https://arxiv.org/abs/2310.11453 and Ma et al. (2024) https://arxiv.org/abs/2402.17764, which estimate energy usage based on the tables provided in Horowitz (2014) https://ieeexplore.ieee.org/document/6757323 and Zhang et al. (2022) https://arxiv.org/abs/2112.00133.

The paper needs to provide concrete evidence of inference speedup to support its acceleration claims.

Addressing the reviewer's concerns about concrete evidence for inference speed-up, we would like to highlight that implementing dedicated hardware to fully leverage linear layers requiring only integer additions is non-trivial, as noted by Ma et al. (2024) https://arxiv.org/abs/2402.17764 in the “New Hardware for 1-bit LLMs” section. To utilize the ternary-weight linear layer, we attempted to use the BitBLAS library https://github.com/microsoft/BitBLAS, which implements a 2-bit kernel, as done by Ma et al. (2024). However, no improvement was observed due to the small sequence size. This is a known issue https://github.com/microsoft/BitBLAS/issues/118, where the overhead of activation dequantization and quantization becomes non-negligible for linear layers with small input and output features dimensions.

We have elaborated on this issue in the conclusions section of the paper, as we agree it is an important point to address.

We thank the reviewer for the valuable feedback and constructive remarks.

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.

ECCT (cited 46 times in two years, which is not negligible for a specialized subdomain) has demonstrated preferable performance to other decoders, with follow-up works building upon it and surpassing classical decoders https://arxiv.org/pdf/2405.04050. However, its adoption in real-world applications has been limited due to runtime and memory constraints. This paper addresses and alleviates these challenges.

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.

Quantization is a well-studied field with many overlapping contributions. However, the challenge of optimizing small networks has been largely overlooked. In this work, we address this gap for the first time, presenting a solution that effectively determines a ternary representation for the model's weights. By employing a dual-scale approach to quantization—one scale based on the weight distribution and the other learnable—we allow the model to adjust its sparsity dynamically. This approach enables the model to explore various sparsity levels, improving the effectiveness of low-precision quantization. The novelty of our method lies in recognizing that, in low-precision quantization, the ability to adapt sparsity is crucial for achieving better performance.

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.

We appreciate the reviewer’s suggestion to demonstrate that AAP generalizes better. Following the review, we compared AAP with abs-mean quantization by Ma et al. (2024) https://arxiv.org/abs/2402.17764, which is the state-of-the-art method for ternary quantization. Abs-mean is also a more natural competitor to AAP than GPTQ, as it is a QAT method.

Our additional experiments include the three models from our ablation study (POLAR(64,48), BCH(31,16), and LDPC(49,24)) as well as TinyBert https://arxiv.org/abs/1909.10351, a 14M parameter model considered very small and suitable for our case. For TinyBert, we used the SST-2 dataset from its original paper and quantized the model with both AAP and abs-mean. The results for the three models from our ablation study are presented in the following table:

The effectiveness of AAP on tiny models, such as ECCT with fewer than 2 million parameters, is evident when compared to the state-of-the-art ternary quantization method. The results for TinyBert https://arxiv.org/abs/1909.10351 are presented below, with the FP32 model achieving an accuracy of 88.7:

Thank you for your detailed review and valuable comments.

Two out of three proposed techniques (HPSA and SPE) may have limited applicability outside of ECC.

As noted in the paper, HPSA can inject any structural information into the self-attention mechanism while also improving efficiency. Given the prevalence of data-structured problems in machine learning, HPSA has broad applicability. For instance, in multi-modal problems, a subgroup of heads within the self-attention mechanism could be allocated to each modality.

Spectral positional encoding (SPE) is widely used in the domain of graph Transformers and, for the first time, is applied to ECC in this work, leveraging the strong connection between ECC and graph structures.

does HPSA and SPE introduce additional implementation complexities on existing hardware? / Discuss any potential challenges or modifications needed to implement HPSA and SPE on current hardware platforms used for error correction.

Thank you for highlighting this point. As stated in the paper at the end of the SPE method section (lines 350–352), SPE does not affect inference, as it remains fixed after training and is merely concatenated to the embedding vectors. Similarly, HPSA introduces no additional implementation overhead, as each group of heads (first-ring and second-ring) performs the standard self-attention mechanism with corresponding masks, requiring only a straightforward concatenation.

Compare the implementation complexity of HPSA and SPE to that of the original ECCT approach.

As noted above, SPE introduces no additional complexity, while HPSA significantly reduces complexity compared to ECCT. For details, please refer to Figure 5 and the complexity analysis in Section 5.

The energy efficiency numbers are estimates, the paper does not provide any real measured numbers on actual hardware.

Regarding the estimation of energy consumption, we follow a line of prior work, including Wang et al. (2023) https://arxiv.org/abs/2310.11453 and Ma et al. (2024) https://arxiv.org/abs/2402.17764, which estimate energy usage based on the tables provided in Horowitz (2014) https://ieeexplore.ieee.org/document/6757323 and Zhang et al. (2022) https://arxiv.org/abs/2112.00133.

If possible, provide estimates or examples of how these techniques might impact hardware design or utilization in practical ECC systems.

Quantization reduces die size, alleviates memory bottlenecks, and simplifies computations to summations. The sparsity introduced by HPSA reduces the quadratic complexity to a level comparable to BP due to its connectivity structure, making the model more efficient. Additionally, model assimilation is generally much easier with AECCT.

Thank you for your thoughtful feedback.

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.

We would like to clarify that our proposed AECCT builds upon the ECCT framework, which is not a "universal decoder" in the sense of employing a single model for all codes. Instead, for each code and its corresponding parity-check matrix, a separate model is specifically trained. This approach ensures that the unique characteristics of each code are effectively captured during training.

A key strength of the ECCT framework, and by extension AECCT, is its ability to leverage a unified and adaptable architecture that can be trained on any given code. This versatility allows it to outperform most classical decoders while maintaining broad applicability. Unlike traditional decoders that are often specialized for specific codes, AECCT's generalizability across diverse code families highlights its practical relevance and effectiveness, even when the particular properties of each code are not explicitly encoded in the parity-check matrix.

One common troubling trend in comparing the performance of neural decoders is the choice of weak classical baselines.

We appreciate your feedback on the selection of classical baselines. For Polar codes, we performed experiments using SCL decoding, and the results are presented in the table below. While SCL decoding is specifically optimized for Polar codes, AECCT demonstrates comparable performance with low L values. This highlights the versatility of AECCT, as it can be trained to decode any code, yet still achieves results on par with an algorithm tailored specifically for Polar codes.

In addition, we added BER and BLER curves to the paper.

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?

Thank you for your comment. However, we would like to clarify that CrossMPT's method uses only first-ring connections and does not incorporate second-ring connections, relying on a design that significantly increases computational and parameter overhead. In contrast, AECCT efficiently utilizes the multi-head attention mechanism by dedicating half the heads to first-ring connections and the other half to second-ring connections. This design substantially reduces computational cost while preserving effectiveness, highlighting a key advantage of our approach.

I believe 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.

Thank you for pointing this out. We have implemented your suggestion by adding BER and BLER curves to the paper for a more intuitive presentation. Negative log BER was employed to remain consistent with a line of previous work in the field of neural decoding, allowing for the compact performance presentation of a multitude of codes.

I see from ablation studies in Table 2 that the performance difference between ECCT + AP vs. ECCT + AAP is very small (0.01 - 0.1).

Achieving a good ternary representation with extremely low precision for the model's weights is a challenging task. In general, prior work https://arxiv.org/abs/2210.17323 in neural network quantization often yields only slight improvements. Therefore, a 0.1 improvement in log scale should not be considered insignificant.

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

We have added a more simplified explanation to the SPE section to improve clarity for readers. To further assist with understanding, please refer to Kreuzer et al. (2021) https://arxiv.org/abs/2106.03893, which provides a clear explanation of spectral positional encoding and the Laplacian eigenspace.