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 能通过其内部的注意力缓存互相对话，为集成式或教师-学生场景提供更好的准确度与更快的推理速度。