Spark Transformer: How Many FLOPs is a Token Worth?

Thank you for taking the time to read and review our paper!

Re: Only a theoretical estimation of the error for the approximate statistical top-k algorithm was provided, without experimental validation.

We appreciate the reviewer's attention to detail. While we provided a theoretical analysis of the error for the approximate statistical top-k algorithm, we also included experimental validation in the original submission.

Figure 1 demonstrates that the percentage of nonzeros resulting from statistical top-k in Spark Transformer closely aligns with the desired sparsity ratio. Our hyperparameter choices aimed for 8% nonzeros in Spark FFN (Figure 1(a)) and at most 256 nonzeros in Spark Attention (Figure 1(b)), which are accurately reflected in the empirical results.

To further address the reviewer's comment, we have added Figure 6 and Figure 7 in the revised manuscript. These figures illustrate the error between the estimated threshold and the true threshold, demonstrating that the estimated threshold (red dotted vertical line) closely approximates the true threshold (blue vertical line). This observation strongly supports the effectiveness and accuracy of the statistical top-k algorithm in practice.

Re: There was no analysis or testing of the gap between the approximate statistical top-k algorithm and the actual top-k selection.

We appreciate the reviewer's inquiry and would like to clarify the potential "gap" referred to in the question.

• Gap in the number of nonzeros: If the reviewer is referring to the difference between the number of nonzeros selected by statistical top-k and the desired sparsity level (k), we would like to direct their attention to Figure 1. This figure demonstrates a close match between the achieved sparsity and the target sparsity.

• Gap in the threshold: If the "gap" refers to the difference between the estimated threshold used in statistical top-k and the true threshold that would select exactly the top-k entries, we have added Figures 6 and 7 to the revised manuscript. These figures illustrate that the estimated threshold closely approximates the true threshold.

• Gap in model quality: If the reviewer is concerned about the potential difference in model quality between using statistical top-k and exact top-k, we acknowledge the computational challenges in conducting a direct comparison. Training with exact top-k is prohibitively expensive (>10x slower), making this ablation study infeasible within the given timeframe. However, the strong agreement between statistical top-k and the desired sparsity level (Figure 1) and the accurate threshold approximation (Figures 6 and 7) suggest that statistical top-k serves as an effective proxy for exact top-k, minimizing any potential quality loss.

Re: The study lacks comparisons with a sufficient number of state-of-the-art sparse attention works.

We appreciate the reviewer's suggestion. While we acknowledge the breadth of research in sparse attention, our work focuses on achieving sparsity through a different mechanism, namely, by inducing activation sparsity within the feedforward network. This approach offers a simpler alternative to complex sparse attention mechanisms, allowing for significant FLOPs reduction without intricate design modifications. We have included a discussion of related sparse attention works in the Related Work section to highlight the broader context of our contribution.

Re: Experiments were conducted only on CPUs, with no results provided for mainstream platforms such as GPUs.

We appreciate the reviewer's observation. While our experiments focused on CPU evaluation, we believe that Spark Transformer can be effectively implemented on GPUs to achieve significant speedups. Recent work like Deja Vu and PowerInfer have demonstrated that dynamic, unstructured activation sparsity, similar to that in Spark Transformer, can be efficiently implemented on GPUs for Transformer architectures, leading to considerable performance gains. We anticipate similar benefits for Spark Transformer given the shared sparsity characteristics. However, our primary focus was on addressing the modeling challenge of achieving high activation sparsity in LLMs without sacrificing quality. As demonstrated in Table 4, Spark Transformer attains higher activation sparsity with lower quality loss compared to existing methods. This contribution lays the groundwork for future efficient GPU implementations, building upon the aforementioned advancements in sparse GPU acceleration.

Re: There were no performance impact studies or experiments on applying the approximate algorithm.

We appreciate the reviewer's interest in the performance impact of the approximate statistical top-k algorithm. While a direct comparison with exact top-k would be ideal, it's important to highlight that training with exact top-k is computationally prohibitive (>10x slower), making this ablation study infeasible within the given timeframe.

However, we provide strong evidence that statistical top-k closely approximates exact top-k. Figures 6 and 7 illustrate the minimal error between the estimated threshold used in statistical top-k and the true threshold for exact top-k selection. The close alignment of these thresholds indicates that statistical top-k selects nearly identical entries as exact top-k, minimizing any potential performance difference.

Re: What is the actual distribution of the scores? How does it deviate from the Gaussian distribution?

We appreciate the reviewer's insightful question. To address this, we have added Figure 6 and Figure 7 in the revised manuscript, which illustrate the actual distribution of scores (blue curve) alongside a fitted Gaussian distribution (red dotted curve). As seen in the figures, the actual distribution closely matches the Gaussian distribution. While some minor discrepancies may exist in certain cases, the estimated threshold still closely aligns with the true threshold. This is because the deviations from the Gaussian distribution tend to cancel each other out, resulting in a sparsity level that closely matches the desired level, as evidenced in Figure 1.

Re: What is the overlap between the selection results of the approximate statistical top-k and the actual top-k selection? How large is the error? What impact does it have on performance?

We appreciate the reviewer's insightful questions regarding the overlap and error between statistical top-k and exact top-k selection.

To clarify the overlap:

• If statistical top-k underestimates the selection threshold, it will include all the actual top-k entries plus some additional ones.
• Conversely, if it overestimates the threshold, it will select the largest values among the actual top-k entries but miss some of the smaller ones.

Regarding the error:

• Figure 1 illustrates the error by comparing the actual number of nonzeros resulting from statistical top-k with the desired number of nonzeros. This provides a direct visualization of the accuracy of our approximation.

While quantifying the precise performance impact of statistical top-k is challenging due to the prohibitive cost of training with exact top-k (>10x slower), Figures 1, 6, and 7 collectively demonstrate that statistical top-k serves as an effective approximation. The close alignment between the achieved sparsity and the desired sparsity, along with the accurate threshold estimation, suggests that any performance difference is likely to be minimal.

Dear Reviewer 1vWh,

Thank you again for your valuable feedback, which has significantly improved our submission. In response to your comments, we have added Figure 6 and Figure 7 in the Appendix, which provides additional justifications for the effectiveness of our statistical top-k.

As the discussion phase is ending soon, we wanted to kindly follow up to confirm if our response has addressed your concerns. Please don't hesitate to let us know and we are happy to answer any additional questions or address any concerns you may have.

Thank you for your time and thoughtful input.

Best regards, The Authors

We thank the reviewer for their thoughtful follow-up and appreciate the continued dialogue on the practical impact of sparsity.

Re: Success of Deja Vu and PowerInfer due to open-source implementations and demonstrated acceleration

We completely agree that real-world acceleration is critical for the success of works like Deja Vu and PowerInfer, as their primary focus is on hardware optimization for GPUs. However, these methods are not applicable to modern Transformers with non-ReLU activations, which lack activation sparsity. Our contribution addresses this gap by introducing a modeling approach that induces strong activation sparsity while preserving task performance. The success of our method in comparison to related work is directly measurable by the degree of sparsity achieved with minimal quality loss, as demonstrated in Table 4.

Re: FLOPs reduction claims without practical value

We share the reviewer’s concern that not all forms of FLOPs reduction translate to practical value on GPUs. However, activation sparsity has repeatedly been shown in recent works, such as Deja Vu and PowerInfer, to directly lead to wall-time reductions as sparsity increases. Thus, maximizing FLOPs reduction through activation sparsity remains a justified and practical goal, with its potential for GPU speedups already validated in prior studies.

Summary

We would like to summarize why the lack of GPU-specific speedup results does not diminish the contribution of our work:

1. Core Contribution: The primary focus of our work is on modeling innovations to induce activation sparsity, not on system optimization.

2. Established Practicality: Prior works like Deja Vu and PowerInfer have already demonstrated that activation sparsity translates to practical GPU speedups, validating the relevance of our sparsity-focused approach.

3. CPU Speedups: We have shown practical acceleration on CPUs, providing further evidence of the real-world benefits of activation sparsity.

4. Broader Impact: Activation sparsity has implications beyond efficiency, informing theoretical studies on generalizability, learnability, and interpretability, as discussed in Section 5.

5. Avoiding the Hardware Lottery: Relying solely on results for today's hardware risks undervaluing ideas that could shape future hardware and software ecosystems.

We thank the reviewer once again for their invaluable feedback, which has significantly enriched the discussion of our work.

Thank you for taking the time to read and review our paper!

Re: Speedup on GPU / TPUs

While our experiments focused on CPU evaluation, we believe that Spark Transformer can be effectively implemented on GPUs to achieve significant speedups. Recent work like Deja Vu and PowerInfer have demonstrated that dynamic, unstructured activation sparsity, similar to that in Spark Transformer, can be efficiently implemented on GPUs for Transformer architectures, leading to considerable performance gains. We anticipate similar benefits for Spark Transformer given the shared sparsity characteristics. However, our primary focus was on addressing the modeling challenge of achieving high activation sparsity in LLMs without sacrificing quality. As demonstrated in Table 4, Spark Transformer attains higher activation sparsity with lower quality loss compared to existing methods. This contribution lays the groundwork for future efficient GPU implementations, building upon the aforementioned advancements in sparse GPU acceleration.

Re: Using the same model as ProSparse for a fair comparison

We acknowledge that comparing with ProSparse/ReGLU using the same baseline model would be ideal. However, ProSparse/ReGLU focus on fine-tuning existing models, while our method introduces a new architecture designed for pre-training from scratch. This fundamental difference makes a direct comparison on Llama2-7B challenging, as the pre-training data required to train a Spark Transformer version of Llama2-7B is not publicly available.

Re: Comparison to token pruning

We have added a short discussion of KV cache pruning methods, such as Heavy-hitter oracle and Scissorhands, to the related work section. These methods drop less-used tokens from the KV cache during decoding. This approach has the drawback that dropped tokens are permanently lost, even if they become useful later in the decoding process. This limitation may explain why these methods only achieve a 5x compression rate. In contrast, our method selectively retrieves information from the entire KV cache. Notably, we showed that Top-256 retrieval is sufficient regardless of the context window length, enabling a much higher compression rate with larger contexts.

Re: explaining L065

We appreciate the reviewer's attention to detail. L065 is explained in the paragraph "FLOPS per Token" in Section 3.1. We have added a pointer to this explanation in footnote 1 in the revised manuscript to clarify how L065 is obtained.

Thank you so much for your time, insightful feedback, and the positive evaluation :-)

Re: table lookup is not friendly to accelerators

We appreciate the reviewer's observation. While our experiments focused on CPU evaluation, we believe that Spark Transformer can be effectively implemented on GPUs to achieve significant speedups. Recent work like Deja Vu and PowerInfer have demonstrated that dynamic, unstructured activation sparsity, similar to that in Spark Transformer, can be efficiently implemented on GPUs for Transformer architectures, leading to considerable performance gains. We anticipate similar benefits for Spark Transformer given the shared sparsity characteristics. However, our primary focus was on addressing the modeling challenge of achieving high activation sparsity in LLMs without sacrificing quality. As demonstrated in Table 4, Spark Transformer attains higher activation sparsity with lower quality loss compared to existing methods. This contribution lays the groundwork for future efficient GPU implementations, building upon the aforementioned advancements in sparse GPU acceleration.

Re: how does each of the two techniques (top-k and predictor) contribute to final quality

We appreciate the reviewer's insightful question. To address this directly, we have added additional ablation studies in Section D.4 of the revised manuscript. Our results demonstrate that applying top-k alone without the low-cost predictor leads to a decrease in model performance. This suggests that the low-cost predictor, while introduced primarily to reduce the FLOPs required for locating nonzero entries, also plays a crucial role in making the model more amenable to sparsification and maintaining high performance.

Re: are the two techniques effective for fine-tuning

Our method requires architectural changes beyond simply switching the activation function as in many previous works, hence cannot be directly applied upon a pretrained model. We recognize that this may be seen as a shortcoming of Spark Transformer as it is not suited for cases where one is seeking to adapt an existing model to having activation sparsity. However, the goal of this paper is to demonstrate that Spark Transformer can be used as a drop-in replacement of the standard Transformer architecture (see line 082), and we advocate adopting Spark Transformer as the architecture in future pretrained models. In that case, the activation sparsity comes entirely for free from future models, which has a big advantage over sparse fine-tuning based approaches which not only incurs training complications but also cannot achieve the same level of FLOPs reduction as ours.

Re: HumanEval of Gemma 2 lower than published ones

Thanks for noting this. We have reached out to the authors of Gemma 2 for understanding this discrepancy and we are looking to resolve it.

Re: statistical topK when activation is sharded

The training of Spark Gemma-2 2B in our paper does not involve sharding the activations upon which statistical top-k is applied to. However, we acknowledge that sharding activations is a valid practical concern and have added a discussion in Appendix F.2. Statistical top-k can be easily extended to this scenario with minimal computational overhead and device communication. Each device can compute the local mean and variance of its activation entries and then broadcast these two numbers to all other devices. This allows each device to efficiently compute the global mean and variance. The communication cost is minimal, as each device only needs to send and receive 2 * #devices floating point numbers per activation vector, which typically contains thousands of entries.

Thank you for taking the time to read and review our paper!

Re: Q1 - Novelty

On novelty compared to [1, 2, 3]: We would like to thank the reviewer for bringing Refs [1, 2, 3] to our attention. While we acknowledge that they are related works, our work differs significantly by applying and adapting statistical top-k to activation sparsity within Transformer architectures. We elaborate on this point below.

• We are the first to use statistical top-k for enforcing activation sparsity in Transformers. Improving Transformer efficiency via activation sparsity has become a very popular research topic (see Section 5), but may have been suffering from a lack of efficient top-k algorithms for enforcing sparsity. Hence, the introduction of statistical top-k may facilitate the development of this area.

• Synergizing statistical top-k into Transformers is nontrivial. Since the method of statistical top-k is based on fitting a statistical distribution to the activation vector, it is necessary to understand the distributions of different activations in order to determine which particular activation vector is suited for the application of statistical top-k and the associated choice of distribution. In our case, we decided that statistical top-k should be applied to the activation before the nonlinear function (for FFN) and before softmax (for Attention) since entries of this vector provably follow a Gaussian distribution at random initialization. We also verify empirically that statistical top-k is still reliable even after initialization (see Figure 1).

• We extend statistical top-k from using the hard-thresholding operator with the estimated statistical threshold to using the soft-thresholding operator. This leads to a continuous optimization landscape that may have facilitated the optimization. Empirically, we found this choice to provide quality benefits for Spark Transformer.

• We provide the first theoretical justification for the correctness of statistical top-k; see Theorem 1.

• We reveal the conceptual connection between statistical top-k and several related top-k operators in the literature; see Section 2.3. Such connections may motivate the development of more powerful top-k algorithms in the future.

On novelty compared to [6]: We also thank the reviewer for pointing us to Ref. [6], which uses an idea similar to the reference (Yang et al., 2024) that we discussed in the introduction. Both [6] and Yang et al. (2024) use a subset of the key dimensions to estimate the important keys, which is similar to Spark Transformer. However, identifying those dimensions requires an arguably complex procedure, and those dimensions are input-dependent, so they need to be recomputed each time. In contrast, we simply use a predetermined set of dimensions that can be used for all inputs.

In the revised manuscript, we have added discussions of Refs [1, 2, 3] on lines 103-105 and an extended discussion in Appendix F. We have also added a reference to Ref [6] on line 76.

Re: Q2 - Practically in terms of speedup over comparing method

We appreciate the reviewer's question regarding the practical speedup of our method. It's important to recognize that the final speedup observed in practice is influenced by two key factors:

1. FLOPs reduction without sacrificing model quality: This is the primary focus of our work. We introduce Spark Transformer as a novel architecture designed to maximize FLOPs reduction while maintaining model quality. To address the reviewer's comment on comparisons with related methods, we have added Table 4 in Appendix F, which lists the FLOPs reduction and quality loss of various approaches. This table demonstrates that Spark Transformer achieves state-of-the-art FLOPs reduction with minimal quality loss.

2. ML system implementation that effectively leverages sparsity: While a comparison of the final speedup would be informative, it's important to note that such comparisons are difficult to make across different papers. The final speedup is highly dependent on various factors such as the specific hardware used (CPU, GPU, TPU), hardware design choices (HBM vs. compute tradeoff, tensor core vs. sparse core), and the specific implementation details. These factors vary significantly across research works, making direct comparisons challenging. In fact, none of the prior works in this area provide such a comparison for this very reason.

FLOPs reduction, on the other hand, offers a fair and consistent comparison across methods that utilize the same form of activation sparsity (as in Table 4). This is because, given the same hardware and implementation, methods with larger FLOPs reduction are likely to exhibit larger speedups. Our focus on maximizing FLOPs reduction while preserving model quality provides a strong foundation for achieving significant speedups in practice when combined with optimized sparse implementations on various hardware platforms.

Dear Reviewer uU7z,

Thank you again for your valuable feedback, which has significantly improved our submission. In response to your comments, we have added several discussions and experimental results in the Appendix, highlighted in red in the revised manuscript.

As the discussion phase is ending soon, we wanted to kindly follow up to confirm if our response has addressed your concerns. Please don't hesitate to let us know and we are happy to answer any additional questions or address any concerns you may have.

Thank you for your time and thoughtful input.

Best regards, The Authors