thu-nics/C2C

[ICLR'26] The official code implementation for "Cache-to-Cache: Direct Semantic Communication Between Large Language Models"

Cache-to-Cache (C2C) – 大型語言模型間的直接語義通訊

它是什麼 – C2C 是一個 Python 函式庫，讓兩個（或多個）LLM 透過轉換並融合其 KV-caches（transformer 注意力機制所使用的鍵值記憶體）來交換資訊，而非來回傳送文字。透過直接操作隱藏狀態表示，該系統能提升回答品質（準確度提升約 8–10%）並將推論延遲降低約 2 倍。

為什麼叫「Rosetta」 – 正如羅塞塔石碑讓學者能解讀埃及象形文字，C2C 的「Rosetta」套件為 KV-cache 張量學習了一種共同的「語言」，讓原本獨立的模型能夠互相理解。

主要功能

功能	作用
KV-Cache 投影與融合	學習輕量級投影器網路，將共享者模型的快取映射到接收者模型的語義空間。
預訓練「融合器」	託管於 Hugging Face，提供多種熱門模型對（例如 Qwen3-0.6B ↔ Qwen2.5-0.5B, Llama-3.2-1B 等）的即用型投影器檢查點。
多共享者支援	將多個教師模型的快取融合到單一接收者中（仍處於實驗階段）。
僅訓練投影器	微調期間僅更新投影器權重；來源與目標 LLM 保持凍結，使訓練成本低廉。
簡單 API	`RosettaModel` 包裝器運作方式如同一般的 `transformers` 模型；您只需傳遞額外的 `kv_cache_index` 張量來告知系統何時應用投影。
演示與腳本	包含 Gradio 演示、即時聊天範例以及完整的訓練/評估腳本。
未來規劃	計劃中的「代理管理 KV-Cache」服務系統 (2026-09)。

典型使用場景

整合推理 – 結合兩個（或多個）LLM 的潛在知識而無需生成中間文字，適用於複雜的問答或多步驟推理。
速度關鍵型推論 – 與單純的「生成後通訊」流程相比，延遲減半。
跨架構知識轉移 – 讓較新、較小的模型受益於較大教師模型的隱藏狀態語義。
LLM 內部研究 – 研究 KV-cache 表示如何編碼資訊以及如何進行轉換。

快速開始（安裝與執行）

# 1. 建立一個新的 conda 環境 (Python 3.10)
conda create -n rosetta python=3.10 && conda activate rosetta

# 2. 安裝套件 (編輯模式)
pip install -e .
# 訓練/評估的額外選項
pip install -e ".[training,evaluation]"

執行預訓練演示

import torch
from huggingface_hub import snapshot_download
from script.playground.inference_example import load_rosetta_model, run_inference_example

ckpt = snapshot_download(
    repo_id="nics-efc/C2C_Fuser",
    allow_patterns=["qwen3_0.6b+qwen2.5_0.5b_Fuser/*"],
)

cfg = {
    "rosetta_config": {
        "base_model": "Qwen/Qwen3-0.6B",
        "teacher_model": "Qwen/Qwen2.5-0.5B-Instruct",
        "checkpoints_dir": f"{ckpt}/qwen3_0.6b+qwen2.5_0.5b_Fuser/final",
    }
}

model, tokenizer = load_rosetta_model(cfg, eval_config={}, device=torch.device("cuda"))
prompt = [{"role": "user", "content": "Say hello in one short sentence."}]
input_text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

# 告知模型在第一個 token 上應用投影器
instr_idx = torch.tensor([1, 0], dtype=torch.long).repeat(inputs['input_ids'].shape[1]-1, 1).unsqueeze(0).to(model.device)
label_idx = torch.tensor([-1, 0], dtype=torch.long).unsqueeze(0).to(model.device)
kv_idx = [instr_idx, label_idx]

with torch.no_grad():
    out = model.generate(**inputs, kv_cache_index=kv_idx, do_sample=False, max_new_tokens=64)
    print(tokenizer.decode(out[0], skip_special_tokens=True))

該腳本會列印出 C2C 生成的答案。

互動式聊天範例

# 單一共享者
python script/playground/live_chat_example.py --checkpoint_dir path/to/checkpoint

# 多個共享者 (列出所有檢查點目錄)
python script/playground/live_chat_example.py --checkpoint_dir ckpt1 ckpt2

將會出現一個簡單的終端聊天介面，展示接收者模型的回答如何被共享者的快取所豐富。

訓練您自己的投影器

建立設定檔 於 recipe/train_recipe/ (參考 C2C_0.6+0.5.json)。

執行：

python script/train/SFT_train.py --config recipe/train_recipe/C2C_0.6+0.5.json   # 單 GPU
# 或多 GPU
torchrun --nproc_per_node=8 script/train/SFT_train.py \
    --config recipe/train_recipe/C2C_0.6+0.5.json

僅更新投影器層；兩個 LLM 保持凍結。

評估

準備一個評估 yaml (例如 recipe/eval_recipe/unified_eval.yaml) 並執行：

python script/evaluation/unified_evaluator.py --config recipe/eval_recipe/unified_eval.yaml

評估器會載入 Rosetta 模型，在選定的基準測試上執行生成，並寫入標準指標（準確度、延遲等）。

擴充框架

新增投影器 – 在 rosetta/model/projector.py 中實作 Projector 的子類別並使用 @register_model 註冊。
新增資料集 – 在 rosetta/train/dataset_adapters.py 中建立 DatasetConfig 並在訓練 JSON 中引用它。
新增基準測試 – 遵循 script/evaluation/unified_evaluator.py 中的模式。

支援的模型對（預訓練）

接收者	共享者	檢查點
Qwen3-0.6B	Qwen2.5-0.5B-Instruct	連結
Qwen3-0.6B	Llama-3.2-1B-Instruct	連結
…	…	…
(完整表格請見 README。)

引用

若您在研究中使用 C2C，請引用 arXiv 論文：

@article{fu2025c2c,
  title={Cache-to-Cache: Direct Semantic Communication Between Large Language Models},
  author={Tianyu Fu and Zihan Min and Hanling Zhang and Jichao Yan and Guohao Dai and Wanli Ouyang and Yu Wang},
  journal={arXiv preprint arXiv:2510.03215},
  year={2025}
}

同一團隊的相關工作

R2R – 用於推理 LLM 的 token 級路由
TaH – 用於推理 LLM 的選擇性潛在思考
FrameFusion – 用於 LVLM 的影片 token 縮減
MoA – 用於 LLM 的稀疏注意力混合

總結： C2C (Rosetta) 提供了一種實用的開源方式，讓 LLM 能透過其內部的注意力快取互相對話，為整合式或教師-學生場景提供更好的準確度與更快的推論速度。