UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression

We are pleased to see your recognition of our work. Regarding the concerns you raised, we would like to offer a detailed clarification, which we believe will help you better understand the core contributions and novelty of our approach.

Weakness 1: Logical soundness and novelty of the method

We understand the reviewer's emphasis on chunk-wise schemes for long-context scenarios. However, there may exist two potential key misunderstandings in this comment, which we wish to clarify here:

1. In the training stage, backpropagation still requires retaining the computational graph and intermediate activations even with the chunk-wise input. Therefore, chunk-wise training cannot reduce peak GPU memory usage. Furthermore, the chunk-wise forward typically leads to a complex computational graph and severe memory fragmentation. This reduces hardware utilization and hinders training speed. In contrast, UniGist's proposed one-pass full-sequence training, combined with the unified sparse gist layout, not only enhances context compression quality but also yields significant training efficiency advantages through better kernel design.

2. UniGist does not mandate full-sequence processing during inference. We fully agree that managing peak memory is a practical challenge during inference, and we wish to emphasize that UniGist fully supports chunk-wise inference and in a more flexible manner. Specifically: (1) During the prefill phase: Users can flexibly set the chunk size and apply sparse attention accordingly. As shown in Section 4.3 and Figure 4, UniGist's gist layout is inherently compatible with different chunk sizes without compromising compression quality. (2) During the decoding phase: UniGist offers a more efficient compression strategy than existing methods. Compression and kv cache eviction can occur after every r generated tokens (where r is the compression ratio, e.g., 4 or 8), rather than waiting for a full chunk (e.g., 2048 tokens) in previous gist methods. This mechanism reduces memory pressure during decoding and boosts throughput.

Therefore, we wish to clarify that the core design principle of UniGist is to provide a more efficient and memory-friendly solution for both training and inference while simultaneously improving compression quality, rather than simply avoiding the challenges of chunk-wise processing. We believe this co-design of methodology and optimizations offers a more fundamental and efficient solution than previous approaches, which constitutes the core contribution and novelty of our work.

Weakness 2: Experimental details in Figure 1c and Table 5

Thank you for raising this point. We are happy to provide more detailed clarifications:

• Regarding Figure 1c: As mentioned in Lines 127–133 of the paper, we evaluate a compression model trained using the traditional chunk-wise method. We randomly sampled 1k sequences of length 32K from the SlimPajama dataset. For each token position in a sequence, we computed the LM loss. To smooth the curve and highlight local trends, each data point in the figure represents the average loss over a sliding window covering the token itself and 16 neighboring tokens on both sides.
• Regarding Table 5: The detailed hardware and model configurations have been provided in Appendix C under “Speed-up Details” (Lines 550–553). We used the standard triton.testing.do_bench function for accurate performance evaluation. This function performs multiple warm-up runs to eliminate initialization overhead, followed by several benchmark runs whose results are averaged to ensure stable and reliable latency measurements.

We will include these implementation details in the final version of the appendix to further enhance the transparency and reproducibility of our work.

Weakness 3: Maximum Context Length and Performance on a Single GPU

To further address your concerns regarding the long-context processing capability on a single GPU, we conducted additional experiments. We benchmarked the peak GPU memory usage of different strategies for processing extremely long contexts on a single 40GB GPU. All experiments were based on the Llama-3.1-8B-Instruct model. The results are presented below:

The results clearly demonstrate that UniGist's KV Cache compression technique significantly extends the maximum context length manageable on a single GPU. Specifically, with a compression ratio of 4, UniGist can process contexts up to 320K on a single 40GB GPU (using 36.7GB of memory), suggesting it can support lengths approaching 340K. When the compression ratio is increased to 8, processing a 512K context requires only 37.7GB of memory, indicating that the supported context length can extend well beyond 512K. In contrast, a standard Full Attention model under the same hardware configuration can only handle a maximum context of approximately 80K. This comparison robustly demonstrates the practical value and substantial advantages of UniGist in extending the long-context processing capabilities of models.

Thank you again for your valuable time and constructive feedback. We hope that our detailed responses and specific revision plan have clearly addressed your concerns. We appreciate your consideration and look forward to your feedback.

Thanks for your insightful feedback. We are pleased that you recognized the motivation and empirical strength of our work. In response to the points you raised, we would like to take this opportunity to offer further clarification:

Weakness 1: Block-wise sparse attention has been studied in numerous previous works [...] it would be better to clarify the unique characteristics of the attention pattern in this paper to better validate the novelty of the work.

We fully understand the concern: given the extensive prior research on sparse attention, does the attention pattern used in UniGist constitute a genuine innovation? We'd like to clarify that our key innovation lies not in the sparse pattern itself, but rather to enable effctive and efficient context compression, which fundamentally differentiates our approach from conventional sparse attention methods.

1. UniGist’s goal differs fundamentally from prior sparse attention works (e.g., Sparse Transformer, Longformer, BigBird). Those methods aim to approximate full attention more efficiently, but still require caching all original tokens' kv cache during inference, thus failing to truly alleviate memory pressure in long-context scenarios. In contrast, UniGist explicitly discards original tokens after compressing them into a small number of learnable gist tokens, enabling sequence-level compression that substantially reduces KV cache size during inference.
2. Moreover, the attention pattern in UniGist is uniquely constructed around connections to the gist tokens, rather than being fixed to blocks or sliding windows. This offers flexibility to support context compression with adjustable chunk sizes as demonstrated in Section 4.3 and Figure 4. This is not feasible with traditional sparse attention designs.

Therefore, the novelty of our work lies not in the geometric shape of the sparse pattern itself, but in its effective and efficient compression characteristics. We will make this fundamental distinction clearer in the revised version of the paper to address your concern.

Weakness 2: The main baselines for Unigist are approaches that adopt a chunk-wise strategy. Although they result in a loss of contextual information, their primary advantages lie in acceleration and memory savings. However, the paper does not compare with them in these two aspects.

We appreciate your insightful observation regarding the efficiency strengths of chunk-wise approaches. However, we would like to clarify a potential misunderstanding: UniGist is fully compatible with chunk-wise compression during inference, and achieves comparable or even better acceleration and memory efficiency.

To comprehensively address this concern, we believe it is important to distinguish between the training and inference stages, as UniGist introduces fundamental improvements in both.

1. Training Stage: UniGist achieves better compression quality and training efficiency. A common misconception is that chunk-wise strategies also save memory or accelerate training. In practice, this is often not the case. Backward propagation requires the computation graph and intermediate activations for the full sequence, meaning chunk-wise training can not reduce peak memory usage. On the contrary, iterative forward passes on the same module usually create complex computation graphs and severe memory fragmentation, which degrades hardware utilization and lowers training throughput. UniGist addresses these issues through its one-pass training paradigm and our custom hardware-aligned kernel, achieving higher compression quality alongside faster training speeds.

2. Inference Stage: UniGist not only supports chunk-wise processing but also surpasses prior methods in both flexibility and efficiency. (1) UniGist is fully compatible with a chunk-wise strategy during the prefilling stage to manage peak memory. This ensures that for the same chunk size, UniGist's peak memory usage is comparable to that of previous chunk-wise methods. More importantly, thanks to our decoupled design, the chunk size can be flexibly configured at inference time without any need for retraining. (2) Prior methods typically perform full attention within each chunk. UniGist, even during chunk-wise prefilling, can leverage our designed sparse attention kernel for acceleration, making it more efficient than baselines even in this step. (3) Previous methods must accumulate to a full chunk (e.g., 2048 tokens) before performing a single compression step. In contrast, UniGist can perform compression and discard the corresponding KV cache immediately after every r new tokens are generated (where r is the compression ratio). This fine-grained compression mechanism reduces peak memory, computation, and memory access during the decoding phase.

To fully address your concern, we have conducted a new end-to-end efficiency experiment. We benchmark UniGist against two representative baseline methods using the Llama-3.1-8B Instruct model in a 64K context (with a chunk size of 2K). As shown in the table below, UniGist achieves the best performance on key metrics such as TTFT, clearly demonstrating its superior overall efficiency:

We will strengthen this analysis in the revised manuscript and include the new quantitative comparison in the appendix to provide clearer evidence of UniGist’s efficiency advantages.

Thank you again for your valuable time and constructive feedback. We hope our detailed responses and the planned revisions have thoroughly addressed your concerns. We appreciate your consideration and look forward to your feedback.

Thanks for your review and constructive feedback. We are very pleased that you recognize the motivation for our work and the strong performance of our method. Regarding the concerns you have raised, we appreciate this opportunity to offer more in-depth clarifications and additional results.

Weakness1: concerns about paper presentation

• On the color annotation of Figure 1(b): Thank you for pointing this out. In the original figure, the color brightness represents the magnitude of the attention weights, where brighter colors indicate higher weights. This is consistent with the standard settings of Python's Matplotlib plotting. We will add a color bar and corresponding legend to Figure 1(b) in the final version to clearly and intuitively depict the distribution of attention.

• On the details of supervised fine-tuning and the synthetic dataset: This might be a misunderstanding. In Appendix A.3 (Training Data Processing), Lines 534-537, we have already mentioned that we primarily use the Magpie dataset and additionally integrated samples from tasks such as LongAlpaca and BookSum, refering to the approach of Activation Beacon. To fully address your concerns and eliminate any ambiguity regarding the composition of our training data, we will make two key additions to Appendix A.3 in the revised version: (1) We will explicitly list the specific number of samples and mixing ratios from each data source. (2) We will use more prominent formatting and a dedicated table to clearly present the data sources, sample counts, and primary task types.

Weakness2: concerns about the ablation study

This is a very insightful comment, and we fully agree that reporting task-level performance is a more intuitive way to demonstrate the importance of each component. That said, we would like to clarify that the manuscript already includes ablation results and discussions related to key components. Our responses are as follows:

• Comparison with Activation Beacon as an ablation for gist attention: As noted, the comparison with Activation Beacon can be interpreted as a holistic ablation. Beacon adopts a traditional chunk-wise training strategy, whereas our method introduces a unified attention mechanism. As shown in Table 1, UniGist significantly outperforms Beacon across multiple tasks. This highlights the superiority of our overall framework.

• Local window ablation is already included: The effect of local window size is illustrated in Figure 6 (Right), which shows its direct impact on performance in the RULER task. To further improve clarity, we will augment this figure with additional performance data from other tasks.

• Training stability is a critical prerequisite for performance: Training stability is essential for learning effective compression. As shown in the left and middle parts of Figure 6, removing sink tokens leads to frequent gradient spikes and unstable convergence, which directly harms downstream performance. To further address your concern, we conduct new experiments by continue pretraining a Llama-3.2-3B-Instruct model for 4B tokens and evaluate its long-context performance. We can observe that introducing sink tokens generally leads to better performance on the final tasks based on the following table. These discussions and new experimental results will be added to the appendix in the revised manuscript.

Weakness3: concerns about efficiency (local window and comparison with previous gist models) We provide responses to both aspects of your concerns below:

Part A: The computational overhead introduced by the local window is negligible

• Theoretical Analysis: The local window size is typically set to a small value (e.g., 32, 64, 128) to achieve the desired performance gains. In our block-wise sparse attention kernel (with block size typically set to 128), this only introduces one additional key-value block computation at the per-query block level. Compared to the entire sequence length, especially in long-context scenarios, this overhead is almost negligible.
• Empirical Validation: To address your concern more thoroughly, we conducted a new experiment measuring the attention forward and backward pass latency with and without the local window for a 64K-length sequence as shown in the following table. The results strongly demonstrate that the overhead introduced by the local window is negligible, whereas its performance benefits are substantial. This validates the superiority of our design.

Part B: Compared to previous GIST methods, UniGist also holds a significant efficiency advantage during inference.

UniGist demonstrates significant inference efficiency gains over previous gist methods in two key phases:

• Faster Prefill Stage: Our method allows for more flexible chunk-wise processing during inference. In previous gist methods, attention within each chunk was computed using a full attention. In contrast, UniGist leverages a custom sparse attention kernel to perform accelerated prefill by attending only to sink, gist, and local tokens within each chunk. This results in faster prefill speeds for the same chunk size.
• More efficient Decoding Stage: Previous gist approaches typically wait until an entire chunk (e.g., 2048 tokens) is generated before performing compression. In contrast, UniGist can insert a new gist token and discard raw tokens every r tokens (where r is the compression ratio, e.g., 4 or 8), enabling fine-grained, more frequent compression. This leads to two major advantages: (1) Lower peak KV cache usage: We can release unused raw token cache earlier, maintaining lower memory usage throughout generation. (2) Reduced computation: Fewer KV are involved in attention computation, reducing both memory access and compute overhead for each next-token prediction.

In summary, the local window introduces negligible overhead, and UniGist achieves superior inference efficiency through accelerated prefill and fine-grained decoding compression compared to prior gist methods. We will include this detailed analysis and new comparative experimental results in the final version of the Appendix.

Question 1: Why did previous methods use chunk-based training? Is it for memory savings? Is your method more memory-intensive than other GIST-based methods?

We would like to clarify a key point: most prior methods adopted chunk-wise training primarily for ease of implementation, not for memory savings. They could reuse existing, highly-optimized kernels like FlashAttention for the full attention within each chunk, thus reducing implementation complexity.

However, this design choice is a compromise, made at the cost of model performance, training efficiency, and inference flexibility. Furthermore, chunk-wise training does not actually save memory during training. The backward propagation requires retaining the computation graph and intermediate activations for all chunks, meaning the GPU memory occupied by a chunk cannot be freed immediately after it is processed.

Specifically, chunk-wise training introduces the following limitations, which our work aims to address:

• Performance Limitations: As illustrated in Figure 1(b), chunk boundaries bring unnecessary obstacles for learning compression. It leads to poor compressed context quality and imbalanced loss within chunks.

• Training Inefficiency: Iteratively performing forward passes on the same Transformer block creates a complex computation graph and leads to significant memory fragmentation, which reduces throughput and hardware efficiency.

• Inference Inflexibility: The chunk size used during inference is limited by the training configuration, making it difficult to adapt to different needs.

Under identical settings, UniGist is more memory-efficient and faster in both training and inference compared to previous gist methods, as detailed in our response to Weakness 3, Part B. In short, we believe UniGist not only delivers significant improvements in model performance (as shown in Table 1) but also offers a more efficient and flexible compression to the previous chunk-wise approach.

Thank you again for your valuable time and constructive feedback. We hope that our detailed responses and specific revision plan have clearly addressed your concerns. We appreciate your consideration and look forward to your feedback.

We are grateful for your thoughtful review and constructive suggestions. We are particularly encouraged that you recognize our strengths. We would like to take this opportunity to offer further clarification and address your concerns:

Weakness 1: The choice of baselines is inadequate.

First, we would like to clarify the core difference between our method and sparse attention methods.

• Our method, UniGist, is a context compression technique. Its central idea is to replace the KV cache of a large number of original tokens using learnable "gist" tokens. This means the model can no longer access the original tokens after compression, which fundamentally reduces memory requirements.
• In contrast, methods like Minference or Deepseek's NSA, which you mentioned, are essentially lossless sparse attention. While they employ various strategies to limit the number of key tokens each query token can attend to, the complete KV cache is still retained in memory. In principle, any original token remains accessible to future tokens.

Therefore, a direct performance comparison between these two types of methods would be unfair and out of scope in some extent. Sparse attention methods naturally have an advantage in detail-recalling tasks, whereas our method excels in memory efficiency.

Despite this fundamental difference, we fully understand your desire to see how our work competes in a broader context. To address your concern and demonstrate the comparative effectiveness of our approach, we conduct new experiments comparing UniGist with MInference as a representative (based on Llama3.1-8B-Instruct):

We plan to add these new experimental results and a detailed discussion to the Appendix. This will allow us to fully address your concern.

Weakness 2: Lines 295-301 only report forward / backward latency in an attention layer, rather than the actual runtime on specific accelerators, or the latency-throughput relationship.

First, we would like to clarify our rationale for initially choosing to evaluate attention latency. As stated in lines 295-296 of our paper, our method's primary improvements are in attention efficiency and GPU memory usage. Therefore, isolating the attention layer for evaluation allows for an objective and clear demonstration of the direct speed-up provided by UniGist, avoiding interference from other model components. We believe this "controlled experiment" is crucial for validating the core contribution of our method.

At the same time, we would like to point out that in our submitted version, we have already provided the raw latency data (in ms) on actual hardware in Appendix C, Table 5, which serves as a foundation for end-to-end performance analysis.

To thoroughly address your concerns and provide a more comprehensive view of UniGist's practical efficiency advantages, we follow your valuable suggestion and conduct additional end-to-end performance evaluations on the Llama-3.1-8B-Instruct model and 64K context (output 128 tokens). Below is an example of the supplementary results:

We plan to add these experiments and discussion to Appendix C in the revised manuscript. This will provide a more complete performance profile for our method.

Minor1: Overall, the introduction would benefit from identifying different phases, their efficiency bottlenecks, and how different strategies (including but not limited to gist-based models) may overcome each of them specifically.

Thank's for your suggestion. We agree that a clearer distinction between the efficiency bottlenecks would significantly strengthen the introduction. In the revised manuscript, we will refine the second paragraph to address this explicitly. The new discussion will differentiate between the compute-bound pre-filling phase and the memory-bandwidth-bound decoding phase. Furthermore, it will clarify that the benefit of reducing the KV cache extends beyond simply saving memory; critically, it enhances decoding throughput by minimizing the data transferred from HBM to SRAM in each step. This enhancement will make the motivation for our work more precise.

Minor2: “Previous empirical study shows […]” (line 102): This statement seems to warrant some citations.

Thank you for pointing this out. The statement is primarily based on the findings from works such as Activation Beacon [54] and the Gist Study [11], which verified that this interleaved gist token insertion is an effective baseline. We will add these citations to the revised manuscript.

Minor3: On the suggestion of using distillation for continued training

This is a very insightful suggestion! In our current work, we adopt the standard continued pre-training approach to maintain consistency with the training paradigm of related works, such as Activation Beacon. We will add a discussion of this strategy in the Appendix or in the future work section. We will acknowledge that distillation could be a more robust training method and frame it as a promising direction for future research.

Minor4: On the lack of detail for the mixed instruction data

This appears to be a misunderstanding. In the current version of the paper, we mention our supervised fine-tuning data composition in Appendix A.3 (Training Data Processing, Lines 534-537), noting that it largely follows the setup of Activation Beacon and providing the corresponding citation. To address your concern and eliminate any ambiguity, we will make two key additions to Appendix A.3 in the revised manuscript: 1. We will explicitly list the specific sample counts and mixing ratios from each data source (e.g., Magpie, LongAlpaca, and BookSum). 2. We will use a dedicated table to clearly present the data sources, sample counts, and primary task types.

Question 1: Regarding the compatibility of position encoding alignment with RoPE

Our position encoding alignment strategy is fully compatible with RoPE. The core of RoPE is to apply a rotation matrix based on a token's absolute position index. Our strategy (as shown in Figure 2c) adjusts these position indices before they are used to compute the RoPE embeddings. Specifically, for an original token $x_i$, we still use its original position index $i$. The relative position difference between any two original tokens $x_i$ and $x_j$ remains $j-i$, perfectly preserving their relative relationship. The gist tokens simply "borrow" the position of a subsequent token without disrupting the structure of the original token sequence. We will clarify this mechanism more explicitly in Section 4.1.

Question 2: Regarding the initialization of Gist Tokens

Thank's for this question about the implementation detail. Gist tokens are learnable virtual tokens. In practice, we randomly initialize their embeddings, typically from a standard normal distribution, which is similar to the initialization of other special tokens like [EOS]. This embedding is then learned and updated end-to-end along with the other model parameters during training. We will add discussion about this detail to Section 4.1.

Thank you again for your valuable time and constructive feedback. We hope that our detailed responses and specific revision plan have clearly addressed your concerns. We appreciate your consideration and look forward to your feedback.