💬 NLP

Compressing Sequences in the Latent Embedding Space: $K$ -Token Merging for Large Language Models

This paper introduces K-Token Merging, a latent-space compression framework that merges contiguous token embeddings via a lightweight encoder to significantly reduce input length and computational costs in Large Language Models while maintaining high performance across diverse tasks.

Original authors: Zihao Xu, John Harvill, Ziwei Fan, Yizhou Sun, Hao Ding, Hao Wang

Published 2026-04-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zihao Xu, John Harvill, Ziwei Fan, Yizhou Sun, Hao Ding, Hao Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to explain a complex story to a friend, but your friend has a very short attention span and a tiny memory. If you try to tell them the whole story word-for-word, they get overwhelmed, and it takes forever.

This is exactly the problem Large Language Models (LLMs) face today. When you ask an AI to read a long document, a huge code file, or a massive history book, the AI has to process every single word (token) individually. The more words you add, the exponentially harder and slower the computer gets to think. It's like trying to carry a library in your backpack; the heavier it gets, the slower you walk.

The paper you shared introduces a clever new trick called K-Token Merging. Here is how it works, explained simply:

The Problem: The "Word-for-Word" Bottleneck

Currently, if you feed an AI 1,000 words, it treats them as 1,000 separate items to juggle.

Old Way: Imagine you are sending a package to a friend. You write every single word of your letter on a separate sticky note. You tape 1,000 sticky notes together. Your friend has to read every single note to understand the message. It's slow and messy.
The Flaw: Many of those sticky notes are redundant. If you say "The cat sat on the mat," the words "The," "cat," "sat," "on," and "the" often appear together so frequently that they form a predictable pattern. The AI doesn't need to see them as five separate things; it just needs to understand the idea of that phrase.

The Solution: The "Smart Summarizer" (K-Token Merging)

The authors propose a new way to send the message. Instead of sending 1,000 sticky notes, they use a Smart Summarizer (a lightweight encoder) to group words together before sending them.

The Analogy: Imagine you group every 4 words (a "K-token" block) and stick them onto a single, super-dense "magic card."
- Instead of sending 1,000 sticky notes, you now only send 250 magic cards.
- Each card contains the essence of those 4 words, compressed into a single, compact signal.
- The AI receives these 250 cards instead of 1,000 notes. It processes them much faster because there are fewer things to juggle.

How Does the AI Understand the Magic Cards?

You might ask: "If we change the words into magic cards, won't the AI get confused?"

That's where the second part of their trick comes in: LoRA Adaptation.

Think of the AI as a student who has learned to read standard sticky notes.
The researchers give this student a special "training course" (using a technique called LoRA) to learn how to read these new "magic cards."
The student learns that one card equals four words. They learn the pattern.
Crucially: When the AI starts writing its answer (generating text), it switches back to normal sticky notes. It doesn't output magic cards; it outputs normal words. This ensures the final answer is still readable by humans.

Why is this better than previous methods?

Previous methods tried to delete words they thought were unimportant (like deleting "the" or "a").

The Risk: If you delete a word in a math problem or a line of code, the whole thing breaks. It's like trying to summarize a recipe by throwing away the ingredients you think are "optional."
The New Way: K-Token Merging doesn't throw anything away. It keeps all the information but packs it tighter. It's like packing a suitcase: instead of leaving clothes on the floor, you fold them tightly. You still have all the clothes, but they take up less space.

The Results: Speed vs. Smarts

The researchers tested this on three different challenges:

Logic Puzzles: Reading a tree diagram to find relationships.
Sentiment Analysis: Reading Amazon reviews to see if they are happy or sad.
Coding: Reading code and fixing a bug.

The Magic Numbers:

They were able to shrink the input size by 75%. (If you had a 1,000-word prompt, the AI only had to process 250 "cards").
The Catch? The AI's performance barely dropped. In some tests, it was almost as smart as the uncompressed version.
The Win: Because the AI had to process 75% fewer items, the computer saved a massive amount of time and energy. In fact, for long tasks, they estimate it could save up to 94% of the computing power needed.

The Big Picture

Think of K-Token Merging as a compression algorithm for human thought. Just like a ZIP file compresses a folder of documents so you can email them faster, this method compresses the "thoughts" inside an AI's brain so it can read longer books, write longer code, and answer more complex questions without getting tired or running out of memory.

It proves that we don't need to delete information to make AI faster; we just need to learn how to pack it more efficiently.

1. Problem Statement

Large Language Models (LLMs) face significant computational and memory bottlenecks when processing long prompts. The standard self-attention mechanism scales quadratically ( $O(N^2)$ ) with input length, making long-context processing expensive.

Existing token compression methods attempt to mitigate this but suffer from specific limitations:

Hard Prompt Compression: Methods that drop or summarize tokens (e.g., SelectiveContext, LLMLingua2) often fail on information-dense tasks (like code or mathematical reasoning) where every token is critical.
Soft Prompt Compression: Methods that learn new "meta-tokens" or adapt parameters (e.g., Gist, LTSC) operate at the token level. They overlook the inherent inefficiency and redundancy within the latent embedding space itself.
Vocabulary Expansion: Simply expanding the vocabulary to include $K$ -grams (sequences of $K$ tokens) causes a "tail-token explosion," leading to prohibitive storage overhead and poor generalization to unseen combinations.

The core insight of this paper is that while the theoretical information content of a token is low (e.g., ~~18 bits for a 150k vocabulary), current embeddings consume significantly more bits (~~28k bits). This gap indicates massive redundancy in the embedding space that existing methods fail to exploit.

2. Methodology: K-Token Merging

The authors propose K-Token Merging, a framework that compresses sequences directly in the latent embedding space rather than the token space.

Core Architecture

Lightweight Encoder ( $f$ ): A small Multi-Layer Perceptron (MLP) is introduced to process the input.
Compression Mechanism: The encoder takes a contiguous block of $K$ $K$ consecutive input token embeddings and merges them into a single compressed embedding ( $C_i$ $C_{i}$ ).
- Input: $K$ original tokens ( $T_{iK+1} \dots T_{iK+K}$ ).
- Output: 1 compressed token embedding ( $C_i$ ).
Two-Stage Workflow:
1. Prefill Stage: The encoder compresses the entire input prompt into $N$ compressed tokens (where original length was $KN$). This reduces the sequence length seen by the LLM's attention mechanism, drastically lowering the $O(N^2)$ cost.
2. Generation Stage: The LLM generates output in the original vocabulary (uncompressed tokens). New tokens are appended to the mixed prefix (compressed input + uncompressed generation) for standard autoregressive processing.

Training Strategy

LoRA Adaptation: The LLM is fine-tuned using Low-Rank Adaptation (LoRA) jointly with the encoder.
Loss Function: The training objective (Negative Log-Likelihood) is calculated only on the positions of the generated (uncompressed) tokens. The model learns to interpret compressed embeddings while maintaining generation quality on the standard vocabulary.
Initialization: To ensure stable convergence, the encoder is initialized to output the mean pooling of the $K$ input embeddings. The encoder is designed as: $Output = \text{Mean}(Inputs) + \text{MLP}(Inputs)$ , where the MLP weights start near zero. This provides a semantically grounded starting point.

3. Key Contributions

Latent-Space Compression: A novel framework that targets redundancy in the embedding space rather than the token sequence, distinguishing it from prior soft/hard compression methods.
Training Recipe: A specific LoRA-based fine-tuning approach and an "Average-Initialized Encoder" strategy that enables LLMs to adapt to compressed inputs efficiently.
Pareto-Optimal Performance: Experimental evidence showing the method lies on the Pareto frontier of performance vs. compression, achieving high compression ratios with minimal accuracy loss.
Generalization: Unlike vocabulary expansion, this method does not increase vocabulary size and generalizes naturally to unseen token combinations.

4. Experimental Results

The method was evaluated on three diverse tasks using a Qwen-2.5 0.5B base model, compared against baselines like SelectiveContext, LLMLingua2, and LTSC.

Dataset / Task	Metric	Key Findings
Textualized Tree (Structural Reasoning)	Accuracy	75% length reduction (K=4) achieved 98.38% accuracy (vs. 99.97% uncompressed). Outperformed baselines by a significant margin in P-L F1 score.
Amazon Reviews (Sentiment Analysis)	Accuracy	75% length reduction (K=4) achieved 91.05% accuracy. The method achieved the best P-L F1 score (0.822), outperforming the strongest baseline by 25.5%.
CommitPackFT (Code Editing)	Perplexity	75% length reduction (K=4) achieved the best P-L F1 score (0.830). The model successfully followed instructions to modify code despite compressed inputs.

Efficiency Gains:

With $K=4$ , input length is reduced by 75%.
Since prefill attention scales as $O(n^2)$ , the compressed prefill requires only $0.25^2 = 6.25\%$ of the original FLOPs.
In scenarios where output length is short relative to input, this results in an estimated 94% reduction in total computation.

5. Significance and Future Directions

Significance:
This work demonstrates that the primary bottleneck in long-context LLMs is not just the number of tokens, but the inefficiency of how those tokens are represented in the embedding space. By merging tokens before they enter the attention mechanism, K-Token Merging offers a path to drastically reduce memory and compute costs without sacrificing the model's ability to handle information-dense tasks like code and reasoning.

Limitations & Future Work:

Generation Phase: Currently, compression is only applied to the input (prefill). Compressing generated tokens during the generation phase is a future direction.
Fixed vs. Adaptive: The current method uses a fixed $K$ for all blocks. Future work could explore adaptive compression (varying $K$ based on token frequency or semantic density).
Model Scale: Experiments were limited to smaller models (0.5B). Validating on larger LLMs is necessary to confirm scalability.

In conclusion, K-Token Merging represents a paradigm shift from token-level to embedding-level compression, offering a highly efficient solution for long-context LLM deployment.