Large Language Model as Token Compressor and Decompressor

1. Problem Statement

Large Language Models (LLMs) face a fundamental bottleneck in handling long contexts due to the quadratic complexity ($O(N^2)$) of the attention mechanism. As sequence lengths grow to tens or hundreds of thousands of tokens, the computational and memory overhead becomes prohibitive for tasks like long-document understanding, multi-hop retrieval, and complex reasoning.

Existing token compression methods suffer from specific limitations:

• Heuristic Methods: Techniques like token merging or pruning based on similarity often discard tokens that carry crucial discourse or pragmatic meaning, leading to information loss.
• Learned Embedding Methods: Approaches like AutoCompressor or ICAE generate fixed-length continuous vectors. These are often non-lexical (hard to interpret), decoupled from the LLM's primary reasoning process, and lack the ability to faithfully reconstruct the original input or adapt to varying information densities.

2. Methodology

The authors propose a novel framework where an off-the-shelf LLM is fine-tuned to function as both a Token Compressor and a Token Decompressor. The core innovation is the use of a self-expressive autoencoding framework that translates natural language into a compact, internal "language" of discrete, variable-length latent codes called Z-tokens.

Core Components

1. LLM Compressor:

   • Variable-Length Generation: Unlike fixed-length vector methods, the compressor generates Z-tokens autoregressively. It dynamically determines the number of tokens ($K$) needed based on semantic density ($K \ll N$). Complex segments receive more Z-tokens; redundant regions are compressed aggressively.
   • Discrete Latent Space: The model is extended with a new vocabulary $V_z$ for Z-tokens. To ensure discreteness while maintaining differentiability for training, the authors employ a Gumbel-Softmax estimator combined with a straight-through estimator.
   • Sliding Window: For extremely long texts, the input is segmented into overlapping windows. A consistency constraint is applied to overlapping regions to ensure semantic coherence across segments.

2. LLM Decompressor:

   • Dual Paradigms: The decompressor supports two modes:
     • Reconstructor: Reconstructs the original surface text from Z-tokens exactly.
     • Inferencer: Performs downstream tasks (e.g., QA, summarization) directly in the Z-token space without reconstructing the surface text first.
   • Lexical Constraints: During decoding, the output logits are restricted to the original pre-trained vocabulary ($V_{base}$). This forces the model to translate its internal abstract "Z-language" back into human-understandable natural language, ensuring semantic fidelity.
   • Training Strategy: The model uses Schedule Sampling to mitigate exposure bias, gradually transitioning from teacher-forcing (using ground truth) to free decoding (using model predictions) to improve robustness.

3. Z-Token Characteristics:

   • Polysemy & Context-Dependency: Z-tokens are not deterministic mappings like standard subwords. They are polysemous; the same Z-token index can represent different concepts depending on the surrounding context and decoding constraints.
   • Semantic Abstraction: Z-tokens act as "morphemes" of a semantic language, folding multiple paraphrases of the same concept into a shared discrete code.

3. Key Contributions

• Novel Insight: Demonstrating that LLMs can learn to compress and decompress text into a discrete, variable-length internal language (Z-tokens) without external auxiliary networks.
• Adaptive Compression: The method achieves content-adaptive compression, allocating more tokens to dense semantic regions and fewer to predictable ones, unlike fixed-ratio compression.
• Interpretability: The Z-tokens are discrete and interpretable. The authors show that identical Z-tokens appear in semantically coherent neighborhoods, proving the compression is structured rather than chaotic.
• Dual-Mode Inference: The framework allows for reasoning directly in the compressed Z-token space, effectively bypassing the need to process redundant surface tokens during high-level reasoning.

4. Experimental Results

The method was evaluated on six large-scale datasets: Wikipedia, CNN/DailyMail, HotpotQA, QuALITY, NarrativeQA, and QASPER.

• Compression Efficiency: Achieved up to 18× token reduction on long-query datasets while preserving reconstruction fidelity.
• Reconstruction Performance: On Wikipedia, the method achieved state-of-the-art BLEU-4 scores (e.g., 99.31 at 4× compression) and low Cross-Entropy loss, outperforming baselines like ICAE, Gist Tokens, and AutoCompressor.
• Downstream Tasks:
  • Summarization (CNN/DailyMail): Maintained or improved ROUGE scores compared to baselines while reducing inference time by 2×.
  • Long-Text QA: Outperformed all baselines on HotpotQA, QuALITY, NarrativeQA, and QASPER, demonstrating the preservation of multi-hop reasoning capabilities even at high compression ratios.
• Efficiency:
  • Inference Speed: Achieved a 4.6× speedup in inference time.
  • Memory: Significantly reduced GPU memory usage (e.g., from 2.81 GB to 2.04 GB for specific batch sizes).
  • Complexity: Reduced computational complexity from $O(N^2)$ to $O(K^2)$, where $K$ is the compressed length.

5. Significance and Future Directions

This work provides a token-efficient pathway for long-context reasoning. By internalizing compression within the LLM's reasoning flow, it transforms the model into a system that can "think" in a compact, abstract language before projecting back to natural language.

Significance:

• It overcomes the quadratic attention bottleneck without sacrificing semantic fidelity.
• It offers a controllable and interpretable alternative to continuous latent embeddings.
• It enables practical applications in prompt compression and long-document analysis on standard hardware.

Future Work:
The authors suggest focusing on:

1. Designing a global coordination mechanism for Z-tokens across sliding windows to handle very long-range dependencies.
2. Making the compression process "task-aware" to preserve specific factual or argumentative structures required for specialized downstream tasks.