thu-nics/C2C

Cache-to-Cache (C2C) – 대규모 언어 모델 간의 직접적인 의미론적 통신

개요 – C2C는 두 개(또는 그 이상)의 LLM이 텍스트를 주고받는 대신 KV-caches(트랜스포머 어텐션에서 사용하는 키-값 메모리)를 변환하고 융합하여 정보를 교환할 수 있게 해주는 Python 라이브러리입니다. 은닉 상태 표현을 직접 조작함으로써 시스템은 답변 품질을 향상시키고(정확도 약 8~10% 향상), 추론 지연 시간을 약 2배 단축할 수 있습니다.

"Rosetta"라는 이름의 의미 – 로제타석이 학자들에게 이집트 상형문자를 읽게 해주었듯이, C2C의 "Rosetta" 패키지는 KV-cache 텐서를 위한 공통의 "언어"를 학습하여 독립적인 모델들이 서로를 이해할 수 있게 합니다.

---

주요 기능

기능	역할
KV-Cache 투영 및 융합	공유자(sharer) 모델의 캐시를 수신자(receiver) 모델의 의미 공간으로 매핑하는 경량 프로젝터 네트워크를 학습합니다.
사전 학습된 "퓨저"	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의 잠재 지식을 결합하여 복잡한 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
   

   프로젝터 레이어만 업데이트되며, 두 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이 내부 어텐션 캐시를 통해 서로 대화할 수 있게 하는 실용적인 오픈 소스 방식을 제공하여, 앙상블 스타일이나 교사-학생 시나리오에서 더 높은 정확도와 더 빠른 추론을 가능하게 합니다.