Add tuned lens as supplementary projection mode for logit lens

Introduces per-layer affine probes (trained to minimise KL divergence with
the final layer) that correct for subspace mismatch in early transformer
layers. Positioned alongside the existing raw logit lens with a frontend
toggle, allowing developers to compare projections and assess whether
observed commitment patterns are robust to projection choice.

New: backend/tuned_lens.py (runtime), scripts/train_tuned_lens.py (training)
Modified: model_service.py (load at startup, tuned lens computation, tuned
commitment summary, modelInfo.tunedLensAvailable, health endpoint)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

backend/model_service.py

@@ -526,6 +526,10 @@ class ModelManager:
 
             logger.info("✅ Model loaded successfully")
 
+            # Load tuned lens probes (optional — falls back to raw logit lens if unavailable)
+            from .tuned_lens import tuned_lens_runtime
+            tuned_lens_runtime.load(self.model_id, self.device, self.dtype)
+
         except Exception as e:
             logger.error(f"Failed to load model: {e}")
             raise
@@ -1238,11 +1242,13 @@ async def root():
 @app.get("/health")
 async def health():
     """Detailed health check - always returns 200 for Docker healthcheck"""
+    from .tuned_lens import tuned_lens_runtime
     return {
         "status": "healthy" if manager.model else "initializing",
         "model_loaded": manager.model is not None,
         "device": str(manager.device) if manager.device else "not set",
         "websocket_clients": len(manager.websocket_clients),
+        "tuned_lens_available": tuned_lens_runtime.available,
         "timestamp": datetime.now().isoformat()
     }
 
@@ -3233,6 +3239,42 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                                     layer_entry["layer_margin"] = layer_margin_val
                                     layer_entry["layer_winner"] = layer_winner_token
                                     layer_entry["layer_runnerup"] = layer_runnerup_token
+
+                                    # Tuned lens: apply per-layer affine correction
+                                    from .tuned_lens import tuned_lens_runtime
+                                    if tuned_lens_runtime.available:
+                                        try:
+                                            corrected = tuned_lens_runtime.apply(layer_idx, hidden_for_lens)
+                                            tuned_normed = normed.__class__  # reuse same LN path
+                                            # Re-apply final LN + lm_head on corrected hidden
+                                            if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
+                                                tuned_normed = manager.model.model.norm(corrected)
+                                                tuned_logits = manager.model.lm_head(tuned_normed)[0]
+                                            elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
+                                                tuned_normed = manager.model.transformer.ln_f(corrected)
+                                                tuned_logits = manager.model.lm_head(tuned_normed)[0]
+                                            else:
+                                                tuned_logits = None
+
+                                            if tuned_logits is not None:
+                                                tuned_probs = torch.softmax(tuned_logits, dim=-1)
+                                                tuned_top_probs, tuned_top_ids = torch.topk(tuned_probs, k=5)
+                                                tuned_entries = []
+                                                for tp, tid in zip(tuned_top_probs.cpu().tolist(), tuned_top_ids.cpu().tolist()):
+                                                    tuned_entries.append({
+                                                        "token": manager.tokenizer.decode([tid], skip_special_tokens=False),
+                                                        "probability": tp
+                                                    })
+                                                layer_entry["tuned_lens_top"] = tuned_entries
+
+                                                tuned_top2_logits, tuned_top2_ids = torch.topk(tuned_logits, k=min(2, len(tuned_logits)))
+                                                tuned_top2_logits_list = tuned_top2_logits.cpu().tolist()
+                                                tuned_top2_ids_list = tuned_top2_ids.cpu().tolist()
+                                                layer_entry["tuned_layer_winner"] = manager.tokenizer.decode([tuned_top2_ids_list[0]], skip_special_tokens=False)
+                                                layer_entry["tuned_layer_runnerup"] = manager.tokenizer.decode([tuned_top2_ids_list[1]], skip_special_tokens=False) if len(tuned_top2_ids_list) > 1 else ""
+                                                layer_entry["tuned_layer_margin"] = (tuned_top2_logits_list[0] - tuned_top2_logits_list[1]) if len(tuned_top2_logits_list) > 1 else 0.0
+                                        except Exception as tuned_err:
+                                            logger.debug(f"Tuned lens error at layer {layer_idx}: {tuned_err}")
                             except Exception as lens_err:
                                 logger.debug(f"Logit lens error at layer {layer_idx}: {lens_err}")
 
@@ -3397,6 +3439,39 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                 "flip_count": flip_count,
             }
 
+            # Tuned lens commitment summary (parallel to raw)
+            tuned_commitment_summary = None
+            from .tuned_lens import tuned_lens_runtime
+            if tuned_lens_runtime.available:
+                tuned_commitment_layers = []
+                tuned_flip_count = 0
+                for step_idx, step_layers in enumerate(layer_data_by_token):
+                    tuned_lens_layers = [l for l in step_layers if l.get("tuned_layer_margin") is not None]
+                    if not tuned_lens_layers:
+                        continue
+                    step_commitment = None
+                    for i, ll in enumerate(tuned_lens_layers):
+                        if ll["tuned_layer_margin"] > 0.3:
+                            stays_positive = all(tuned_lens_layers[j]["tuned_layer_margin"] > 0 for j in range(i, len(tuned_lens_layers)))
+                            if stays_positive:
+                                step_commitment = ll["layer_idx"]
+                                break
+                    if step_commitment is not None:
+                        tuned_commitment_layers.append(step_commitment)
+                    for i in range(1, len(tuned_lens_layers)):
+                        prev_w = (tuned_lens_layers[i-1].get("tuned_layer_winner") or "").strip()
+                        curr_w = (tuned_lens_layers[i].get("tuned_layer_winner") or "").strip()
+                        if prev_w and curr_w and prev_w != curr_w:
+                            tuned_flip_count += 1
+
+                tuned_avg = sum(tuned_commitment_layers) / len(tuned_commitment_layers) if tuned_commitment_layers else n_layers
+                tuned_late = sum(1 for cl in tuned_commitment_layers if cl > late_threshold)
+                tuned_commitment_summary = {
+                    "avg_commitment_layer": round(tuned_avg, 1),
+                    "late_commitment_count": tuned_late,
+                    "flip_count": tuned_flip_count,
+                }
+
             # Build response
             response = {
                 "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
@@ -3413,12 +3488,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     "numHeads": n_heads,
                     "modelDimension": d_model,
                     "headDim": head_dim,
-                    "vocabSize": manager.model.config.vocab_size
+                    "vocabSize": manager.model.config.vocab_size,
+                    "tunedLensAvailable": tuned_lens_runtime.available,
                 },
                 "generationTime": generation_time,
                 "numTokensGenerated": len(generated_tokens),
                 "marginStats": margin_stats,
                 "commitmentSummary": commitment_summary,
+                **({"tunedCommitmentSummary": tuned_commitment_summary} if tuned_commitment_summary else {}),
             }
 
             # Estimate response size

backend/tuned_lens.py

@@ -0,0 +1,130 @@
+"""
+Tuned Lens Runtime — load and apply per-layer affine probes for improved
+intermediate-layer predictions.
+
+Each probe applies a learned linear correction A_l(x) = x @ W_l^T + b_l
+(initialised to identity + zero during training) that is trained to minimise
+KL divergence between the corrected layer's predictions and the model's
+final-layer predictions.
+
+See scripts/train_tuned_lens.py for the training pipeline.
+"""
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+TUNED_LENS_DIR = os.environ.get("TUNED_LENS_DIR", "./tuned_lens_weights")
+
+
+class TunedLensRuntime:
+    """Load, cache, and apply per-layer affine probes at inference time."""
+
+    def __init__(self):
+        self._probes: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}
+        self._metadata: Optional[dict] = None
+        self._available = False
+
+    @property
+    def available(self) -> bool:
+        return self._available
+
+    def load(self, model_id: str, device: torch.device, dtype: torch.dtype,
+             weights_dir: Optional[str] = None) -> bool:
+        """Load tuned lens checkpoint for *model_id*.
+
+        Returns True if weights were loaded successfully, False otherwise.
+        Failure is non-fatal — the system falls back to raw logit lens.
+        """
+        base_dir = Path(weights_dir or TUNED_LENS_DIR)
+        model_dir = base_dir / model_id
+
+        if not model_dir.exists():
+            logger.info(f"Tuned lens: no weights directory for {model_id} at {model_dir}")
+            return False
+
+        # Find the checkpoint — pick the first .pt file
+        pt_files = sorted(model_dir.glob("tuned_lens_*.pt"))
+        if not pt_files:
+            logger.info(f"Tuned lens: no .pt checkpoint found in {model_dir}")
+            return False
+
+        checkpoint_path = pt_files[0]
+        metadata_path = model_dir / "metadata.json"
+
+        try:
+            # Load metadata
+            if metadata_path.exists():
+                with open(metadata_path, "r") as f:
+                    self._metadata = json.load(f)
+                logger.info(f"Tuned lens: metadata loaded — {self._metadata.get('n_layers')} layers, "
+                            f"d_model={self._metadata.get('d_model')}")
+            else:
+                self._metadata = {}
+
+            # Load state dict
+            state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+
+            # Parse layer_N.weight / layer_N.bias entries
+            self._probes = {}
+            layer_indices = set()
+            for key in state_dict:
+                parts = key.split(".")
+                if len(parts) == 2 and parts[0].startswith("layer_"):
+                    idx = int(parts[0].split("_")[1])
+                    layer_indices.add(idx)
+
+            for idx in sorted(layer_indices):
+                w_key = f"layer_{idx}.weight"
+                b_key = f"layer_{idx}.bias"
+                if w_key in state_dict and b_key in state_dict:
+                    weight = state_dict[w_key].to(device=device, dtype=dtype)
+                    bias = state_dict[b_key].to(device=device, dtype=dtype)
+                    self._probes[idx] = (weight, bias)
+
+            if not self._probes:
+                logger.warning(f"Tuned lens: checkpoint loaded but no layer probes found")
+                return False
+
+            self._available = True
+            logger.info(f"Tuned lens: loaded {len(self._probes)} layer probes from {checkpoint_path} "
+                        f"(device={device}, dtype={dtype})")
+            return True
+
+        except Exception as e:
+            logger.warning(f"Tuned lens: failed to load checkpoint — {e}")
+            self._probes = {}
+            self._metadata = None
+            self._available = False
+            return False
+
+    def apply(self, layer_idx: int, hidden_state: torch.Tensor) -> torch.Tensor:
+        """Apply the affine probe for *layer_idx*: hidden @ W^T + b.
+
+        If no probe exists for this layer, returns the hidden state unchanged
+        (identity fallback).
+        """
+        if layer_idx not in self._probes:
+            return hidden_state
+        weight, bias = self._probes[layer_idx]
+        return hidden_state @ weight.T + bias
+
+    def get_info(self) -> dict:
+        """Return metadata dict for health/debug endpoints."""
+        return {
+            "available": self._available,
+            "num_probes": len(self._probes),
+            "layer_indices": sorted(self._probes.keys()),
+            "metadata": self._metadata or {},
+        }
+
+
+# Global singleton
+tuned_lens_runtime = TunedLensRuntime()

scripts/train_tuned_lens.py

@@ -0,0 +1,312 @@
+"""
+Train Tuned Lens probes — per-layer affine corrections that minimise
+KL divergence between an intermediate layer's corrected predictions and
+the model's final-layer predictions.
+
+Usage:
+    python -m scripts.train_tuned_lens \
+        --model-id codegen-350m \
+        --corpus-file calibration_data.txt \
+        --output-dir ./tuned_lens_weights/ \
+        --max-samples 2000 --epochs 5
+
+Each probe is a simple affine map  A_l(x) = x @ W_l^T + b_l
+initialised to identity + zero so that the untrained probe reproduces
+the raw logit lens exactly.
+"""
+
+import argparse
+import hashlib
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# AffineProbe
+# ---------------------------------------------------------------------------
+
+class AffineProbe(nn.Module):
+    """Per-layer affine correction initialised to identity."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.eye(d_model))
+        self.bias = nn.Parameter(torch.zeros(d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x @ self.weight.T + self.bias
+
+
+# ---------------------------------------------------------------------------
+# Architecture detection — mirrors model_service.py
+# ---------------------------------------------------------------------------
+
+def get_final_ln_and_lm_head(model):
+    """Return (final_layer_norm, lm_head) for the loaded model."""
+    # Mistral / LLaMA / CodeGen-style
+    if hasattr(model, "model") and hasattr(model.model, "norm"):
+        return model.model.norm, model.lm_head
+    # GPT-style
+    if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
+        return model.transformer.ln_f, model.lm_head
+    raise RuntimeError(
+        "Cannot detect final layer norm — model architecture not recognised. "
+        "Supported: Mistral/LLaMA (.model.norm), GPT (.transformer.ln_f)"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Model hash — ties checkpoint to exact model weights
+# ---------------------------------------------------------------------------
+
+def compute_model_hash(model, n_tensors: int = 20) -> str:
+    """SHA-256 of the first *n_tensors* parameter tensors' bytes."""
+    h = hashlib.sha256()
+    for i, (_, param) in enumerate(model.named_parameters()):
+        if i >= n_tensors:
+            break
+        h.update(param.data.cpu().numpy().tobytes())
+    return h.hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# Corpus loader
+# ---------------------------------------------------------------------------
+
+def load_corpus(path: str, max_samples: int, max_seq_len: int, tokenizer) -> list:
+    """Load and tokenize a plain-text corpus (one sample per line or paragraph)."""
+    texts = []
+    with open(path, "r", encoding="utf-8") as f:
+        buf = []
+        for line in f:
+            line = line.rstrip("\n")
+            if line.strip() == "" and buf:
+                texts.append("\n".join(buf))
+                buf = []
+                if len(texts) >= max_samples:
+                    break
+            else:
+                buf.append(line)
+        if buf and len(texts) < max_samples:
+            texts.append("\n".join(buf))
+
+    # Tokenize
+    samples = []
+    for text in texts[:max_samples]:
+        ids = tokenizer.encode(text, add_special_tokens=False, truncation=True,
+                               max_length=max_seq_len)
+        if len(ids) >= 8:  # skip very short sequences
+            samples.append(torch.tensor(ids, dtype=torch.long))
+    logger.info(f"Loaded {len(samples)} samples from {path} (max_seq_len={max_seq_len})")
+    return samples
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def train_tuned_lens(
+    model,
+    tokenizer,
+    samples: list,
+    device: torch.device,
+    lr: float = 1e-3,
+    l2_weight: float = 1e-4,
+    epochs: int = 5,
+):
+    """Train one AffineProbe per layer, streaming hidden states (no disk storage)."""
+    final_ln, lm_head = get_final_ln_and_lm_head(model)
+    config = model.config
+    d_model = getattr(config, "hidden_size", None) or getattr(config, "n_embd")
+    n_layers = getattr(config, "num_hidden_layers", None) or getattr(config, "n_layer")
+
+    # Create probes + optimizers
+    probes = {}
+    optimizers = {}
+    for l in range(n_layers):
+        probe = AffineProbe(d_model).to(device)
+        probes[l] = probe
+        optimizers[l] = torch.optim.AdamW(probe.parameters(), lr=lr, weight_decay=0.0)
+
+    logger.info(f"Training {n_layers} probes (d_model={d_model}, {len(samples)} samples, {epochs} epochs)")
+
+    for epoch in range(epochs):
+        epoch_losses = {l: 0.0 for l in range(n_layers)}
+        epoch_count = 0
+
+        for si, sample_ids in enumerate(samples):
+            input_ids = sample_ids.unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                outputs = model(input_ids, output_hidden_states=True)
+                hidden_states = outputs.hidden_states  # tuple of (n_layers+1) tensors
+
+                # Reference distribution from final layer
+                ref_hidden = hidden_states[-1]
+                ref_normed = final_ln(ref_hidden)
+                ref_logits = lm_head(ref_normed)
+                ref_log_probs = F.log_softmax(ref_logits, dim=-1).detach()
+
+            # Train each layer's probe independently
+            for l in range(n_layers):
+                probe = probes[l]
+                optimizer = optimizers[l]
+
+                # hidden_states[0] = embedding, hidden_states[l+1] = after layer l
+                h = hidden_states[l + 1].detach()
+
+                corrected = probe(h)
+                corrected_normed = final_ln(corrected)
+                probe_logits = lm_head(corrected_normed)
+                probe_log_probs = F.log_softmax(probe_logits, dim=-1)
+
+                # KL(ref || probe) — ref is the target distribution
+                kl = F.kl_div(probe_log_probs, ref_log_probs.exp(), reduction="batchmean", log_target=False)
+
+                # L2 regularisation toward identity: ||W - I||^2 + ||b||^2
+                identity = torch.eye(d_model, device=device, dtype=probe.weight.dtype)
+                l2_reg = ((probe.weight - identity) ** 2).sum() + (probe.bias ** 2).sum()
+
+                loss = kl + l2_weight * l2_reg
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                epoch_losses[l] += loss.item()
+
+            epoch_count += 1
+
+            # Free memory
+            del outputs, hidden_states, ref_hidden, ref_normed, ref_logits, ref_log_probs
+
+            if (si + 1) % 100 == 0:
+                avg_loss = sum(epoch_losses[l] for l in range(n_layers)) / (n_layers * epoch_count)
+                logger.info(f"  Epoch {epoch+1}, sample {si+1}/{len(samples)}, avg loss: {avg_loss:.4f}")
+
+        avg_epoch_loss = sum(epoch_losses[l] for l in range(n_layers)) / (n_layers * max(epoch_count, 1))
+        logger.info(f"Epoch {epoch+1}/{epochs} complete — avg loss: {avg_epoch_loss:.4f}")
+
+    return probes
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint saving
+# ---------------------------------------------------------------------------
+
+def save_checkpoint(probes: dict, model, model_id: str, output_dir: str,
+                    training_config: dict):
+    """Save probe state dicts and metadata."""
+    model_hash = compute_model_hash(model)
+    config = model.config
+    d_model = getattr(config, "hidden_size", None) or getattr(config, "n_embd")
+    n_layers = getattr(config, "num_hidden_layers", None) or getattr(config, "n_layer")
+
+    save_dir = Path(output_dir) / model_id
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    # Build combined state dict
+    state_dict = {}
+    for layer_idx, probe in probes.items():
+        state_dict[f"layer_{layer_idx}.weight"] = probe.weight.data.cpu()
+        state_dict[f"layer_{layer_idx}.bias"] = probe.bias.data.cpu()
+
+    checkpoint_path = save_dir / f"tuned_lens_{model_hash[:16]}.pt"
+    torch.save(state_dict, checkpoint_path)
+
+    metadata = {
+        "model_id": model_id,
+        "model_hash": model_hash,
+        "n_layers": n_layers,
+        "d_model": d_model,
+        "training_config": training_config,
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+    }
+    metadata_path = save_dir / "metadata.json"
+    with open(metadata_path, "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    logger.info(f"Saved checkpoint to {checkpoint_path} ({checkpoint_path.stat().st_size / 1024 / 1024:.1f}MB)")
+    logger.info(f"Saved metadata to {metadata_path}")
+    return checkpoint_path
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Train tuned lens probes for a model")
+    parser.add_argument("--model-id", required=True, help="Model identifier (e.g. codegen-350m)")
+    parser.add_argument("--model-name", default=None,
+                        help="HuggingFace model name (defaults to model-id)")
+    parser.add_argument("--corpus-file", required=True, help="Plain-text calibration corpus")
+    parser.add_argument("--output-dir", default="./tuned_lens_weights/",
+                        help="Output directory for checkpoints")
+    parser.add_argument("--max-samples", type=int, default=2000)
+    parser.add_argument("--max-seq-len", type=int, default=512)
+    parser.add_argument("--epochs", type=int, default=5)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--l2-weight", type=float, default=1e-4)
+    parser.add_argument("--device", default=None, help="Device (auto-detected if omitted)")
+    parser.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
+    args = parser.parse_args()
+
+    model_name = args.model_name or args.model_id
+    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+    dtype = dtype_map[args.dtype]
+
+    if args.device:
+        device = torch.device(args.device)
+    elif torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    logger.info(f"Device: {device}, dtype: {dtype}")
+    logger.info(f"Loading model: {model_name}")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=dtype).to(device)
+    model.eval()
+
+    samples = load_corpus(args.corpus_file, args.max_samples, args.max_seq_len, tokenizer)
+    if not samples:
+        logger.error("No valid samples loaded — aborting")
+        sys.exit(1)
+
+    training_config = {
+        "lr": args.lr,
+        "l2_weight": args.l2_weight,
+        "epochs": args.epochs,
+        "max_samples": args.max_samples,
+        "max_seq_len": args.max_seq_len,
+        "dtype": args.dtype,
+        "num_samples_used": len(samples),
+    }
+
+    probes = train_tuned_lens(
+        model, tokenizer, samples, device,
+        lr=args.lr, l2_weight=args.l2_weight, epochs=args.epochs,
+    )
+
+    save_checkpoint(probes, model, args.model_id, args.output_dir, training_config)
+    logger.info("Done.")
+
+
+if __name__ == "__main__":
+    main()