UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
Authors: Chenlong Deng
Zhisong Zhang
Kelong Mao
Shuaiyi Li
Tianqing Fang
Hongming Zhang
Haitao Mi
Dong Yu
Zhicheng Dou
Venue: NeurIPS 2025 poster ()
Keywords: Long context compression
sparse attention
Primary area: deep_learning
TL;DR: We introduce UniGist, a unified gist token-based long context compression method without chunk-wise training, which significantly enhances long context retention and efficiency through a hardward-aligned design.

Abstract: Large language models are increasingly capable of handling long-context inputs, but the memory overhead of KV cache remains a major bottleneck for general-purpose deployment. While many compression strategies have been explored, sequence-level compression is particularly challenging due to its tendency to lose important details. We present UniGist, a gist token-based long context compression framework that removes the need for chunk-wise training, enabling the model to learn how to compress and utilize long-range context during training. To fully exploit the sparsity, we introduce a gist shift trick that transforms the attention layout into a right-aligned block structure and develop a block-table-free sparse attention kernel based on it. UniGist further supports one-pass training and flexible chunk sizes during inference, allowing efficient and adaptive context processing. Experiments across multiple long-context tasks show that UniGist significantly improves compression quality, with especially strong performance in recalling details and long-range dependency modeling.

UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
Chenlong Deng" lf,g Zhisong Zhang Kelong Mao', Shuaiyi Li², Tianqing Fang", Hongming Zhang", Haitao Mi², Dong Yu², Zhicheng Dou
Renmmn University of China Tencent Al "City University of Hong Kong fdengchenlong, dou) Cruc. edu. cn, zhisong. zhw.gecittng... edu.
Abstract
Large language models are increasingly capable of handling long-context inputs, but the memory overhead of key-value (KV) cache remains a major bottleneck for general-purpose deployment. While various compression strategies have been explored, sequence-level compression, which drops the full KV caches for certain tokens, is particularly challenging it can lead the loss of important contextual information. To address this, we introduce UniGist, sequence-level long- compression framework that efficiently preserves context information replacing raw tokens with special compression tokens (gists) in a fine-grained manner. We adopt chunk-free training strategy and design efficient kernel with gist shift trick, enabling optimized GPU training. Our scheme also supports flexible inference by allowing the actual removal of compressed tokens, resulting in real- time memory savings. Experiments across multiple long-context tasks demonstrate that UniGist significantly improves compression quality, with especially strong performance in detail- detail-recalling tasks and long-range dependency modeling.
Introduction
As large language models (LLMs) are applied to more sophisticated and demanding applications, the ability to process long-range context has emerged as a fundamental requirement [29, 59]. A wide range of real-world applications, such as understanding, require models to retain and reason over input sequences with extensive lengths [16, 60]. This has spurred growing interest in scaling up context windows from a few thousand tokens to hundreds of thousands or even millions, enabling models to incorporate more history, maintain coherence, and ground responses in broader contexts [23, 35, 46, 14].
generation and long-form document
Despite this progress, long-contest modeling still remains a challenge due to the inherent compute and memory bottlenecks of the underlying Transformer architecture [37]. In particular, the cost of self- attention scales quadratically with sequence length, making both training and inference increasingly resource-intensive the context grows [2, 50]. More critically, key-value (KV) caching, which is essential for efficient autoregressive decoding during inference, has also become major source of memory consumption [18, 36]. In long-contes settings, the memory required store KV caches can even exceed that of the model's parameters. These limitations highlight the need for effective KV compression techniques that can reduce computational and memory overhead without sacrificing the model's 's ability to retain and reason over essential contextual information.
This work focuses on sequence-leve compression, which directly reduces the number of token repre- sentations along the sequence dimension [43, 58). A partieularly promising and general sequence-
"Tii work done during internship Tencent Al Lab. Corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

level strategy is the gist token-based approach [27, 17, 26, 30, 6, 54, 11], which segments long sequences into chunks and inserts learnable virtual tokens (i.e., gists) to represent the original tokens. In principle, this method has potential for dramatically reducing sequence length, as gist tokens can be much fewer than raw tokens.
However, we observe two major drawbacks in existing gist-based methods. First, they suffer from notable degradation when required to recall information from distant chunks. We hypothesize that this is due to the prevalent chunk-wise training scheme, which introduces a shortcut letting tokens to minimize loss by attending only to nearby uncompressed tokens within the same chunk [6, 54]. As result, the model is less inclined to learn to integrate information across chunk boundaries, limiting its ability to reason over long-range dependencies. Second, the chunk-wiss design suffers from fragmented memory usage and complex computational graphs, which hampers overall training efficiency. In this training scheme, chunks are processed in an iterative style and separate memory allocations are required for each chunk, leading to potential under-utilization of hardware resources.
In this work, we address the above challenge by introducing UniGist, a unitied gist-based framework for context compression. Unlike prior approaches that rely on chunk-wise training, we remove this constraint and enable the model learn compression more effectively over the entire sequence through a unified sparse gist layout. To further enhance the efficiency of this layout, we propose gist shift trick, which is a hardware- aligned kernel design that transforms the irregular, vertically sparse attention pattern into hight-aligned block structure. This transformation aligns with the block-wise execution patterns of modern GPU architectures, significantly enhancing training throughput as well as supporting fast autoregressive inference. Experiments over wide range of long-context tasks illustrate the effeciveness and efficiency of the proposed approach.
Our main contributions are as follows:
1) We propose UniGist, a sequence- compression method with unified sparse gist layout, enabling effective and efficient long context modeling without chunk-wiss training.
2) We introduce the gist shift trick to align the proposed sparse attention layout with Prunnfiindly block structures, significantly improving training throughput and supporting efficient inference.
3) With experiments on variety of long-context tasks, we demonstrate that UniGist can provide improvements for context compression, with especially strong performance in detail-recallnng tasks and long-range dependency modeling.
2 Related Work
KV Cache Compression. Transformer-bases large language models mostly employ key-value caching to avoid repetitive attention computations during inference. However, as real-world appli- cations demand ever-increasing context lengths, the associated memory and bandwidth overhead has quickly become primary bottleneck for real-time systems. To tackle this, KV compression techniques are widely explored along four orthogonal dimensions (i.e, layer, head, token sequence, and numerical precision). Layer-level strategies typically skip certain transformer layers or share activations across layers [33, 34, 3, 32]. At the head level, multiple query heads can share a single key head, or the long-term cache of specific heads can be eliminated [1. 31]. Recent low-rank approaches further introduce a promising direction for this line [10, 56). Precisson-focusee methods apply custom quantization and control the floating-poin error they introduce to maintain the model's ability [24]. Moreover, some techniques within the layer, head, and precision dimensions demon- strate ncar-lossless performance on general long context tasks, making them broadly applicable in cutting-edge models.
On the other hand, sequence-level compression falls into two broad categories. Token eviction dis- cards older token activations in real time once a budget is exceeded. For example, StreamingL [43] retains only the initial and recent tokens, which irreversibly lose the details in evicted tokens. To miti- gate this, most of the following methods [58, 22, 47, 4. 12, 13] have to rely on hybrid eviction of layer or head level, or even resort to question heuristics rather than general compression. Token merging aggregates multiple tokens to reduce the overall cache size [57, 38, 40]. Models that introduce new "gist" tokens and adopt more training often achieve more stable compression [27, 6, 54]. Nonetheless, prior study shows that gist-based compression still struggles to retain highly complex details [1 l].

---

Lopeleeeinng Connet Teokookk
Salikkey dhweee al
120
Conperaso of Compremtin
Oepauiio Nokc
Mkde
D
180
Compreed Doves
Cenwuher caee
Irpatii! Damndax
code colyy steak oxqpoosed auee Clcck
Lest www.paaiio...ooooomm ac
The axepessdd ssedoradly rece sos lage itadii wights. mal woggh are atcaue unsundd verurin..
D
BDEE DDDG
136cx pwidic With 00ch dhachh (kg sesdec
Cuatted Soquste Iapa Cal Cueggoosko Archiecuure
(c) Altatian Weigh Dise Bankes Aualysi
66 Distiaaki wilkis Clasks
Figure 1: Analysis of previous architectures. (a) Previous approaches adopt chunk- wise strategy, which accumulates gist representations iteratively. (b) Our attention analysis reveals that tokens tend to ignore compressed contexts outside the current chunk. (c) Our token-position analysis indicates that tokens near the chunk ends may rely only on uncompressed tokens to minimize the target loss.
Our method addresses the key limitations of gist-baser architectures, achieving high-quality and general-purpose compression using sequence-level alone.
Sparse Attention. Sparse attention has been widely adopted to mitigate the quadratic complexity of self-attention by restricting computation to a subset of token interactions. Empirical studies show that attention maps in language models are inherently sparse, enabling full attention can be approximated closely. Early methods focus on designing fixed sparsity patterns and training models to adapt to them [2,51, 51, 50]. With large language models scaling, researchers find that models trained with full attention could preserve performance by simply adopting dynamically searched sparsity patterns without extra tuning [20, 42, 41, 44, 53]. More recent works train dynamic attention patterns from scratch [15, 25, 49], in some cases even outperforming full attention. However, to ensure generality across diverse inputs, state-of-the-en sparse attention methods still require the cache for each token to maintain reachability, Our method leverages the intrinsic sparsity of the gist token-based architecture, eliminating the need for distant original token caches and achieving substantial acceleration during both training and inference.
3
Why Does Typical Gist-based Compression Fail in Details?
3.1 Preliminary
Previous empirical study [11] shows that the most effective variant of gist token-based compression inserts learnable gist tokens in an interleaved manner within each fixed-length chunk, and retains the KV cache of gist tokens across all transformer layers. This section introduces this variant as the representative architecture for our analysis.
Figure I(a) illustrates the overall architecture. Given long sequence X = [x], T2........................... 2T]. a chunk length L. and a compression ratio r, the sequence is first divided into N = T fixed-length chunks. Within each chunk, one gist token is inserted after every raw tokens to compress the sequence. As result, each chunk contains L raw tokens and n L/r interleaved gist tokens. The i-th chunk's content is given by:
10
Inl,
r << L,
the 1ist gist unit
the n-th gist unit
where g denotes the inserted gist tokens. Each group consisting of raw tokens and the subsequent gist token is defined as a gist unit. For simplicity, all inserted gist tokens share the same embedding.
For the i-th chunk, its input to the transformer blocks contains not only the current (+),, but also the
accumulated KV cache from all preceding chunks' gist tokens Ge: as the prefix, thereby enabling information flow across chunks. To maintain consistency with the inference stage, the same chunking
Here assumes that I' can be divided by L and r for clarity.
3

---

Reserved KV Cache (all loyers)
Position
For Language Modeling (Boo hiiilen
Recuval Disaase
Relocttoo
annguage Model (One-pass Forward and Backward)
MEREEEOOM Reconstruciaa Long Sequenoe Input
Sink
(a) Model Architecture
(b) Attention Pattem
(c) Revised Position Encodings
Figure 2: Overall architecture of UniGist. (a) We adopt a chunk-free training strategy that allows one-pass sequence processing. (b) We design gist-enhanced attention pattern effectively aggregate information. (c) We revise positional encoding to maintain the relevant distances between raw tokens.
and iterative chunk-size forward strategy used during training. The final training objective follows the standard cross-entropy language modeling loss over raw tokens.
3.2 Analysis
We analyze why previous gist-based models struggle to fully utilize compressed context from two perspectives by experimenting on a well-trained gist-based model:
Attention Patterns. We examine the attention patterns on two datasets with long-range dependen- cies (Figure 1(b)). In natural data such textbooks, most tokens exhibit little attention across chunk boundaries. For structured tasks like RULER MultiKey-2 [19], which follow a parallel needle-in-a- haystack pattern, attention also remains mostly localized within the current chunk. In addition, many attention weights concentrate at the beginning of each chunk, even though this region may contain irrelevant key-value pairs. In contrast, the actual related piece of context receives little attention.
Token Positions. We perform continue-trainin training with a well-trained compressed model on 32K- length sequences and measure the average loss at each relative position within chunks, excluding the first chunk that has no previous contexts. chunk consistently exhibit higher loss, while those closer to the end show loss patterns similar to full attention models. This suggests that tokens near the end can rely on nearby uncompressed tokens to minimize the language modeling loss. As result, previous contexts may not be effectively used for compression learning, which limits overall efficiency of training.
shown Figure I(c), tokens the beginning of each
These observations suggest that chunk-wiss training is primary bottleneck limiting compression quality. Previous approaches rely on chunk-wis. training ensure consistency between training and inference, but this strategy prevents models from learning cross-chunk dependencies and limits compression performance, calling for better designs for training strategies.
4 Method
4.1
Unified Gist Context-basedd Language Modeling
As illustrated in Section 3.2, the chunk-wise training leads to an under-attention problem, that is, tokens in each block tend to ignore earlier compressed contexts. To address this, we eliminate the coarse-grained chunking strategy and introduce a unified attention pattern to enable the language model learn more effective context compression. Figure 2 provides an overview of our strategy. Given token sequence [I1, T2......... .r)), we construct the augmented input with the following steps:
1) Attention Sink Prefix. We prepend fixed number of sink tokens [81,82, ... 8s] at the beginning of the entire sequence. They serve as fixed "attention sinks" [43] to prevent mode collapse that can occur when initial raw tokens are compressed and then removed.

---

Saser Loep
Osate
Glai Stall Trick
soiees
soians
siizes
Outer Leep
Attention Patien Transformation
(b) Block Teihle-froe Kernel Design
We adopt the gist shift trick by moving all gist tokens to the rightmost position of the sequence to facilitate GPU block processing. (b) With fixed compression rate, non-visible blocks can skipped by pre-calculating block indexes.
Figure 3: Kernel design of our efficient sparse gist attention. (a)
2) Gist Token Insertion. We interleave learnable gist tokens into the token sequence with compression ratio (i.e., one gist token for every r raw tokens). We uniformly apply this insertion operation over all raw tokens. After this operation, the new sequence will be:
Z = [1
E1, Ir, 91, 21 2/r) T' ==sa
22r 92, VT, 9T/rl [T,9T/1
T/r
3) Position Encoding Alignment. Inserting gist tokens shifts the position indices of later tokens, which changes the original distances between raw tokens and may harm language coherence. To mitigate this, assign each gist token the same position id the raw token right after it, so that all raw-token distances remain unchanged.
With this augmeneeett input, we define unified sparse autoregressive attention pattern that is applied to all tokens. For each E Z, the set of its visible tokens is defined as:
A(x.) = (S U SSUGUWA n {zylū < (),
where the intersection with {æylj j < t} enforces the autoregressive constraint, restricting each token to attend only to itself prior tokens. S and G sink and gist tokens, respectively, which can provide conpressed contexts. The local window W, consists the current gist unit that contains 2,, well the previous & gist units (each containing raw tokens followed by gist token). This helps the model retain its basic language modeling ability with local contexts. The final training objective remains the standard autoregressive cross-entropy loss
tokens 24 E X Z.
4.2 Hardware-aligned Kernel Design
Sparse attention patterns in gist-token models are inherently efficient. Nevertheless, to fully realize such efficiency GPU devices, need suitable hardware-aligne aligned kernels. Because of the sparsely inserted gist tokens, our attention patterns are incompatibility with FLASH_ATTN ATTN kernels [8]. Therefore, design custom kernel that supports efficient processing.
Since our attention pattern evenly distributes the gist tokens in a fine-grained way, conventional GPU kernels cannot fully exploit the advantages of sparse attention. The main reason is that GPU kernels are optimized for processing data in blocks (typically of size 64 or 128). In our setting (with compression ratio of 4 or 8), gist tokens are always attended to and are present in every block. As result, no blocks can be skipped during attention computation in training, even though the raw tokens in previous blocks outside the local window are not visible to later tokens and can be skipped.
We address this issue with the gist shift trick, which moves all gist tokens to the rightmost position of the sequence. This transformation converts the attention pattern into dense, right- aligned scheme that aligns well with standard GPU block processing, as shown in Figure 3. this way, the gist tokens are gathered together, and each token will only need to attend to the compacted gist blocks in addition to sink and local window blocks. This leads to a scheme that can fully exploit the sparsity benefits by skipping non attended blocks. Note that with fixed compression ratio, our attention
5

---

Prefill Stage (chunk sizem4)
Settings: sink sizc-2, compression ratio-2, local gist unit-1
Sink takens
Ram toleess
Gast tolems
(1) Truncute the sequence input and apply forward
(2) Progressively compress the sequence
Compression completed
10
4
6
sccond chusk
Cadhcalier compresiio
Comproeed
Docoding Stage
(1) Forward with the comprewsed cache
(2) Reach the next gist token
(3) Compress and continue generating
10
10
10 10
11
10
10
12
12
Cumproocti KVV cache
New tubes irpua
Poot KVV cachc
saken impa.
Camprosao KV axche
kolken tepui
Figure 4: Illustration of flexible inference. Since the raw tokens outside the local window will never be attended to by future tokens, these invisible raw tokens can be safely discarded at each compression point. For prefilling, we can re-adopt the chunk-wise scheme and encode each chunk of the input with our efficent kernel, with nonvisible tokens dropped after each chunk. For decoding, nonvisible raw tokens can be discarded after each new gist is inserted.
layout is deterministic, and there is no need to build or store any block index tables. Instead, the indexes of the visible blocks for each token can be directly pre-calculated.
This trick brings a side effect that the original token indexes are distorted, which leads the difficulty of calculation in-block attention masks. To mitigate this issue, we pass an auxiliary input of the original token indexes to the kernel. With this information, we can easily calculate the attention masks according to Equation 3.
4.3 One-pass Training and Flexible Inference
One-pass with Full Sequence. With our unified attention pattern and custom kernel, we can process long sequences in single pass without chunk wise iterative processing during training. This eliminates memory fragmentation and reduces the complexity associated with the chunk-based scheme, leading to faster and more -eemory-effcceent training.
Training:
Inference: Flexible Decoding with Low Memory Usage. For prefilling, which naturally follows the iterative manner, we can re-adopt the chunk-wis scheme to encode each chunk with our kernel to reduce peak memory usage.2 The overall procedure illustrated in Figure During inference, each token only attends to tokens defined by the autoregressive pattern. Raw tokens beyond the local window are no longer required and can be dropped on the fly during generation. This also allows compatibility with efficient decoding techniques like FlashDecoding [9].
5
Experiments
5.1 Experimental Setup
Baselines. We compare UniGist primarily with methods that perform sequence-leves context compression, as compression techniques along other dimensions are often orthogonal and not directly comparable. Our baselines include three categories of methods: (1) Vanilla Full Attention Models, including the original model and its supervised fine-tuned variant using the same data as ours to control for variables. (2) Training-free Context Compression, where include StreamingI.LM [43]. SnapKV [22], PyramidKV / [4], and AdaKV [12] (based on SnapK V) as popular eviction methods. (3) Gist Token-based Compression, where we evaluate AutoCompressors [6] and Activation Beacon [54] as state-of-the-art methods. Both use chunk-wise compression to enable general-purpose compression.
Implementation Details. We use Llama3. Instrutttt and Llama-3.2-3P1 -Istruutt as the base models to evaluate performance across different model sizes [35]. Continued pretraining is conducted on 16B tokens of 32K-length samples drawn from Prolong's [14] mixed dataset, followed by super- vised tuning on 1B tokens of mixed instruction data. Cross-document masking is applied during training to block attention across document boundaries. All custom attention kernels are implemented in Triton. For all methods, we adopt question agnnosic compression setup to ensure generality and
The chunk size can be configured any integer multiple of the compression ratio during prefilling.

---

Method
RAG
Rerank
LongQA
ICL
Synthetic
Summ.
Average
Llama-3.1 -ama-3.1-8B-Inaa 3--8-hrrrrt
Full Attention
74.5
55.1
43.5
81.8
99.3
28.9
63.9
Full Attention (FT)
72.7
48.4
43.8
79.5
97.4
26.8
61.4
LLreamingLLM
58.5
28.4
33.1
36.0
18.6
12.2
31.1
SnapKV
60.3
11.8
40.4
29.6
19.8
12.6
29.1
PyramidKV
60.9
9.3
41.3
22.4
21.7
12.7
28.0
AdaKV
61.1
14.0
40.5
37.6
31.6
32.9
AutoCompressors
64.7
20.9
36.4
40.1
23.2
16.5
33.6
Activation Beacon
67.9
34.4
40.6
79.2
61.8
21.0
51.8
UniGist (Ours)
71.3
45.5
45.5
85.8
91.3
22.9
60.4
Llama-3.2-2 3--tuc--ntt
Full Attention
68.6
17.5
38.6
79.6
78.9
51.9
Full Attention (FT)
66.3
16.9
39.1
77.1
85.6
27.7
52.1
StreamingLLM
52.3
10.9
13.2
5.8
15.1
11.5
18.1
SnapKV
52.8
3,0
29.4
8.6
16.2
12.3
20.4
PyramidKV
52.5
3.2
31.5
7.8
12.6
11.6
20.0
AdaKV
53.8
3.1
30.5
6.6
28.6
12.4
22.5
AutoCompressors
55.7
7.2
26.3
31.2
7.9
15.2
23.9
Activation Beacon
58.4
11.4
36.5
57.4
47.1
16.8
37.9
UniGist (Ours)
63.6
15.6
41.7
82.3
47.8
Table 1: Evaluation results of various context compression methods under two base models across long-context tasks. "FT" denotes to the fine-tuned full attention model. Bolded numbers indicate the best results among compression methods.
difficulty. Greedy decoding is used for all tasks. For UniGist, the sink size is set to 128, and the local window corresponds to 128 raw tokens. The main results reported use a compression ratio of 4. All training and inference experiments are conducted using the Huggingface framework. More details about data composition and training setup are introduced in Appendix A.
5.2
Main Results
Method
MMLU-Pro
GSM8K
HellaSwag
Long Context Evaluation. We evaluate mod- els' long-context understanding using the HEL- MET benchmark [48], which includes a broad range of long-context tasks, including retrieval- augmented generation (RAG), summarization, and synthetic recall. HELMET incorporates several representative datasets such as Inf- Bench [55] and RULER [19], offering com- prehensive assessment of model performance across diverse long-context scenarios. For each task, we report the average score over datasets. Detailed task configurations are provided in Appendix B.
Full
47.7
83.6
60.2
Full-FT
47.1
84.3
59.8
Beacon
47.3
84.2
60.1
UniGist
47.6
83.9
60.1
Table 2: Performance three widely-used short- context datasets. The "Full" and "Fi" denote to full attention and fine-tuning respectively.
Table presents the overall evaluation results. We highlight three key observations: (1) UniGist achieves the best performance among all compression methods across both model sizes. On the synthetic tasks such as RULER MK-3, the only method that performs close full attention. (2) UniGist shows clear improvements over prior methods with similar architectures (e.g.. Beacon). This suggests that our unified attention pattern design plays a key role helping the model read compressed context effectively. (3) The effect of compression varies by task. In tasks like RAG, many methods still perform reasonably well. In contrast, many-shot ICL is much more sensitive, especially for training-free approaches. This may be because compression disrupts the parallel structure required for in-context learning. Gist taken methods show better robustness under these conditions.

---

Loss (Steps 50-130)
Gradiert Noom (Steps 50-100)
Performance Vs. Local Window Size
24
vid 128 skace
wat 130 sdlar
24
RAIi
-Bi wilbia dill toloee
74
24
32
mm
Ca
18
63
Gloed
Glohal
Lead Window
Figure 6: Left & Middle: Loss and gradient norm curves of -lama-3.1-8B-Instte during training steps 50-100. Using 128 sink tokens leads more stable training, with lower loss and smaller gradient fluctuations. Right: Performance of two long-context tasks Llama-3.2 B.ttma--.----Intttuue after 4B-token continual pretraining and then supervised finetuning under different local window sizes.
Basic Short-context Capability. To assess whether context compression compromises the model's core capabilities, we further evaluate all methods on a set of standard short-context benchmarks, including MMLU-Pro [39] (knowledge and reasoning). GSM8K [7] (math), and HellaSwag [52] (commonsense inference) with 3-shot demonstration, as shown in Table 2. These tasks do not involve long contexts and thus serve proxy to measure whether compression-specee undesirable side effects on general performance. We find that all gist roken-based methods maintain performance comparable the full attention baseline across all tasks. Differences between methods are within the range of natural variance, and no consistent degradation is observed. This suggests that continued pretraining with the compression target does not impair the model's ability to perform basic tasks. The gains in long-context understanding come without sacrificing short-conted capabilities.
training introduces
Boundary Effect Test. As shown in Section 3.2, chunk- based training tends to introduce boundary effects, where a token's 's performance depends on its position within the chunk. We evaluate UniGist under the same setting and compare with full attention. Figure 5 shows that UniGist maintains nearly uniform perplexity across po- sitions, closely matching full attention. This confirms that UniGist avoids position-related bias caused by the wise training scheme.
1.70
Averoge Loos
2.89
166
Attatio
Telmn prettion warh chunk
Figure 5: Boundary Effect Test. The compression model exhibits a similar trend to full attention, showing degra- dation near chunk boundaries.
5.3 Ablation Study
The core mechanism of UniGist is the unified attention layout, which is designed for learning better compression. Its effectiveness is already evident in the main results through the performance comparison with Beacon's 's chunk- training. In this section, we focus two additional components that contribute to UniGist's performance: the use of sink tokens and the choice of local window size.
Sink tokens for Stable Learning. To understand their impact, we compare training behaviors between models with and without sink tokens during the first 50 to 100 training steps. As shown in the left of Figure 6, the model with sink tokens achieves consistently lower loss and smoother gradient norms. In contrast, the model without sink tokens shows frequent gradient spikes, which slow down convergence. One possible reason is that, without sink tokens, the model fails to consistently attend to the initial tokens in the sequence because of progressive compression. This may break the autoregressive pattern and weaken early-stage learning signals.
Effect of Local Window Size. We further examine how the size of the local attention window affects the trade-off between language modeling and compression learning. A small window limits the model's sss access to raw tokens, forcing it to rely more heavily on gist tokens. This encourages the model to learn to extract compressed representations, but also increases the difficulty of language modeling due to reduced contextual continuity. In contrast, large window gives the model easy
8

---

Alcition
Unitiist (Comprecsios rulle Forward Spord-up
Uniiid (Compreesio ralio-5)
Peak Mentory Usage
Backward Spoed-up
45
6
Peak Mencry Usuge (G3)
4.5
35-
3.5+
up
up
25-
20
16K
32K
6-K
128K 16k
32K
64
124K
16K
32K
5/K
126K
Lengh
Lengh
Langth
Figure 7: Comparison of efficiency between UniGist and standard full attention. Left: Peak memory usage increases steeply with context length for full attention, while UniGist maintains significantly lower memory consumption. Middle & Right: UniGist delivers substantial forward and backward attention speed-up, with higher gains under longer sequences and larger compression ratios.
access to raw tokens, which weakens the pressure use gist tokens and may cause the model to bypass the compression path altogether. To quantify this trade-off, we evaluate UniGist with different window sizes on two long context tasks. shown on the right of Figure 6, performance on both the RAG and synthetic recall tasks exhibits non-monotonic pattern, initially improving and then declining. Notably, both tasks achieve peak performance when the window size is set to 128. This suggests that moderate window provides sufficient local context for stable language modeling, while still encouraging the model rely on gist tokens for global understanding.
5.4 Efficiency Comparison
UniGist is primarily designed to reduce the memory occupation of the KV cache, while also providing substantial speed-ups during both training and inference. We evaluate the quantitative benefits from two perspectives: memory usage and computation speed.
Memory Usage.
Long-context modeling places heavy demands on memory resources, often becoming a bottleneck for deployment. We evaluate peak GPU memory usage under compression ratios of 4 and 8, comparing UniGist with full attention on the Llama-3, -.------Instree model. As shown in the left of Figure 7, UniGist achieves a substantial reduction in memory usage in both settings, with more pronounced savings under ratio=8. These results demonstrate that UniGist can effectively alleviate memory bottlenecks while preserving modeling performance. The benefits grow with longer contexts, making UniGist a strong choice for scaling to ultra-long sequences.
Speed-up Evaluation. To isolate the speed-up benefit of the attention mechanism itself, measure the forward and backward pass latency of a single attention layer. This avoids confounding factors from other components of the model. As shown in the right of Figure 7, the speed-up achieved by UniGist grows with context length and gradually approaches the theoretical upper bound, which equals the compression ratio (e.g., 4x speed-up under compression ratio=4)². This confirms that the sparse attention path introduced by UniGist leads to real acceleration in real scenarios, especially in long sequences where the quadratic cost of full attention becomes prohibitive.
Conclusion
This work presents UniGist, sequence-level compression framework for long-context language modeling without relying on hybrid strategies. By removing chunk-wise training and introducing unified gist attention layout, UniGist enables effective learning over compressed sequences. We design a gist shift trick and block-table-free sparse attention kernel that improve training efficiency and are fully reusable during inference. UniGist also supports flexible chunking at inference time for adaptive context processing. Experiments across a broad range of tasks demonstrate that UniGist
We provide the original data and experimental details in Appendix C.

---

achieves solid performance across range of tasks, and proves particularly effective at retaining details in long context scenarios. It also improves efficiency over full attention, with faster runtime and lower memory usage during inference.
7
Acknowledgments
This work was supported by Beijing Municipal Science and Technology Project No. 10001033330009 National Natural Science Foundation of China No. 62272467, Beijing Natural Science Foundation No. L.230088. The work was partially done at the Engineering Research Center of Gext-Generation Intelligent Search and Recommendation, MOE.
References
[1] Joshua Ainslie, James Lee- Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pages 4895- 4901. Association for Computational Linguistics, 2023.
[2] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long- transformer. CoRR, abs/2004.05150, 2020.
William Brandon, Mayank Mishra, Aniruddaa Nrusimha, Rameswar Panda, and Jonathan Kelle- Reducing transformer key-value cache size with cross-layer attention. Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, UIrich Paquet, Jakub M. Tomczak, and Cheng Zhang, editors, Advances in Neural Information Processing Systems 38; Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang. Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic KV cache compression based on pyramidal information funneling. CoRR, abs/2406 02069, 2024.
[5] Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Elficient fine-tuning of cong-context large language models. In The Twelfth International Conference on Learning Reprexentations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.ned 2024.
[6] Alexis Chevalier, Alexander Wettig, Anirudh Ajith, Danqi Chen. Adapting language models to compress contexts. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processng, EMNLP 2023, Singapore, December 6-10, 2023, pages 3829-3846. Association for Computational Linguistics. 2023.
[7] Karl Cobbe. Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/21 14,,, 2021.
[8] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
[9] Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flashdeonding. httpe://erfa.stantfor. stanford. html.
[10] Deepseck-AI. Deepseck-v2: strong, economical, and efficient mixture-of-expprts language model. CoRR, abs/2405.004344 2024.
[11] Chenlong Deng, Zhisong Zhang, Kelong Mao, Shuaiyi Li, Xinting Huang, Dong Yu, and Zhicheng Dou. A silver bullet or compromise for full attention? A comprehensive study of gist token-based context compression. CoRR, abs/2412.174833, 2024.
[12] Yuan Feng, Junlin Lx, Yukun Cao, Xike Xie, and S. Kevin Zhou. Ada-kv: Optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference. CoRR, abs/2407. 11550. 2024.
[13] Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. Not all heads matter: A head-level KV cache compression method with integrated retrieval and reasoning. CoRR, abs/2410.19258 2024.
10

---

[14] Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively), 2025.
[15] Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang. and Mao Yang. Seerattention: Learning intrinsic sparse attention in your Ilms. CoRR, 132/2410.13276, 2024.
[16] Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun. Qianyu Guo, Meng Wang. and Haofen Wang. Retrieval-augmente generation for large language models: A survey. CoRR, atbs/231..1099 10997, 2023.
[17] Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model. In The Twelfih International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
[18] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/2312.00752, 2023.
[19] Cheng- Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
[20] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang I Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin- Yew Lin, Yuqing Yang. and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context Iims via dynamic sparse attention. In Amir Globersons. Lester Mackey, Danielle Belgrave. Angela Fan. Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, Advances Neurul Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NewrIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[21] Wojeiech Kryscinski, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. BOOK- SUM: A collection of datasets for long form narrative summarization. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Findings of the A Associatin for Computational Linguistics: EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pages 6536-6558. Association for Computational Linguistics, 2022.
[22] Yuhong Li, Yinghing Huang, Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: L.M knows what you are looking for hefore generation. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024. NeurlPS 2024, Vancouver, BC, Canada, December 10 15. 2024, 2024.
[23] Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang. Chenchen Zhang, Ge Zhang, Jiebin Zhang, Yuanxing Zhang, Zhuo Chen, Hangyu Guo, Shilong Li, Ziqiang Liu, Yong Shan, Yifan Song. Jiayi Tian, Wenhao Wu, Zhejian Zhou, Ruijie Zhu, Junlan Feng. Yang Gao, Shizhu He, Zhoujun Li. Tianyu Liu, Fanyu Meng. Wenbo Su. Yingshui Tan. Zili Wang, Jian Yang, Wei Ye, Bo Zheng, Wangchunshu Zhou, Wenhao Huang, Sujian Li, and Zhaoxiang Zhang. A comprehensive survey on long context language modeling, 2025.
[24] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: tuning-free free asymmetric 2bit quantization for KY cache. In Forty- International Conference on Machine Learning ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
[25] Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulum Du, Tao Jiang. Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou, Mingxing Zhang, and Jiezhong Qiu. Moba: Mixture of block attention for long-context 1ms.. CoRR, abs/2502..13888 189, 2025.
[26] Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. CoRR, abs/2305. 16300, 2023.
[27] Jesse MU, Xiang Li, and Noah D. Goodman. Learning to compress prompts with gist tokens. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023.
[28] NVIDIA. Kypress. http://stttt/////ttttoooooott//t
[29] OpenAL. GPT-4 technical report. CoRR, abs/2303.08777777 2023.

---

[30] Guanghui Qin, Corby Rosset, Ethan C. Chau, Nikhil Rao, and Benjamin Van Durme. Nugget 2d: Dynamic contextual compression for scaling decoder-only language models. CoRR, abs/2310.02409, 2023.
[31] Noam Shazeer. Fast transformer decoding: One write-head is all you need. CoRR, abs/1911.02 2019.
[32] Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomczak, and Cheng Zhang. editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NewrIPS 2024, Vancouver BC, Canada, December 10 15, 2024, 2024.
[33] Gemma Team. Gemma 2: Improving open language models a practical size. CoRR, abs/2408.001 18, 2024.
[34] Gemma Team. Gemma 3 technical report. CoRR, abs/2503.19786, 2025.
[35] Llama 3 Team, The laaaa 3 herd of models, 2024.
[36] MiniMax Team, Minimax-01: Scaling foundation models with lightning attention. CoRR, ahs/2501.08313, 2025.
[37] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakoh Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polnsukhin. Attention is all you need. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.. 2017.
[38] Zhongwei Wan, Xinjian Wu, Yu Zhang, Yi Xin, Chaofan Tao, Zhihong Zhu, Xin Wang. Siqi Luo, Jing Xiong, and Mi Zhang. D20; dynamic discriminative operations for efficient generative inference of large language models. CoRR, abs/2406. 13035, 2024.
[39] Yubo Wang. Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj. Xuan He. Ziyan Jiang, Tianle Li, Max Ku, Kai Wang. Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: A robust and challenging multi-task language understanding benchmark. Amir Globersons. Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomozak, and Cheng Zhang, editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver BC, Canada, December 10- 15, 2024, 2024.
[40] Zheng Wang, Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. Model tells you where to merge: Adaptive KV cache merging for Ilms on long-contest tasks. CoRR, 0as/2407.08454 2024.
[41] Wei Wu, Zhuoshi Pan, Chao Wang, Liyi Chen, Yunchu Bai, Kun Fu, Zheng Wang, and Hui Xiong. Tokenselect: Efficient long-context inference and length extrapolation for llms via dynamic token-level KV cache selection. CORR, abs/241 11 1,0286,, 2024.
[42] Chaojun Xiao, Pengle Zhang. Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. Infllm: Training-free free contecontext extrapolation for llms with efficient context memory. In Amir Globersons. Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet. Jakub M. Tomezak. and Cheng Zhang, editors, Advances Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15. 2024, 2024.
[43] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
[44] Ruyi Xu, Guangxuan Xiao, Haofeng Huang. Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. CoRR, Abs/2503.16422, 2025.
[45] Zhangchen Xu, Fengqing Jiang, Luyao Niu, Yuntian Deng, Radha Poovendran, Yejin Choi, and Bill Yuchen Lin. Magpie: Alignment data synthesis from scratch by prompting aligned 1lms with nothing. In The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28. 2025. OpenReview.net, 2025.
12

---

[46] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui. Bo Zheng, Bowen Yu, Chang Gao. Chengen Huang. Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge. Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang. Jianxin Yang, Jiaxi Yang. Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang. Keqin Bao, Kexin Yang. Le Yu, Lianghao Deng, Mci Li, Mingfeng Xue, Mingze Li. Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo. Tianhao Li, Tianyi Tang. Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang. Yinger Zhang, Yu Wan, Yugiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025.
[47] Dongjie Yang. Xiaodong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid KV cache compression for high-throughpus LLM inference. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024, pages 3258 -270. Association for Computational Linguistics, 2024.
[48] Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izsak, Moshe Wasserblat, and Danqi Chen. Helmet: How to evaluate long-context language models effectively and thoroughly, 2025,
[49] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuging Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng- Native sparse attention: Hardware-aligned and natively trainable sparse attention. CORR, abs/2502. 10099, 2025.
[50] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[51] Manzil Zaheer. Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañóón. Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[52] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Anna Korhonen, David R. Traum, and Lluís Mårquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Papers, pages 4791 4800. Association for Computational Linguistics, 2019.
[53] Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattm: Accurate sparse attention accelerating any model inference. CoRR, abs/2502.18133 2025.
[54] Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Long context compression with activation beacon. arXiv preprint arXiv:2401.034.
2024.
[55] Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang. Zhiyuan Liu, and Maosong Sun. ocbench: Extending long context evaluation beyond 100k tokens, 2024.
[56] Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi-Chih Yao. Tensor product attention is all you need. CoRR, abs/2501 106425, 2025.
[57] Yuxin Zhang, Yuxuan Du, Gen Luo, Yunshan Zhong, Zhenyu Zhang, Shiwei Liu, and Rongrong Ji. Cam: Cache merging for memory-efficient Ilms inference. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21 27- 2024. OpenReview.net, 2024.
[58] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H20: heavy-hitter oracle for efficient generative inference large language models. Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[59] Wayne Xin Zhao, Kun Zhou, Junyi Li. Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang. Zican Dong, Yifan Du, Chen Yang. Yushuo Chen, Zhipeng Chen, Jinhao Jiang. Ruiyang Ren, Yifan Li, Xinyu Tang. Zikang Liu, Peiyu Liu, Jian- Yun Nie, and Ji Rong Wen. A survey of large language models. CoRR, bts//333.18223,. 2023.
13

---

[60] Yutao Zhu, Huaying Yuan, Shuting Wang. Jiongnan Liu, Wenhan Liu, Chenlong Deng. Zhicheng Dou, and Ji-Rong Wen. Large language models for information retrieval: A survey. CoRR, 2023.
14

---

Category
Tasks
Metrics
NQ TriviaQA PopQA
SubEM
SubEM
RAG
SubEM
HotpotQA
SumEM
Rerank
MS Marco
NDCG@5
0c Bench QA
ROUGE Recall
Long-doc QA
coBench MC
Accuracy
TREC Coarse
Accuracy
TREC Fine
Accuracy
Many-shot ICL
NLU BANKING77
Accuracy
Accuracy
CLINIC150
Accuracy
JSON KV
SubEM
RULER MK Needle
SubEM
Synthetic recall
RULER MK UUID
SubEM
RULER MV
SubEM
x=Bench Sum Multi-LexSum
ROUGE-Sum FI ROUGE-Sum FI
Summ.
Table 3: Data Composition and metrics of the used HELMET benchmark.
A Model and Training Details
A.1 Baseline Configuration
For training-freee compression methods, we adopt the implementation based on the KVPress frame- work [28] ensure consistency. For chunk-wise training methods, we divide sequences into 2K- length chunks and continue pretraining and fine-tuning on 32K-length samples. To ensure faimess, all trainable methods are fine-tuned with full- pal-meareetter updates. As for vanilla LLaMA models, given their extensive pretraining on approximately 15T tokens, we skip additional pretraining and apply supervised fine-tuning directly.
A.2 Hyper-parameters
For continued pretraining, we use a batch size of 2M tokens and set the learning rate to 1e-5. The learning rate is warmed up linearly from 0 over 256 steps and then decayed to 50% of its peak using cosine scheduling. The AdamW optimizer is used for our experiments. The compression ratio and local window size remain fixed throughout training. During fine-tuning, we retain most hyperparameters but reduce the warm-up steps to 128.
A.3 Training Data Processing
For continued pretraining, we use the 64K-length dataset from Prolong [14] and restructure it into 32K-Iength samples. Specifically, we split all 64K-length samples into two 32K segments. The original data includes two types of samples: those consisting of single long document, and those formed by concatenating multiple shorter documents. For single-document samples, we apply no additional processing beyond the split. For multi-document samples, pad each document individually so that its length is divisible by the compression ratio, ensuring compatibility with our attention pattern.
For supervised fine-tuning, we use the Magpie-Llama-33, -Po.:.1----------0---0--- main source. Then, we refer to the Beacon's 's [54] setup to further augment it with samples from LongAlpaca [5], BookSum [21], and 500 RULER-like synthetic data designed for long-context tasks. In total, the finettunin data comprises approximately IB tokens.
dataset [45] as the
15

---

Direction
Method
Ratio
16K
32K
64K
128K
Causal
13.8
52.7
206.9
863.9
Forward
UniGist
4
6.45
20.5
72.9
271.8
UniGist
8
4.10
11.6
37.3
130.9
Causal
57.4
225.6
871.5
3586.8
Backward
UniGist
4
24.3
80.9
294.7
1116.8
UniGist
8
14.9
44.5
147.9
538.9
Table Forward and backward time (ms) under different context lengths. UniGist is evaluated compression ratios (CR) of 4 and 8.
B Evaluation Details
Long Context Evaluation. We evaluate on six tasks from the HELMET benchmark. The datasets and metrics are listed in Table 3. We chat template for all prompts. For the LongQA task, use a 2-shot demonstration to ensure an exact match in multi-choice questions. Since most models trained with 32K context length, use the corresponding 32K configuration files. For tasks like Rerank, we append a brief reminder after the question to reinforce the expected output format and task objective.
C More Experiments
Speed-up Details We report the raw measurements of speed-up under controlled setting with batch size 1, 32 attention heads, and head dimension of 128. All methods are implemented with Triton to ensure fair comparison. Each result is averaged over five runs, as shown in Table 4.
D Limitations
Model Scale and Context Length Our experiments are limited to models with at most 8B pa- rameters and context length of 32K due computational constraints. Although this setup enables controlled and efficient experiments, it does not fully capture the potential of larger models to learn and benefit from compression. Understanding how model scale and context length influence compression quality remains an open question for future research.
Training Protocol In this work, we adapt existing language models our attention pattern through continued pretraining. While we do not train from scratch, the UniGist architecture is fully compatible with end-to-end training. Training from scratch may allow the model to better internalize the structure and inductive biases of our compression pattern, without being influenced by pre-existing attention mechanisms. However, such training typically requires hundreds of billions of tokens yield meaningful insights, which is beyond our current computational budget. We leave comprehensive investigation of seratch training for future work, especially understand how compression-awee architectures behave during early-stage pertaining.
E
Ethical Discussion
UniGist aims to reduce the computational and memory overhead of large language models by introducing efficient context compression. This brings tangible benefits: it lowers the cost of inference, reduces carbon footprint, and facilitates the deployment of long-context models in resource- constrained environments. However, compression may inevitably alters model behavior. Compared to full attention baselines, compressed models may yield less accurate or misleading outputs in certain edge cases, especially when critical information is omitted or distorted during compression. This poses a risk in high-stakes applications where incorrect responses can have real-world consequences. We urge to carefully evaluate reliability when applying UniGist in downstream tasks and to consider human oversight when appropriate.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We claim our contributions and scope in these two sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We include limitation statement in Appendix D.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
17

---

Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide details in Section 5.1, Appendix A and Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
18

---

Answer:
[No]
Justification:
we plan to release our code in the future.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
Possiicy)))) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips.cooc wwwppppwlpgpppiieeeeeeeeooooepee http//uwd////eee CCodeSubmiieeeeece
Stt.d///ggiiiee//eeeuubbbiiiiiii iaatemmsssiooooooli (eionnolccy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide details in Section 5.1, Appendix A and Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification:
We report the average scores for most of our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
19

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We include this details in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https: / http///eerrppp.eeeee cc htcp/pulcc/thiiccudeeeeeee..eee
didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer:
[Yes]
Justification:
We don't volate the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss these potential impacts in Appendix E.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
20

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
We discuss these in Appendix E.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The assets are explicitly mentioned and respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
21

---

Answer: [No]
Justification: We don't publish new assets until submission, but we plan to release in accordance with regulations.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for rescarch with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is an important, original, or non-standard component of the methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is required.
Answer: [NA]
22

---

Justification: the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy lattps:/nneuriss /eearrpps.ccCCCofeeeeeeeee for what should or should not be described.
23


Reviews:
[{"review_id": "AO1dPod7Vw", "reviewer": "Reviewer_XQBu", "summary": "This paper overcomes some of the limitations of efficient Transformers based on gist tokens. After finding that attention remains localised within chunks and that early tokens within chunks exhibit higher loss, the Authors propose a chunk-free variant of gist tokens. This is achieved through a special mask (while taking care of sink tokens and positional encoding). To this end, they design a special attention kernel efficiently supporting the required masking by restructuring the mask matrix. As a result, they report minor degradation in long-context performance, while peak memory and latency are substantially reduced.", "questions": "- How does position encoding alignment work with relative positions (e.g., RoPE)?\n- How are gist tokens initialised in practice?", "limitations": "The main limitation, which is not properly discussed, is the overhead of (continued) training a model, which may be prohibitively expensive. In addition, no discussion of broader societal or environmental implications is present in the main text.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- The analyses on the behaviour of current gist-based models and their limitations (section 3.2) are quite insightful.\n- The improvements over gist-based models are meaningful and allow for better accessing gist tokens beyond the current chunk.\n- The results \n\n**Weaknesses** \n- The choice of baselines is inadequate: training-free methods do not represent the state of the art in sparse attention / KV cache compression. Based on large-scale experiments in Sparse Frontier (Nawrot et al. 2025), these are rather MInference for pre-filling and Quest or cartridges (Eyuboglu et al. 2025) for decoding during inference. As for training-based baselines, again, state-of-the-art methods beyond gist-based ones are ignored, e.g., Deepseek’s NSA. To understand if the proposed method is competitive, it must be compared with additional baselines.\n- Considering that one of the main purposes of KV cache compression is to speed up inference, I was a bit surprised by how little of the paper is dedicated to properly measuring this property. 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 for different batch sizes and sequence lengths.\n\n**Minor**\n- The problem statement in the introduction could be improved: while “the cost of self-attention” is relevant for both training and inference, in the latter it matters mostly only during pre-filling, but not as much for decoding, as this phase is mostly memory-bound. This should be clarified a bit better (also to explain why reducing KV cache size is not only related to VRAM memory load but also to throughput). 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.\n- “Previous empirical study shows […]” (line 102): This statement seems to warrant some citations.\n- For continued training, you may consider distillation from the vanilla model rather than cross-entropy with the data (Łańcucki et al. 2025, Dynamic Memory Sparsification), as Llama’s data mix is unavailable and hence Prolong may cause a distribution shift.\n- Which mixed instruction data was used for continued training (line 210)? This detail is important for reproducibility.", "was_revised": true, "final_justification": "Based on the discussion on the choice of baselines and the new end-to-end performance evaluations, I maintain my positive assessment of this paper.", "rebuttal": "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:\n\n> **Weakness 1:** The choice of baselines is inadequate.\n\nFirst, we would like to clarify the core difference between our method and sparse attention methods.\n- **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.\n- 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.\n\nTherefore, 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.\n\nDespite 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):\n\n| Method | RAG | Rerank | LongQA | ICL | Synthetic | Summ. |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|\n| Full Attention | 61.9 | 51.4 | 43.5 | 81.8 | 99.3 | 28.9 |\n| MInference | 60.8 | 50.9 | 46.1 | 36.6 | 95 | 28.5 |\n| UniGist | 59.5 | 46.7 | 44.7 | 83.4 | 98.3 | 26.7 |\n\nWe plan to add these new experimental results and a detailed discussion to the Appendix. This will allow us to fully address your concern.\n\n> **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.\n\nFirst, 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. \n\nAt 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.\n\nTo 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:\n\n| Method | Compression Ratio | TTFT (ms) | TPOT (ms/token) | End-to-End Latency (ms) |\n| :---: | :---: | :---: | :---: | :---: |\n| Full Attention | - | 9484.7 | 50.5 | 15900.9 |\n| UniGist | 4 | 8587.4 | 37.0 | 13284.0 |\n| UniGist | 8 | 6581.6 | 33.2 | 10795.2 |\n\nWe 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.\n\n> **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.\n\nThank'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.\n\n> **Minor2:** “Previous empirical study shows […]” (line 102): This statement seems to warrant some citations.\n\nThank 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.\n\n> **Minor3:** On the suggestion of using distillation for continued training\n\nThis 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.\n\n> **Minor4:** On the lack of detail for the mixed instruction data\n\nThis 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.\n\n> **Question 1:** Regarding the compatibility of position encoding alignment with RoPE\n\n**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.\n\n> **Question 2:** Regarding the initialization of Gist Tokens\n\nThank'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.\n\nThank 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."}, {"review_id": "gREbxkGKUL", "reviewer": "Reviewer_ZixF", "summary": "This paper introduces UniGist, a sequence-level long context compression framework for large language models. The proposed approach moves away from the prevalent chunk-wise gist training strategies, instead enabling a chunk-free, unified attention pattern with learnable gist tokens interleaved across the entire sequence. The authors further introduce a “gist shift” trick—an efficient hardware-aligned kernel design—to exploit GPU architecture, accelerating both training and inference. Experiments on long-context tasks demonstrate that UniGist significantly outperforms previous gist-based compression methods.", "questions": "Please refer to the weaknesses.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n* Clear motivation: The limitations of prior chunk-wise gist-based approaches are clearly identified and experimentally analyzed.\n* Performance:Empirical evaluations demonstrate that UniGist approaches full attention performance on long-context tasks, particularly outperforming other compression methods by a large margin. Also, the gist shift trick is an engineering contribution that aligns the sparse attention pattern with GPU-friendly block processing.\n\n**Weaknesses**\n* Block-wise sparse attention has been studied in numerous previous works [1], and various variants of attention patterns have been proposed recently. It would be better to clarify the unique characteristics of the attention pattern in this paper to better validate the novelty of the work. \n* 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.\n\n[1] Generating Long Sequences with Sparse Transformers", "was_revised": true, "final_justification": "In the rebuttal, the authors addressed the core differences between Unigist and previous works, and explained why this approach achieves more effective compression. Therefore, I have raised my score.", "rebuttal": "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:\n\n> **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.\n\nWe 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.\n\n1. **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.\n2. **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.\n\nTherefore, 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.\n\n> **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.\n\nWe 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.**\n\nTo comprehensively address this concern, we believe it is important to distinguish between the **training and inference stages**, as UniGist introduces fundamental improvements in both.\n\n1. **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.\n\n2. **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.\n\n**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:\n\n| Method | Compression Ratio | TTFT (ms) | TPOT (ms/token) | End-to-End Latency (ms) |\n| :---: | :---: | :---: | :---: | :---: |\n| Full Attention | - | 9484.7 | 50.5 | 15900.9 |\n| Beacon | 4 | 8749.6 | 37.9 | 12561.5 |\n| UniGist | 4 | 8587.4 | 37.0 | 13284.0 |\n\nWe 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.\n\nThank 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."}, {"review_id": "aDjmM2BJ9Q", "reviewer": "Reviewer_upn7", "summary": "This paper proposes a new gist-based sequence-level compression method that efficient long context modeling without chunk-wise training. In order to improve throughput for the sparse attention, a gist shift trick (moving all gist tokens to the rightmost position of the sequence) is proposed. The experiments on several long-context tasks demonstrate the good performance of this proposed method.", "questions": "1. It is unclear why chunk-based training is used in the previous method. Is it to save memory? If it is, is your method more memory-consuming for other GIST-based method?", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. Some analysis on previous GIST architectures:  ignore compressed contexts and high loss on the start part of each chunk. This clearly show the issues of the previous GIST architectures.\n\n2. A new gist-based sequence-level compression method is proposed. Some new techniques ( (attention sink prefix, position encoding alignment, local window) ) are used in the new method. \n\n3. The proposed method shows the consistent improvements on several long-context tasks. Two different-size models (Llama3.1-8B-Instruct and Llama-3.2-3B-Instruct) comparison are also shown. \n\n\nWeaknesses:\n\n1. Some presentation are not clear. In Figure 1(b), there is no any annotation what are the attention weight range for each color. What are the supervised tuning on 1B tokens of mixed instruction data? What are the synthetic dataset? Otherwise, it is difficult to evaluate the generation ability of the proposed method.\n\n2. In the proposed methods, there are two main added techniques (sink tokens and local window). In section 5.3 some ablation is done. However, ablation is only on training curves (loss, and gradient norm). There is no any study on the effect of the performance.  These can help to know how importance of each part.\n\n3. It is not surprised that local window with raw tokens can help the performance. However, it could hurt the efficiency. In Figure 7, there is  a comparison of efficiency between UniGist and standard full attention. There is no comparison with previous GIST method.", "was_revised": true, "final_justification": "The rebuttal addresses my concerns on the presentation and contributions. So I keep this score as 4.", "rebuttal": "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.\n\n> **Weakness1:** concerns about paper presentation\n\n- **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.\n\n- **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.\n\n> **Weakness2:** concerns about the ablation study\n\nThis 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:\n\n- **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.\n\n- **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.\n\n- **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.\n\n| Sink size | RAG | Rerank | LongQA | ICL | Synthetic | Summ. |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|\n| sink=0 | 51.5 | 2.9 | 37.9 | 71.4 | 68.0 | 21.4 |\n| sink=128 | 53.3 | 3.7 | 38.4 | 72.0 | 72.1 | 21.4 |\n\n> **Weakness3:** concerns about efficiency (local window and comparison with previous gist models)\nWe provide responses to both aspects of your concerns below:\n\n**Part A: The computational overhead introduced by the local window is negligible**\n\n- **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.\n- **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.\n\n| Local window size | Forward (ms) | Backward (ms) |\n| :---: | :---: | :---: |\n| 0 | 22.1 | 141.5 |\n| 128 | 21.6 | 137.9 |\n\n**Part B: Compared to previous GIST methods, UniGist also holds a significant efficiency advantage during inference.**\n\nUniGist demonstrates significant inference efficiency gains over previous gist methods in two key phases:\n\n- **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.\n- **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.\n\nIn 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.\n\n> **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?\n\nWe 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.\n\nHowever, 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.\n\nSpecifically, chunk-wise training introduces the following limitations, which our work aims to address:\n\n- **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.\n\n- **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.\n\n- **Inference Inflexibility:** The chunk size used during inference is limited by the training configuration, making it difficult to adapt to different needs.\n\n**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.\n\nThank 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."}, {"review_id": "g18jyY4hoW", "reviewer": "Reviewer_7dFr", "summary": "To address the accuracy challenge of the gist token-based KV cache compression approach, this work proposes a unified gist-based framework for effective compression, named UniGist. By discarding the chunk-wise compression scheme, the LLMs can learn the compression more effectively over the entire sequence. Extensive experiments validate that UniGist outperforms the existing gist token-based compression methods in terms of algorithm accuracy. By implementing the GPU kernel specifically, the proposed method also achieves higher inference throughput.", "questions": "See above.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "[Strengths]:\n1. The motivation of the proposed method is demonstrated clearly, and the paper is well-written.\n2. The GPU kernel for UniGist is implemented and achieves actual throughput improvement.\n\n\n[Weaknesses]:\n1. The proposed scheme lacks logical rationality. Unlike the full sequence processing, the chunk-wise processing is proposed to satisfy the memory constraint throughout the full inference process in the case of long text. Adopting full sequence-based scheme to avoid the challenges of chunk-wise processing lacks novelty.\n2. How was the experiment in Figure 1c and Table 5 conducted?\n3. Under the same token-level compression ratio, how long of a context can the proposed method process on a single machine, and how is its performance compared with other works?", "was_revised": true, "final_justification": "After reading the rebuttal, I think the author has adequately addressed my main concerns. Specifically, the author further clarifies the superiority and flexibility of the proposed UniGist compared to conventional chunk-wise compression methods. The supported maximum context length on a Single GPU was also evaluated and supplemented. I believe the main contributions are solid and reasonably supported, and the paper may be of interest to the community. Thus, I am raising my rating to 'Borderline Accept'.", "rebuttal": "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.\n\n> **Weakness 1:** Logical soundness and novelty of the method\n\nWe 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:\n\n1. 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.\n\n2. **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.\n\nTherefore, 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.\n\n> **Weakness 2:** Experimental details in Figure 1c and Table 5\n\nThank you for raising this point. We are happy to provide more detailed clarifications:\n\n- **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.\n- **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.\n\nWe will include these implementation details in the final version of the appendix to further enhance the transparency and reproducibility of our work.\n\n> **Weakness 3:** Maximum Context Length and Performance on a Single GPU\n\nTo 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:\n\n| Method | Compression Ratio | 128K | 256K | 320K | 512K |\n| :---: | :---: | :---: | :---: | :---: | :---: |\n| UniGist | 4 | 24.3GB | 32.6GB | 36.7GB | OOM |\n| UniGist | 8 | 20.2GB | 26.4GB | 30.5GB | 37.7GB |\n\n**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.\n\nThank 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."}]

Author rebuttal:

Meta-review: This paper introduces UniGist, a sequence-level long context compression framework designed to be both general across tasks and hardware-aligned for efficient deployment. Reviewers agreed that the problem addressed is important and timely, and they appreciated the technical idea of combining general compression with hardware efficiency considerations. Strengths noted include clear motivation, practical relevance, and evidence that the method offers competitive performance with meaningful efficiency gains. At the same time, reviews were cautious, with concerns about limited novelty relative to prior compression methods, and questions on whether the empirical evaluation fully support broad claims. The rebuttal helped clarify design choices and provided additional results, partially alleviating concerns. Given that the reviewers’ reservations were mostly about scope rather than fundamental flaws, and considering the strengths in problem importance, methodological soundness, and demonstrated potential, I recommend acceptance.
Decision: Accept (poster)