---
library_name: transformers
license: mit
tags:
- qwen-compress
- kv-cache
- compression
- qwen
- model-merging
- efficiency
datasets:
- gsm8k
- ifeval
---

# Qwen-Compress：可學習加權平均實現 KV Cache 壓縮

**一種即插即用的 KV cache 壓縮方法，在每隔 4 層的 Transformer 層之間透過可學習的加權平均選擇性壓縮 hidden states，在不重新訓練基礎模型的情況下，將 prefill KV 大小減少最高約 4 倍。**

## 概述

標準的自迴歸 Transformer 需要**與 prefill 總長度成正比的 KV cache 成本**。Qwen-Compress 在特定層插入小型 **`HiddenStateCompressor`** 模組來緩解這個問題。每個 compressor 執行 **可學習的加權平均**：透過 learned linear projection 為每個 token 計算重要性分數，然後在連續 `ratio=4` 個 token 的滑動視窗上做 softmax，加權求和後將 4 個 token 壓縮成 1 個。同時保留最後 `window_size=128` 個 token 不做壓縮（tail retention）。

關鍵洞察：在 hidden states **進入 compressor 層的 attention QKV projection 之前**進行壓縮。下游層透過正確的旋轉位置編碼（RoPE）索引來關注壓縮後的 KV 位置，避免位置錯配。

### 靈感來源

本專案受 **DeepSeek 的壓縮注意力機制**（DeepSeek-V4 使用）啟發。Qwen-Compress 將此概念適配到 Qwen 架構，並加入了逐序列壓縮排程與解碼階段 buffer 管理。

## 架構

### Compressor 模組

一個 `HiddenStateCompressor` 是輕量級的可學習加權層：

```
輸入:  (B, 4, D)  — 4 個連續的 hidden states
          │
    ┌─────┴─────┐
    │  wgate    │  Linear(D, 1) → 每個 token 的重要性分數
    │  + ape    │  學習到的位置偏置 (4, 1)
    └─────┬─────┘
          │ softmax over 4 tokens → 歸一化權重
          │
    加權求和 → (B, 1, D)  — 1 個壓縮後的 token
```

**參數量**: ~D（一個 Linear 層 + 4 個學習到的標量）— 相較於基礎模型可忽略不計。

### 放置位置

Compressor 被插入在 `layer_idx % 4 == 3` 的層（每 4 層）。非 compressor 層的 attention 保持不變。

### 壓縮排程（Window + Ratio）

對於 compressor 層的每個序列：

```
true_len   = 序列長度（左側填充後）
window     = min(window_size=128, true_len)   ← 完整保留
compressible = true_len - window
to_compress  = compressible - (compressible % ratio=4)
compressed   = to_compress / 4
residual     = compressible % ratio            ← 完整保留（0–3 個 token）
tail         = window                          ← 最後 128 個 token 完整保留
```

KV 位置會逐 token 追蹤，以確保正確的 RoPE 索引。批次中經零填充的條目透過 `kv_valid` 遮罩在生成過程中正確屏蔽。

### 解碼階段的 Buffer 壓縮

在自迴歸解碼過程中，新產生的 token 會累積在 **uncompressed buffer**（`buf_k`、`buf_v`、`buf_h`）中。當 buffer 達到 `window_size + ratio = 132` 個 token 時，最舊的 `ratio=4` 個 token 會被壓縮成 1 個並合併到持久 KV cache 中，使 buffer 維持在 `window_size` 個 token。這將解碼階段的 KV 開銷限制在每個 compressor 層最多 `window` 個 token。

## 效能

測試於 **Qwen3.6-27B (NF4，64 層)**。壓縮應用在每隔 4 層（i % 4 == 3），共 16 個 compressor。

### GSM8K（5-shot）

| Filter              | 基準     | 壓縮後   | 變化    |
|---------------------|----------|----------|--------|
| flexible-extract    | 0.8749   | 0.7331   | −16.2% |
| strict-match        | 0.8658   | 0.5785   | −33.2% |

### IFEval（0-shot）

| Filter                     | 基準     | 壓縮後   | 變化    |
|----------------------------|----------|----------|--------|
| 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%  |

### 關鍵結論

- **IFEval 幾乎不受影響** — 指令遵循、格式化和約束滿足能力被完整保留。
- **GSM8K 出現衰退** — 多步驟數學推理在壓縮下受到影響，尤其是嚴格的輸出格式匹配（`#### N`）。
- **壓縮比**：301 個 token 的 prefill 被壓縮至約 172 個 KV 條目（1.75 倍）。有效壓縮比隨提示長度變化（較短的提示收益較小，非常長的提示可接近 4 倍）。

## 使用方式

參考實作為一個 FastAPI 伺服器（`server.py`），啟動後會自動為 Qwen 模型打上 compressor patch 並提供 OpenAI 相容 API。

```bash
# 安裝依賴
pip install torch transformers accelerate bitsandbytes

# 啟動伺服器（載入模型、patch compressors、listen on :8001）
python server.py

# 透過 curl 或任何 OpenAI client 發送請求
curl http://localhost:8001/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "qwen-compress", "messages": [{"role": "user", "content": "1+1="}]}'
```

Patching 邏輯在 `build_franken_qwen()`（`server.py:344`）：每第 4 層的 attention forward 被替換為 `get_custom_qwen_forward(compressor, window_size=128, ratio=4)`，compressor 權重從檢查點載入。

## 限制

- **推理能力衰退**：複雜的多步驟推理（GSM8K）準確率下降 16–33%。即便在關閉思考模式的情況下，compressor 仍可能破壞算術推理所需的 intermediate representations。
- **無聯合訓練恢復**：compressor 檢查點是獨立訓練的，未與基礎模型聯合微調。透過 post-training 或 LoRA 恢復可能縮小差距。
- **上下文無關的壓縮**：加權平均 compressor 獨立處理每個 4-token 窗口。跨窗口或雙向設計可能更好地保留整體推理結構。