K-Token Merging: Compressing Sequences in Latent Embedding Space for Efficient Inference of Large Language Models

K-Token Merging is an innovative prompt compression method that merges consecutive token blocks in the latent embedding space. It significantly reduces input sequence length while maintaining model performance, opening up a new path for efficient inference of large language models.

Large language models face a fundamental challenge when processing long texts: the computational cost of the prefill stage is linearly related to the input sequence length. When users submit a long document containing tens of thousands of tokens, the model needs to consume a lot of computational resources to process these inputs before it can start generating responses. This 'long context penalty' severely limits the application efficiency of LLMs in scenarios such as document analysis, code understanding, and multi-turn dialogues.

Traditional solutions include sparse attention mechanisms, sliding windows, and hierarchical processing, but these methods often require modifications at the model architecture level. K-Token Merging takes a different approach, proposing a new idea of prompt compression in the latent embedding space.

The core insight of K-Token Merging is that natural language contains a lot of redundant information, and adjacent tokens often carry similar semantic information. If these redundant tokens are merged at the embedding level, the sequence length can be significantly shortened without significant loss of information.

Specifically, this method treats every K consecutive input tokens as a block, and merges the embeddings of these K tokens into a single latent embedding through a lightweight encoder. This compressed prefix is then fed into the large language model for prefill, while the generation stage still takes place in the original token space.

The workflow of K-Token Merging is divided into two distinct stages:

In the prefill stage, encoder f receives every K consecutive input tokens and generates a single compressed token embedding. The specific process is as follows:

1. Tokenize the input prompt
2. Retrieve the token embeddings of the base model from the cached embedding table
3. Split the prompt embeddings into consecutive blocks of size K
4. Merge each block using a lightweight encoder (the encoder is initialized to behave like mean pooling, then jointly trained with LoRA adapters)
5. Feed the compressed prefix into the base LLM

In the generation stage, the LLM outputs original uncompressed tokens. Each newly generated token is appended to the mixed compressed/uncompressed prefix, followed by standard autoregressive generation. This design ensures that generation quality is not affected while enjoying the benefits of prefill acceleration.

K-Token Merging uses an ingenious encoder initialization strategy. The encoder is initialized to behave like mean pooling, then trained end-to-end with LoRA adapters. This design brings several advantages:

• Stability: Mean initialization provides a reasonable starting point, avoiding gradient instability in the early stages of training
• Flexibility: LoRA adapters allow learning compression strategies while keeping the base model frozen
• Efficiency: The lightweight encoder has minimal computational overhead, which does not offset the gains from compression

The research team validated the effectiveness of K-Token Merging on three benchmark tests: Textualized Tree, Amazon Reviews, and CommitPackFT.

Taking the Textualized Tree benchmark as an example, when the merging factor K=4, this method achieved:

• 75% reduction in input length: The original sequence was compressed to 1/4 of its original size
• Only 1.59% drop in accuracy: Maintained extremely high performance while significantly shortening the sequence

This result proves that K-Token Merging successfully leverages redundancy in the latent embedding space while preserving most of the model's reasoning capabilities.