thu-nics/C2C

Cache-to-Cache (C2C) – 大規模言語モデル間の直接的な意味通信

概要 – C2C は、2つ（またはそれ以上）の LLM がテキストをやり取りする代わりに、KV-caches（Transformer のアテンションで使用されるキー・バリューメモリ）を変換・融合することで情報を交換できるようにする Python ライブラリです。隠れ状態の表現を直接操作することで、回答の品質を向上させ（精度が約 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)。

---

代表的なユースケース

• アンサンブル推論 – 中間テキストを生成することなく、2つ（またはそれ以上）の LLM の潜在的な知識を組み合わせます。複雑な QA や多段階推論に役立ちます。
• 速度重視の推論 – 単純な「生成してから通信する」パイプラインと比較して、レイテンシを半分に短縮します。
• クロスアーキテクチャ知識転送 – より新しく小さなモデルが、より大きな教師モデルの隠れ状態の意味論から恩恵を受けられるようにします。
• 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)

# 最初のトークンでプロジェクターを適用するようにモデルに指示
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

シンプルなターミナルチャットが表示され、受信者モデルの応答が共有者のキャッシュによってどのように強化されるかを確認できます。

---

独自のプロジェクターの学習

1. recipe/train_recipe/ に設定ファイルを作成します (例：C2C_0.6+0.5.json)。
2. 実行：

   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
   

   プロジェクター層のみが更新され、2つの 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}
}

---

詳細情報

• プロジェクトページ: https://fuvty.github.io/C2C_Project_Page/
• 論文: https://arxiv.org/abs/2510.03215
• Hugging Face モデルハブ: https://huggingface.co/nics-efc/C2C_Fuser
• ライブデモ (Gradio): https://huggingface.co/spaces/nics-efc/C2C_demo

---

同グループによる関連研究

• R2R – 推論 LLM のためのトークンレベルルーティング
• TaH – 推論 LLM のための選択的潜在思考
• FrameFusion – LVLM のためのビデオトークン削減
• MoA – LLM のためのスパースアテンション混合

---

結論: C2C (Rosetta) は、LLM が内部のアテンションキャッシュを通じて互いに通信できるようにする実用的でオープンソースな方法を提供し、アンサンブルスタイルや教師・生徒シナリオにおいて、より高い精度と高速な推論を実現します。