From 6c25d80bce7445eb6ceae8d65a5b8b73a3ebbeed Mon Sep 17 00:00:00 2001 From: Stanislav Date: Thu, 19 Mar 2026 01:23:30 +0000 Subject: [PATCH 1/2] Add activation-space norm fix for mid-layer cross-model injection Previously mid-layer injection projected to embedding-space norm (~20-40) but injected at layer 21/28 where activation norms are ~500-2000. This 50-100x norm mismatch caused the model to barely notice the injection, explaining the 45% GSM8K result (vs 77% for layer-0 rosetta). Fix: compute target model's activation norm at the injection layer via calibration prompts, then renormalize the projected vector to match. Co-Authored-By: Claude Opus 4.6 (1M context) --- benchmarks/gsm8k_2agent/pipeline_mid_layer.py | 255 ++++++++++++++++ benchmarks/gsm8k_2agent/run_gsm8k_2agent.py | 163 +++++++++- src/avp/rosetta/mid_layer.py | 284 ++++++++++++++++++ 3 files changed, 697 insertions(+), 5 deletions(-) create mode 100644 benchmarks/gsm8k_2agent/pipeline_mid_layer.py create mode 100644 src/avp/rosetta/mid_layer.py diff --git a/benchmarks/gsm8k_2agent/pipeline_mid_layer.py b/benchmarks/gsm8k_2agent/pipeline_mid_layer.py new file mode 100644 index 0000000..6475571 --- /dev/null +++ b/benchmarks/gsm8k_2agent/pipeline_mid_layer.py @@ -0,0 +1,255 @@ +"""Mid-layer injection pipeline: 2-agent chain with intermediate layer injection. + +Researcher runs on model A (latent steps -> extract hidden state -> project). +Solver runs on model B (inject at layer ~75% depth via forward hook -> generate). + +Unlike rosetta (injects projected embedding at layer 0 via inputs_embeds), +mid-layer injects at an intermediate layer, operating directly in the +semantic representation space. +""" + +import time +from typing import Any, Dict, List + +import torch + +from benchmarks.shared.generation import generate_text, render_prompt, tokenize_prompt +from benchmarks.shared.kv_utils import get_past_length +from benchmarks.shared.metrics import gpu_memory_tracker +from .agents import AGENTS, build_latent_prompt +from .evaluate import extract_gold, extract_gsm8k_answer, check_correct + + +def run_mid_layer_pipeline( + conn_a: Any, + model_a: Any, + tokenizer_a: Any, + identity_a: Any, + model_b: Any, + tokenizer_b: Any, + device: str, + avp_map: Any, + question: str, + gold_solution: str, + latent_steps: int = 10, + max_new_tokens: int = 512, + temperature: float = 0.7, + top_p: float = 0.95, + verbose: bool = False, + depth_ratio: float = 0.75, +) -> Dict: + """Run the 2-agent cross-model pipeline with mid-layer injection. + + Researcher (model A): latent steps -> extract hidden state -> project + Solver (model B): inject at intermediate layer via forward hook -> generate + """ + from avp.rosetta.mid_layer import ( + compute_activation_norm, + compute_injection_layer, + mid_layer_injection_hook, + renormalize_to_activation_space, + ) + + with gpu_memory_tracker(device) as mem: + t0 = time.perf_counter() + agent_traces: List[Dict] = [] + total_prompt_tokens = 0 + total_latent_steps = 0 + total_output_tokens = 0 + + researcher = AGENTS[0] + solver = AGENTS[1] + + # --- Agent 1: Researcher on model A (latent steps) --- + messages = build_latent_prompt(researcher.role, question) + prompt_text = render_prompt(tokenizer_a, messages) + input_ids, attention_mask = tokenize_prompt(tokenizer_a, prompt_text, device) + + agent_t0 = time.perf_counter() + prompt_tokens = int(input_ids.shape[-1]) + total_prompt_tokens += prompt_tokens + total_latent_steps += latent_steps + + # Collect hidden states from all latent steps + past_kv, hidden_states = conn_a.generate_latent_steps( + input_ids, latent_steps=latent_steps, attention_mask=attention_mask, + collect_hidden_states=True, + ) + + # Use last hidden state for projection + last_hidden = hidden_states[-1].unsqueeze(0) # [1, D_src] + + # Project to target model space + proj_t0 = time.perf_counter() + projected, proj_metrics = conn_a.project_hidden_for_cross_model( + last_hidden, avp_map, return_metrics=True, + ) + projection_ms = (time.perf_counter() - proj_t0) * 1000 + + # Ensure [1, D] shape for injection + if projected.dim() == 1: + projected = projected.unsqueeze(0) + if projected.dim() == 3: + projected = projected.squeeze(0)[-1:, :] + + wire_bytes = projected.nelement() * projected.element_size() + agent_time_ms = (time.perf_counter() - agent_t0) * 1000 + + # Compute injection layer + target_num_layers = model_b.config.num_hidden_layers + injection_layer = compute_injection_layer(target_num_layers, depth_ratio) + + # Renormalize from embedding-space norm to activation-space norm + activation_norm = compute_activation_norm( + model_b, tokenizer_b, injection_layer, + ) + projected = renormalize_to_activation_space(projected, activation_norm) + + if verbose: + print(f" [{researcher.name}] activation_norm at layer {injection_layer}: " + f"{activation_norm:.1f}") + + agent_traces.append({ + "name": researcher.name, + "role": researcher.role, + "prompt_tokens": prompt_tokens, + "latent_steps": latent_steps, + "projection_ms": projection_ms, + "wire_bytes": wire_bytes, + "agent_time_ms": agent_time_ms, + "injection_layer": injection_layer, + "target_num_layers": target_num_layers, + "output": "", + }) + + if verbose: + print(f" [{researcher.name}] latent steps={latent_steps}, " + f"projection={projection_ms:.1f}ms, " + f"inject at layer {injection_layer}/{target_num_layers} " + f"({100*injection_layer/target_num_layers:.0f}% depth)") + + # Free model A KV-cache + del past_kv, hidden_states + if device == "cuda": + torch.cuda.empty_cache() + + # --- Agent 2: Solver on model B (mid-layer injection + generate) --- + messages = build_latent_prompt(solver.role, question) + prompt_text = render_prompt(tokenizer_b, messages) + input_ids, attention_mask = tokenize_prompt(tokenizer_b, prompt_text, device) + + agent_t0 = time.perf_counter() + prompt_tokens = int(input_ids.shape[-1]) + total_prompt_tokens += prompt_tokens + + # Generate with mid-layer injection hook + inject_hidden = projected.to(device).to(model_b.dtype) + + with mid_layer_injection_hook(model_b, injection_layer, inject_hidden): + text, _ = generate_text( + model_b, tokenizer_b, input_ids, attention_mask, device, + past_key_values=None, # No KV-cache priming + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + ) + + output_encoded = tokenizer_b(text, add_special_tokens=False) + output_tokens = len(output_encoded["input_ids"]) + total_output_tokens += output_tokens + agent_time_ms = (time.perf_counter() - agent_t0) * 1000 + + agent_traces.append({ + "name": solver.name, + "role": solver.role, + "prompt_tokens": prompt_tokens, + "output_tokens": output_tokens, + "agent_time_ms": agent_time_ms, + "output": text, + }) + + if verbose: + print(f" [{solver.name}] output ({len(text)} chars): {text[:200]}...") + + wall_time = time.perf_counter() - t0 + + total_tokens = total_prompt_tokens + total_latent_steps + total_output_tokens + tokens_per_sec = total_tokens / wall_time if wall_time > 0 else 0 + + gold = extract_gold(gold_solution) + prediction = extract_gsm8k_answer(agent_traces[-1]["output"]) + correct = check_correct(prediction, gold) + + return { + "question": question, + "gold": gold, + "prediction": prediction, + "raw_output": agent_traces[-1]["output"], + "correct": correct, + "wall_time": wall_time, + "total_prompt_tokens": total_prompt_tokens, + "total_latent_steps": total_latent_steps, + "total_output_tokens": total_output_tokens, + "total_tokens": total_tokens, + "tokens_per_sec": tokens_per_sec, + "peak_memory_mb": mem["peak_memory_mb"], + "projection_overhead_ms": projection_ms, + "projection_wire_bytes": wire_bytes, + "injection_layer": injection_layer, + "depth_ratio": depth_ratio, + "hidden_state_norm": float(proj_metrics["hidden_state_norm"].mean()) if "hidden_state_norm" in proj_metrics else None, + "nearest_cos_sim": float(proj_metrics["nearest_cos_sim"].mean()) if "nearest_cos_sim" in proj_metrics else None, + "agents": agent_traces, + "mode": "mid_layer", + } + + +def run_mid_layer_benchmark( + conn_a: Any, + model_a: Any, + tokenizer_a: Any, + identity_a: Any, + model_b: Any, + tokenizer_b: Any, + device: str, + avp_map: Any, + dataset: List[Dict], + latent_steps: int = 10, + max_new_tokens: int = 512, + temperature: float = 0.7, + top_p: float = 0.95, + verbose: bool = False, + depth_ratio: float = 0.75, +) -> List[Dict]: + """Run mid-layer pipeline on a list of GSM8K samples.""" + results = [] + for i, sample in enumerate(dataset): + if verbose: + print(f"\n[Mid-Layer] Sample {i + 1}/{len(dataset)}: " + f"{sample['question'][:80]}...") + + result = run_mid_layer_pipeline( + conn_a, model_a, tokenizer_a, identity_a, + model_b, tokenizer_b, device, avp_map, + question=sample["question"], + gold_solution=sample["answer"], + latent_steps=latent_steps, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + verbose=verbose, + depth_ratio=depth_ratio, + ) + results.append(result) + + if verbose: + status = "CORRECT" if result["correct"] else "WRONG" + print(f" => {status} (pred={result['prediction']}, gold={result['gold']}, " + f"time={result['wall_time']:.1f}s)") + else: + correct = sum(1 for r in results if r["correct"]) + print(f" [Mid-Layer d={depth_ratio}] {i + 1}/{len(dataset)} " + f"({correct}/{i + 1} correct, {result['wall_time']:.1f}s)", + flush=True) + + return results diff --git a/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py b/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py index f077d52..93d7352 100644 --- a/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py +++ b/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py @@ -39,7 +39,8 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--mode", - choices=["latent", "text", "direct", "rosetta", "text_cross_model", "both", "all"], + choices=["latent", "text", "direct", "rosetta", "logit_guided", + "mid_layer", "trained", "text_cross_model", "both", "all"], default="all", help="Pipeline(s) to run (default: all)", ) @@ -64,6 +65,10 @@ def parse_args() -> argparse.Namespace: help="Softmax temperature for cross-model projection (default: 1.0)") parser.add_argument("--num_transfer_states", type=int, default=1, help="Number of hidden states to transfer in rosetta mode (default: 1)") + parser.add_argument("--logit_bias_alpha", type=float, default=0.5, + help="Logit bias scaling factor for logit_guided mode (default: 0.5)") + parser.add_argument("--logit_bias_confidence_threshold", type=float, default=0.8, + help="Confidence threshold for logit bias gating (default: 0.8)") return parser.parse_args() @@ -104,6 +109,8 @@ def run_benchmark(config: dict) -> dict: output_dir = config.get("output_dir") projection_temperature = config.get("projection_temperature", 1.0) num_transfer_states = config.get("num_transfer_states", 1) + logit_bias_alpha = config.get("logit_bias_alpha", 0.5) + logit_bias_confidence_threshold = config.get("logit_bias_confidence_threshold", 0.8) model_b_name = config.get("model_b", "Qwen/Qwen2.5-0.5B-Instruct") @@ -111,12 +118,15 @@ def run_benchmark(config: dict) -> dict: run_latent = mode in ("latent", "both", "all") run_text = mode in ("text", "both", "all") run_rosetta = mode in ("rosetta", "all") + run_logit_guided = mode in ("logit_guided", "all") + run_mid_layer = mode in ("mid_layer", "all") + run_trained = mode in ("trained",) # not in "all" — requires training run_text_cross_model = mode in ("text_cross_model", "all") print(f"Device: {device}") print(f"Mode: {mode}") print(f"Model A: {model_name}") - if run_rosetta or run_text_cross_model: + if run_rosetta or run_logit_guided or run_mid_layer or run_trained or run_text_cross_model: print(f"Model B: {model_b_name}") print(f"Samples: {max_samples}") print(f"Latent steps: {latent_steps}") @@ -124,7 +134,9 @@ def run_benchmark(config: dict) -> dict: print(f"Temperature: {temperature}") print(f"Seed: {seed}") print(f"Pipelines: direct={run_direct}, text={run_text}, latent={run_latent}, " - f"rosetta={run_rosetta}, text_cross_model={run_text_cross_model}") + f"rosetta={run_rosetta}, logit_guided={run_logit_guided}, " + f"mid_layer={run_mid_layer}, trained={run_trained}, " + f"text_cross_model={run_text_cross_model}") print() dataset = load_dataset(max_samples) @@ -134,6 +146,9 @@ def run_benchmark(config: dict) -> dict: latent_results = None text_results = None rosetta_results = None + logit_guided_results = None + mid_layer_results = None + trained_results = None text_cross_model_results = None if run_direct: @@ -182,7 +197,7 @@ def run_benchmark(config: dict) -> dict: # Load model B if needed for cross-model modes model_b = tokenizer_b = connector_b = identity_b = None - if run_rosetta or run_text_cross_model: + if run_rosetta or run_logit_guided or run_mid_layer or run_trained or run_text_cross_model: model_b, tokenizer_b, connector_b, identity_b = load_model(model_b_name, device) if run_text_cross_model: @@ -235,6 +250,117 @@ def run_benchmark(config: dict) -> dict: num_transfer_states=num_transfer_states, ) + if run_logit_guided: + from benchmarks.gsm8k_2agent.pipeline_logit_guided import run_logit_guided_benchmark + from avp.rosetta.calibrate import calibrate + + print("\n" + "=" * 50) + print("Running LOGIT-GUIDED (cross-model logit bias) pipeline...") + print(f" Model A (Researcher): {model_name}") + print(f" Model B (Solver): {model_b_name}") + print(f" Alpha: {logit_bias_alpha}") + print(f" Confidence threshold: {logit_bias_confidence_threshold}") + print("=" * 50) + set_seed(seed) + + # Calibrate (reuse if already done for rosetta) + if 'avp_map' not in dir() or avp_map is None: + print("Calibrating Rosetta Stone projection...") + avp_map = calibrate( + source_model=model, target_model=model_b, + source_tokenizer=tokenizer, target_tokenizer=tokenizer_b, + device=device, + ) + print(f" Method: {avp_map.method.value}, " + f"validation_score: {avp_map.validation_score:.4f}, " + f"{avp_map.source_dim}d → {avp_map.target_dim}d") + + logit_guided_results = run_logit_guided_benchmark( + conn_a=connector, model_a=model, tokenizer_a=tokenizer, + identity_a=identity, model_b=model_b, tokenizer_b=tokenizer_b, + device=device, avp_map=avp_map, dataset=dataset, + latent_steps=latent_steps, max_new_tokens=max_new_tokens, + temperature=temperature, top_p=top_p, verbose=verbose, + logit_bias_alpha=logit_bias_alpha, + logit_bias_confidence_threshold=logit_bias_confidence_threshold, + ) + + if run_mid_layer: + from benchmarks.gsm8k_2agent.pipeline_mid_layer import run_mid_layer_benchmark + from avp.rosetta.calibrate import calibrate + + print("\n" + "=" * 50) + print("Running MID-LAYER (cross-model mid-layer injection) pipeline...") + print(f" Model A (Researcher): {model_name}") + print(f" Model B (Solver): {model_b_name}") + print(f" Depth ratio: 0.75") + print("=" * 50) + set_seed(seed) + + # Calibrate (reuse if already done) + if 'avp_map' not in dir() or avp_map is None: + print("Calibrating Rosetta Stone projection...") + avp_map = calibrate( + source_model=model, target_model=model_b, + source_tokenizer=tokenizer, target_tokenizer=tokenizer_b, + device=device, + ) + print(f" Method: {avp_map.method.value}, " + f"validation_score: {avp_map.validation_score:.4f}, " + f"{avp_map.source_dim}d -> {avp_map.target_dim}d") + + mid_layer_results = run_mid_layer_benchmark( + conn_a=connector, model_a=model, tokenizer_a=tokenizer, + identity_a=identity, model_b=model_b, tokenizer_b=tokenizer_b, + device=device, avp_map=avp_map, dataset=dataset, + latent_steps=latent_steps, max_new_tokens=max_new_tokens, + temperature=temperature, top_p=top_p, verbose=verbose, + ) + + if run_trained: + from benchmarks.gsm8k_2agent.pipeline_trained import run_trained_benchmark + from avp.rosetta.train import train_projector, TrainConfig + + print("\n" + "=" * 50) + print("Running TRAINED (per-layer learned projection) pipeline...") + print(f" Model A (Researcher): {model_name}") + print(f" Model B (Solver): {model_b_name}") + print("=" * 50) + + # Training phase + train_config = TrainConfig( + num_samples=config.get("train_samples", 2000), + batch_size=config.get("train_batch_size", 4), + num_epochs=config.get("train_epochs", 2), + learning_rate=config.get("train_lr", 1e-4), + gate_init=config.get("train_gate_init", -5.0), + gate_reg_weight=config.get("train_gate_reg", 0.01), + ) + print(f"Training projector: {train_config.num_samples} samples, " + f"{train_config.num_epochs} epochs...") + + trained_map = train_projector( + source_model=model, + target_model=model_b, + source_tokenizer=tokenizer, + target_tokenizer=tokenizer_b, + device=device, + config=train_config, + ) + active = [i for i, g in enumerate(trained_map.layer_gates) if g > 0.01] + print(f"Training complete. Active layers: {len(active)}/{len(trained_map.layer_gates)}") + print(f"Validation score: {trained_map.validation_score:.4f}") + + set_seed(seed) + + trained_results = run_trained_benchmark( + conn_a=connector, model_a=model, tokenizer_a=tokenizer, + identity_a=identity, model_b=model_b, tokenizer_b=tokenizer_b, + device=device, avp_map=trained_map, dataset=dataset, + latent_steps=latent_steps, max_new_tokens=max_new_tokens, + temperature=temperature, top_p=top_p, verbose=verbose, + ) + # Free model B to reclaim GPU memory if model_b is not None: del model_b, tokenizer_b, connector_b, identity_b @@ -254,6 +380,12 @@ def run_benchmark(config: dict) -> dict: modes.append(("Text", 13, text_results)) if rosetta_results is not None: modes.append(("Rosetta", 13, rosetta_results)) + if logit_guided_results is not None: + modes.append(("Logit-Guided", 13, logit_guided_results)) + if mid_layer_results is not None: + modes.append(("Mid-Layer", 13, mid_layer_results)) + if trained_results is not None: + modes.append(("Trained", 13, trained_results)) if text_cross_model_results is not None: modes.append(("Text Cross-Model", 16, text_cross_model_results)) @@ -267,6 +399,12 @@ def run_benchmark(config: dict) -> dict: available["latent"] = latent_results if rosetta_results is not None: available["rosetta"] = rosetta_results + if logit_guided_results is not None: + available["logit_guided"] = logit_guided_results + if mid_layer_results is not None: + available["mid_layer"] = mid_layer_results + if trained_results is not None: + available["trained"] = trained_results if text_cross_model_results is not None: available["text_cross_model"] = text_cross_model_results agreement_data = compute_agreement(available) if len(available) > 1 else None @@ -289,7 +427,7 @@ def run_benchmark(config: dict) -> dict: "config": { "benchmark": "gsm8k_2agent", "model_a": model_name, - "model_b": model_b_name if (run_rosetta or run_text_cross_model) else None, + "model_b": model_b_name if (run_rosetta or run_logit_guided or run_mid_layer or run_trained or run_text_cross_model) else None, "device": device, "mode": mode, "max_samples": max_samples, @@ -320,6 +458,21 @@ def run_benchmark(config: dict) -> dict: "summary": compute_stats(rosetta_results), "samples": rosetta_results, } + if logit_guided_results is not None: + output_data["logit_guided"] = { + "summary": compute_stats(logit_guided_results), + "samples": logit_guided_results, + } + if mid_layer_results is not None: + output_data["mid_layer"] = { + "summary": compute_stats(mid_layer_results), + "samples": mid_layer_results, + } + if trained_results is not None: + output_data["trained"] = { + "summary": compute_stats(trained_results), + "samples": trained_results, + } if text_cross_model_results is not None: output_data["text_cross_model"] = { "summary": compute_stats(text_cross_model_results), diff --git a/src/avp/rosetta/mid_layer.py b/src/avp/rosetta/mid_layer.py new file mode 100644 index 0000000..909f049 --- /dev/null +++ b/src/avp/rosetta/mid_layer.py @@ -0,0 +1,284 @@ +"""Mid-layer injection for cross-model latent transfer. + +Instead of injecting projected hidden states at layer 0 (via inputs_embeds), +injects at an intermediate layer (~75% depth) using a forward hook. This +bypasses the early embedding/position-encoding layers and operates directly +in the semantic representation space. + +Based on: +- Ramesh & Li (2501.14082): Cross-model hidden state injection at intermediate + layers, up to 27% improvement over text, cross-family confirmed. +- Proportional depth mapping (2504.08775): Layer L_a/N_a maps to L_b/N_b + across architectures (p < 0.005 for 24 LLMs from 1B-70B). + +Key design decisions: +- REPLACE, not sum/mean (Ramesh & Li found sum/mean produce OOD norms) +- Proportional depth mapping: injection_layer = int(N_tgt * extraction_ratio) +- Forward hook scoped to prefill only (removed after first forward pass) +""" + +import logging +from contextlib import contextmanager +from typing import Any, List, Optional, Tuple + +import torch + +logger = logging.getLogger(__name__) + +# Default extraction/injection depth ratio (validated at 0.75 = ~75% depth) +DEFAULT_DEPTH_RATIO = 0.75 + + +def compute_extraction_layer(num_layers: int, depth_ratio: float = DEFAULT_DEPTH_RATIO) -> int: + """Compute the layer index to extract hidden states from. + + Args: + num_layers: Total number of transformer layers in the model. + depth_ratio: Fraction of depth to extract from (0.0=first, 1.0=last). + + Returns: + Layer index (0-indexed). + """ + layer = int(num_layers * depth_ratio) + return min(layer, num_layers - 1) + + +def compute_injection_layer(num_layers: int, depth_ratio: float = DEFAULT_DEPTH_RATIO) -> int: + """Compute the layer index to inject hidden states into. + + Uses proportional depth mapping: if source extracted from 75% depth, + inject at 75% depth of target model (even if different number of layers). + + Args: + num_layers: Total number of transformer layers in target model. + depth_ratio: Fraction of depth to inject at. + + Returns: + Layer index (0-indexed). + """ + layer = int(num_layers * depth_ratio) + return min(layer, num_layers - 1) + + +def extract_mid_layer_hidden( + model_outputs: Any, + extraction_layer: int, +) -> Any: + """Extract hidden state from an intermediate layer of model outputs. + + Args: + model_outputs: Model output dict with hidden_states. + extraction_layer: Layer index to extract from. + + Returns: + Hidden state tensor [B, D] from the specified layer's last token. + """ + # hidden_states is a tuple of (num_layers + 1) tensors, each [B, seq, D] + # Index 0 = embedding output, index i = output of layer i + hidden_states = model_outputs.hidden_states + if extraction_layer + 1 >= len(hidden_states): + extraction_layer = len(hidden_states) - 2 # -1 is last layer output + # +1 because index 0 is embedding layer output, index 1 is layer 0 output + layer_hidden = hidden_states[extraction_layer + 1] + return layer_hidden[:, -1, :] # [B, D] — last token only + + +def _get_decoder_layers(model: Any): + """Get the list of decoder layers from a HuggingFace model. + + Handles different model architectures (Llama, Qwen, GPT-2, etc.). + """ + # Try common attribute paths + inner = getattr(model, "model", None) + if inner is not None: + layers = getattr(inner, "layers", None) + if layers is not None: + return layers + + # GPT-2 style + transformer = getattr(model, "transformer", None) + if transformer is not None: + h = getattr(transformer, "h", None) + if h is not None: + return h + + raise AttributeError( + f"Cannot find decoder layers in model {type(model).__name__}. " + "Expected model.model.layers or model.transformer.h" + ) + + +CALIBRATION_PROMPTS = [ + "Solve step by step: What is 24 * 17 + 3?", + "Write a Python function that checks if a number is prime.", + "Explain the difference between a stack and a queue.", + "The quick brown fox jumps over the lazy dog.", + "Analyze the following: renewable energy costs have decreased by 90% since 2010.", + "What are the main causes of the French Revolution?", + "Debug this code: def add(a, b): return a - b", + "Summarize: Machine learning models learn patterns from data.", +] + + +def compute_activation_norm( + model: Any, + tokenizer: Any, + layer_index: int, + prompts: Optional[List[str]] = None, +) -> float: + """Compute mean L2 norm of activations at a specific layer. + + Runs calibration prompts through the model and records the + last-token hidden state norm at the target layer. + + Args: + model: HuggingFace model. + tokenizer: HuggingFace tokenizer. + layer_index: Which layer to measure (0-indexed). + prompts: Calibration prompts. Uses defaults if None. + + Returns: + Mean L2 norm (float) of the last-token hidden state at that layer. + """ + if prompts is None: + prompts = CALIBRATION_PROMPTS + + norms = [] + model.eval() + with torch.no_grad(): + for text in prompts: + inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256) + inputs = {k: v.to(model.device) for k, v in inputs.items()} + outputs = model(**inputs, output_hidden_states=True) + # hidden_states[0] = embedding output, hidden_states[i+1] = layer i output + layer_hidden = outputs.hidden_states[layer_index + 1] + norm = layer_hidden[:, -1, :].float().norm(dim=-1).item() + norms.append(norm) + + mean_norm = sum(norms) / len(norms) + logger.info( + "Activation norm at layer %d: %.1f (from %d prompts)", + layer_index, mean_norm, len(norms), + ) + return mean_norm + + +def renormalize_to_activation_space( + projected: torch.Tensor, + activation_norm: float, +) -> torch.Tensor: + """Renormalize a projected vector from embedding-space norm to activation-space norm. + + Args: + projected: Projected tensor [..., D] (currently at embedding-space norm). + activation_norm: Target L2 norm for the injection layer. + + Returns: + Renormalized tensor with L2 norm matching the activation space. + """ + current_norm = projected.float().norm(dim=-1, keepdim=True).clamp_min(1e-6) + return projected * (activation_norm / current_norm) + + +@contextmanager +def mid_layer_injection_hook( + model: Any, + injection_layer: int, + projected_hidden: Any, +): + """Context manager that installs a forward hook to replace hidden states + at a specific layer during the first forward pass (prefill). + + The hook fires once and then removes itself, so it only affects the + initial prefill pass, not subsequent autoregressive generation steps. + + Args: + model: HuggingFace model to hook into. + injection_layer: Layer index to inject at. + projected_hidden: Tensor [1, D] or [B, D] to replace the last token's + hidden state with. + + Yields: + None. The hook is active during the context. + """ + import torch + + layers = _get_decoder_layers(model) + target_layer = layers[injection_layer] + + fired = [False] # mutable flag for closure + + def hook_fn(module, input, output): + if fired[0]: + return output + + fired[0] = True + + # Decoder layer output is a tuple: (hidden_states, ...) or just hidden_states + if isinstance(output, tuple): + hidden = output[0] # [B, seq_len, D] + else: + hidden = output + + # Replace last token's hidden state with projected source hidden state + injection = projected_hidden.to(device=hidden.device, dtype=hidden.dtype) + if injection.dim() == 1: + injection = injection.unsqueeze(0) # [D] -> [1, D] + + # Clone to avoid in-place modification + modified = hidden.clone() + modified[:, -1, :] = injection # Replace last position + + if isinstance(output, tuple): + return (modified,) + output[1:] + return modified + + handle = target_layer.register_forward_hook(hook_fn) + try: + yield + finally: + handle.remove() + + +def project_for_mid_layer( + source_hidden: Any, + avp_map: Any, + source_model: Any, + target_model: Any, + target_num_layers: int, + injection_depth_ratio: float = DEFAULT_DEPTH_RATIO, +) -> Tuple[Any, int]: + """Project source hidden state for mid-layer injection. + + Unlike rosetta (which projects to target embedding space for layer-0 inputs_embeds), + mid-layer projects to the target's intermediate representation space. Since we don't + have a direct map between intermediate spaces, we use the same vocab-mediated/overlap + projection but normalize to the target layer's activation norm instead of the + embedding norm. + + Args: + source_hidden: Source hidden state [1, D_src] or [D_src]. + avp_map: AVPMap with projection data. + source_model: Source HuggingFace model. + target_model: Target HuggingFace model. + target_num_layers: Number of layers in target model. + injection_depth_ratio: Depth ratio for injection point. + + Returns: + Tuple of (projected_hidden [1, D_tgt], injection_layer_index). + """ + import torch + from .project import apply_cross_model_projection + + injection_layer = compute_injection_layer(target_num_layers, injection_depth_ratio) + + # Use standard vocab-mediated/overlap projection + projected = apply_cross_model_projection( + source_hidden, avp_map, source_model, target_model, + ) + + # Ensure correct shape [1, D] + if projected.dim() == 1: + projected = projected.unsqueeze(0) + + return projected, injection_layer From 04bbcd822e4ce84a36ffd203a1165ef7646128de Mon Sep 17 00:00:00 2001 From: Stanislav Date: Thu, 19 Mar 2026 02:10:36 +0000 Subject: [PATCH 2/2] Add multi-vector mid-layer injection support Extract N evenly-spaced hidden states from latent steps, project each through vocabulary, prepend N pad tokens as placeholders, inject all N at mid-layer. Tests whether multi-vector through vocabulary works better at mid-layer than at layer 0 (where it was validated negative #2). Co-Authored-By: Claude Opus 4.6 (1M context) --- benchmarks/gsm8k_2agent/pipeline_mid_layer.py | 52 ++++++++++++++----- benchmarks/gsm8k_2agent/run_gsm8k_2agent.py | 2 + src/avp/rosetta/mid_layer.py | 26 ++++++---- 3 files changed, 58 insertions(+), 22 deletions(-) diff --git a/benchmarks/gsm8k_2agent/pipeline_mid_layer.py b/benchmarks/gsm8k_2agent/pipeline_mid_layer.py index 6475571..41e980d 100644 --- a/benchmarks/gsm8k_2agent/pipeline_mid_layer.py +++ b/benchmarks/gsm8k_2agent/pipeline_mid_layer.py @@ -37,6 +37,7 @@ def run_mid_layer_pipeline( top_p: float = 0.95, verbose: bool = False, depth_ratio: float = 0.75, + num_vectors: int = 1, ) -> Dict: """Run the 2-agent cross-model pipeline with mid-layer injection. @@ -76,22 +77,36 @@ def run_mid_layer_pipeline( collect_hidden_states=True, ) - # Use last hidden state for projection - last_hidden = hidden_states[-1].unsqueeze(0) # [1, D_src] + # Select hidden states for projection + if num_vectors == 1: + selected = [hidden_states[-1].unsqueeze(0)] # [1, D_src] + else: + # Evenly spaced from the latent steps + n_available = len(hidden_states) + indices = [int(i * n_available / num_vectors) for i in range(num_vectors)] + indices[-1] = n_available - 1 # always include last + selected = [hidden_states[i].unsqueeze(0) for i in indices] - # Project to target model space + # Project each hidden state to target model space proj_t0 = time.perf_counter() - projected, proj_metrics = conn_a.project_hidden_for_cross_model( - last_hidden, avp_map, return_metrics=True, - ) + projected_list = [] + proj_metrics = {} + for h in selected: + p, m = conn_a.project_hidden_for_cross_model( + h, avp_map, return_metrics=True, + ) + if p.dim() == 1: + p = p.unsqueeze(0) + if p.dim() == 3: + p = p.squeeze(0)[-1:, :] + projected_list.append(p) + if not proj_metrics: + proj_metrics = m + + # Stack into [N, D] for multi-vector or [1, D] for single + projected = torch.cat(projected_list, dim=0) # [N, D] projection_ms = (time.perf_counter() - proj_t0) * 1000 - # Ensure [1, D] shape for injection - if projected.dim() == 1: - projected = projected.unsqueeze(0) - if projected.dim() == 3: - projected = projected.squeeze(0)[-1:, :] - wire_bytes = projected.nelement() * projected.element_size() agent_time_ms = (time.perf_counter() - agent_t0) * 1000 @@ -138,6 +153,14 @@ def run_mid_layer_pipeline( prompt_text = render_prompt(tokenizer_b, messages) input_ids, attention_mask = tokenize_prompt(tokenizer_b, prompt_text, device) + # For multi-vector: prepend N pad tokens as placeholders for injection + if num_vectors > 1: + pad_id = tokenizer_b.pad_token_id or tokenizer_b.eos_token_id + pad_ids = torch.full((1, num_vectors), pad_id, dtype=input_ids.dtype, device=device) + pad_mask = torch.ones((1, num_vectors), dtype=attention_mask.dtype, device=device) + input_ids = torch.cat([pad_ids, input_ids], dim=-1) + attention_mask = torch.cat([pad_mask, attention_mask], dim=-1) + agent_t0 = time.perf_counter() prompt_tokens = int(input_ids.shape[-1]) total_prompt_tokens += prompt_tokens @@ -145,7 +168,7 @@ def run_mid_layer_pipeline( # Generate with mid-layer injection hook inject_hidden = projected.to(device).to(model_b.dtype) - with mid_layer_injection_hook(model_b, injection_layer, inject_hidden): + with mid_layer_injection_hook(model_b, injection_layer, inject_hidden, num_vectors=num_vectors): text, _ = generate_text( model_b, tokenizer_b, input_ids, attention_mask, device, past_key_values=None, # No KV-cache priming @@ -201,6 +224,7 @@ def run_mid_layer_pipeline( "nearest_cos_sim": float(proj_metrics["nearest_cos_sim"].mean()) if "nearest_cos_sim" in proj_metrics else None, "agents": agent_traces, "mode": "mid_layer", + "num_vectors": num_vectors, } @@ -220,6 +244,7 @@ def run_mid_layer_benchmark( top_p: float = 0.95, verbose: bool = False, depth_ratio: float = 0.75, + num_vectors: int = 1, ) -> List[Dict]: """Run mid-layer pipeline on a list of GSM8K samples.""" results = [] @@ -239,6 +264,7 @@ def run_mid_layer_benchmark( top_p=top_p, verbose=verbose, depth_ratio=depth_ratio, + num_vectors=num_vectors, ) results.append(result) diff --git a/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py b/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py index 93d7352..596464a 100644 --- a/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py +++ b/benchmarks/gsm8k_2agent/run_gsm8k_2agent.py @@ -309,12 +309,14 @@ def run_benchmark(config: dict) -> dict: f"validation_score: {avp_map.validation_score:.4f}, " f"{avp_map.source_dim}d -> {avp_map.target_dim}d") + num_vectors = config.get("num_vectors", 1) mid_layer_results = run_mid_layer_benchmark( conn_a=connector, model_a=model, tokenizer_a=tokenizer, identity_a=identity, model_b=model_b, tokenizer_b=tokenizer_b, device=device, avp_map=avp_map, dataset=dataset, latent_steps=latent_steps, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, verbose=verbose, + num_vectors=num_vectors, ) if run_trained: diff --git a/src/avp/rosetta/mid_layer.py b/src/avp/rosetta/mid_layer.py index 909f049..7dffd1e 100644 --- a/src/avp/rosetta/mid_layer.py +++ b/src/avp/rosetta/mid_layer.py @@ -185,6 +185,7 @@ def mid_layer_injection_hook( model: Any, injection_layer: int, projected_hidden: Any, + num_vectors: int = 1, ): """Context manager that installs a forward hook to replace hidden states at a specific layer during the first forward pass (prefill). @@ -195,14 +196,15 @@ def mid_layer_injection_hook( Args: model: HuggingFace model to hook into. injection_layer: Layer index to inject at. - projected_hidden: Tensor [1, D] or [B, D] to replace the last token's - hidden state with. + projected_hidden: Tensor [1, D], [N, D], or [B, N, D] to inject. + For N>1, replaces the first N positions (prepended pad tokens). + For N=1, replaces the last position (original behavior). + num_vectors: Number of vectors being injected. Must match + projected_hidden's N dimension. Yields: None. The hook is active during the context. """ - import torch - layers = _get_decoder_layers(model) target_layer = layers[injection_layer] @@ -220,14 +222,20 @@ def hook_fn(module, input, output): else: hidden = output - # Replace last token's hidden state with projected source hidden state injection = projected_hidden.to(device=hidden.device, dtype=hidden.dtype) - if injection.dim() == 1: - injection = injection.unsqueeze(0) # [D] -> [1, D] - # Clone to avoid in-place modification modified = hidden.clone() - modified[:, -1, :] = injection # Replace last position + + if num_vectors == 1: + # Single vector: replace last position (original behavior) + if injection.dim() == 1: + injection = injection.unsqueeze(0) + modified[:, -1, :] = injection + else: + # Multi-vector: replace first N positions (prepended pad tokens) + if injection.dim() == 2: + injection = injection.unsqueeze(0) # [N, D] -> [1, N, D] + modified[:, :num_vectors, :] = injection if isinstance(output, tuple): return (modified,) + output[1:]