카트리지

Cartridges는 압축된 장문맥 표현을 매개변수화된 KV-cache 접두사로 저장하는 프롬프트 학습 방법입니다. 핵심 아이디어는 Cartridges: Lightweight and general-purpose long context representations via self-study 논문에서 유래했습니다.

고수준 개요와 동기에 대해서는 블로그 포스트 Cartridges: Storing long contexts in tiny caches with self-study를 참고하세요.

Cartridges와 Prefix Tuning의 차이

Prefix Tuning과 Cartridges 모두 past_key_values(prefix KV 캐시)를 기본 모델에 주입하여 서빙됩니다.

• Prefix Tuning은 가상 토큰 임베딩(및 선택적으로 MLP 투영)을 학습하여 KV 접두사를 생성합니다.
• Cartridges는 p개의 가상 토큰에 대한 레이어별 키/밸류 벡터인 KV 프리픽스 자체를 직접 학습하며, 실제 prefill KV(예: 코퍼스/시스템 프롬프트의 첫 p개 토큰)에서 초기화되도록 설계되었습니다.

이 논문은 안정성을 위해 첫 번째 토큰을 attention sink로 고정하는 것을 권장합니다(num_frozen_tokens=1이 기본값입니다).

사용법 (추론)

학습된 CARTRIDGE 어댑터를 로드하고 생성을 실행합니다:

from transformers import AutoModelForCausalLM, AutoTokenizer

from peft import PeftModel

model_id = "Qwen/Qwen2.5-0.5B-Instruct"
adapter_path = "path/to/cartridge_adapter"

base = AutoModelForCausalLM.from_pretrained(model_id)
model = PeftModel.from_pretrained(base, adapter_path)

tok = AutoTokenizer.from_pretrained(model_id)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

out = model.generate(**tok("Question about the corpus:", return_tensors="pt"), max_new_tokens=64)
print(tok.decode(out[0], skip_special_tokens=True))

학습 전에 카트리지를 생성하고 초기화해야 하는 경우, 아래 초기화 옵션을 참조하세요.

초기화 옵션

이 논문에서는 몇 가지 실용적인 초기화 전략을 다룹니다:

• Random KV(기본값): CartridgeConfig를 생성하고 학습을 시작합니다. 이렇게 하면 KV 접두사가 무작위로 초기화됩니다.
• 프롬프트/코퍼스에서 첫 번째 토큰의 KV: initialize_kv_prefix_from_text(model, tokenizer, text=...)을 사용합니다. 이 함수는 text에 대해 prefill을 실행하고, 생성된 KV 캐시의 첫 num_virtual_tokens를 어댑터로 복사합니다.
• 기존 캐시에서 KV: 베이스 모델의 prefill에서 이미 past_key_values 객체를 가지고 있다면 initialize_kv_prefix_from_past_key_values(model, past_key_values=...)를 사용합니다.

학습

Cartridges 논문은 SELF-STUDY 증류 목적 함수를 제안합니다(동결된 베이스 모델이 교사 로짓을 제공하고, CARTRIDGE 어댑터는 학생 모델이 대상 세그먼트에서 교사의 다음 토큰 분포와 일치하도록 학습됩니다). PEFT는 학습 로직을 핵심 라이브러리에서 분리합니다. 참고 워크플로우는 https://github.com/huggingface/peft/tree/main/examples/cartridge_self_study를 참조하세요. 예제 스크립트는 동결된 베이스 모델을 교사로, 어댑트된 모델을 학생으로 사용하므로, 두 모델은 동일한 기반 체크포인트를 공유합니다.

구성

독립적으로 학습된 카트리지를 단일 어댑터로 연결하려면 compose_cartridge_adapters(...)를 사용하세요.

API

CartridgeConfig

CartridgeEncoder

initialize_kv_prefix_from_past_key_values

prompt_embeddings_from_past_key_values