| --- |
| library_name: transformers |
| license: mit |
| tags: |
| - qwen-compress |
| - kv-cache |
| - compression |
| - qwen |
| - model-merging |
| - efficiency |
| datasets: |
| - gsm8k |
| - ifeval |
| --- |
| (有提供繁體中文版readme) |
| # Qwen-Compress: Learned Weighted Pooling for KV Cache Compression |
|
|
| **A drop-in KV cache compression method that selectively compresses hidden states via learned weighted pooling between every 4th transformer layer, reducing prefill KV size by up to ~1.75× without retraining the base model.** |
|
|
| ## Overview |
|
|
| Standard autoregressive transformers pay a **KV cache cost proportional to total prefill length**. Qwen-Compress mitigates this by inserting small **`HiddenStateCompressor`** modules at selected layers. Each compressor performs **learned weighted pooling**: it computes an importance score per token via a learned linear projection, then applies softmax over a sliding window of `ratio=4` consecutive hidden states and takes the weighted sum, producing 1 compressed token from 4. This shrinks the KV cache while preserving the final `window_size=128` tokens unchanged (tail retention). |
|
|
| The key insight: compress hidden states **before** they enter the attention QKV projection at the compressor layer. Downstream layers attend to compressed KV positions via correct rotary position embedding (RoPE) indices, avoiding position mismatch. |
|
|
| ### Inspiration |
|
|
| This work is inspired by **DeepSeek's compressed attention mechanism** (used in DeepSeek-V4), where inter-layer compression reduces KV cache footprint. Qwen-Compress adapts this concept to the Qwen architecture with a per-sequence compression schedule and decode-phase buffer management. |
|
|
| ## Architecture |
|
|
| ### Compressor Module |
|
|
| A single `HiddenStateCompressor` is a lightweight learned weight layer: |
|
|
| ``` |
| Input: (B, 4, D) — 4 consecutive hidden states |
| │ |
| ┌─────┴─────┐ |
| │ wgate │ Linear(D, 1) → per-token importance score |
| │ + ape │ Learned position bias (4, 1) |
| └─────┬─────┘ |
| │ softmax over 4 tokens → normalized weights |
| │ |
| Weighted sum → (B, 1, D) — 1 compressed token |
| ``` |
|
|
| **Parameters**: ~D (one Linear layer + 4 learned scalars) — negligible vs. base model. No latent bottleneck or QKV attention involved. |
|
|
| ### Placement |
|
|
| Compressors are inserted at layers where `layer_idx % 4 == 3` (every 4th layer). At non-compressor layers, attention runs unchanged. |
|
|
| ### Compression Schedule (Window + Ratio) |
|
|
| For each sequence at a compressor layer: |
|
|
| ``` |
| true_len = sequence length (after left-padding) |
| window = min(window_size=128, true_len) ← kept verbatim |
| compressible = true_len - window |
| to_compress = compressible - (compressible % ratio=4) |
| compressed = to_compress / 4 |
| residual = compressible % ratio ← kept verbatim (0–3 tokens) |
| tail = window ← last 128 tokens verbatim |
| ``` |
|
|
| KV positions are tracked per-token for correct RoPE indexing. Zero-padded batch entries are masked via `kv_valid` throughout generation. |
|
|
| ### Decode-Phase Buffer Compression |
|
|
| During autoregressive decode, new tokens accumulate in an **uncompressed buffer** (`buf_k`, `buf_v`, `buf_h`). Once the buffer reaches `window_size + ratio = 132` tokens, the oldest `ratio=4` tokens are compressed into 1 and merged into the persistent KV cache, keeping the buffer at `window_size` tokens. This bounds the decode-phase KV overhead at `window` tokens per compressor layer. |
|
|
| ## Performance |
|
|
| Tested on **Qwen3.6-27B (NF4, 64 layers)**. Compression applied at every 4th layer (i % 4 == 3), totaling 16 compressors. |
|
|
| ### GSM8K (5-shot) |
|
|
| | Filter | Baseline | Compressed | Δ | |
| |---------------------|----------|------------|--------| |
| | flexible-extract | 0.8749 | 0.7331 | −16.2% | |
| | strict-match | 0.8658 | 0.5785 | −33.2% | |
|
|
| ### IFEval (0-shot) |
|
|
| | Filter | Baseline | Compressed | Δ | |
| |----------------------------|----------|------------|--------| |
| | inst_level_loose_acc | 0.8849 | 0.8897 | +0.5% | |
| | inst_level_strict_acc | 0.8525 | 0.8513 | −0.1% | |
| | prompt_level_loose_acc | 0.8299 | 0.8355 | +0.7% | |
| | prompt_level_strict_acc | 0.7856 | 0.7856 | ±0.0% | |
|
|
| ### Key Takeaways |
|
|
| - **IFEval is virtually unaffected** — instruction following, formatting, and constraint satisfaction are preserved. |
| - **GSM8K shows degradation** — multi-step mathematical reasoning suffers under compression, especially strict format matching (`#### N`). |
| - **Compression ratio**: A 301-token prefill is compressed to ~172 KV entries (1.75×). Effective ratio varies with prompt length (shorter prompts see less benefit; very long prompts approach 4×). |
|
|
| ## Usage |
|
|
| The reference implementation is a FastAPI server (`server.py`) that patches a Qwen model with compressors and serves an OpenAI-compatible API. |
|
|
| ```bash |
| # Install dependencies |
| pip install torch transformers accelerate bitsandbytes |
| |
| # Start the server (loads model, patches compressors, serves on :8001) |
| python server.py |
| |
| # Send requests via curl or any OpenAI client |
| curl http://localhost:8001/v1/chat/completions \ |
| -H "Content-Type: application/json" \ |
| -d '{"model": "qwen-compress", "messages": [{"role": "user", "content": "1+1="}]}' |
| ``` |
|
|
| The patching happens in `build_franken_qwen()` (`server.py:344`): every 4th layer's attention forward is replaced with `get_custom_qwen_forward(compressor, window_size=128, ratio=4)`, and the compressor weights are loaded from a checkpoint. |
|
|
| ## Limitations |
|
|
| - **Reasoning degradation**: Complex multi-step reasoning (GSM8K) shows 16–33% accuracy drop. The compressor likely disrupts intermediate representations critical for arithmetic reasoning, even with thinking mode disabled. |
| - **No training recovery**: The compressor checkpoint is trained in isolation, not fine-tuned jointly with the base model. Post-training or LoRA-based recovery may close the gap. |
| - **Context-independent compression**: The weighted pooling treats each 4-token window independently. A cross-window or bidirectional design could better preserve global reasoning structure. |
|
|
