2 | 
3 |
4 |
15 |
16 | Cache-to-Cache (C2C) enables Large Language Models to communicate directly through their KV-Caches, bypassing text generation. By projecting and fusing KV-Caches between models, C2C achieves 8.5–10.5% higher accuracy than individual models and 3.0–5.0% better performance than text-based communication, with 2.0× speedup in latency.
17 |
18 | Feel free to star the repo or cite the paper if you find it interesting.
19 |
20 | ```bibtex
21 | @article{fu2025c2c,
22 | title={Cache-to-Cache: Direct Semantic Communication Between Large Language Models},
23 | author={Tianyu Fu and Zihan Min and Hanling Zhang and Jichao Yan and Guohao Dai and Wanli Ouyang and Yu Wang},
24 | journal={arXiv preprint arXiv:2510.03215},
25 | year={2025},
26 | }
27 | ```
28 |
29 | > **Why "Rosetta"?** The Python package is named after the **Rosetta Stone**, the ancient artefact that unlocked the translation of Egyptian hieroglyphs by presenting the same text in multiple scripts. Likewise, C2C translates KV-cache representations between otherwise independent LLMs, allowing them to speak a common language in a richer and more direct way.
30 |
31 |
32 | ## News
33 |
34 | [2025/12] 🧪 Multi-sharer support is now available! Fuse KV-caches from multiple sharer models to a single receiver. This feature is in preliminary stages and we are still actively working on it. See `live_chat_example.py` for usage.
35 |
36 | [2025/11] 🚀 Thank you for the enthusiasm from the community! [Live demo](https://huggingface.co/spaces/nics-efc/C2C_demo) is now available! Try C2C in action with side-by-side model comparison.
37 |
38 | [2025/10] 🤗 Our paper was featured as the **#1 Paper of the Day** on [Hugging Face Daily Papers](https://huggingface.co/papers/2510.03215)
39 |
40 | ## Demo
41 |
42 |
3 |
4 | Cache-to-Cache
5 |Direct Semantic Communication Between Large Language Models
6 | 7 |8 | 🌐 Project Page • 9 | 📑 Paper • 10 | 🤗 HuggingFace • 11 | 🚀 Live Demo 12 |
13 | 14 |
43 |
50 |
51 | The demo can be reproduced with `script/playground/gradio_demo.py`.
52 |
53 | ## Environment Setup
54 |
55 | Create a new environment:
56 |
57 | ```bash
58 | conda create -n rosetta python=3.10
59 | conda activate rosetta
60 | ```
61 |
62 | Install the package:
63 |
64 | ```bash
65 | pip install -e .
66 | ```
67 |
68 | For training and evaluation, install additional dependencies:
69 |
70 | ```bash
71 | pip install -e ".[training,evaluation]"
72 | ```
73 |
74 | ## How to
75 |
76 | ### Use Hugging Face weights
77 |
78 | Minimal example to load published C2C weights from the Hugging Face collection and run the provided inference script:
79 |
80 | ```python
81 | import torch
82 | from huggingface_hub import snapshot_download
83 | from script.playground.inference_example import load_rosetta_model, run_inference_example
84 |
85 | checkpoint_dir = snapshot_download(
86 | repo_id="nics-efc/C2C_Fuser",
87 | allow_patterns=["qwen3_0.6b+qwen2.5_0.5b_Fuser/*"],
88 | )
89 |
90 | model_config = {
91 | "rosetta_config": {
92 | "base_model": "Qwen/Qwen3-0.6B",
93 | "teacher_model": "Qwen/Qwen2.5-0.5B-Instruct",
94 | "checkpoints_dir": f"{checkpoint_dir}/qwen3_0.6b+qwen2.5_0.5b_Fuser/final",
95 | }
96 | }
97 |
98 | rosetta_model, tokenizer = load_rosetta_model(model_config, eval_config={}, device=torch.device("cuda"))
99 | device = rosetta_model.device
100 |
101 | prompt = [{"role": "user", "content": "Say hello in one short sentence."}]
102 | input_text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True, enable_thinking=False)
103 | inputs = tokenizer(input_text, return_tensors="pt").to(device)
104 |
105 | instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(inputs['input_ids'].shape[1] - 1, 1).unsqueeze(0).to(device)
106 | label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(device)
107 | kv_cache_index = [instruction_index, label_index]
108 |
109 | with torch.no_grad():
110 | sampling_params = {
111 | 'do_sample': False,
112 | 'max_new_tokens': 256
113 | }
114 | outputs = rosetta_model.generate(**inputs, kv_cache_index=kv_cache_index, **sampling_params)
115 | output_text = tokenizer.decode(outputs[0, instruction_index.shape[1] + 1:], skip_special_tokens=True)
116 | print(f"C2C output text: {output_text}")
117 | ```
118 |
119 | ### Run an example
120 |
121 | We provide an interactive chat example to demonstrate cache-to-cache communication with pre-trained projectors in `script/playground/live_chat_example.py`.
122 |
123 | ```bash
124 | # Single sharer
125 | python script/playground/live_chat_example.py --checkpoint_dir path/to/checkpoint
126 |
127 | # Multiple sharers (teacher models read from each checkpoint's config.json)
128 | python script/playground/live_chat_example.py --checkpoint_dir path/to/ckpt1 path/to/ckpt2
129 | ```
130 |
131 | ### Apply Cache-to-Cache
132 |
133 | You can apply C2C to your own models with a few lines of code. Here is an example:
134 |
135 | ```python
136 | import torch
137 | from transformers import AutoModelForCausalLM
138 | from rosetta.model.wrapper import RosettaModel
139 | from rosetta.model.projector import C2CProjector
140 |
141 | # Load target (receiver) and source (sharer) models
142 | target_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
143 | source_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
144 |
145 | # Create C2C projector for KV-Cache transformation
146 | projector_list = []
147 | for i in range(target_model.config.num_hidden_layers):
148 | projector = C2CProjector(
149 | source_dim=128, target_dim=128,
150 | source_num_heads=8, target_num_heads=8,
151 | hidden_dim=1024, num_layers=3
152 | )
153 | projector_list.append(projector)
154 | # If you want to use a pretrained projector, you can load it from the checkpoint directory
155 |
156 | # Wrap with RosettaModel for cache-to-cache communication
157 | c2c_model = RosettaModel(
158 | model_list=[target_model, source_model],
159 | base_model_idx=0,
160 | projector_list=projector_list
161 | )
162 |
163 | # Configure layer-wise projection mappings
164 | for idx, layer_idx in enumerate(range(target_model.config.num_hidden_layers)):
165 | c2c_model.set_projector_config(
166 | source_model_idx=1, source_model_layer_idx=layer_idx,
167 | target_model_idx=0, target_model_layer_idx=layer_idx,
168 | projector_idx=idx
169 | )
170 |
171 | # Generate: kv_cache_index controls when to apply C2C projection
172 | # [1, 0] = apply projection from sharer 1, [-1, 0] = no projection
173 | seq_len = input_ids.shape[1]
174 | instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(seq_len-1, 1)[None, :, :]
175 | response_index = torch.tensor([[-1, 0]], dtype=torch.long)[None, :, :]
176 | outputs = c2c_model.generate(
177 | kv_cache_index=[instruction_index, response_index],
178 | input_ids=inputs.input_ids,
179 | )
180 | ```
181 |
182 | ### Train C2C Projectors
183 |
184 | Prepare a training configuration file in `recipe/train_recipe/`. Specify the base model, teacher model, projector type and parameters, training hyperparameters, dataset, and output directory. See `recipe/train_recipe/C2C_0.6+0.5.json` for a complete example.
185 |
186 | Run training:
187 |
188 | ```bash
189 | # Single GPU
190 | python script/train/SFT_train.py --config recipe/train_recipe/C2C_0.6+0.5.json
191 |
192 | # Multi-GPU
193 | export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
194 | torchrun --nproc_per_node=8 script/train/SFT_train.py \
195 | --config recipe/train_recipe/C2C_0.6+0.5.json
196 | ```
197 |
198 | During training, only the C2C projector parameters are updated while both source and target models remain frozen.
199 |
200 | ### Evaluate C2C
201 |
202 | Prepare an evaluation configuration file in `recipe/eval_recipe/`. Specify the model configuration with base model, teacher model, checkpoint directory, generation config, evaluation dataset, and output directory. See `recipe/eval_recipe/unified_eval.yaml` for a complete example.
203 |
204 | Run evaluation:
205 |
206 | ```bash
207 | python script/evaluation/unified_evaluator.py --config recipe/eval_recipe/unified_eval.yaml
208 | ```
209 |
210 | ## Understanding the Code
211 |
212 | ### Code Structure
213 |
214 | - `rosetta/`: The main package for Cache-to-Cache.
215 | - `model/`: Core model components.
216 | - `train/`: Training utilities.
217 | - `baseline/`: Baseline implementations.
218 | - `utils/`: Utility functions for evaluation and model registry.
219 | - `script/`: Scripts for running experiments.
220 | - `train/`: Training scripts including `SFT_train.py`.
221 | - `evaluation/`: Evaluation scripts including `unified_evaluator.py`.
222 | - `dataset/`: Dataset preparation scripts.
223 | - `examples/`: Usage examples.
224 | - `recipe/`: Configuration files.
225 | - `train_recipe/`: Training configurations (e.g., `C2C_0.6+0.5.json`).
226 | - `eval_recipe/`: Evaluation configurations (e.g., `unified_eval.yaml`).
227 |
228 | ### Adding Projector
229 |
230 | Add a new projector architecture in `rosetta/model/projector.py`. The projector transforms source model's KV-Cache to target model's semantic space.
231 |
232 | > Key components: projection networks (MLP/concat-based), gating mechanism for layer-wise selection, and temperature-annealed training. See `C2CProjector` in `rosetta/model/projector.py` as an example.
233 |
234 | ```python
235 | from rosetta.utils.registry import register_model, capture_init_args
236 |
237 | @register_model
238 | @capture_init_args
239 | class MyProjector(Projector):
240 | def __init__(self, source_dim, target_dim, **kwargs):
241 | super().__init__()
242 | # Your architecture
243 | self.projection = nn.Linear(source_dim, target_dim)
244 | self.gate = nn.Parameter(torch.tensor(0.0))
245 | def forward(self, source_kv, target_kv):
246 | # Project and fuse KV-caches
247 | return projected_kv
248 | ```
249 |
250 | Register in configuration: `{"projector": {"type": "MyProjector", "params": {...}}}`
251 |
252 | ### Adding Dataset
253 |
254 | Add a new dataset in `rosetta/train/dataset_adapters.py` for training with your data.
255 |
256 | ```python
257 | @dataclass
258 | class MyDatasetConfig(DatasetConfig):
259 | dataset_name: str = "my_dataset"
260 | def load(self):
261 | return load_dataset("path/to/dataset")
262 |
263 | def my_formatting_func(examples):
264 | return {"text": [f"Q: {q}\nA: {a}" for q, a in zip(...)]}
265 |
266 | DATASET_CONFIGS["MyDataset"] = MyDatasetConfig
267 | FORMATTING_FUNCS["MyDataset"] = my_formatting_func
268 | ```
269 |
270 | Use in configuration: `{"data": {"type": "MyDataset"}}`
271 |
272 | ### Adding Benchmark
273 |
274 | Add evaluation logic in `script/evaluation/` following the pattern in `unified_evaluator.py`. The evaluator loads models, runs inference, and computes metrics for your benchmark dataset.
275 |
276 | ## Supported Model Pairs
277 |
278 | ### Qwen Family
279 |
280 | * Qwen3-0.6B + Qwen2.5-0.5B-Instruct
281 | * Qwen3-0.6B + Llama-3.2-1B-Instruct
282 | * Qwen3-0.6B + Qwen3-4B-Base
283 |
284 | ### Other Configurations
285 |
286 | C2C supports arbitrary model pairs. The framework automatically handles:
287 | - Different hidden dimensions
288 | - Different number of layers
289 | - Different attention head configurations
290 | - Different tokenizers
291 |
292 | To use custom model pairs, simply specify them in your configurations.
293 |
--------------------------------------------------------------------------------
/script/playground/inference_example.py:
--------------------------------------------------------------------------------
1 | """
2 | Example inference using RosettaModel with Qwen3-0.6B and Qwen3-1.7B models and MLP projector
3 | """
4 |
5 | import torch
6 | import sys
7 | import os
8 | from pathlib import Path
9 |
10 | # Add the project root to the path
11 | sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
12 |
13 | from transformers import AutoModelForCausalLM, AutoTokenizer
14 | from rosetta.model.wrapper import RosettaModel
15 | from rosetta.model.aligner import TokenAligner, AlignmentStrategy
16 | from rosetta.model.projector import AllInOneProjector
17 | from rosetta.train.dataset_adapters import generate_kv_cache_index
18 | from typing import Dict, Any, List, Tuple, Optional
19 | from transformers import AutoModelForCausalLM, AutoTokenizer
20 |
21 | from rosetta.model.projector import load_projector
22 | from rosetta.model.wrapper import RosettaModel
23 | from rosetta.utils.evaluate import set_default_chat_template
24 | import re
25 |
26 | def test_token_aligner(slm_tokenizer: AutoTokenizer, llm_tokenizer: AutoTokenizer):
27 | """Test the TokenAligner functionality
28 | Args:
29 | slm_tokenizer: SLM tokenizer
30 | llm_tokenizer: LLM tokenizer
31 | """
32 | print("\n" + "="*80)
33 | print("Testing TokenAligner")
34 | print("="*80)
35 |
36 | # Test with FIRST strategy
37 | aligner_first = TokenAligner(
38 | slm_tokenizer=slm_tokenizer,
39 | llm_tokenizer=llm_tokenizer,
40 | strategy=AlignmentStrategy.FIRST,
41 | verbose=True
42 | )
43 |
44 | # Test with LONGEST strategy
45 | aligner_longest = TokenAligner(
46 | slm_tokenizer=slm_tokenizer,
47 | llm_tokenizer=llm_tokenizer,
48 | strategy=AlignmentStrategy.LONGEST,
49 | verbose=True
50 | )
51 |
52 | # Test text samples
53 | test_texts = [
54 | "Hello world!",
55 | "The future of artificial intelligence is",
56 | "北京是中国的首都", # Chinese text
57 | "🚀 Emojis and special characters!",
58 | ]
59 |
60 | for text in test_texts:
61 | print(f"\nTest text: '{text}'")
62 | print("-" * 40)
63 |
64 | # Test FIRST strategy
65 | print("\nFIRST Strategy:")
66 | aligner_first.visualize_alignment(text)
67 |
68 | # Test LONGEST strategy
69 | print("\nLONGEST Strategy:")
70 | aligner_longest.visualize_alignment(text)
71 |
72 | # Test alignment without visualization
73 | sample_text = "This is a test."
74 | slm_tokens, aligned_llm_tokens = aligner_first.align_sequence(sample_text)
75 | print(f"\nQuick alignment test for: '{sample_text}'")
76 | print(f"SLM tokens: {slm_tokens}")
77 | print(f"Aligned LLM tokens: {aligned_llm_tokens}")
78 |
79 | print("\n✅ TokenAligner test completed")
80 |
81 |
82 | def run_inference_example(rosetta_model: RosettaModel, tokenizer: AutoTokenizer, prompt: str):
83 | """Run inference example with RosettaModel
84 | Args:
85 | rosetta_model: RosettaModel
86 | tokenizer: AutoTokenizer
87 | prompt: str
88 | """
89 | print("Running inference example...")
90 |
91 | device = rosetta_model.device
92 |
93 | # Prepare input
94 |
95 |
96 | input_text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True, enable_thinking=False)
97 | print(f"Input text: {input_text}")
98 | inputs = tokenizer(input_text, return_tensors="pt").to(device)
99 | print(f"Input tokens: {inputs['input_ids']}")
100 | instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(inputs['input_ids'].shape[1] - 1, 1).unsqueeze(0).to(device)
101 | label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(device)
102 | kv_cache_index = [instruction_index, label_index]
103 | # slm_tokenizer = tokenizer
104 | # llm_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
105 | # strategy = "first"
106 | # aligner = TokenAligner(slm_tokenizer=slm_tokenizer, llm_tokenizer=llm_tokenizer, strategy=AlignmentStrategy(strategy))
107 | # messages = [
108 | # {"role": "user", "content": prompt}
109 | # ]
110 | # details = aligner.align_chat_messages(messages, add_generation_prompt=True, return_details=True)
111 | # slm_ids = torch.tensor(details['slm_ids_padded']).unsqueeze(0)
112 | # llm_ids = torch.tensor(details['llm_ids_padded']).unsqueeze(0)
113 |
114 | # slm_pad_mask = torch.tensor(details['slm_padding_mask']).unsqueeze(0)
115 | # llm_pad_mask = torch.tensor(details['llm_padding_mask']).unsqueeze(0)
116 |
117 | # slm_attention_mask = (~slm_pad_mask).float()
118 | # llm_attention_mask = (~llm_pad_mask).float()
119 |
120 | # message_mask = torch.tensor(details['message_mask'])
121 | # kv_cache_index = generate_kv_cache_index(slm_ids.shape[1], slm_ids.shape[1])
122 | # kv_cache_index[~message_mask] = torch.tensor([[-1,0]])
123 |
124 | # kv_idx = kv_cache_index
125 | # change_points = [0]
126 | # for i in range(1, kv_idx.size(0)):
127 | # if not torch.equal(kv_idx[i], kv_idx[i - 1]):
128 | # change_points.append(i)
129 | # change_points.append(kv_idx.size(0))
130 |
131 | # kv_cache_list = []
132 |
133 | # for i in range(len(change_points) - 1):
134 | # start = change_points[i]
135 | # end = change_points[i + 1]
136 | # kv_cache_list.append(kv_idx[start:end, :].unsqueeze(0).to(device))
137 | # prefill_kv_cache_list = kv_cache_list[:-1]
138 | # print(f"Input prompt: '{prompt}'")
139 | # print(f"Input shape: {slm_ids.shape}")
140 | # print(f"Device: {device}")
141 |
142 | # slm_ids = slm_ids.to(device)
143 | # llm_ids = llm_ids.to(device)
144 | # slm_attention_mask = slm_attention_mask.to(device)
145 | # llm_attention_mask = llm_attention_mask.to(device)
146 |
147 | # Run inference
148 | # with torch.no_grad():
149 | # # outputs = rosetta_model.forward(
150 | # # input_ids=[slm_ids, llm_ids],
151 | # # attention_mask=[slm_attention_mask, llm_attention_mask],
152 | # # kv_cache_index=kv_cache_list,
153 | # # position_ids=torch.arange(slm_ids.shape[1]).unsqueeze(0).to(device),
154 | # # use_cache=True,
155 | # # output_attentions=False,
156 | # # output_hidden_states=False,
157 | # # sample=False,
158 | # # )
159 | # outputs = rosetta_model(**inputs, kv_cache_index=kv_cache_index)
160 |
161 | # # Get logits and generate next token
162 | # logits = outputs.logits
163 | # next_token_logits = logits[0, -1, :]
164 | # next_token_id = torch.argmax(next_token_logits, dim=-1)
165 | # next_token = tokenizer.decode(next_token_id)
166 |
167 | # print(f"Output logits shape: {logits.shape}")
168 | # print(f"Next predicted token: '{next_token}'")
169 | # print("✅ Inference completed successfully")
170 |
171 | # Run generation
172 | with torch.no_grad():
173 | # outputs = rosetta_model.generate(
174 | # prefill_kv_cache_index=prefill_kv_cache_list,
175 | # input_ids=[slm_ids, llm_ids],
176 | # attention_mask=[slm_attention_mask, llm_attention_mask],
177 | # use_cache=True,
178 | # output_attentions=False,
179 | # output_hidden_states=False,
180 | # max_new_tokens=256,
181 | # do_sample=False,
182 | # )
183 | sampling_params = {
184 | 'do_sample': True,
185 | 'temperature': 0.7,
186 | 'top_p': 0.8,
187 | 'top_k': 20,
188 | 'min_p': 0.0,
189 | 'repetition_penalty': 1.2,
190 | 'max_new_tokens': 1024
191 | }
192 | outputs = rosetta_model.generate(**inputs, kv_cache_index=kv_cache_index, **sampling_params)
193 | output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
194 | print(f"Rosetta output text: {output_text}")
195 |
196 | with torch.no_grad():
197 | sampling_params = {
198 | 'do_sample': True,
199 | 'temperature': 0.7,
200 | 'top_p': 0.8,
201 | 'top_k': 20,
202 | 'min_p': 0.0,
203 | 'repetition_penalty': 1.2,
204 | 'max_new_tokens': 1024
205 | }
206 | slm_model = rosetta_model.model_list[0]
207 | outputs = slm_model.generate(**inputs, **sampling_params)
208 | output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
209 | print(f"SLM output text: {output_text}")
210 |
211 | with torch.no_grad():
212 | sampling_params = {
213 | 'do_sample': True,
214 | 'temperature': 0.7,
215 | 'top_p': 0.8,
216 | 'top_k': 20,
217 | 'min_p': 0.0,
218 | 'repetition_penalty': 1.2,
219 | 'max_new_tokens': 1024
220 | }
221 | llm_model = rosetta_model.model_list[1]
222 | outputs = llm_model.generate(**inputs, **sampling_params)
223 | output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
224 | print(f"LLM output text: {output_text}")
225 |
226 | def load_rosetta_model(model_config: Dict[str, Any], eval_config: Dict[str, Any],
227 | device: torch.device) -> Tuple[Any, Any]:
228 | """
229 | Load Rosetta model with projectors and aggregators.
230 |
231 | Args:
232 | model_config: Model configuration dict
233 | eval_config: Evaluation configuration dict
234 | device: Device to load model on
235 |
236 | Returns:
237 | Tuple of (rosetta_model, tokenizer)
238 | """
239 | # Prefer checkpoints_dir under model.rosetta_config; fall back to eval config for backward compatibility
240 | rosetta_config = model_config["rosetta_config"]
241 | checkpoint_dir = rosetta_config.get("checkpoints_dir", eval_config.get("checkpoints_dir"))
242 | if checkpoint_dir is None:
243 | raise KeyError("checkpoints_dir must be provided under model.rosetta_config (preferred) or eval config (legacy)")
244 | slm_model_path = rosetta_config["base_model"]
245 | llm_model_path = rosetta_config["teacher_model"]
246 |
247 | # Load tokenizer
248 | slm_tokenizer = AutoTokenizer.from_pretrained(str(slm_model_path))
249 | set_default_chat_template(slm_tokenizer, slm_model_path)
250 |
251 | # Load models
252 | slm_model = AutoModelForCausalLM.from_pretrained(
253 | str(slm_model_path),
254 | torch_dtype=torch.bfloat16,
255 | device_map={"": device}
256 | ).eval()
257 |
258 | llm_model = AutoModelForCausalLM.from_pretrained(
259 | str(llm_model_path),
260 | torch_dtype=torch.bfloat16,
261 | device_map={"": device}
262 | ).eval()
263 |
264 | # Load projectors
265 | num_projectors = len([f for f in os.listdir(checkpoint_dir) if re.match(r"projector_\d+\.pt", f)])
266 | projector_list = []
267 | for t in range(num_projectors):
268 | json_cfg = os.path.join(checkpoint_dir, f"projector_{t}.json")
269 | proj = load_projector(json_cfg)
270 | proj = proj.to(device)
271 | pt_path = os.path.join(checkpoint_dir, f"projector_{t}.pt")
272 | if os.path.exists(pt_path):
273 | state_dict = torch.load(pt_path, map_location=device)
274 | proj.load_state_dict(state_dict, strict=False)
275 | projector_list.append(proj)
276 |
277 | # Load aggregators
278 | num_aggregators = len([f for f in os.listdir(checkpoint_dir) if re.match(r"aggregator_\d+\.pt", f)])
279 | aggregator_list = []
280 | for t in range(num_aggregators):
281 | json_cfg = os.path.join(checkpoint_dir, f"aggregator_{t}.json")
282 | agg_path = os.path.join(checkpoint_dir, f"aggregator_{t}.pt")
283 | agg = load_aggregator(json_cfg)
284 | if os.path.exists(agg_path):
285 | sd = torch.load(agg_path, map_location="cpu")
286 | agg.load_state_dict(sd, strict=False)
287 | agg = agg.to(device)
288 | aggregator_list.append(agg)
289 |
290 | # Initialize Rosetta model
291 | rosetta_model = RosettaModel(
292 | model_list=[slm_model, llm_model],
293 | base_model_idx=0,
294 | projector_list=projector_list,
295 | aggregator_list=aggregator_list,
296 | ).to(device).eval()
297 |
298 | # Load projector/aggregator mapping configs
299 | proj_cfg_path = os.path.join(checkpoint_dir, "projector_config.json")
300 | agg_cfg_path = os.path.join(checkpoint_dir, "aggregator_config.json")
301 | rosetta_model.load_projector_config(proj_cfg_path)
302 | rosetta_model.load_aggregator_config(agg_cfg_path)
303 |
304 | return rosetta_model, slm_tokenizer
305 |
306 | def main():
307 | """Main function to run the inference example"""
308 |
309 | rosetta_model, slm_tokenizer = load_rosetta_model(
310 | model_config={
311 | "rosetta_config": {
312 | "base_model": "Qwen/Qwen3-0.6B",
313 | "teacher_model": "Qwen/Qwen3-4B",
314 | "checkpoints_dir": "local/checkpoints/0.6B_4B_general/final"
315 | }
316 | },
317 | eval_config={},
318 | device=torch.device("cuda")
319 | )
320 |
321 | # Test token aligner
322 | # test_token_aligner(slm_tokenizer, llm_tokenizer)
323 |
324 | # Run inference
325 | prompt = [{
326 | "role": "user",
327 | "content": "Accurately answer the following question:\n\nStatement 1 | If T: V -> W is a linear transformation and dim(V ) < dim(W) < 1, then T must be injective. Statement 2 | Let dim(V) = n and suppose that T: V -> V is linear. If T is injective, then it is a bijection.\n\nAre these statements correct? Let's think step by step and then answer the question starting with Answer:"
328 | }]
329 | run_inference_example(rosetta_model, slm_tokenizer, prompt)
330 | # run_inference_example(rosetta_model, slm_tokenizer, "从美国向北进入加拿大时，您会看到北星（北极星）越来越")
331 |
332 |
333 | if __name__ == "__main__":
334 | # import debugpy
335 | # debugpy.listen(("0.0.0.0", 5678))
336 | # print("Waiting for debugger attach...")
337 | # debugpy.wait_for_client()
338 | # print("Debugger attached, running...")
339 | main()
340 |
--------------------------------------------------------------------------------
/script/analysis/scaling/batch_evaluate_checkpoints.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | """
3 | 批量评测已训练的模型检查点脚本
4 | 自动遍历scaling_MMLU_15k目录下的所有checkpoint，生成评测配置并运行评测
5 | """
6 |
7 | import os
8 | import sys
9 | import json
10 | import yaml
11 | import subprocess
12 | from pathlib import Path
13 | from typing import Dict, List, Tuple, Optional
14 | import logging
15 | from datetime import datetime
16 |
17 | # 设置日志
18 | logging.basicConfig(
19 | level=logging.INFO,
20 | format='%(asctime)s - %(levelname)s - %(message)s',
21 | handlers=[
22 | logging.FileHandler('batch_evaluate.log'),
23 | logging.StreamHandler()
24 | ]
25 | )
26 | logger = logging.getLogger(__name__)
27 |
28 | class BatchCheckpointEvaluator:
29 | """批量checkpoint评测器"""
30 |
31 | def __init__(self):
32 | # 路径配置
33 | self.checkpoints_base_dir = Path("local/checkpoints/scaling_MMLU_15k")
34 | self.eval_template_path = "eval_recipe/test_eval.yaml"
35 | self.training_configs_dir = Path("recipe")
36 | self.evaluation_script = "script/evaluation/unified_evaluator.py"
37 |
38 | # 模型家族定义（与auto_scaling_experiment.py保持一致）
39 | self.family1 = {
40 | "tiny": "Qwen/Qwen3-0.6B",
41 | "small": "Qwen/Qwen3-1.7B",
42 | "medium": "Qwen/Qwen3-4B",
43 | "large": "Qwen/Qwen3-8B",
44 | "xlarge": "Qwen/Qwen3-14B"
45 | }
46 |
47 | self.family2 = {
48 | "tiny": "Qwen/Qwen2.5-0.5B-Instruct",
49 | "small": "Qwen/Qwen2.5-1.5B-Instruct",
50 | "medium": "Qwen/Qwen2.5-3B-Instruct",
51 | "large": "Qwen/Qwen2.5-7B-Instruct",
52 | "xlarge": "Qwen/Qwen2.5-14B-Instruct"
53 | }
54 |
55 | # 尺寸映射
56 | self.size_name_to_key = {
57 | "0.6B": "tiny", "1.7B": "small", "4B": "medium", "8B": "large", "14B": "xlarge",
58 | "0.5B": "tiny", "1.5B": "small", "3B": "medium", "7B": "large"
59 | }
60 |
61 | def find_checkpoints(self) -> List[Tuple[str, Path]]:
62 | """
63 | 找到所有可用的checkpoint目录
64 |
65 | Returns:
66 | List of (experiment_name, checkpoint_path) tuples
67 | """
68 | checkpoints = []
69 |
70 | if not self.checkpoints_base_dir.exists():
71 | logger.error(f"Checkpoints base directory does not exist: {self.checkpoints_base_dir}")
72 | return checkpoints
73 |
74 | for item in self.checkpoints_base_dir.iterdir():
75 | if item.is_dir() and item.name.startswith("scaling_"):
76 | # 检查是否存在final子目录
77 | final_dir = item / "final"
78 | if final_dir.exists() and final_dir.is_dir():
79 | experiment_name = item.name.replace("scaling_", "")
80 | checkpoints.append((experiment_name, final_dir))
81 | logger.info(f"Found checkpoint: {experiment_name} -> {final_dir}")
82 | else:
83 | logger.warning(f"No 'final' directory found in {item}")
84 |
85 | logger.info(f"Found {len(checkpoints)} valid checkpoints")
86 | return checkpoints
87 |
88 | def parse_experiment_name(self, experiment_name: str) -> Optional[Tuple[str, str]]:
89 | """
90 | 解析实验名称，提取base模型和teacher模型信息
91 |
92 | Args:
93 | experiment_name: 实验名称，如 "0.6B+0.5B", "1.7B+3B"
94 |
95 | Returns:
96 | (base_size_key, teacher_size_key) 或 None
97 | """
98 | if "+" not in experiment_name:
99 | logger.error(f"Invalid experiment name format: {experiment_name}")
100 | return None
101 |
102 | try:
103 | base_size, teacher_size = experiment_name.split("+")
104 |
105 | # 查找对应的尺寸键
106 | base_key = self.size_name_to_key.get(base_size)
107 | teacher_key = self.size_name_to_key.get(teacher_size)
108 |
109 | if base_key is None or teacher_key is None:
110 | logger.error(f"Unknown model sizes in experiment: {experiment_name}")
111 | logger.error(f" Base size '{base_size}' -> {base_key}")
112 | logger.error(f" Teacher size '{teacher_size}' -> {teacher_key}")
113 | return None
114 |
115 | return base_key, teacher_key
116 |
117 | except ValueError:
118 | logger.error(f"Failed to parse experiment name: {experiment_name}")
119 | return None
120 |
121 | def load_eval_template(self) -> Dict:
122 | """加载评测模板配置"""
123 | try:
124 | with open(self.eval_template_path, 'r', encoding='utf-8') as f:
125 | return yaml.safe_load(f)
126 | except Exception as e:
127 | logger.error(f"Failed to load eval template {self.eval_template_path}: {e}")
128 | raise
129 |
130 | def create_eval_config(self, experiment_name: str, checkpoint_path: Path,
131 | base_model_path: str, teacher_model_path: str) -> str:
132 | """
133 | 创建评测配置文件
134 |
135 | Args:
136 | experiment_name: 实验名称
137 | checkpoint_path: checkpoint路径
138 | base_model_path: base模型路径
139 | teacher_model_path: teacher模型路径
140 |
141 | Returns:
142 | 生成的配置文件路径
143 | """
144 | # 加载模板
145 | config = self.load_eval_template()
146 |
147 | # 修改配置
148 | config["model"]["model_name"] = "Rosetta"
149 | config["model"]["rosetta_config"]["base_model"] = base_model_path
150 | config["model"]["rosetta_config"]["teacher_model"] = teacher_model_path
151 | config["model"]["rosetta_config"]["checkpoints_dir"] = str(checkpoint_path)
152 |
153 | # 设置输出目录
154 | output_dir = f"local/scaling_new_prompt_results/{experiment_name}_mmlu-redux"
155 | config["output"]["output_dir"] = output_dir
156 |
157 | # 保存配置文件
158 | config_filename = f"eval_recipe/scaling_new_prompt_{experiment_name}_eval.yaml"
159 | with open(config_filename, 'w', encoding='utf-8') as f:
160 | yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
161 |
162 | logger.info(f"Created eval config: {config_filename}")
163 | logger.info(f" Base model: {base_model_path}")
164 | logger.info(f" Teacher model: {teacher_model_path}")
165 | logger.info(f" Checkpoint: {checkpoint_path}")
166 | logger.info(f" Output dir: {output_dir}")
167 |
168 | return config_filename
169 |
170 | def run_evaluation(self, config_path: str, experiment_name: str) -> bool:
171 | """
172 | 运行单个实验的评测
173 |
174 | Args:
175 | config_path: 评测配置文件路径
176 | experiment_name: 实验名称
177 |
178 | Returns:
179 | 是否评测成功
180 | """
181 | logger.info(f"Starting evaluation for experiment: {experiment_name}")
182 |
183 | try:
184 | # 构建评测命令
185 | cmd = [
186 | "python", self.evaluation_script,
187 | "--config", config_path
188 | ]
189 |
190 | logger.info(f"Running command: {' '.join(cmd)}")
191 |
192 | # 运行评测（实时输出）
193 | process = subprocess.Popen(
194 | cmd,
195 | cwd=os.getcwd(),
196 | stdout=subprocess.PIPE,
197 | stderr=subprocess.STDOUT,
198 | text=True,
199 | bufsize=1,
200 | )
201 |
202 | # 实时输出日志
203 | if process.stdout:
204 | for line in process.stdout:
205 | sys.stdout.write(line)
206 | sys.stdout.flush()
207 |
208 | process.wait()
209 |
210 | if process.returncode == 0:
211 | logger.info(f"Evaluation completed successfully for {experiment_name}")
212 | return True
213 | else:
214 | logger.error(f"Evaluation failed for {experiment_name}")
215 | logger.error(f"Process returned code {process.returncode}")
216 | return False
217 |
218 | except Exception as e:
219 | logger.error(f"Exception during evaluation {experiment_name}: {e}")
220 | return False
221 |
222 | def cleanup_temp_configs(self, temp_configs: List[str]):
223 | """清理临时生成的配置文件"""
224 | for config_path in temp_configs:
225 | try:
226 | if os.path.exists(config_path):
227 | os.remove(config_path)
228 | logger.info(f"Removed temporary config: {config_path}")
229 | except Exception as e:
230 | logger.warning(f"Failed to remove temporary config {config_path}: {e}")
231 |
232 | def check_existing_results(self, experiment_name: str) -> bool:
233 | """检查实验结果是否已存在"""
234 | output_dir = Path(f"local/scaling_new_prompt_results/{experiment_name}_mmlu-redux")
235 | if output_dir.exists() and any(output_dir.iterdir()):
236 | logger.info(f"Results already exist for {experiment_name}, skipping")
237 | return True
238 | return False
239 |
240 | def run_batch_evaluation(self, skip_existing: bool = True, dry_run: bool = False):
241 | """
242 | 运行批量评测
243 |
244 | Args:
245 | skip_existing: 是否跳过已有结果的实验
246 | dry_run: 是否为试运行（不实际执行评测）
247 | """
248 | logger.info("=" * 80)
249 | logger.info("STARTING BATCH CHECKPOINT EVALUATION")
250 | logger.info("=" * 80)
251 |
252 | # 查找所有checkpoint
253 | checkpoints = self.find_checkpoints()
254 |
255 | if not checkpoints:
256 | logger.error("No valid checkpoints found!")
257 | return
258 |
259 | logger.info(f"Found {len(checkpoints)} checkpoints to evaluate:")
260 | for exp_name, checkpoint_path in checkpoints:
261 | logger.info(f" ✓ {exp_name} -> {checkpoint_path}")
262 |
263 | if dry_run:
264 | logger.info("DRY RUN MODE - No actual evaluations will be performed")
265 |
266 | successful_evaluations = 0
267 | failed_evaluations = 0
268 | skipped_evaluations = 0
269 | temp_configs = []
270 |
271 | for i, (experiment_name, checkpoint_path) in enumerate(checkpoints, 1):
272 | logger.info(f"\n[{i}/{len(checkpoints)}] Processing experiment: {experiment_name}")
273 |
274 | # 检查是否已有结果
275 | if skip_existing and self.check_existing_results(experiment_name):
276 | skipped_evaluations += 1
277 | continue
278 |
279 | # 解析实验名称
280 | parsed = self.parse_experiment_name(experiment_name)
281 | if parsed is None:
282 | logger.error(f"Failed to parse experiment name: {experiment_name}")
283 | failed_evaluations += 1
284 | continue
285 |
286 | base_key, teacher_key = parsed
287 | base_model_path = self.family1[base_key]
288 | teacher_model_path = self.family2[teacher_key]
289 |
290 | try:
291 | # 创建评测配置
292 | config_path = self.create_eval_config(
293 | experiment_name, checkpoint_path,
294 | base_model_path, teacher_model_path
295 | )
296 | temp_configs.append(config_path)
297 |
298 | if dry_run:
299 | logger.info(f"DRY RUN: Would evaluate {experiment_name}")
300 | successful_evaluations += 1
301 | else:
302 | # 运行评测
303 | success = self.run_evaluation(config_path, experiment_name)
304 |
305 | if success:
306 | successful_evaluations += 1
307 | else:
308 | failed_evaluations += 1
309 |
310 | except Exception as e:
311 | logger.error(f"Failed to process {experiment_name}: {e}")
312 | failed_evaluations += 1
313 |
314 | # 清理临时配置文件
315 | if not dry_run:
316 | self.cleanup_temp_configs(temp_configs)
317 |
318 | # 输出最终结果
319 | logger.info(f"\n" + "=" * 80)
320 | logger.info("BATCH EVALUATION SUMMARY")
321 | logger.info(f"=" * 80)
322 | logger.info(f"Total checkpoints found: {len(checkpoints)}")
323 | logger.info(f"Successful evaluations: {successful_evaluations}")
324 | logger.info(f"Failed evaluations: {failed_evaluations}")
325 | logger.info(f"Skipped evaluations: {skipped_evaluations}")
326 |
327 | if dry_run:
328 | logger.info("\nThis was a DRY RUN. No actual evaluations were performed.")
329 | logger.info("Run with --no-dry-run to perform actual evaluations.")
330 |
331 |
332 | def main():
333 | """主函数"""
334 | import argparse
335 |
336 | parser = argparse.ArgumentParser(description='Batch evaluate trained checkpoints')
337 | parser.add_argument('--skip-existing', action='store_true', default=True,
338 | help='Skip experiments that already have results (default: True)')
339 | parser.add_argument('--no-skip-existing', action='store_false', dest='skip_existing',
340 | help='Evaluate all experiments, even if results exist')
341 | parser.add_argument('--dry-run', action='store_true', default=False,
342 | help='Show what would be evaluated without running actual evaluations')
343 | parser.add_argument('--no-dry-run', action='store_false', dest='dry_run',
344 | help='Perform actual evaluations (default)')
345 |
346 | args = parser.parse_args()
347 |
348 | # 创建评测器并运行
349 | evaluator = BatchCheckpointEvaluator()
350 | evaluator.run_batch_evaluation(
351 | skip_existing=args.skip_existing,
352 | dry_run=args.dry_run
353 | )
354 |
355 |
356 | if __name__ == "__main__":
357 | main()
358 |
--------------------------------------------------------------------------------
/rosetta/model/ablation_projector.py:
--------------------------------------------------------------------------------
1 | """
2 | Ablation Projector: A configurable projector for ablation studies based on C2CProjector.
3 | Allows gradual removal of components to study their individual contributions.
4 | """
5 |
6 | import torch
7 | import torch.nn as nn
8 | from torch import Tensor
9 | from typing import Optional, Tuple, Literal
10 |
11 | from rosetta.utils.registry import register_model, capture_init_args
12 | from rosetta.model.projector import Projector
13 | from rosetta.model.projector import RegularMLP
14 |
15 |
16 | @register_model
17 | @capture_init_args
18 | class AblationProjector(Projector):
19 | """
20 | Ablation study projector based on C2CProjector with configurable component removal.
21 |
22 | Ablation levels:
23 | 0. Full C2C (baseline)
24 | 1. Remove scalar weights (set to 1.0)
25 | 2. Remove gates (set to 1.0)
26 | 3. Remove target contribution (only use source)
27 | 4. Remove gates only (gates=1.0), keep scalars and target
28 |
29 | Each level builds on the previous one, allowing gradual degradation study.
30 | """
31 |
32 | def __init__(
33 | self,
34 | source_dim: int,
35 | target_dim: int,
36 | source_num_heads: int = 1,
37 | target_num_heads: int = 1,
38 | intermediate_dim: int = 1024,
39 | hidden_dim: int = 1024,
40 | num_layers: int = 3,
41 | dropout: float = 0.1,
42 | initial_temperature: float = 1.0,
43 | final_temperature: float = 0.001,
44 | anneal_steps: int = 1929,
45 | dtype: torch.dtype = torch.float32,
46 |
47 | # Ablation configuration
48 | ablation_level: int = 0, # 0=full, 1=no_scalar, 2=no_gate+no_scalar, 3=no_target, 4=no_gate_only
49 | use_scalar_weights: bool = True, # Can be overridden by ablation_level
50 | use_gates: bool = True, # Can be overridden by ablation_level
51 | use_target: bool = True, # Can be overridden by ablation_level
52 | ):
53 | super().__init__()
54 |
55 | assert 0 <= ablation_level <= 4, "ablation_level must be 0, 1, 2, 3, or 4"
56 |
57 | # Dimensions
58 | self.source_dim = source_dim
59 | self.target_dim = target_dim
60 | self.source_num_heads = source_num_heads
61 | self.target_num_heads = target_num_heads
62 | self.ablation_level = ablation_level
63 |
64 | # Override component usage based on ablation level
65 | if ablation_level == 4:
66 | # Special case: disable gates only, keep scalars and target
67 | use_scalar_weights = True
68 | use_gates = False
69 | use_target = True
70 | else:
71 | if ablation_level >= 1:
72 | use_scalar_weights = False
73 | if ablation_level >= 2:
74 | use_gates = False
75 | if ablation_level >= 3:
76 | use_target = False
77 |
78 | self.use_scalar_weights = use_scalar_weights
79 | self.use_gates = use_gates
80 | self.use_target = use_target
81 |
82 | # Sizes
83 | in_dim = source_dim * source_num_heads
84 | out_dim = target_dim * target_num_heads
85 |
86 | # 1) concat(source_X, target_X) then project to hidden_dim
87 | # If not using target, only use source features
88 | if self.use_target:
89 | self.key_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
90 | self.value_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
91 | else:
92 | # Only use source features
93 | self.key_in = nn.Linear(in_dim, hidden_dim, bias=True, dtype=dtype)
94 | self.value_in = nn.Linear(in_dim, hidden_dim, bias=True, dtype=dtype)
95 |
96 | # 2) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
97 | self.key_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
98 | self.value_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
99 |
100 | # 3a) intermediate representation → (L-2)-layer MLP for weights → project to head dim
101 | # Only build if using scalar weights
102 | if self.use_scalar_weights:
103 | self.key_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
104 | self.value_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
105 | self.key_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
106 | self.value_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
107 |
108 | # 3b) intermediate representation → (L-2)-layer MLP for projected_X → finally project hidden_dim → out_dim
109 | self.key_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
110 | self.value_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
111 | self.key_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
112 | self.value_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
113 |
114 | # Scalar key/value gate parameters and temperature schedule
115 | # Only build if using gates
116 | if self.use_gates:
117 | self.key_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
118 | self.value_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
119 | self.use_gumbel = True
120 | self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype))
121 | self.initial_temperature = initial_temperature
122 | self.final_temperature = final_temperature
123 | self.anneal_steps = anneal_steps
124 |
125 | # Temperature for weight normalization
126 | self.scalar_temperature = 1.0
127 |
128 | def update_temperature(self, step: int):
129 | """Update temperature using exponential annealing schedule for gates."""
130 | if self.use_gates:
131 | ratio = min(step / self.anneal_steps, 1.0)
132 | temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** ratio
133 | self.gate_temperature.fill_(temp)
134 |
135 | def forward(
136 | self,
137 | source_kv: Tuple[Tensor, Tensor],
138 | target_kv: Tuple[Tensor, Tensor],
139 | position_ids: Optional[Tensor] = None,
140 | max_pos: Optional[Tensor] = None,
141 | ) -> Tuple[Tensor, Tensor]:
142 | source_key, source_value = source_kv
143 | target_key, target_value = target_kv
144 |
145 | B, Hs, N, Ds = source_key.shape
146 | _, Ht, _, Dt = target_key.shape
147 |
148 | # Flatten heads
149 | source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
150 | source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
151 | target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
152 | target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
153 |
154 | # 1) Prepare input features based on ablation level
155 | if self.use_target:
156 | # Full C2C: concat source and target features
157 | key_cat = torch.cat([source_key_flat, target_key_flat], dim=-1)
158 | value_cat = torch.cat([source_value_flat, target_value_flat], dim=-1)
159 | else:
160 | # Ablation level 3: only use source features
161 | key_cat = source_key_flat
162 | value_cat = source_value_flat
163 |
164 | # 2) project to hidden dim
165 | key_hidden = self.key_in(key_cat)
166 | value_hidden = self.value_in(value_cat)
167 |
168 | # 3) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
169 | key_hidden = self.key_mlp1(key_hidden)
170 | value_hidden = self.value_mlp1(value_hidden)
171 |
172 | # 4b) intermediate representation -> projected feature path
173 | key_proj_hidden = self.key_proj_out(self.key_proj_mlp2(key_hidden)) # (B, N, Ht * Dt)
174 | value_proj_hidden = self.value_proj_out(self.value_proj_mlp2(value_hidden)) # (B, N, Ht * Dt)
175 | projected_key = key_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
176 | projected_value = value_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
177 |
178 | # 4a) intermediate representation -> scalar path (if using scalar weights)
179 | if self.use_scalar_weights:
180 | key_scalar = self.key_scalar_head(self.key_scalar_mlp2(key_hidden)) # (B, N, Ht)
181 | value_scalar = self.value_scalar_head(self.value_scalar_mlp2(value_hidden)) # (B, N, Ht)
182 | key_scalar = key_scalar.permute(0, 2, 1).unsqueeze(-1) # (B, Ht, N, 1)
183 | value_scalar = value_scalar.permute(0, 2, 1).unsqueeze(-1) # (B, Ht, N, 1)
184 | # Normalize scalars
185 | norm_key_scalar = torch.sigmoid(key_scalar)
186 | norm_value_scalar = torch.sigmoid(value_scalar)
187 | else:
188 | # Ablation level 1+: set scalar weights to 1.0
189 | norm_key_scalar = torch.ones(B, Ht, N, 1, device=projected_key.device, dtype=projected_key.dtype)
190 | norm_value_scalar = torch.ones(B, Ht, N, 1, device=projected_value.device, dtype=projected_value.dtype)
191 |
192 | # Key/value gates (if using gates)
193 | if self.use_gates:
194 | key_gate_logit = self.key_gate_logit.view(1, 1, 1, 1)
195 | value_gate_logit = self.value_gate_logit.view(1, 1, 1, 1)
196 | if self.training and self.use_gumbel:
197 | u1 = torch.rand(B, Ht, N, 1, device=key_gate_logit.device, dtype=key_gate_logit.dtype)
198 | u2 = torch.rand(B, Ht, N, 1, device=value_gate_logit.device, dtype=value_gate_logit.dtype)
199 | g1 = -torch.log(-torch.log(u1 + 1e-20) + 1e-20)
200 | g2 = -torch.log(-torch.log(u2 + 1e-20) + 1e-20)
201 | key_gate = torch.sigmoid((key_gate_logit + g1) / self.gate_temperature)
202 | value_gate = torch.sigmoid((value_gate_logit + g2) / self.gate_temperature)
203 | else:
204 | key_gate = (key_gate_logit > 0).float()
205 | value_gate = (value_gate_logit > 0).float()
206 | else:
207 | # Gates disabled: set gates to 1.0 (always open)
208 | key_gate = torch.ones(B, Ht, N, 1, device=projected_key.device, dtype=projected_key.dtype)
209 | value_gate = torch.ones(B, Ht, N, 1, device=projected_value.device, dtype=projected_value.dtype)
210 |
211 | # Compute projected contribution
212 | projected_key_term = key_gate * norm_key_scalar * projected_key
213 | projected_value_term = value_gate * norm_value_scalar * projected_value
214 |
215 | # Compute target contribution (if using target)
216 | if self.use_target:
217 | # Full C2C: add target with projected
218 | output_key = target_key + projected_key_term
219 | output_value = target_value + projected_value_term
220 | else:
221 | # Ablation level 3: only use projected (no target)
222 | output_key = projected_key_term
223 | output_value = projected_value_term
224 |
225 | return output_key, output_value
226 |
227 | def get_ablation_info(self) -> dict:
228 | """Return information about current ablation configuration."""
229 | return {
230 | 'ablation_level': self.ablation_level,
231 | 'use_scalar_weights': self.use_scalar_weights,
232 | 'use_gates': self.use_gates,
233 | 'use_target': self.use_target,
234 | 'description': self._get_ablation_description()
235 | }
236 |
237 | def _get_ablation_description(self) -> str:
238 | """Get human-readable description of current ablation level."""
239 | descriptions = {
240 | 0: "Full C2C (baseline)",
241 | 1: "No scalar weights (scalars=1.0)",
242 | 2: "No gates (gates=1.0) + No scalar weights",
243 | 3: "No target (source-only) + No gates + No scalar weights",
244 | 4: "No gates (gates=1.0), keep scalars and target"
245 | }
246 | return descriptions.get(self.ablation_level, "Unknown ablation level")
247 |
248 |
249 | # Convenience functions for creating specific ablation levels
250 | def create_ablation_projector(
251 | source_dim: int,
252 | target_dim: int,
253 | source_num_heads: int = 1,
254 | target_num_heads: int = 1,
255 | ablation_level: int = 0,
256 | **kwargs
257 | ) -> AblationProjector:
258 | """Create an AblationProjector with specified ablation level."""
259 | return AblationProjector(
260 | source_dim=source_dim,
261 | target_dim=target_dim,
262 | source_num_heads=source_num_heads,
263 | target_num_heads=target_num_heads,
264 | ablation_level=ablation_level,
265 | **kwargs
266 | )
267 |
268 |
269 | def create_full_c2c_projector(**kwargs) -> AblationProjector:
270 | """Create full C2C projector (ablation level 0)."""
271 | return create_ablation_projector(ablation_level=0, **kwargs)
272 |
273 |
274 | def create_no_scalar_projector(**kwargs) -> AblationProjector:
275 | """Create projector without scalar weights (ablation level 1)."""
276 | return create_ablation_projector(ablation_level=1, **kwargs)
277 |
278 |
279 | def create_no_gate_projector(**kwargs) -> AblationProjector:
280 | """Create projector without gates (ablation level 2)."""
281 | return create_ablation_projector(ablation_level=2, **kwargs)
282 |
283 |
284 | def create_source_only_projector(**kwargs) -> AblationProjector:
285 | """Create source-only projector (ablation level 3)."""
286 | return create_ablation_projector(ablation_level=3, **kwargs)
287 |
288 |
289 | def create_no_gate_only_projector(**kwargs) -> AblationProjector:
290 | """Create projector without gates but with scalar weights and target (ablation level 4)."""
291 | return create_ablation_projector(ablation_level=4, **kwargs)
292 |
--------------------------------------------------------------------------------