Compressing Sequences in the Latent Embedding Space: K-Token Merging for Large Language Models Zihao Xu♢ * , John Harvill♡ * , Ziwei Fan♠ , Yizhou Sun♣† , Hao Ding♠ , Hao Wang♢♠ , ♢ Rutgers University, ♠ AWS AI Labs, ♣ Amazon, ♡ Mistral AI, † University of California Los Angeles [email protected]
K-Token Merging vs Uncompressed
arXiv:2604.15153v1 [cs.CL] 16 Apr 2026
Abstract Large Language Models (LLMs) incur significant computational and memory costs when processing long prompts, as full self-attention scales quadratically with input length. Token compression aims to address this challenge by reducing the number of tokens representing inputs. However, existing promptcompression approaches primarily operate in token space and overlook inefficiencies in the latent embedding space. In this paper, we propose K-Token Merging, a latent-space compression framework that merges each contiguous block of K token embeddings into a single embedding via a lightweight encoder. The compressed sequence is processed by a LoRAadapted LLM, while generation remains in the original vocabulary. Experiments on structural reasoning (Textualized Tree), sentiment classification (Amazon Reviews), and code editing (CommitPackFT) show that K-Token Merging lies on the Pareto frontier of performance vs. compression, achieving up to 75% input length reduction with minimal performance degradation.
1
Introduction
We have witnessed tremendous progress in Large Language Models (LLMs) (Zhao et al., 2023). Since the advent of ChatGPT (Brown et al., 2020), LLMs have become deeply integrated into many aspects of daily life – from AI-enhanced customer support (Shareef, 2024; Scatolin and Pedrini, 2026) and game NPCs (Cox and Ooi, 2023; Christiansen et al., 2024) to agentic systems capable of data analysis (Tang et al., 2025) and complex software generation (Joel et al., 2024; Jiang et al., 2026). These increasingly sophisticated applications demand ever-longer input contexts, which sharply escalate both computational and memory costs: with softmax attention (Vaswani et al., 2017), the * Work done at AWS AI Labs.
Uncompressed K-Token Merging (K = 4) 100
100.00
Small Accuracy Gap (1.59%)
99.97 98.38
Large Length Reduction
80
(75.00%)
60 40 20 0
25.00
Input Length (%)
Accuracy (%)
Figure 1: Our K-Token Merging method (K = 4) achieves a 75% reduction in input length with only a 1.59% drop in accuracy on the Textualized Tree benchmark, demonstrating that it exploits redundancy in the latent embedding space while preserving high performance. See the “Experiments” section for details.
memory and compute requirements of LLMs grow quadratically with input length. A natural direction to mitigate this challenge is token compression (Li et al., 2025). Token compression aims to reduce the input length by representing prompts with fewer tokens. Existing approaches largely fall into two categories. The first category, hard prompt compression (Li et al., 2023; Pan et al., 2024; Jiang et al., 2024; Liskavets et al., 2025), decreases token count by dropping or summarizing tokens. While effective for tasks where key information is sparse, these methods often fail on information-dense tasks such as document revision, mathematical reasoning, symbolic computation, or code translation– settings in which every token may carry essential information. The second category, soft prompt compression (Mu et al., 2023; Harvill et al., 2025a; Gao et al., 2024), draws inspiration from soft
Prefill …
𝐶1
𝐶2
Encoder 𝒇
Encoder 𝒇
𝑇1
𝑇2
𝑇3
Generation
𝑇4
𝐺1
𝐶𝑁 Encoder 𝒇
… 𝑇2𝑁−1
𝐺2
𝐺3
LLM with LoRA
𝑇2𝑁 𝐶1
𝐶2
…
𝐶𝑁
𝐺1
𝐺2
…
Figure 2: Model Structure for K-Token Merging Model (Case K = 2). Left: During the prefill stage, the encoder f takes each K consecutive input tokens and produces a single compressed token embedding. Here, Ti denotes the original input tokens and Ci denotes the resulting compressed tokens. Right: During the generation stage, the LLM outputs original (uncompressed) tokens. Each newly generated token is appended to the mixed compressed/uncompressed prefix, after which standard auto-regressive generation continues. Here, Gi denotes the generated uncompressed tokens.
prompts (Lester et al., 2021) and learns new tokens or adapts model parameters to produce more compact representations. These methods have gained increasing attention because, empirically, they preserve more information than hard methods while still offering substantial compression (Harvill et al., 2025a; Gao et al., 2024). However, existing methods focus almost exclusively on reducing the number of tokens, overlooking a major source of redundancy: the embedding space itself. Consider the Q WEN 2.5 model (Qwen et al., 2024), which has a vocabulary size of 151,936. Each token consumes approximately 896 × 32 = 28,672 bits, whereas the theoretical minimum needed to identify one token from its vocabulary is only log2 (151,936) ≈ 18 bits. For a K-gram, the minimum grows to merely log2 (151,936K ) ≈ 18K bits. This enormous gap highlights significant inefficiency in current embeddings and indicates substantial room for improving input representation. In this paper, we propose K-Token Merging, a new soft prompt compression method that directly targets redundancy in the latent embedding space. Our key idea is to use a learned encoder to compress a sequence of K consecutive tokens (Kgrams) into a single embedding. In other words, we represent multiple tokens with one embedding on the input side only. We then finetune the LLM with LoRA to adapt to these compressed embeddings, enabling strong downstream performance. As illustrated in Figure 1, our approach achieves high compression ratios with minimal performance degradation (see the “Experiments” section for details). Importantly, our method is fundamentally differ-
ent from expanding the vocabulary with K-grams. Vocabulary expansion requires representing multitoken sequences as single tokens on both the input and output sides. This leads to an explosion of low-frequency “tail” tokens, which severely complicates optimization and becomes quickly infeasible for large K. For example, the number of unique 4-grams can exceed the size of the original vocabulary by 29× (Yu et al., 2025), resulting in prohibitive storage overhead. Furthermore, since each K-gram is assigned a fixed embedding, such approaches cannot generalize to unseen K-grams during inference, sharply limiting their applicability. In contrast, our encoder-based approach avoids tail-token explosion, supports large K values without increasing vocabulary size, and generalizes naturally to unseen token combinations. Our main contributions can be summarized as follows: • We introduce K-Token Merging, a novel framework for latent-space token sequence compression. • We present a training recipe for adapting LLMs to such compressed inputs. • We show experimentally that our approach achieves up to 75% token compression with minimal performance degradation across diverse tasks. In scenarios where output length is much shorter than input length, this corresponds to an estimated 94%1 reduction in computation. 1 With K=4, input length is reduced by 75%. Prefill attention cost scales as O(n²), so compressed prefill requires only 0.252 = 6.25% of the original FLOPs. When output length M ≪ input length KN , total FLOPs ≈ 6.25% of uncompressed, yielding ∼ 94% reduction.
Textualized Tree
Amazon Reviews
Indentation Tree: └A ├B ├C └D └E └F
Reviews: I’m really happy with this purchase! The fabric feels soft and comfortable, and the fit is true to size. After a couple of washes it still …
Q: Parent - Child Relationship of B and C?
Q: Positive / Negative Review?
A: NO
CommitPackFT Code Before Update: import math import matplotlib.pyplot as plt import numpy as np def compute_sin(x): return np.sin(x) … Instruct: Plot the “cos” function instead of the “sin” function.
Generated Code: … def compute_cos(x): return np.cos(x) …
A: Positive
Figure 3: Datasets & Tasks. Left: Textualized Tree. Given a textualized indentation tree, the LLM determines whether two nodes have a parent - child relationship. Middle: Amazon Reviews. The LLM performs sentiment classification to judge whether a product review is positive or negative. Right: CommitPackFT. Given code and an update instruction, the LLM needs to output the modified code that follows the instruction.
2
Related Work
Token compression (Li et al., 2025) has been extensively studied since the introduction of the Transformer architecture (Vaswani et al., 2017). Existing approaches can be grouped into hard prompt compression and soft prompt compression. Hard prompt compression. Hard compression methods (Li et al., 2023; Pan et al., 2024; Chuang et al., 2024; Jiang et al., 2024; Jung and Kim, 2024; Shandilya et al., 2025; Liu et al., 2023; Liskavets et al., 2025) reduce input length by dropping or summarizing tokens. For example, SelectiveContext (Li et al., 2023) removes tokens deemed uninformative based on self-information, and LLMLingua2 (Pan et al., 2024) trains a classifier to identify redundant tokens. However, these methods often suffer significant performance degradation on information-dense tasks, as even a small number of omitted tokens may contain essential context. In contrast, our method never removes tokens; it re-embeds all tokens into compressed embeddings. This enables substantially higher compression rates while mitigating information loss. Furthermore, our technique is compatible with hard compression and can be applied subsequently, making it complementary rather than competing. Soft prompt compression. Soft compression methods (Bolya et al., 2022; Mu et al., 2023; Chevalier et al., 2023; Ge et al., 2023; Cheng et al., 2024; Harvill et al., 2025a; Gao et al., 2024) learn new to-
kens or adapt LLM parameters to produce compact input representations. For instance, Gist (Mu et al., 2023) compresses instructions into a small set of meta-tokens prepended to the input, enforcing restricted attention patterns to ensure that later tokens cannot be attended beyond the meta-tokens. LTSC (Harvill et al., 2025a) discovers frequently occurring token subsequences and replaces them with shorter learned meta-tokens, analogous to classical file-compression algorithms. However, these methods operate entirely at the token level. In contrast, our approach performs compression directly in the latent embedding space, targeting inefficiencies in token embeddings themselves rather than the token sequence. This distinction allows our framework – though still a form of soft compression – to achieve a substantially more favorable balance between performance and compression ratio than prior approaches.
3
Method
3.1
Model Structure
Our insight is to exploit inefficiency in token-level representations. We introduce a lightweight encoder f that compresses every contiguous block of K tokens into a single compressed token embedding. The overall architecture is illustrated in Fig. 2, using the 2-Token Merging model as an example. 1
Results come from Harvill et al. (2025b).
Table 1: Accuracy (%), Length Reduction Ratio (%), and P–L F1 on Textualized Tree. We bold the best result and underline the second-best result for each metric. Our method achieves both the best and second-best results on Length Reduction Ratio and P–L F1 , while achieving the second-best Accuracy, surpassed only by the Uncompressed baseline.
Method
Uncompressed 99.97 0.0 0.000
Accuracy Length Reduction P–L F1
SelectiveContext 90.43 52.5 0.664
For our K-Token Merging model, we assume the total number of input tokens be KN . If necessary, we append padding tokens to ensure that the input length is always a multiple of K. We describe the model’s workflow in two stages: the prefill stage and the generation stage. Prefill. In the prefill stage, for each block {TiK+j }K j=1 , i ∈ {0, . . . , N − 1}, the encoder f takes as input their original embeddings from the model’s embedding table Emb and outputs a single compressed embedding Ci : Ci = f (Emb(TiK+1 ), . . . , Emb(TiK+K )) . (1) This compression is computationally inexpensive because: (i) the encoder is a small multi-layer perceptron (MLP) of roughly 50 MB, and (ii) compressed embeddings for frequently occurring Kgrams may be cached. Generation. During generation, the model always outputs original (uncompressed) tokens. These newly produced tokens are appended to the mixed compressed/uncompressed prefix and autoregressively processed by the LLM. 3.2
Objective Function
We finetune the LLM using a LoRA adapter jointly with the encoder f . Importantly, the training objective is evaluated only on positions corresponding to uncompressed (original) tokens – e.g., the generated tokens G1 , G2 , G3 , . . . in Fig. 2. Let the full sequence after inserting compressed tokens be X = (C1 , . . . CN , G1 , . . . GM ), where M is the number of generated tokens. Let U = {N + 1, . . . , N + M } denote the index set of uncompressed token positions. The LLM defines next-token probabilities pθ (· | X<t ). The training loss is the negative log-likelihood restricted to uncompressed targets: X L(θ, f ) = − log pθ (Xt | X<t ) . (2) t∈U
LLMLingua21 82.17 27.0 0.406
LTSC1 99.68 27.1 0.426
K-Token Merging (Ours) 2-Token
3-Token
4-Token
99.91 50.0 0.666
98.63 66.7 0.796
98.38 75.0 0.851
This objective ensures that the model learns to interpret compressed embeddings while maintaining generation quality on the original vocabulary. 3.3
Compressed Embedding Initialization
The initialization strategy for compressed embeddings has a significant impact on the convergence speed of training. Empirically, we find that initializing the embedding as the average (mean pooling) of the K original token embeddings leads to substantially faster convergence compared to random initialization. We design a dedicated encoder architecture that naturally produces such an initialization. The encoder computes a residual combination of (i) the mean of the K original embeddings and (ii) the output of a small MLP. The MLP weights are initialized close to zero so that, at the beginning of training, its contribution is negligible. Consequently, the encoder initially outputs a compressed embedding that closely matches the average-pooled embedding of the K tokens, providing a stable and semantically grounded starting point for optimization. The pseudo code for this Average-Initialized Encoder is shown in Alg. 1.
4
Experiments
In this section, we demonstrate our method’s effectiveness across three tasks and present both an ablation and a case study. 4.1
Baselines
We select two representative hard prompt compression methods, SelectiveContext (Li et al., 2023) and LLMLingua2 (Pan et al., 2024), along with a recent soft prompt compression approach, LTSC (Harvill et al., 2025b), as our baselines. For comparison, we also report results from the uncompressed model.
Accuracy vs. Length Reduction on Textualized Tree
Accuracy (%)
2-Token
LTSC
3-Token
95.0 92.5
SelectiveContext
90.0 87.5
Uncompressed LTSC
93.5
4-Token
1.30 Uncompressed
93.0
LLMLingua2
82.5
92.0
0
10
20
30
40
50
Length Reduction (%)
(a) Textualized Tree
60
70
91.0
LTSC
1.32 2-Token 3-Token
92.5
1.34 2-Token
1.36
LLMLingua2
91.5
85.0
Perplexity vs. Length Reduction on CommitPackFT
Perplexity
Uncompressed
97.5
Accuracy (%)
100.0
Accuracy vs. Length Reduction on Amazon Reviews
SelectiveContext 4-Token
0
10
20
30
40
50
Length Reduction (%)
60
(b) Amazon Reviews
70
SelectiveContext 3-Token
1.38
4-Token
LLMLingua2
0
10
20
30
40
50
Length Reduction (%)
60
70
(c) CommitPackFT
Figure 4: Performance Score (Accuracy / Perplexity) vs. Length Reduction Ratio on three datasets: (a) Textualized Tree, (b) Amazon Reviews, and (c) CommitPackFT. A higher Performance Score / Length Reduction Ratio indicates better performance; therefore, points located toward the upper-right region of the plot are preferred. Pareto-optimal points are marked with hollow pink circle markers (◦). Our method, K-Token Merging with K ∈ {2, 3, 4}, is highlighted in red. As shown, our approach lies on the Pareto-optimal frontier across all three datasets. For consistency, panel (c) (CommitPackFT) uses a y-axis (perplexity) that increases from top to bottom.
4.2
Datasets & Tasks
We use three datasets, each designed for a unique task. Textualized Tree (Harvill et al., 2025b) is a synthetic dataset consisting of textualized indentation trees, as shown in Fig. 3 (left). We design a classification task on this dataset to evaluate whether our method preserves most of the information after prompt compression. Specifically, we randomly select a pair of nodes in the tree and ask the LLM to classify whether a parent-child relationship exists between them. The model outputs “true” if such a relationship exists, and “false” otherwise. Since any node pair can be selected, the compressed token embeddings must preserve complete structural information to maintain high classification accuracy. This dataset includes 2.45 million training samples and 50 thousand test samples, with each tree containing up to 150 nodes. Amazon Reviews (Hou et al., 2024) is a dataset comprising user reviews, item metadata, and useritem interactions from Amazon. We formulate a sentiment classification task to test whether the compressed embeddings perform well on natural language, as shown in Fig. 3 (middle). Given a user review (e.g., “This legging is amazing”), the model is expected to classify the sentiment as either “positive” or “negative”. Reviews with a rating above 3.5 are labeled as positive, and those below or equal to 3.5 as negative. We use the “Amazon Fashion” category, consisting of 2.45 million training samples and 50 thousand test samples. CommitPackFT (Muennighoff et al., 2023) is a
Git commit dataset covering 350 programming languages. Each sample contains the code before an update, the corresponding commit message, and the updated code. We define a code update task where the model receives the code before the update and the commit message as input, and is expected to generate the updated code, which is shown in Fig. 3 (right). We use the CommitPackFT subset, which consists of high-quality filtered data from the original CommitPack dataset. We further focus our experiments on the Python category, which includes 50.4 thousand training samples and 5 thousand test samples. 4.3
Evaluation Metrics
For each method, we report both the model performance with compressed prompts and the average length reduction ratio. For the Textualized Tree and Amazon Reviews datasets, we use classification accuracy as the performance metric, while for CommitPackFT, we use perplexity. In practice, different methods may trade off performance and length reduction differently; a method with stronger compression may exhibit slightly lower performance, and vice versa. To quantitatively evaluate this trade-off, we propose the Performance–Length Reduction F1 score (P– L F1 ), defined as the harmonic mean of performance and length reduction, analogous to the classical F1 score. The P–L F1 score is computed as: P–L F1 =
2P L , P +L
(3)
Table 2: Accuracy (%), Length Reduction Ratio (%), and P–L F1 on Amazon Reviews. We bold the best result and underline the second-best result for each metric. Our method achieves the best and second-best results on P–L F1 , outperforming the strongest baseline by 25.5%. Method
Uncompressed
Accuracy Length Reduction P–L F1
SelectiveContext
93.54 0.0 0.000
91.30 44.6 0.599
LLMLingua2 91.49 51.0 0.655
where P ∈ [0, 1] denotes the normalized performance score and L denotes the length reduction ratio. For the Textualized Tree and Amazon Reviews datasets, we directly use classification accuracy as P . For CommitPackFT, we use a relative perplexity ratio. Let ppli be the perplexity of method i and pplmin = min ppli . The normalized performance
3-Token
4-Token
92.51 50.0 0.649
92.23 66.7 0.774
91.05 75.0 0.822
def forward(self, x: matrix): // Input shape:(batch_size, embedding_dim, k) mean ← Mean(x, axis = -1); x_flat ← Reshape(x, shape = (batch_size, embedding_dim * k)); output ← mean + self.net(x_flat); return output;
(4)
which ensures that lower perplexity corresponds to higher normalized performance. 4.4
2-Token
def init(self, embedding_dim: integer, k: integer): self.net ← MLP(input_dim = k * embedding_dim, output_dim = embedding_dim);
i
pplmin , ppli
93.48 0.1 0.002
K-Token Merging (Ours)
// Encoder compressing k embeddings into 1 embedding class Average-Initialized Encoder:
Pi ∈ [0, 1] for each method i is given by: Pi =
LTSC
Implementation Details
We use Qwen-2.5 0.5B (Qwen et al., 2024) as the base model for all baselines. We apply LoRA finetuning to this base model separately on the compressed prompts produced by different baselines, using the same LoRA configuration across all experiments (rank r = 4, LoRA scaling factor α = 16, and dropout rate 0.05. We apply LoRA adaptation across all layers, including both the selfattention modules and the MLP components). For the Tree Classification task, we adopt a curriculum learning strategy: the model is first trained on smaller trees and progresses to larger trees only after its performance surpasses a predefined threshold. This stage uses approximately 19.6–34.6 million auxiliary training samples. For all other tasks, the model is trained directly without curriculum learning. The encoder in our model is implemented as a three-layer MLP. All models are trained using the AdamW (Loshchilov and Hutter, 2017) optimizer with a learning rate of 1 × 10−4 . For all methods, we run each experiment three times and report the best result.
Algorithm 1: Pseudo Code for AverageInitialized Encoder. Since the MLP is initialized near 0, this encoder at first outputs a compressed embedding that closely matches the averagepooled embedding of the K tokens, giving training a stable starting point.
4.5
Result Discussion
Textualized Tree. Table 1 reports results for the Textualized Tree classification task, also visualized in Fig. 4(a). Our K-Token merging models achieve accuracies of 99.91%, 98.63%, and 98.38%. The 4Token model, with a 75% Length Reduction Ratio, attains the best P–L F1 score (0.851), outperforming the strongest baseline by 28.2%. Hard prompt compression methods suffer from information loss. For instance, LLMLingua2 reduces input length by only 27%, and its accuracy drops to 82.17%, yielding a low P–L F1 score of 0.406. The soft compression method LTSC preserves content well but does not exploit inefficiencies in the latent embedding space, resulting in a relatively low Length Reduction Ratio of 27.1%.
Table 3: Perplexity, Length Reduction Ratio (%), and P–L F1 on CommitPackFT. We bold the best result and underline the second-best result for each metric. On P–L F1 , our method obtains both the best and the second-best performance. Method Perplexity Length Reduction P–L F1
Uncompressed 1.293 0.0 0.000
SelectiveContext 1.380 39.9 0.560
The uncompressed baseline achieves the highest accuracy (99.97%). The small gap between our KToken models and this baseline suggests that our method preserves most of the input information while achieving substantial length reduction. Amazon Reviews. Table 2 and Fig. 4 (b) show results for the semantic analysis task on the Amazon Reviews dataset. For this natural language dataset, the soft prompt compression method LTSC struggles to identify structured, repeatable patterns in the corpus, resulting in a length reduction ratio of only 0.1%. Hard prompt compression methods are more suitable in this setting due to the sparsity of key information; both achieve over 40% length reduction while maintaining accuracy above 90%. However, all baselines are inferior to our method in terms of P–L F1 score. Our K-Token Merging models maintain high accuracy, with the 2-Token model performing best at 92.51%, only 1.03% lower than the uncompressed upper bound while achieving a 50% input length reduction. The 4Token model, which attains the highest compression ratio (75.0%), achieves the best P–L F1 score of 0.822, improving over the strongest baseline by 25.5%. CommitPackFT. Table 3 and Fig. 4(c) present the results on the CommitPackFT dataset. Because performance is evaluated using perplexity, lower values indicate better language modeling quality. Among the baselines, SelectiveContext and LLMLingua2 achieve moderate length reductions (39.9% and 30.0%), with P–L F1 scores of 0.560 and 0.454, respectively. LTSC largely preserves perplexity (1.296) but compresses only 17.2%, resulting in a relatively low P–L F1 score of 0.293. In contrast, our method achieves both the best and second-best P–L F1 scores. Under a 50.0% length reduction ratio, the 2-Token model still outperforms SelectiveContext and LLMLingua2 while maintaining a perplexity of 1.343. The 4-Token model achieves the highest Length Reduction Ratio (75.0%) and the best P–L F1 score (0.830), indicat-
LLMLingua2 1.381 30.0 0.454
LTSC 1.296 17.2 0.293
K-Token Merging (Ours) 2-Token
3-Token
4-Token
1.343 50.0 0.658
1.382 66.7 0.779
1.391 75.0 0.830
ing that for the coding task, our approach achieves a stronger compression–performance trade-off. 4.6
Ablation Study
As shown in Fig. 5, we present an ablation study comparing different initialization strategies for the compressed embeddings. Our average-based initialization achieves 97% classification accuracy in just 8 epochs – one epoch faster than random initialization. This indicates that the average-based approach enables faster convergence and serves as a more effective initialization method than random initialization. 4.7
Case Study
We include a case study of our 2-Token Merging model trained on the CommitPackFT dataset to demonstrate that, even with compressed inputs, the model can still preserve key information and follow instructions to generate correct outputs. Listing 1: Code input used for our model’s case study. import math import matplotlib . pyplot as plt import numpy as np # Define the function def compute_sin (x): return np . sin (x) # Generate input values x_vals = np . linspace ( -2 * np .pi , 2 * np . pi , 500) y_vals = compute_sin ( x_vals ) # Plotting plt . figure ( figsize =(8 , 4) ) plt . plot ( x_vals , y_vals , label = ’ sin (x) ’) plt . show ()
In this example, we provide the original code in Listing 1 and instruct the model to “plot the ‘cos’ function at the same time.” The generated code is shown in Listing 2. Two observations are noteworthy. First, the model successfully generates an additional function to plot the cosine curve. Second, rather than introducing a new variable, it reuses the previously defined
Accuracy (%)
100
97% Accuracy
90
Listing 2: Output code. Our model follows the instructions to add an additional plot of cosine function. It reuse the previously defined variable “x_vals” (highlighted in yellow) as the function domain.
80
import math import matplotlib . pyplot as plt import numpy as np # Define the function def compute_sin (x): return np . sin (x)
70
# Generate input values x_vals = np . linspace ( -2 * np .pi , 2 * np . pi , 500) y_vals = compute_sin ( x_vals )
60
2
Average-based Initialization Random Initialization 4 6 8
Epoch
Figure 5: Ablation Study on Embedding Initialization Strategies. We evaluate both initialization strategies on the Textualized Tree classification task using smaller trees (only 5 nodes). All experiments use our 4-Token Merging Model. The results show that the model with average-based initializationconverges one epoch faster than the model with random initialization when reaching 97% accuracy.
variable “x_vals” (highlighted in yellow) as the function domain, ensuring that both plots share the same range. These two points indicate that, despite compressing both the original code and the instruction, the model can still recover most relevant information and generate the correct implementation. This suggests our method effectively leverages redundancy in the embedding space while preserving model performance.
5
performance.
Future Directions
Compression During Generation: Currently, we do not apply token compression to generated tokens for simplicity. Exploring how to extend our compression technique to the generation phase could lead to significant memory and efficiency gains. Recursive Encoding: In our setup, we apply the encoder once to compress the input tokens. A natural extension is to apply the encoder recursively – e.g., compress 8 tokens into 4, then compress those 4 into 2 – enabling progressively shorter representations. This hierarchical compression approach may further improve efficiency while preserving
# Plotting plt . figure ( figsize =(8 , 4) ) plt . plot ( x_vals , y_vals , label = ’ sin (x) ’) plt . show () # Plotting the " cos " function plt . figure ( figsize =(8 , 4) ) plt . plot ( x_vals , np . cos ( x_vals ) , label = ’ cos (x) ’) plt . legend () plt . show ()
Adaptive Compression: In this paper, we apply uniform compression over every fixed group of K tokens, regardless of the statistical properties of different K-grams. An adaptive strategy – such as compressing 2 tokens in high-frequency regions and 16 in low-frequency ones – could further improve the compression ratio while maintaining or even boosting performance.
6
Conclusions
In this paper, we introduced K-Token Merging, a latent-space compression framework that merges groups of tokens into single embeddings to reduce the effective input length of LLMs. By combining a lightweight encoder with LoRA adaptation, our method enables LLMs to process compressed inputs while preserving standard generation. Experiments across synthetic, natural language, and code tasks demonstrate that our approach achieves up to 75% input length reduction with minimal performance degradation, offering a favorable compression-performance trade-off. These results suggest that exploiting inefficiency in the embedding space is a promising direction for improving the efficiency of LLMs, particularly for long-context inputs.
Limitations Our work has the following limitations that we leave for future research: No compression during generation. Our method compresses only input tokens and does not compress generated tokens, limiting efficiency gains as the output sequence grows. Fixed compression ratio. We apply uniform compression to every group of K tokens regardless of corpus statistics. Adaptive strategies (e.g., smaller K in dense regions and larger K in sparse ones) could further improve compression while maintaining performance. Limited evaluation on LLMs. Due to limited computational resources, we evaluated our method only on relatively small models such as Qwen 2.5 0.5B. It would be valuable to investigate whether our conclusions generalize to larger models.
References Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. 2022. Token merging: Your vit but faster. arXiv preprint arXiv:2210.09461. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, and 1 others. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901. Xin Cheng, Xun Wang, Xingxing Zhang, Tao Ge, SiQing Chen, Furu Wei, Huishuai Zhang, and Dongyan Zhao. 2024. xrag: Extreme context compression for retrieval-augmented generation with one token. Advances in Neural Information Processing Systems, 37:109487–109516. Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. 2023. Adapting language models to compress contexts. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 3829–3846. Frederik Roland Christiansen, Linus Nørgaard Hollensberg, Niko Bach Jensen, Kristian Julsgaard, Kristian Nyborg Jespersen, and Ivan Nikolov. 2024. Exploring presence in interactions with llm-driven npcs: a comparative study of speech recognition and dialogue options. In Proceedings of the 30th ACM Symposium on Virtual Reality Software and Technology, pages 1–11. Yu-Neng Chuang, Tianwei Xing, Chia-Yuan Chang, Zirui Liu, Xun Chen, and Xia Hu. 2024. Learning to compress prompt in natural language formats. In Proceedings of the 2024 Conference of the North
American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 7756–7767. Samuel Rhys Cox and Wei Tsang Ooi. 2023. Conversational interactions with npcs in llm-driven gaming: Guidelines from a content analysis of player feedback. In International Workshop on Chatbot Research and Design, pages 167–184. Springer. Jun Gao, Qi Lv, Zili Wang, Tianxiang Wu, Ziqiang Cao, and Wenjie Li. 2024. Uniicl: An efficient unified framework unifying compression, selection, and generation. arXiv preprint arXiv:2405.17062. Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. 2023. In-context autoencoder for context compression in a large language model. arXiv preprint arXiv:2307.06945. John Harvill, Ziwei Fan, Hao Wang, Yizhou Sun, Hao Ding, Luke Huan, and Anoop Deoras. 2025a. Lossless token sequence compression via meta-tokens. arXiv preprint arXiv:2506.00307. John Harvill, Ziwei Fan, Hao Wang, Yizhou Sun, Hao Ding, Luke Huan, and Anoop Deoras. 2025b. Lossless token sequence compression via meta-tokens. arXiv preprint arXiv:2506.00307. Yupeng Hou, Jiacheng Li, Zhankui He, An Yan, Xiusi Chen, and Julian McAuley. 2024. Bridging language and items for retrieval and recommendation. arXiv preprint arXiv:2403.03952. Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2024. Longllmlingua: Accelerating and enhancing llms in long context scenarios via prompt compression. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1658–1677. Juyong Jiang, Fan Wang, Jiasi Shen, Sungju Kim, and Sunghun Kim. 2026. A survey on large language models for code generation. ACM Transactions on Software Engineering and Methodology, 35(2):1–72. Sathvik Joel, Jie Wu, and Fatemeh Fard. 2024. A survey on llm-based code generation for low-resource and domain-specific programming languages. ACM Transactions on Software Engineering and Methodology. Hoyoun Jung and Kyung-Joong Kim. 2024. Discrete prompt compression with reinforcement learning. IEEE Access, 12:72578–72587. Brian Lester, Rami Al-Rfou, and Noah Constant. 2021. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 conference on empirical methods in natural language processing, pages 3045–3059.
Yucheng Li, Bo Dong, Frank Guerin, and Chenghua Lin. 2023. Compressing context to enhance inference efficiency of large language models. In Proceedings of the 2023 conference on empirical methods in natural language processing, pages 6342–6353. Zongqian Li, Yinhong Liu, Yixuan Su, and Nigel Collier. 2025. Prompt compression for large language models: A survey. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 7182–7195. Barys Liskavets, Maxim Ushakov, Shuvendu Roy, Mark Klibanov, Ali Etemad, and Shane K Luke. 2025. Prompt compression with context-aware sentence encoding for fast and improved llm inference. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 24595–24604. Junyi Liu, Liangzhi Li, Tong Xiang, Bowen Wang, and Yiming Qian. 2023. Tcra-llm: Token compression retrieval augmented large language model for inference cost reduction. In Findings of the association for computational linguistics: EMNLP 2023, pages 9796–9810. Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. Jesse Mu, Xiang Li, and Noah Goodman. 2023. Learning to compress prompts with gist tokens. Advances in Neural Information Processing Systems, 36:19327– 19352. Niklas Muennighoff, Qian Liu, Armel Zebaze, Qinkai Zheng, Binyuan Hui, Terry Yue Zhuo, Swayam Singh, Xiangru Tang, Leandro Von Werra, and Shayne Longpre. 2023. Octopack: Instruction tuning code large language models. In NeurIPS 2023 workshop on instruction tuning and instruction following. Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Menglin Xia, Xufang Luo, Jue Zhang, Qingwei Lin, Victor Rühle, Yuqing Yang, Chin-Yew Lin, and 1 others. 2024. Llmlingua-2: Data distillation for efficient and faithful task-agnostic prompt compression. arXiv preprint arXiv:2403.12968. A Yang Qwen, Baosong Yang, B Zhang, B Hui, B Zheng, B Yu, Chengpeng Li, D Liu, F Huang, H Wei, and 1 others. 2024. Qwen2. 5 technical report. arXiv preprint. Matheus Ferracciú Scatolin and Helio Pedrini. 2026. Stellar: A structured, trustworthy, and explainable llm-led architecture for reliable customer support. Journal of the Brazilian Computer Society, 32(1):128– 144. Shivam Shandilya, Menglin Xia, Supriyo Ghosh, Huiqiang Jiang, Jue Zhang, Qianhui Wu, Victor Rühle, and Saravan Rajmohan. 2025. Taco-rl: Task
aware prompt compression optimization with reinforcement learning. In Findings of the Association for Computational Linguistics: ACL 2025, pages 1582–1597. Farooq Shareef. 2024. Enhancing conversational ai with llms for customer support automation. In 2024 2nd International Conference on Self Sustainable Artificial Intelligence Systems (ICSSAS), pages 239– 244. IEEE. Zirui Tang, Weizheng Wang, Zihang Zhou, Yang Jiao, Bangrui Xu, Boyu Niu, Dayou Zhou, Xuanhe Zhou, Guoliang Li, Yeye He, and 1 others. 2025. Llm/agent-as-data-analyst: A survey. arXiv preprint arXiv:2509.23988. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30. Da Yu, Edith Cohen, Badih Ghazi, Yangsibo Huang, Pritish Kamath, Ravi Kumar, Daogao Liu, and Chiyuan Zhang. 2025. Scaling embedding layers in language models. arXiv preprint arXiv:2502.01637. Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, and 1 others. 2023. A survey of large language models. arXiv preprint arXiv:2303.18223, 1(2):1–124.