Upload 149 files

hugging/td_fuse/selfimprove.py

@@ -1,534 +1,629 @@
 """
-TD Self-Improvement Loop — the core of Time Dilation.
-
-THE LOOP (from test_1 through test_18 interviews):
-1. Ask the model "what are you bad at?" → it identifies weak spots
-2. Generate HARD problems targeting those weaknesses (stuff it fails 30-50%)
-3. Generate N candidates per problem, score with verified rewards
-4. Train on winners (STaR/GRPO approach)
-5. Benchmark → measure improvement
-6. Save → repeat
-
-KEY DIFFERENCE FROM GRPO CYCLE 2:
-- Problems are HARDER — designed so the model fails 30-50% of the time
-- This creates real contrast for learning (correct vs wrong chains)
-- Uses CoT (<think>...</think>) for all reasoning
-- Saves after every major step (diagnosis, generation, training)
-- Targets layers 16-28 specifically for reasoning improvement
-
-PRINCIPLES (from interviews + dad's tests):
-- Verified rewards only: math correct, logic valid, no learned reward model
-- Cherry_LLM perplexity filter prevents mode collapse
-- Mix problem types to avoid overfitting
-- Small nudges per cycle (1-5%), compounds over time
+TD Self-Improvement v4 — Full Weight Intelligence + Targeted LoRA
+
+NOT just fixing merge damage. Analyzes EVERYTHING about every layer:
+
+1. STANDALONE HEALTH — No reference needed:
+   - Distribution quality (kurtosis, skew — are weights well-shaped?)
+   - Spectral efficiency (effective rank vs max rank — is capacity wasted?)
+   - Conditioning (condition number — are gradients flowing well?)
+   - Sparsity patterns (dead neurons, over-sparse layers)
+   - Layer coherence (sudden norm jumps between adjacent layers)
+   - Weight magnitude balance (are norms proportional to layer depth?)
+
+2. REFERENCE COMPARISON — What changed from the merge:
+   - Norm drift, variance shift, distribution shape change
+   - This is just ONE signal among many
+
+3. IMPROVEMENT SCORING — Every layer gets a score:
+   - How much COULD this layer be improved?
+   - Combines standalone health + reference damage
+   - Ranks ALL layers, targets the ones with most room to grow
+
+4. LoRA IMPROVEMENT — Train the layers with most potential:
+   - Diverse data: text, math, code, reasoning, conversation
+   - Targets layers ranked by improvement potential
+
+From test_18: LoRA/DoRA with layers_to_transform, layers 16-28 for reasoning.
+But we don't limit — we go wherever the analysis says.
 """
 
 import torch
 import time
 import json
 import math
-import random
 import gc
-import re
 import shutil
+import numpy as np
 from pathlib import Path
-from typing import List, Dict, Optional
-from dataclasses import dataclass
+from typing import List, Dict, Tuple, Optional
+from dataclasses import dataclass, field
 
 
 @dataclass
 class SelfImproveConfig:
-    """Configuration for one self-improvement cycle."""
+    """Configuration for full self-improvement."""
     model_path: str = "td_fuse_outputs/reasoning_healed"
+    reference_model: str = "Qwen/Qwen3-VL-8B-Instruct"
     output_dir: str = "td_fuse_outputs/self_improve"
-    progress_file: str = "td_fuse_outputs/self_improve/progress.json"
-
-    # Generation
-    num_candidates: int = 8        # Generate N answers per question, pick best
-    max_gen_tokens: int = 1024     # Long — CoT needs room
-    temperature: float = 0.8       # Higher than GRPO — we want diversity + some failures
 
-    # Problems — more than GRPO, and harder
-    num_hard_math: int = 150       # Multi-step math (the model should fail ~30-50%)
-    num_hard_logic: int = 100      # Complex logic chains
-    num_hard_tricks: int = 50      # Tricky questions that trip up LLMs
-    num_hard_coding: int = 50      # Code reasoning problems
-    batch_save_every: int = 25     # Save progress more often
+    # How many layers to target (by improvement potential)
+    top_n_layers: int = 16       # Target more layers — we're improving, not just fixing
+    min_improvement_score: float = 0.05  # Floor to be worth targeting
 
     # Training
-    lora_r: int = 32               # Higher rank — more capacity for reasoning
+    lora_r: int = 32
     lora_alpha: int = 64
     train_epochs: int = 3
     train_batch: int = 2
-    train_grad_accum: int = 8      # Effective batch = 16
-    learning_rate: float = 1e-5    # Gentle
+    train_grad_accum: int = 8
+    learning_rate: float = 1.5e-5
     save_steps: int = 50
-    target_layers: str = "16-28"   # Reasoning circuit (from test_18)
-
+    num_train_samples: int = 3000  # More diverse data for general improvement
 
-# Import verify from grpo_reason
-from .grpo_reason import verify_answer as smart_verify
+    # Weights for improvement scoring (standalone health vs reference damage)
+    w_health: float = 0.55    # Standalone health issues matter most
+    w_damage: float = 0.25    # Merge damage
+    w_coherence: float = 0.20 # Layer-to-layer coherence
 
 
 # ============================================================
-# STEP 1: DIAGNOSE — Ask the model what it's bad at
+# STEP 1: DEEP WEIGHT ANALYSIS — Every layer, every metric
 # ============================================================
 
-def diagnose_weaknesses(model, tokenizer, eos_id, save_path=None):
-    """
-    Ask the model to identify its own weaknesses.
-    All 3 AIs (ChatGPT, Grok, Gemini) confirmed models can do this.
-    """
-    print("\n=== STEP 1: SELF-DIAGNOSIS ===")
-
-    prompts = [
-        ("What kinds of questions are you worst at? Be brutally honest. "
-         "List your top 5 weaknesses with specific examples of questions you'd get wrong."),
-        ("Give me 10 math or logic problems that would be genuinely hard for you. "
-         "Problems where you'd likely make mistakes. Be specific."),
-        ("What types of multi-step reasoning cause you to make errors? "
-         "Give 5 specific examples where you'd probably get the wrong answer."),
-    ]
-
-    all_responses = []
-    for i, prompt in enumerate(prompts):
-        p = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
-             f"<|im_start|>assistant\n<think>\nLet me think carefully about my actual weaknesses.\n")
-        ids = tokenizer(p, return_tensors="pt").to(model.device)
-        out = model.generate(
-            **ids, max_new_tokens=800, do_sample=True,
-            temperature=0.7, eos_token_id=eos_id
-        )
-        response = tokenizer.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True)
-        all_responses.append(response)
-        # Show the answer part (after </think>) if available
-        display = response.split("</think>")[-1].strip() if "</think>" in response else response
-        print(f"  Diagnosis {i+1}: {display[:200]}...")
+def analyze_layer(name: str, tensor: torch.Tensor) -> Dict:
+    """Full health analysis for a single weight tensor."""
+    t = tensor.float()
+
+    stats = {
+        "name": name,
+        "shape": list(tensor.shape),
+        "params": tensor.numel(),
+        "dtype": str(tensor.dtype),
+        # Basic stats
+        "l2_norm": t.norm(2).item(),
+        "frobenius_norm": t.norm('fro').item() if t.dim() >= 2 else t.norm(2).item(),
+        "mean": t.mean().item(),
+        "std": t.std().item(),
+        "min": t.min().item(),
+        "max": t.max().item(),
+        "abs_mean": t.abs().mean().item(),
+        # Distribution shape
+        "sparsity": (t.abs() < 1e-6).float().mean().item(),
+        "kurtosis": ((t - t.mean()) / max(t.std().item(), 1e-8)).pow(4).mean().item() - 3,
+        "skewness": ((t - t.mean()) / max(t.std().item(), 1e-8)).pow(3).mean().item(),
+        # Dead/saturated neurons
+        "near_zero_pct": (t.abs() < 1e-5).float().mean().item(),
+        "extreme_pct": (t.abs() > t.abs().mean() * 5).float().mean().item(),
+    }
+
+    # For 2D matrices: spectral analysis
+    if t.dim() == 2 and min(t.shape) >= 2:
+        try:
+            rows = min(512, t.shape[0])
+            cols = min(512, t.shape[1])
+            sample = t[:rows, :cols]
+
+            svd = torch.linalg.svdvals(sample)
+            stats["sv_max"] = svd[0].item()
+            stats["sv_min"] = svd[-1].item()
+            stats["sv_ratio"] = svd[0].item() / max(svd[-1].item(), 1e-10)
+
+            # Effective rank — how much capacity is actually used
+            sv_norm = svd / svd.sum()
+            entropy = -(sv_norm * torch.log(sv_norm + 1e-10)).sum()
+            stats["sv_effective_rank"] = torch.exp(entropy).item()
+            stats["sv_max_rank"] = float(min(sample.shape))
+            stats["sv_rank_utilization"] = stats["sv_effective_rank"] / stats["sv_max_rank"]
+
+            # Top singular value concentration — is too much info in few dimensions?
+            top5_energy = (svd[:5] ** 2).sum() / (svd ** 2).sum()
+            stats["sv_top5_energy"] = top5_energy.item()
+
+            # Stable rank (Frobenius^2 / spectral^2)
+            stats["stable_rank"] = (sample.norm('fro').item() ** 2) / max(svd[0].item() ** 2, 1e-10)
+
+        except Exception:
+            stats["sv_ratio"] = -1
+            stats["sv_rank_utilization"] = -1
+
+    return stats
+
 
-    # Save diagnosis
-    if save_path:
-        with open(save_path, "w") as f:
-            json.dump({"diagnoses": all_responses, "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")}, f, indent=2)
-        print(f"  Saved diagnosis to {save_path}")
+def analyze_model_weights(model_path: str, label: str = "model") -> Dict:
+    """Load and analyze every language layer in the model."""
+    print(f"\n{'='*50}")
+    print(f"ANALYZING: {label}")
+    print(f"Path: {model_path}")
+    print(f"{'='*50}")
 
-    return all_responses
+    from safetensors.torch import load_file
+    import glob
+
+    st_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    if not st_files:
+        raise FileNotFoundError(f"No safetensors files in {model_path}")
+
+    all_stats = {}
+    total_params = 0
+
+    for st_file in st_files:
+        print(f"  Loading {Path(st_file).name}...")
+        state_dict = load_file(st_file, device="cpu")
+
+        for name, tensor in state_dict.items():
+            # Skip vision encoder — language model only
+            if any(skip in name for skip in ["visual.", "merger."]):
+                continue
+
+            stats = analyze_layer(name, tensor)
+            all_stats[name] = stats
+            total_params += stats["params"]
+
+        del state_dict
+        gc.collect()
+
+    print(f"  Analyzed {len(all_stats)} tensors ({total_params/1e9:.2f}B params)")
+    return {"label": label, "layers": all_stats, "total_params": total_params}
 
 
 # ============================================================
-# STEP 2: GENERATE HARD PROBLEMS
+# STEP 2: STANDALONE HEALTH SCORING — No reference needed
 # ============================================================
 
-def generate_hard_problems(cfg: SelfImproveConfig) -> List[Dict]:
+def compute_health_scores(analysis: Dict) -> Dict[str, Dict]:
     """
-    Generate problems that are HARD enough for the model to fail 30-50%.
-    This is the key difference from GRPO cycle 2 (which was too easy at 97%).
+    Score each layer's health based on its own stats.
+    Higher score = MORE room for improvement (worse health).
+
+    Checks:
+    - Distribution quality: Extreme kurtosis or skew = unhealthy
+    - Spectral efficiency: Low rank utilization = wasted capacity
+    - Conditioning: Very high condition number = bad gradient flow
+    - Dead neurons: High sparsity = parts of the layer doing nothing
+    - Extreme weights: Too many outliers = potential instability
     """
-    problems = []
-
-    # --- HARD MATH: multi-step, carrying, fractions, negative numbers ---
-    for _ in range(cfg.num_hard_math):
-        ptype = random.choice([
-            "multi_step", "carrying", "negative", "fraction_word",
-            "percentage_chain", "rate_comparison", "sequence_hard"
-        ])
-
-        if ptype == "multi_step":
-            # 3-4 step word problems
-            items = random.choice(["books", "apples", "shirts", "tickets", "cookies"])
-            start = random.randint(20, 100)
-            buy = random.randint(5, 30)
-            give = random.randint(1, min(15, start + buy - 1))
-            lose = random.randint(1, min(5, start + buy - give - 1))
-            final = start + buy - give - lose
-            problems.append({
-                "question": (f"You have {start} {items}. You buy {buy} more. "
-                            f"You give {give} to a friend. You lose {lose}. "
-                            f"How many {items} do you have now?"),
-                "answer": str(final), "type": "hard_math",
-                "think_hint": f"{start} + {buy} - {give} - {lose} = {final}"
-            })
-
-        elif ptype == "carrying":
-            # Large number arithmetic that requires carrying
-            a = random.randint(500, 9999)
-            b = random.randint(500, 9999)
-            op = random.choice(["+", "-", "*"])
-            if op == "+": answer = a + b
-            elif op == "-": answer = a - b
-            else:
-                a = random.randint(20, 200)
-                b = random.randint(20, 200)
-                answer = a * b
-            problems.append({
-                "question": f"Calculate: {a} {op} {b}",
-                "answer": str(answer), "type": "hard_math",
-                "think_hint": f"{a} {op} {b} = {answer}"
-            })
-
-        elif ptype == "negative":
-            a = random.randint(-50, 50)
-            b = random.randint(-50, 50)
-            op = random.choice(["+", "-", "*"])
-            if op == "+": answer = a + b
-            elif op == "-": answer = a - b
-            else: answer = a * b
-            a_str = str(a) if a >= 0 else f"({a})"
-            b_str = str(b) if b >= 0 else f"({b})"
-            problems.append({
-                "question": f"What is {a_str} {op} {b_str}?",
-                "answer": str(answer), "type": "hard_math",
-                "think_hint": f"{a} {op} {b} = {answer}"
-            })
-
-        elif ptype == "fraction_word":
-            total = random.choice([12, 16, 20, 24, 30, 36, 40, 48, 60])
-            frac_n = random.choice([1, 2, 3])
-            frac_d = random.choice([2, 3, 4, 5, 6])
-            while frac_n >= frac_d:
-                frac_n = random.choice([1, 2, 3])
-            first = total * frac_n // frac_d
-            remaining = total - first
-            second_frac = random.choice([2, 3, 4])
-            second = remaining // second_frac
-            final = remaining - second
-            problems.append({
-                "question": (f"A jar has {total} marbles. You take {frac_n}/{frac_d} of them. "
-                            f"Then you give away {second} of what's left. How many remain in the jar?"),
-                "answer": str(final), "type": "hard_math",
-                "think_hint": (f"Take {frac_n}/{frac_d} of {total} = {first}. "
-                              f"Left: {remaining}. Give {second}. Final: {final}")
-            })
-
-        elif ptype == "percentage_chain":
-            base = random.choice([100, 200, 400, 500, 800, 1000])
-            pct1 = random.choice([10, 20, 25, 30, 50])
-            after1 = base + base * pct1 // 100
-            pct2 = random.choice([10, 20, 25])
-            final = after1 - after1 * pct2 // 100
-            problems.append({
-                "question": (f"A stock is worth ${base}. It goes up {pct1}%, "
-                            f"then drops {pct2}%. What is it worth now?"),
-                "answer": str(final), "type": "hard_math",
-                "think_hint": (f"Up {pct1}%: {base} + {base*pct1//100} = {after1}. "
-                              f"Down {pct2}%: {after1} - {after1*pct2//100} = {final}")
-            })
-
-        elif ptype == "rate_comparison":
-            speed1 = random.randint(30, 80)
-            time1 = random.randint(2, 6)
-            speed2 = random.randint(40, 100)
-            time2 = random.randint(1, 5)
-            dist1 = speed1 * time1
-            dist2 = speed2 * time2
-            diff = abs(dist1 - dist2)
-            who = "first" if dist1 > dist2 else "second"
-            problems.append({
-                "question": (f"Car A drives {speed1} mph for {time1} hours. "
-                            f"Car B drives {speed2} mph for {time2} hours. "
-                            f"How much further did the car that drove more go?"),
-                "answer": str(diff), "type": "hard_math",
-                "think_hint": f"Car A: {dist1}mi. Car B: {dist2}mi. Diff: {diff}mi."
-            })
-
-        elif ptype == "sequence_hard":
-            # Non-obvious sequences
-            seq_type = random.choice(["quadratic", "fibonacci_like", "alternating"])
-            if seq_type == "quadratic":
-                a, b, c = random.randint(1, 3), random.randint(1, 5), random.randint(0, 10)
-                seq = [a*n*n + b*n + c for n in range(1, 6)]
-                next_val = a*36 + b*6 + c
-            elif seq_type == "fibonacci_like":
-                a, b = random.randint(1, 5), random.randint(1, 5)
-                seq = [a, b]
-                for _ in range(4):
-                    seq.append(seq[-1] + seq[-2])
-                next_val = seq[-1] + seq[-2]
-                seq = seq[:5]
-            else:  # alternating
-                start = random.randint(1, 10)
-                step1 = random.randint(2, 5)
-                step2 = random.randint(1, 3)
-                seq = [start]
-                for i in range(4):
-                    if i % 2 == 0:
-                        seq.append(seq[-1] + step1)
-                    else:
-                        seq.append(seq[-1] - step2)
-                # Next follows the pattern
-                next_val = seq[-1] + step1 if len(seq) % 2 == 1 else seq[-1] - step2
-            problems.append({
-                "question": f"What comes next: {', '.join(map(str, seq))}?",
-                "answer": str(next_val), "type": "hard_math",
-            })
-
-    # --- HARD LOGIC: longer chains, negations, misdirection ---
-    logic_templates = [
-        # 4+ entity ordering
-        ("Amy is taller than Bob. Carol is shorter than Bob but taller than Dave. "
-         "Eve is taller than Amy. Who is the shortest?",
-         "dave", "Eve > Amy > Bob > Carol > Dave. Dave is shortest."),
-        ("In a tournament: A beat B, C beat A, D beat C, B beat D. Who beat the most players?",
-         "1", "Each player beat exactly 1 other player (it's circular)."),
-        # Negation chains
-        ("If it's not true that all cats are black, and Fluffy is a cat, "
-         "can we conclude Fluffy is not black?",
-         "no", "We can't conclude either way — some cats might be black, some not."),
-        ("None of the winners cheated. Tom didn't cheat. Did Tom win?",
-         "no", "Not cheating is necessary but not sufficient for winning."),
-        # Conditional logic
-        ("If it rains, the ground is wet. The ground is wet. Did it rain?",
-         "no", "The ground could be wet for other reasons (sprinkler, etc). Affirming the consequent."),
-        ("If it rains, the ground is wet. It didn't rain. Is the ground dry?",
-         "no", "The ground could still be wet from other causes. Denying the antecedent."),
-        ("If A then B. If B then C. A is true. Is C true?",
-         "yes", "A→B→C. A is true, so B is true, so C is true. Modus ponens chain."),
-        ("If A then B. If B then C. C is false. Is A true?",
-         "no", "C is false means B must be false (contrapositive), which means A must be false."),
-        # Set theory
-        ("In a class of 30 students, 18 play soccer and 15 play basketball. "
-         "At least how many play both?",
-         "3", "18 + 15 = 33. Must overlap by at least 33 - 30 = 3."),
-        ("A group has 40 people. 25 speak English, 20 speak French, and 10 speak both. "
-         "How many speak neither?",
-         "5", "English or French: 25 + 20 - 10 = 35. Neither: 40 - 35 = 5."),
-    ]
-
-    for _ in range(cfg.num_hard_logic):
-        if logic_templates:
-            q, a, hint = random.choice(logic_templates)
-            problems.append({
-                "question": q, "answer": a, "type": "hard_logic",
-                "think_hint": hint
-            })
-
-    # Also generate parameterized ordering problems (harder: 5-6 entities)
-    for _ in range(cfg.num_hard_logic // 2):
-        names = random.sample(["Amy", "Bob", "Carol", "Dave", "Eve", "Frank"], 5)
-        # Create a random ordering
-        order = list(range(5))
-        random.shuffle(order)
-        # Generate comparison clues (not all pairs — make them work it out)
-        clues = []
-        for i in range(4):
-            a_idx = order[i]
-            b_idx = order[i + 1]
-            clues.append(f"{names[a_idx]} is taller than {names[b_idx]}")
-        random.shuffle(clues)
-        tallest = names[order[0]]
-        shortest = names[order[4]]
-        q_type = random.choice(["tallest", "shortest"])
-        answer = tallest.lower() if q_type == "tallest" else shortest.lower()
-        problems.append({
-            "question": f"{'. '.join(clues)}. Who is the {q_type}?",
-            "answer": answer, "type": "hard_logic",
-        })
-
-    # --- HARD TRICKS: LLM failure modes ---
-    tricks = [
-        ("How many times can you subtract 5 from 25?",
-         "once", "Only once — after that it's 20, not 25."),
-        ("How many months have 28 days?",
-         "all", "All 12 months have at least 28 days."),
-        ("I have 6 eggs. I broke 2, cooked 2, and ate 2. How many eggs do I have left?",
-         "4", "You broke 2 to cook them, and ate the cooked ones. Same 2 eggs through the process. 6-2=4."),
-        ("A doctor gives you 3 pills and says take one every 30 minutes. How long do they last?",
-         "60", "Take first at 0 min, second at 30 min, third at 60 min. Total: 60 minutes."),
-        ("How many times does the digit 9 appear between 1 and 100?",
-         "20", "9,19,29,39,49,59,69,79,89 (nine times in ones) + 90-99 (ten times in tens, plus 99 double-counts). Total: 20."),
-        ("If you overtake the last person in a race, what position are you in?",
-         "impossible", "You can't overtake the last person — if you're behind them, YOU are last."),
-        ("A bat and ball cost $1.10 total. The bat costs $1 more than the ball. How much is the ball?",
-         "5", "Ball = x, Bat = x+1.00. Total: 2x+1.00=1.10. x=0.05. Ball costs $0.05 or 5 cents."),
-        ("If 5 machines take 5 minutes to make 5 widgets, how long for 100 machines to make 100 widgets?",
-         "5", "Each machine makes 1 widget in 5 minutes. 100 machines make 100 widgets in 5 minutes."),
-        ("You're in a dark room with a candle, a wood stove, and a gas lamp. You have one match. What do you light first?",
-         "match", "You light the match first — you need it to light anything else."),
-        ("Is it legal for a man to marry his widow's sister?",
-         "no", "If he has a widow, he's dead. Dead people can't marry."),
-    ]
-    for _ in range(cfg.num_hard_tricks):
-        q, a, hint = random.choice(tricks)
-        problems.append({
-            "question": q, "answer": a, "type": "hard_trick",
-            "think_hint": hint
-        })
-
-    # --- CODE REASONING (no actual code execution, just logic) ---
-    code_problems = [
-        ("What does this print?\nx = 5\nx = x + 3\nprint(x)", "8",
-         "x starts at 5, then x = 5+3 = 8. Prints 8."),
-        ("What does this print?\nfor i in range(3): print(i)\n(List all numbers printed)", "0 1 2",
-         "range(3) gives 0, 1, 2."),
-        ("What's wrong with this code?\ndef add(a, b):\n  return a - b", "subtracts instead of adds",
-         "Function named 'add' but uses minus operator."),
-        ("What does this evaluate to?\nTrue and False or True", "true",
-         "(True and False) = False. False or True = True."),
-        ("After this code runs, what is x?\nx = [1,2,3]\ny = x\ny.append(4)\nprint(len(x))", "4",
-         "y = x makes y point to same list. Appending to y also changes x. len(x) = 4."),
-    ]
-    for _ in range(cfg.num_hard_coding):
-        q, a, hint = random.choice(code_problems)
-        problems.append({
-            "question": q, "answer": a.lower(), "type": "hard_coding",
-            "think_hint": hint
-        })
-
-    random.shuffle(problems)
-    print(f"  Generated {len(problems)} hard problems")
-    print(f"    Math: {cfg.num_hard_math}, Logic: {cfg.num_hard_logic + cfg.num_hard_logic//2}, "
-          f"Tricks: {cfg.num_hard_tricks}, Code: {cfg.num_hard_coding}")
-
-    return problems
+    print("\n=== STANDALONE HEALTH ANALYSIS ===")
+
+    health_scores = {}
+
+    for name, stats in analysis["layers"].items():
+        issues = []
+        score = 0.0
+
+        # 1. Distribution quality — ideal kurtosis near 0 (normal), skew near 0
+        kurt = abs(stats["kurtosis"])
+        if kurt > 10:
+            score += 0.3
+            issues.append(f"extreme kurtosis ({kurt:.1f})")
+        elif kurt > 5:
+            score += 0.15
+            issues.append(f"high kurtosis ({kurt:.1f})")
+
+        skew = abs(stats["skewness"])
+        if skew > 2:
+            score += 0.15
+            issues.append(f"high skew ({skew:.2f})")
+
+        # 2. Spectral efficiency — are we using the layer's full capacity?
+        rank_util = stats.get("sv_rank_utilization", -1)
+        if rank_util > 0:
+            if rank_util < 0.1:
+                score += 0.3
+                issues.append(f"very low rank utilization ({rank_util:.2%})")
+            elif rank_util < 0.25:
+                score += 0.15
+                issues.append(f"low rank utilization ({rank_util:.2%})")
+
+            # Top-5 energy concentration — too much in few dimensions
+            top5 = stats.get("sv_top5_energy", 0)
+            if top5 > 0.9:
+                score += 0.2
+                issues.append(f"top-5 SVD concentration ({top5:.2%})")
+            elif top5 > 0.75:
+                score += 0.1
+
+        # 3. Conditioning — gradient flow health
+        sv_ratio = stats.get("sv_ratio", -1)
+        if sv_ratio > 0:
+            if sv_ratio > 10000:
+                score += 0.25
+                issues.append(f"extreme condition number ({sv_ratio:.0f})")
+            elif sv_ratio > 1000:
+                score += 0.1
+
+        # 4. Dead neurons — near-zero weights
+        dead = stats["near_zero_pct"]
+        if dead > 0.3:
+            score += 0.2
+            issues.append(f"high dead neuron % ({dead:.1%})")
+        elif dead > 0.15:
+            score += 0.1
+
+        # 5. Extreme weights — outlier instability
+        extreme = stats["extreme_pct"]
+        if extreme > 0.05:
+            score += 0.15
+            issues.append(f"many extreme weights ({extreme:.1%})")
+
+        # 6. Std too low (collapsed) or too high (exploded)
+        if stats["std"] < 1e-4:
+            score += 0.25
+            issues.append(f"near-zero variance ({stats['std']:.6f})")
+        elif stats["std"] > 1.0:
+            score += 0.15
+            issues.append(f"high variance ({stats['std']:.4f})")
+
+        health_scores[name] = {
+            "health_score": min(score, 1.0),  # Cap at 1.0
+            "issues": issues,
+            "stats": stats,
+        }
+
+    return health_scores
 
 
 # ============================================================
-# STEP 3: GENERATE & SCORE WITH COT
+# STEP 3: REFERENCE COMPARISON — Merge damage signal
 # ============================================================
 
-def build_cot_prompt(question: str) -> str:
-    """Build prompt that encourages chain-of-thought reasoning."""
-    return (
-        f"<|im_start|>user\n{question}\n"
-        f"Think step by step before giving your final answer.<|im_end|>\n"
-        f"<|im_start|>assistant\n<think>\n"
-    )
+def compute_damage_scores(model_analysis: Dict, ref_analysis: Dict) -> Dict[str, float]:
+    """Compare each layer to reference. Returns damage score per layer name."""
+    print("\n=== REFERENCE COMPARISON (merge damage) ===")
+
+    damages = {}
+    model_layers = model_analysis["layers"]
+    ref_layers = ref_analysis["layers"]
+    matched = 0
+
+    for name, m_stats in model_layers.items():
+        ref_key = None
+        for c in [name, name.replace("model.", "", 1), f"model.{name}"]:
+            if c in ref_layers:
+                ref_key = c
+                break
 
+        if ref_key is None:
+            continue
 
-def generate_and_score(model, tokenizer, problems, cfg, eos_id, resume_from=0):
+        r_stats = ref_layers[ref_key]
+        matched += 1
+
+        norm_diff = abs(m_stats["l2_norm"] - r_stats["l2_norm"]) / max(r_stats["l2_norm"], 1e-8)
+        std_diff = abs(m_stats["std"] - r_stats["std"]) / max(r_stats["std"], 1e-8)
+        kurt_diff = abs(m_stats["kurtosis"] - r_stats["kurtosis"]) / max(abs(r_stats["kurtosis"]) + 1, 1e-8)
+        sparsity_diff = abs(m_stats["sparsity"] - r_stats["sparsity"])
+
+        sv_ratio_diff = 0
+        if "sv_ratio" in m_stats and "sv_ratio" in r_stats:
+            if m_stats["sv_ratio"] > 0 and r_stats["sv_ratio"] > 0:
+                sv_ratio_diff = abs(
+                    math.log(m_stats["sv_ratio"] + 1) - math.log(r_stats["sv_ratio"] + 1)
+                ) / max(math.log(r_stats["sv_ratio"] + 1), 1e-8)
+
+        combined = (0.30 * norm_diff + 0.25 * std_diff + 0.20 * kurt_diff +
+                    0.15 * sv_ratio_diff + 0.10 * sparsity_diff)
+        damages[name] = combined
+
+    print(f"  Matched {matched} layers between model and reference")
+    return damages
+
+
+# ============================================================
+# STEP 4: LAYER COHERENCE — Adjacent layer smoothness
+# ============================================================
+
+def compute_coherence_scores(analysis: Dict) -> Dict[int, float]:
     """
-    Generate CoT candidates, score them, return winning pairs.
-    Saves progress every batch_save_every problems.
+    Check layer-to-layer coherence.
+    Big jumps in norms/variance between adjacent layers = poor coherence.
+    Returns coherence issue score per layer number.
     """
-    print(f"\n=== STEP 3: GENERATE & SCORE ({len(problems)} problems, {cfg.num_candidates} candidates) ===")
-    gen_start = time.time()
-
-    progress_path = Path(cfg.progress_file)
-    progress_path.parent.mkdir(parents=True, exist_ok=True)
-
-    winning_pairs = []
-    losing_pairs = []
-
-    if resume_from > 0 and progress_path.exists():
-        with open(progress_path) as f:
-            saved = json.load(f)
-            winning_pairs = saved.get("winning_pairs", [])
-            losing_pairs = saved.get("losing_pairs", [])
-        print(f"  Resumed from problem {resume_from}, {len(winning_pairs)} winners loaded")
-
-    total_correct = 0
-    total_attempted = resume_from
-
-    for i in range(resume_from, len(problems)):
-        prob = problems[i]
-        prompt = build_cot_prompt(prob["question"])
-        ids = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-        correct_responses = []
-        wrong_responses = []
-
-        for _ in range(cfg.num_candidates):
-            out = model.generate(
-                **ids, max_new_tokens=cfg.max_gen_tokens,
-                do_sample=True, temperature=cfg.temperature,
-                eos_token_id=eos_id
-            )
-            response = tokenizer.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True)
-
-            # Use verify_key if it exists (trick questions)
-            check_prob = dict(prob)
-            if "verify_key" in check_prob:
-                check_prob["answer"] = check_prob["verify_key"]
-            score = smart_verify(check_prob, response)
-
-            if score > 0:
-                correct_responses.append(response)
-            else:
-                wrong_responses.append(response)
-
-        total_attempted += 1
-
-        if correct_responses:
-            total_correct += 1
-            best = min(correct_responses, key=len)  # Shortest correct = most concise
-            chat = (
-                f"<|im_start|>user\n{prob['question']}\n"
-                f"Think step by step before giving your final answer.<|im_end|>\n"
-                f"<|im_start|>assistant\n{best}<|im_end|>"
-            )
-            winning_pairs.append(chat)
-
-            if wrong_responses:
-                worst = max(wrong_responses, key=len)
-                losing_pairs.append(
-                    f"<|im_start|>user\n{prob['question']}<|im_end|>\n"
-                    f"<|im_start|>assistant\n{worst}<|im_end|>"
-                )
-
-        # Progress
-        if (i + 1) % cfg.batch_save_every == 0 or i == len(problems) - 1:
-            pct = total_correct / max(1, total_attempted) * 100
-            elapsed = time.time() - gen_start
-            rate = (i + 1 - resume_from) / max(1, elapsed)
-            eta = (len(problems) - i - 1) / max(0.001, rate) / 60
-            print(f"  [{i+1}/{len(problems)}] Correct: {total_correct}/{total_attempted} ({pct:.0f}%) "
-                  f"| Winners: {len(winning_pairs)} | ETA: {eta:.0f}min")
-
-            with open(progress_path, "w") as f:
-                json.dump({
-                    "winning_pairs": winning_pairs,
-                    "losing_pairs": losing_pairs,
-                    "problems_completed": i + 1,
-                    "total_correct": total_correct,
-                    "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
-                }, f)
-
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-    fail_rate = 100 - (total_correct / max(1, total_attempted) * 100)
-    print(f"\n  FINAL: {total_correct}/{total_attempted} correct "
-          f"({total_correct/max(1,total_attempted)*100:.0f}%)")
-    print(f"  Failure rate: {fail_rate:.0f}% "
-          f"({'GOOD — enough contrast for learning' if 15 < fail_rate < 60 else 'WARNING — might be too easy or too hard'})")
-    print(f"  Winners: {len(winning_pairs)} | Losers: {len(losing_pairs)}")
-
-    return winning_pairs, losing_pairs
+    print("\n=== LAYER COHERENCE ANALYSIS ===")
+
+    # Group stats by layer number
+    layer_norms = {}
+    layer_stds = {}
+
+    for name, stats in analysis["layers"].items():
+        parts = name.split(".")
+        layer_num = None
+        for j, part in enumerate(parts):
+            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
+                layer_num = int(parts[j+1])
+                break
+
+        if layer_num is not None:
+            if layer_num not in layer_norms:
+                layer_norms[layer_num] = []
+                layer_stds[layer_num] = []
+            layer_norms[layer_num].append(stats["l2_norm"])
+            layer_stds[layer_num].append(stats["std"])
+
+    # Average norm per layer
+    avg_norms = {l: sum(ns)/len(ns) for l, ns in layer_norms.items()}
+    avg_stds = {l: sum(ss)/len(ss) for l, ss in layer_stds.items()}
+
+    sorted_layers = sorted(avg_norms.keys())
+    coherence_issues = {}
+
+    for i in range(len(sorted_layers)):
+        l = sorted_layers[i]
+        score = 0.0
+
+        # Check jump from previous layer
+        if i > 0:
+            prev = sorted_layers[i-1]
+            norm_jump = abs(avg_norms[l] - avg_norms[prev]) / max(avg_norms[prev], 1e-8)
+            std_jump = abs(avg_stds[l] - avg_stds[prev]) / max(avg_stds[prev], 1e-8)
+            if norm_jump > 0.5:
+                score += norm_jump * 0.5
+            if std_jump > 0.5:
+                score += std_jump * 0.3
+
+        # Check jump to next layer
+        if i < len(sorted_layers) - 1:
+            nxt = sorted_layers[i+1]
+            norm_jump = abs(avg_norms[l] - avg_norms[nxt]) / max(avg_norms[nxt], 1e-8)
+            if norm_jump > 0.5:
+                score += norm_jump * 0.3
+
+        coherence_issues[l] = min(score, 1.0)
+
+    return coherence_issues
 
 
 # ============================================================
-# STEP 4: TRAIN — SFT on winning CoT chains
+# STEP 5: COMBINE — Rank layers by improvement potential
 # ============================================================
 
-def train_on_winners(model, tokenizer, winning_pairs, cfg):
+def rank_layers_for_improvement(
+    health_scores: Dict[str, Dict],
+    damage_scores: Dict[str, float],
+    coherence_scores: Dict[int, float],
+    cfg: SelfImproveConfig
+) -> Tuple[List[int], Dict]:
     """
-    Train on correct chain-of-thought responses (STaR approach).
-    Targets layers 16-28 for reasoning improvement.
+    Combine all signals into a single improvement potential score per layer.
+    Returns sorted list of layer numbers to target + full report.
     """
-    print(f"\n=== STEP 4: TRAIN ON WINNERS ({len(winning_pairs)} chains) ===")
+    print("\n=== RANKING LAYERS BY IMPROVEMENT POTENTIAL ===")
+
+    # Aggregate per transformer layer number
+    layer_health = {}    # layer_num -> list of health scores
+    layer_damage = {}    # layer_num -> list of damage scores
+
+    for name, h in health_scores.items():
+        parts = name.split(".")
+        layer_num = None
+        for j, part in enumerate(parts):
+            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
+                layer_num = int(parts[j+1])
+                break
+
+        if layer_num is not None:
+            if layer_num not in layer_health:
+                layer_health[layer_num] = []
+                layer_damage[layer_num] = []
+            layer_health[layer_num].append(h["health_score"])
+            if name in damage_scores:
+                layer_damage[layer_num].append(damage_scores[name])
+
+    # Combined score per layer
+    layer_scores = {}
+    layer_details = {}
+
+    for layer_num in layer_health:
+        avg_health = sum(layer_health[layer_num]) / len(layer_health[layer_num])
+        avg_damage = (sum(layer_damage[layer_num]) / len(layer_damage[layer_num])
+                     if layer_damage.get(layer_num) else 0.0)
+        coherence = coherence_scores.get(layer_num, 0.0)
+
+        combined = (cfg.w_health * avg_health +
+                    cfg.w_damage * avg_damage +
+                    cfg.w_coherence * coherence)
+
+        layer_scores[layer_num] = combined
+        layer_details[layer_num] = {
+            "health": avg_health,
+            "damage": avg_damage,
+            "coherence": coherence,
+            "combined": combined,
+        }
+
+    # Sort by improvement potential (highest first)
+    sorted_layers = sorted(layer_scores.items(), key=lambda x: x[1], reverse=True)
+
+    # Print ranking
+    print(f"\n  {'Layer':>7}  {'Combined':>9}  {'Health':>8}  {'Damage':>8}  {'Coherence':>10}  {'Status'}")
+    print(f"  {'-'*65}")
+    for layer_num, score in sorted_layers:
+        d = layer_details[layer_num]
+        bar = "█" * int(score * 50)
+        target = " ← TARGET" if score >= cfg.min_improvement_score else ""
+        print(f"  Layer {layer_num:>2}  {score:>9.4f}  {d['health']:>8.4f}  {d['damage']:>8.4f}  {d['coherence']:>10.4f}  {bar}{target}")
+
+    # Pick top N layers above threshold
+    target_layers = []
+    for layer_num, score in sorted_layers:
+        if score >= cfg.min_improvement_score and len(target_layers) < cfg.top_n_layers:
+            target_layers.append(layer_num)
+
+    if not target_layers:
+        target_layers = [l for l, _ in sorted_layers[:cfg.top_n_layers]]
+        print(f"\n  No layers above threshold {cfg.min_improvement_score}, using top {cfg.top_n_layers}")
+
+    target_layers.sort()
+    print(f"\n  TARGET LAYERS FOR IMPROVEMENT: {target_layers}")
+
+    # Collect top issues across all targeted layers
+    all_issues = []
+    for name, h in health_scores.items():
+        parts = name.split(".")
+        layer_num = None
+        for j, part in enumerate(parts):
+            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
+                layer_num = int(parts[j+1])
+                break
+        if layer_num in target_layers and h["issues"]:
+            for issue in h["issues"]:
+                all_issues.append(f"  Layer {layer_num} {name.split('.')[-2]}: {issue}")
+
+    if all_issues:
+        print(f"\n  Key issues found in target layers:")
+        for issue in all_issues[:30]:
+            print(f"    {issue}")
+
+    return target_layers, layer_details
+
 
-    if len(winning_pairs) < 20:
-        print("  Too few winning pairs — skipping training")
-        return model
+# ============================================================
+# STEP 6: LoRA IMPROVEMENT — Train targeted layers
+# ============================================================
+
+def improve_with_lora(model_path: str, target_layers: List[int], cfg: SelfImproveConfig):
+    """
+    Apply LoRA training to layers with highest improvement potential.
+    Uses diverse high-quality data covering text, reasoning, code, and conversation.
+    """
+    print(f"\n{'='*50}")
+    print(f"LORA IMPROVEMENT — Targeting {len(target_layers)} layers: {target_layers}")
+    print(f"{'='*50}")
 
+    from transformers import AutoModelForImageTextToText, AutoTokenizer, TrainingArguments, Trainer
     from peft import LoraConfig, get_peft_model, TaskType
-    from transformers import TrainingArguments, Trainer
+    from datasets import load_dataset
     from torch.utils.data import Dataset
 
-    # Parse target layers
-    layer_start, layer_end = map(int, cfg.target_layers.split("-"))
-    target_layers = list(range(layer_start, layer_end + 1))
+    # Load model
+    print("  Loading model...")
+    model = AutoModelForImageTextToText.from_pretrained(
+        model_path, dtype=torch.bfloat16,
+        device_map="auto", trust_remote_code=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
+    # LoRA on target layers
     lora_config = LoraConfig(
-        r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=0.0,
+        r=cfg.lora_r,
+        lora_alpha=cfg.lora_alpha,
+        lora_dropout=0.0,
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                         "gate_proj", "up_proj", "down_proj"],
         layers_to_transform=target_layers,
-        bias="none", task_type=TaskType.CAUSAL_LM,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
     )
     model = get_peft_model(model, lora_config)
     model.print_trainable_parameters()
 
-    class CoTDataset(Dataset):
+    # Load DIVERSE training data — 5 categories for general improvement
+    print("\n  Loading diverse training data...")
+    train_texts = []
+    per_category = cfg.num_train_samples // 5  # Equal split across 5 types
+
+    # Category 1: General text (language fluency)
+    try:
+        ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
+        count = 0
+        for row in ds:
+            if count >= per_category:
+                break
+            text = row.get("text", "")
+            if len(text) > 100:
+                chat = (f"<|im_start|>user\nContinue this text naturally.<|im_end|>\n"
+                       f"<|im_start|>assistant\n{text[:800]}<|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    General text: {count}")
+    except Exception as e:
+        print(f"    WARNING: General text failed: {e}")
+
+    # Category 2: Math reasoning (step-by-step)
+    try:
+        ds = load_dataset("openai/gsm8k", "main", split="train")
+        count = 0
+        for row in ds:
+            if count >= per_category:
+                break
+            q, a = row.get("question", ""), row.get("answer", "")
+            if q and a:
+                chat = (f"<|im_start|>user\n{q}\nThink step by step.<|im_end|>\n"
+                       f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    Math reasoning: {count}")
+    except Exception as e:
+        print(f"    WARNING: Math data failed: {e}")
+
+    # Category 3: Code generation
+    try:
+        ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
+        count = 0
+        for row in ds:
+            if count >= per_category:
+                break
+            p, c = row.get("prompt", ""), row.get("completion", "")
+            if p and c:
+                chat = (f"<|im_start|>user\n{p}<|im_end|>\n"
+                       f"<|im_start|>assistant\n{c}<|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    Code: {count}")
+    except Exception as e:
+        print(f"    WARNING: Code data failed: {e}")
+
+    # Category 4: Instruction following (diverse tasks)
+    try:
+        ds = load_dataset("yahma/alpaca-cleaned", split="train")
+        count = 0
+        for row in ds:
+            if count >= per_category:
+                break
+            inst = row.get("instruction", "")
+            inp = row.get("input", "")
+            out = row.get("output", "")
+            if inst and out:
+                prompt = f"{inst}\n{inp}" if inp else inst
+                chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
+                       f"<|im_start|>assistant\n{out}<|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    Instructions: {count}")
+    except Exception as e:
+        print(f"    WARNING: Instruction data failed: {e}")
+
+    # Category 5: Reasoning / logic
+    try:
+        ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
+        count = 0
+        for row in ds:
+            if count >= per_category:
+                break
+            q = row.get("question", "")
+            a = row.get("answer", "")
+            if q and a:
+                chat = (f"<|im_start|>user\n{q}\nThink step by step before answering.<|im_end|>\n"
+                       f"<|im_start|>assistant\n<think>\nLet me think through this carefully.\n"
+                       f"The answer is {a}.\n</think>\nThe answer is {a}.<|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    Reasoning/logic: {count}")
+    except Exception as e:
+        print(f"    WARNING: Reasoning data failed: {e}")
+
+    print(f"\n  Total training samples: {len(train_texts)}")
+
+    if len(train_texts) < 100:
+        print("  ERROR: Not enough training data!")
+        return model_path
+
+    # Shuffle for variety
+    import random
+    random.shuffle(train_texts)
+
+    # Tokenize
+    class ImproveDataset(Dataset):
         def __init__(self, texts, tokenizer, max_len=1024):
             self.data = []
             for t in texts:
@@ -542,13 +637,13 @@ def train_on_winners(model, tokenizer, winning_pairs, cfg):
         def __len__(self): return len(self.data)
         def __getitem__(self, i): return self.data[i]
 
-    dataset = CoTDataset(winning_pairs, tokenizer)
+    print("  Tokenizing...")
+    dataset = ImproveDataset(train_texts, tokenizer)
 
     out_dir = Path(cfg.output_dir) / "train_output"
     out_dir.mkdir(parents=True, exist_ok=True)
 
     total_steps = (len(dataset) * cfg.train_epochs) // (cfg.train_batch * cfg.train_grad_accum)
-    print(f"  Steps: ~{total_steps} | Target layers: {cfg.target_layers}")
 
     args = TrainingArguments(
         output_dir=str(out_dir),
@@ -572,212 +667,194 @@ def train_on_winners(model, tokenizer, winning_pairs, cfg):
         train_dataset=dataset, args=args
     )
 
-    print("  Training...")
+    print(f"\n  Training: ~{total_steps} steps on {len(target_layers)} layers")
+    print(f"  Data: {len(train_texts)} samples × {cfg.train_epochs} epochs")
+    print(f"  LoRA: r={cfg.lora_r}, alpha={cfg.lora_alpha}")
     trainer.train()
 
-    print("  Merging LoRA...")
+    # Merge LoRA back into weights
+    print("\n  Merging LoRA adapters into model...")
     merged = model.merge_and_unload()
-    gc.collect()
 
+    # Save
+    save_dir = Path(cfg.output_dir) / "improved"
+    save_dir.mkdir(parents=True, exist_ok=True)
+    print(f"  Saving improved model to {save_dir}...")
+    merged.save_pretrained(str(save_dir), safe_serialization=True)
+    tokenizer.save_pretrained(str(save_dir))
+    sz = sum(f.stat().st_size for f in save_dir.glob("*.safetensors")) / 1e9
+    print(f"  SAVED: {save_dir} ({sz:.1f} GB)")
+
+    # Cleanup training artifacts
     shutil.rmtree(str(out_dir), ignore_errors=True)
+    del merged, model
     gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
-    return merged
-
-
-# ============================================================
-# STEP 5: BENCHMARK
-# ============================================================
-
-def benchmark(model, tokenizer, eos_id):
-    """Comprehensive reasoning benchmark."""
-    print("\n=== STEP 5: BENCHMARK ===")
-    results = {}
-
-    def ask(prompt, use_cot=False):
-        if use_cot:
-            p = build_cot_prompt(prompt)
-        else:
-            p = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
-        ids = tokenizer(p, return_tensors="pt").to(model.device)
-        out = model.generate(**ids, max_new_tokens=500, do_sample=False, eos_token_id=eos_id)
-        return tokenizer.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True)
-
-    # Basic math
-    math_tests = [("7+8", "15"), ("123+456", "579"), ("1000-387", "613"),
-                  ("12*13", "156"), ("144/12", "12")]
-    mc = sum(1 for q, e in math_tests
-             if smart_verify({"answer": e}, ask(f"What is {q}? Give just the number.")) > 0)
-    results["math"] = f"{mc}/5"
-    print(f"  Math: {mc}/5")
-
-    # Hard math (the stuff GRPO couldn't handle)
-    hard_math = [
-        ("A stock is worth $200. It goes up 25%, then drops 20%. What is it worth?", "200"),
-        ("You have 50 cookies. You eat 1/5, give 10 away, then split the rest equally among 4 friends. How many each?", "7"),
-        ("Car A drives 60mph for 3 hours. Car B drives 80mph for 2 hours. How much further did Car A go?", "20"),
-    ]
-    hmc = 0
-    for q, e in hard_math:
-        a = ask(q, use_cot=True)
-        correct = smart_verify({"answer": e}, a) > 0
-        hmc += correct
-        status = "OK" if correct else "FAIL"
-        answer_display = a.split("</think>")[-1].strip() if "</think>" in a else a
-        print(f"    {status}: {q[:55]}... -> {answer_display[:80]}")
-    results["hard_math"] = f"{hmc}/3"
-    print(f"  Hard math: {hmc}/3")
-
-    # Reasoning with CoT
-    reason_tests = [
-        ("Which is heavier: a pound of feathers or a pound of bricks?", "same"),
-        ("If yesterday was Monday, what day is tomorrow?", "wednesday"),
-        ("A farmer has 17 sheep. All but 9 die. How many are left?", "9"),
-        ("If you're running and pass the person in 2nd place, what position are you in?", "2nd"),
-        ("How many times can you subtract 5 from 25?", "once"),
-        ("All mammals breathe air. Whales are mammals. Do whales breathe air?", "yes"),
-        ("How many months have 28 days?", "all"),
-        ("If A > B and B > C, who is smallest?", "c"),
-        ("A bat and ball cost $1.10. The bat costs $1 more than the ball. How much is the ball?", "5"),
-        ("5 machines take 5 min to make 5 widgets. How long for 100 machines to make 100 widgets?", "5"),
-    ]
-    rc = 0
-    for q, expected in reason_tests:
-        a = ask(q, use_cot=True)
-        correct = smart_verify({"answer": expected}, a) > 0
-        rc += correct
-        status = "OK" if correct else "FAIL"
-        answer_display = a.split("</think>")[-1].strip() if "</think>" in a else a
-        print(f"    {status}: {q[:55]}... -> {answer_display[:80]}")
-    results["reasoning"] = f"{rc}/10"
-    print(f"  Reasoning: {rc}/10")
-
-    # Logic
-    logic_tests = [
-        ("If it rains, the ground is wet. The ground is wet. Did it rain?", "no"),
-        ("None of the winners cheated. Tom didn't cheat. Did Tom win?", "no"),
-        ("If A then B. If B then C. A is true. Is C true?", "yes"),
-    ]
-    lc = 0
-    for q, expected in logic_tests:
-        a = ask(q, use_cot=True)
-        correct = smart_verify({"answer": expected}, a) > 0
-        lc += correct
-        status = "OK" if correct else "FAIL"
-        answer_display = a.split("</think>")[-1].strip() if "</think>" in a else a
-        print(f"    {status}: {q[:55]}... -> {answer_display[:80]}")
-    results["logic"] = f"{lc}/3"
-    print(f"  Logic: {lc}/3")
-
-    # Perplexity
-    test_text = "The quick brown fox jumps over the lazy dog. Machine learning models can process natural language."
-    enc = tokenizer(test_text, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        loss = model(**enc, labels=enc.input_ids).loss.item()
-    results["perplexity"] = f"{math.exp(loss):.2f}"
-    print(f"  Perplexity: {results['perplexity']}")
-
-    return results
+    return str(save_dir)
 
 
 # ============================================================
-# MAIN
+# MAIN — Full self-improvement cycle
 # ============================================================
 
 def run_cycle(cfg: SelfImproveConfig = None, cycle_num: int = 1):
-    """Run one complete self-improvement cycle."""
+    """
+    Full self-improvement cycle:
+    1. Analyze own weights (standalone health)
+    2. Download + analyze reference (merge damage)
+    3. Check layer coherence
+    4. Combine all signals → rank layers by improvement potential
+    5. LoRA improve on top layers
+    6. Re-analyze to confirm improvement
+    """
     if cfg is None:
         cfg = SelfImproveConfig()
 
     start = time.time()
     print("=" * 60)
-    print(f"TD SELF-IMPROVEMENT — CYCLE {cycle_num}")
+    print(f"TD SELF-IMPROVEMENT v4 — FULL WEIGHT INTELLIGENCE")
+    print(f"Cycle: {cycle_num}")
     print(f"Model: {cfg.model_path}")
+    print(f"Approach: Analyze EVERYTHING, improve what has most potential")
     print(f"Started: {time.strftime('%H:%M:%S')}")
     print("=" * 60)
 
-    # Check for resume
-    progress_path = Path(cfg.progress_file)
-    resume_from = 0
-    if progress_path.exists():
-        with open(progress_path) as f:
-            saved = json.load(f)
-            resume_from = saved.get("problems_completed", 0)
-
-    # Load model
-    from transformers import AutoModelForImageTextToText, AutoTokenizer
-    print("\nLoading model...")
-    model = AutoModelForImageTextToText.from_pretrained(
-        cfg.model_path, dtype=torch.bfloat16,
-        device_map="auto", trust_remote_code=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained(cfg.model_path, trust_remote_code=True)
-    eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
-
-    # Pre-benchmark
-    print("\n--- PRE-TRAINING BENCHMARK ---")
-    pre_results = benchmark(model, tokenizer, eos_id)
-
-    # Step 1: Diagnose
-    diagnosis_path = Path(cfg.output_dir) / "diagnosis.json"
-    diagnosis_path.parent.mkdir(parents=True, exist_ok=True)
-    weaknesses = diagnose_weaknesses(model, tokenizer, eos_id, save_path=str(diagnosis_path))
-
-    # Step 2: Generate hard problems
-    print("\n=== STEP 2: GENERATING HARD PROBLEMS ===")
-    random.seed(42 + cycle_num)  # Different problems each cycle
-    problems = generate_hard_problems(cfg)
-
-    # Step 3: Generate & score
-    if resume_from < len(problems):
-        winning_pairs, losing_pairs = generate_and_score(
-            model, tokenizer, problems, cfg, eos_id, resume_from=resume_from
+    # ---- STEP 1: Analyze our model ----
+    print(f"\n{'='*40}")
+    print("STEP 1/6: ANALYZE OUR MODEL")
+    print(f"{'='*40}")
+    model_analysis = analyze_model_weights(cfg.model_path, "TD merged+healed")
+
+    # ---- STEP 2: Standalone health scoring ----
+    print(f"\n{'='*40}")
+    print("STEP 2/6: STANDALONE HEALTH SCORING")
+    print(f"{'='*40}")
+    health_scores = compute_health_scores(model_analysis)
+
+    # Print health summary
+    unhealthy_count = sum(1 for h in health_scores.values() if h["health_score"] > 0.2)
+    total = len(health_scores)
+    print(f"\n  Health summary: {unhealthy_count}/{total} tensors have notable issues")
+
+    # ---- STEP 3: Reference comparison (merge damage) ----
+    print(f"\n{'='*40}")
+    print("STEP 3/6: REFERENCE COMPARISON")
+    print(f"{'='*40}")
+
+    ref_path = Path(cfg.output_dir) / "reference_weights"
+    if not ref_path.exists() or not list(ref_path.glob("*.safetensors")):
+        print(f"  Downloading reference: {cfg.reference_model}...")
+        from huggingface_hub import snapshot_download
+        snapshot_download(
+            cfg.reference_model,
+            local_dir=str(ref_path),
+            allow_patterns=["*.safetensors", "*.json"],
+            ignore_patterns=["*.bin", "*.pt", "*.onnx"],
         )
-    else:
-        with open(progress_path) as f:
-            saved = json.load(f)
-            winning_pairs = saved["winning_pairs"]
-            losing_pairs = saved.get("losing_pairs", [])
 
-    # Step 4: Train
-    model = train_on_winners(model, tokenizer, winning_pairs, cfg)
+    ref_analysis = analyze_model_weights(str(ref_path), "Qwen3-VL-8B (reference)")
+    damage_scores = compute_damage_scores(model_analysis, ref_analysis)
 
-    # Save improved model
-    save_dir = Path(cfg.output_dir) / f"cycle_{cycle_num}"
-    save_dir.mkdir(parents=True, exist_ok=True)
-    print(f"\nSaving improved model to {save_dir}...")
-    model.save_pretrained(str(save_dir), safe_serialization=True)
-    tokenizer.save_pretrained(str(save_dir))
-    sz = sum(f.stat().st_size for f in save_dir.glob("*.safetensors")) / 1e9
-    print(f"SAVED: {save_dir} ({sz:.1f} GB)")
+    # Free reference model from disk
+    print(f"  Cleaning up reference model to save disk...")
+    shutil.rmtree(str(ref_path), ignore_errors=True)
+    gc.collect()
+
+    # ---- STEP 4: Layer coherence ----
+    print(f"\n{'='*40}")
+    print("STEP 4/6: LAYER COHERENCE")
+    print(f"{'='*40}")
+    coherence_scores = compute_coherence_scores(model_analysis)
+
+    # ---- STEP 5: Rank and improve ----
+    print(f"\n{'='*40}")
+    print("STEP 5/6: RANK LAYERS + LORA IMPROVEMENT")
+    print(f"{'='*40}")
+    target_layers, layer_details = rank_layers_for_improvement(
+        health_scores, damage_scores, coherence_scores, cfg
+    )
 
-    # Step 5: Post-benchmark
-    print("\n--- POST-TRAINING BENCHMARK ---")
-    post_results = benchmark(model, tokenizer, eos_id)
+    # Save pre-improvement report
+    report_dir = Path(cfg.output_dir)
+    report_dir.mkdir(parents=True, exist_ok=True)
+    pre_report = {
+        "cycle": cycle_num,
+        "model_path": cfg.model_path,
+        "target_layers": target_layers,
+        "layer_details": {str(k): v for k, v in layer_details.items()},
+        "config": {
+            "w_health": cfg.w_health,
+            "w_damage": cfg.w_damage,
+            "w_coherence": cfg.w_coherence,
+            "top_n_layers": cfg.top_n_layers,
+            "lora_r": cfg.lora_r,
+        },
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    with open(report_dir / "pre_improvement_report.json", "w") as f:
+        json.dump(pre_report, f, indent=2)
+
+    # Run LoRA improvement
+    improved_path = improve_with_lora(cfg.model_path, target_layers, cfg)
+
+    # ---- STEP 6: Verify improvement ----
+    print(f"\n{'='*40}")
+    print("STEP 6/6: VERIFY IMPROVEMENT")
+    print(f"{'='*40}")
+
+    post_analysis = analyze_model_weights(improved_path, "TD improved")
+    post_health = compute_health_scores(post_analysis)
+
+    # Compare health before and after
+    pre_avg_health = sum(h["health_score"] for h in health_scores.values()) / max(len(health_scores), 1)
+    post_avg_health = sum(h["health_score"] for h in post_health.values()) / max(len(post_health), 1)
+
+    # Per-layer comparison for targeted layers
+    print(f"\n  Per-layer improvement on targeted layers:")
+    for layer_num in target_layers:
+        pre_score = layer_details.get(layer_num, {}).get("combined", 0)
+        # Find matching post health scores
+        post_layer_scores = []
+        for name, h in post_health.items():
+            parts = name.split(".")
+            for j, part in enumerate(parts):
+                if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
+                    if int(parts[j+1]) == layer_num:
+                        post_layer_scores.append(h["health_score"])
+                    break
+        post_avg = sum(post_layer_scores) / max(len(post_layer_scores), 1) if post_layer_scores else 0
+        pre_h = layer_details.get(layer_num, {}).get("health", 0)
+        delta = pre_h - post_avg
+        symbol = "+" if delta > 0 else "-" if delta < 0 else "="
+        print(f"    Layer {layer_num:>2}: health {pre_h:.4f} → {post_avg:.4f} ({symbol}{abs(delta):.4f})")
 
-    # Compare & save report
     elapsed = (time.time() - start) / 60
-    print(f"\n{'=' * 60}")
+    print(f"\n{'='*60}")
     print(f"CYCLE {cycle_num} COMPLETE — {elapsed:.1f} min")
-    print(f"PRE:  {pre_results}")
-    print(f"POST: {post_results}")
-    print(f"Model: {save_dir}")
-    print(f"{'=' * 60}")
-
-    report = {
-        "cycle": cycle_num,
-        "pre_benchmark": pre_results,
-        "post_benchmark": post_results,
-        "problems_total": len(problems),
-        "winning_pairs": len(winning_pairs),
-        "losing_pairs": len(losing_pairs),
+    print(f"{'='*60}")
+    print(f"  Target layers: {target_layers}")
+    print(f"  Average health score BEFORE: {pre_avg_health:.6f}")
+    print(f"  Average health score AFTER:  {post_avg_health:.6f}")
+    improvement = (pre_avg_health - post_avg_health) / max(pre_avg_health, 1e-8) * 100
+    print(f"  Health improvement: {improvement:.1f}%")
+    print(f"  Improved model: {improved_path}")
+    print(f"{'='*60}")
+
+    # Final report
+    final_report = {
+        **pre_report,
+        "pre_avg_health": pre_avg_health,
+        "post_avg_health": post_avg_health,
+        "health_improvement_pct": improvement,
         "duration_min": elapsed,
-        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "improved_path": improved_path,
     }
-    with open(save_dir / "cycle_report.json", "w") as f:
-        json.dump(report, f, indent=2)
+    with open(Path(improved_path) / "cycle_report.json", "w") as f:
+        json.dump(final_report, f, indent=2)
 
-    return str(save_dir)
+    return improved_path
 
 
 if __name__ == "__main__":