Adaptive Computation Pruning for the Forgetting Transformer

Below we show FLOP savings, attention runtime reduction, training throughput improvement provided by ACP for FoX trained with 48B tokens. We also show results for models trained with 16B tokens for comparison. The results are as follows:

As shown in the results, the improvements for 48B-token models are slightly less but still signficant.

Furthermore, the original FoX paper provided evaluation on LongBench, which is not included here, missing the way for a more comprehensive and fair comparison.

LongBench is a generation-based benchmark for instruction-tuned models, so the results can have high variance for small base models that have not been instruction-tuned. Therefore, it is less suitable for our purposes. To demonstrate this, we have repeated three runs for the 760M-param/16B-token/4k-context length setting with unpruned FoX:

Note the large variance across seeds especially for some tasks such as Qasper and Samsum.

For completeness, we have run LongBench evaluation for the 760M-param models trained with 16B tokens with different context lengths. The results are follows:

As shown in the table, the results for models trained with 16B tokens with and without ACP remain similar, with some variations in some tasks. As discussed above, this is likely due to variance across training runs. Note that LongBench results for models trained with 48B tokens are also included in the response to your last concern regarding dataset size.

We hope our response addresses your concerns. If so, we would greatly appreciate it if you could raise your score. Feel free to let us know if you have more questions!

The experimental section lacks comparative baselines. The results only compare the proposed method to unpruned FoX, making it difficult to assess the overall competitiveness or practical advantage of ACP. It's possible that its performance remains below that of existing pruning methods.

Our method is designed specifically for FoX so it is natural to measure how much it speeds up (unpruned) FoX. As mentioned before, most "pruning" methods either deviate from standard softmax attention or are only designed for inference of already pretrained models, so they are not comparable to our method which is applicable during pretraining. Furthermore, our method is guaranteed to be virtually lossless (note our default $\varepsilon=e^{-10}\approx 0.000045$ is even below the roughly $0.001$ to $0.01$ relative error introduced by the BF16 mixed training) and only requires minimal modifications to FlashAttention, so in principle it could be combined with any other (lossy) pruning methods. In this sense, ACP can be seen as orthogonal to all other pruning methods.

Some reported results in this paper appear questionable. The FoX (Pro) result on HellaSwag is around 33 in this paper, while the original FoX paper reports a score of 38.39. It's unclear why a different dataset size was used.

As you mentioned, the results are different because we use 16B tokens for the 760M-parameter models, while the main results in the FoX paper use 48B tokens. This choice is mainly due to the large number of experiments we need to run (combinations of different model sizes, context lengths, with and without ACP) and our limited computational resources. Also note that our set of configurations (125M-param/2.7B-token, 360M-param/7.5B-token, 760M-param/16B-token) are from the ablation experiments in the FoX paper which aligns with the Chinchilla scaling law.

For completeness, we have also trained a 760M FoX model for 48B tokens with ACP. We compare the results with the official results as follows. In the following results marked "FoX (Pro) w/o ACP" are either taken from the FoX paper or are reproduced using their official checkpoint.

Per-token loss at different token positions:

Language Model Evaluation Harness:

LongBench

As shown in these results, the per-token language modeling loss almost match exactly and the downstream task results are similar. As mentioned in our paper and also will be discussed later, downstream task performance can vary greatly across training runs so it is not possible to get exactly the same results. We also run the needle-in-a-haystack test and FoX with ACP achieves perfect accuracy, similar to FoX without ACP.

There is no discussion of how choices of block sizes (Bq, Bk) and resulting grid dimensions (M, N) affect the granularity of pruning versus the overhead of boundary detection .

Typical values of $B_q$ and $B_k$ are 32, 64, 128 (chosen based on profiling), which are limited by GPU's on-chip shared memory size. This is very small compared to typical sequence lengths such as 4k and 16k, so their impact on pruning granularity is negligible. As mentioned before the complexity of boundary index search is $O(\max(M, N))$ where $M = L/B_q$ and $N=L/B_k$, which is negligible compared to the $O(Ld^2)$ complexity the attention computation regardless $B_q$ and $B_k$. In the following tables we show the FLOP savings (which depends on pruning granularity) and the fraction of runtime spent on the boundary index search algorithm within the attention kernel for different ($B_q$, $B_k$), for a 760M-parameter model trained with a context length of 4k tokens. Note that we separate the forward pass and backward pass because they use different block sizes.

Forward pass results are shown as follows:

Backward pass results are shown as follows. In our implementation the gradient computation for queries and keys/values use different kernels and have different block sizes, so there are two pairs of $(B_q, B_k)$ in each column (one for the query grad kernel and one for key/value grad kernel).

As can be seen in these results, the block sizes have little impact on boundary search overheads and FLOP savings.

We hope our response addresses your concerns. If so, we would greatly appreciate it if you could raise your score. Feel free to let us know if you have more questions!

Can you quantify end-to-end speed-ups and memory savings during pre-fill and autoregressive decoding when ACP is applied to KV-cache eviction, and how do these gains compare with the training-time numbers?

In our implementation, prefilling is a simple forward pass, so the improvements are comparable to those for pretraining (e.g., 34% improvement in end-to-end prefilling speed compared to the 43% pretraining throughput improvement).

For decoding, note the percentage of memory IO savings provided by ACP is the same as the percentage of FLOP savings because the kernel does not need to load pruned KV blocks. In our preliminary results with our simple FlashAttention decoding implementation, we observe that memory IO savings and attention kernel runtime during decoding are comparable to those during pretraining. However, we do not observe a clear improvement in end-to-end decoding throughput. Our analysis shows that this is due to the large CPU overheads during decoding, which most likely come from Triton/CUDA kernel launch costs. These overheads overlap with the attention kernel execution, so even if the runtime of the attention kernel is greatly reduced, the CPU overheads still persist (as seen from the low GPU utilization), and thus there would be no throughput improvement.

We also point out that even though our implementation saves memory IO, it does not perform explicit KV-cache eviction so it does not reduce GPU memory consumption. KV-cache eviction is technically possible but is very tricky to implement with our current implementation.

Below we report the memory IO/FLOP savings and attention kernel runtime reduction during decoding for different model sizes and context lengths $L$. The prefill-length is $(7/8)L$ and we decode the last $L/8$ tokens. For 125M models we use 256k tokens per batch and for 360M and 760M models we use 128k tokens per batch.

Memory IO/FLOP savings are shown in the following table. The savings are generally larger than those for pretraining likely because decoding happens at the end of the sequence so there are more KV entries that could pruned at the decoding stage.

Attention runtime reduction are shown in the following table. Note we fix the number of tokens per batch and the runtime is for a single decoding step, so the runtime stays roughly constant across different context lengths.

We emphasize that these are preliminary results. For future work, reducing overheads such as Triton kernel launch costs or simply using a larger model (so that we won't be bottlenecked by CPU overheads) should allow us to translate the speedup of the attention kernel to throughput improvement.

scalability is only tested up to 760 M parameters on a single 4 × L40S node, so behavior on multi-billion-parameter models and distributed settings remains unknown.

This is indeed a limitation due to our limited computational resources. However, given that the improvements are fairly consistent across the three small scales we tested, it is reasonable to expect similar results for larger scales.

What fraction of total attention runtime is spent on the boundary-search algorithm at sequence length 16 k on L40S versus A100 GPUs

Below we report the fraction of time spent on boundary search within the attention kernel for different context lengths and model sizes (i.e., boundary index search time divided by the total runtime of the attention kernel)

As shown in the table, the time spent on boundary index search is always below 6%. The result for A100 is similar (always below 7%). As mentioned in the General Response, the complexity of boundary index search is $O(\max(M, N))$ as opposed to $O(MN)$, where $M = L/B_q$ and $N=L/B_k$. In other words, the complexity is linear w.r.t. the sequence length $L$ (note in Algorithm 1 we never reset $l$ to 1 in the for loop). This is negligible compared to the quadratic $O(L^2 d)$ costs of the actual attention computation.

We hope our response addresses your concerns. If so, we would greatly appreciate it if you could raise your score. Feel free to let us know if you have more questions!