v7

improve_gainlora/_add_c5_flags.py

@@ -1,54 +0,0 @@
-"""Add --run_single False and --n_batches_c5 100 to task 2+ in gen_scripts."""
-import re, os
-
-BASE = '/Users/nnminh322/Desktop/personal/Continual/improve_gainlora'
-
-T5_SCRIPTS = [
-    os.path.join(BASE, 'gen_script_superni_order1_t5_specroute.sh'),
-    os.path.join(BASE, 'gen_script_superni_order2_t5_specroute.sh'),
-    os.path.join(BASE, 'gen_script_long_order3_t5_specroute.sh'),
-    os.path.join(BASE, 'gen_script_long_order4_t5_specroute.sh'),
-]
-
-for name in T5_SCRIPTS:
-    if not os.path.exists(name):
-        print(f'SKIP (not found): {name}')
-        continue
-    
-    with open(name) as f:
-        lines = f.readlines()
-    
-    # Process line by line to find and modify task 2+ blocks
-    # Task 2+ blocks have "--per_device_train_batch_size $BSZ" but NOT "--run_single True"
-    # We need to add --run_single False and --n_batches_c5 100 before --data_replay_freq
-    
-    result = []
-    in_task_ge_2 = False
-    for i, line in enumerate(lines):
-        # Detect task 1: has "--run_single True"
-        if '--run_single True' in line:
-            in_task_ge_2 = False
-            result.append(line)
-        # Detect task 2+: previous task block ended, and this block has "--output_dir" with task number > 1
-        elif 'outputs/' in line and re.search(r'/(\d+)-', line):
-            task_num = int(re.search(r'/(\d+)-', line).group(1))
-            in_task_ge_2 = task_num > 1
-            result.append(line)
-        # Add flags before --data_replay_freq for task 2+ 
-        elif in_task_ge_2 and line.strip() == '--data_replay_freq -1 \\':
-            # Check if we need to add the flags
-            # Look back to see if --run_single or --n_batches_c5 already present in recent lines
-            recent = ''.join(lines[max(0, i-20):i])
-            if '--run_single' not in recent and '--n_batches_c5' not in recent:
-                result.append('   --run_single False \\\n')
-                result.append('   --n_batches_c5 100 \\\n')
-            result.append(line)
-        else:
-            result.append(line)
-    
-    with open(name, 'w') as f:
-        f.writelines(result)
-    
-    print(f'{os.path.basename(name)}: added C5 flags')
-
-print('Done.')

improve_gainlora/gen_script_long_order3_t5_specroute.sh

@@ -42,7 +42,7 @@ elif [ "$IS_T4" -eq 1 ]; then
     GPU_IDS="${1:-0}"
     FP16_FLAG="--gradient_checkpointing"
     echo "[GPU] Strategy: 1x T4 + fp32 + gradient_checkpointing"
-elif [ "$GPU_MEM" -gt 16000 ]; then
+elif [ "$GPU_MEM" -le 17000 ]; then
     GPU_MODE="p100"
     GPU_IDS="${1:-0}"
     FP16_FLAG="--gradient_checkpointing"
@@ -63,7 +63,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=8; GA=4; EVAL_BSZ=128
 fi
@@ -118,7 +118,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -175,7 +175,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -232,7 +232,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -289,7 +289,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -346,7 +346,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -403,7 +403,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -460,7 +460,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -517,7 +517,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -574,7 +574,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -631,7 +631,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -688,7 +688,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -745,7 +745,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -802,7 +802,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -859,7 +859,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi

improve_gainlora/gen_script_long_order4_t5_specroute.sh

@@ -42,7 +42,7 @@ elif [ "$IS_T4" -eq 1 ]; then
     GPU_IDS="${1:-0}"
     FP16_FLAG="--gradient_checkpointing"
     echo "[GPU] Strategy: 1x T4 + fp32 + gradient_checkpointing"
-elif [ "$GPU_MEM" -gt 16000 ]; then
+elif [ "$GPU_MEM" -le 17000 ]; then
     GPU_MODE="p100"
     GPU_IDS="${1:-0}"
     FP16_FLAG="--gradient_checkpointing"
@@ -63,7 +63,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=8; GA=4; EVAL_BSZ=128
 fi
@@ -118,7 +118,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -175,7 +175,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -232,7 +232,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -289,7 +289,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -346,7 +346,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -403,7 +403,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -460,7 +460,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -517,7 +517,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -574,7 +574,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -631,7 +631,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -688,7 +688,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -745,7 +745,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -802,7 +802,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi
@@ -859,7 +859,7 @@ if [ "$GPU_MODE" = "t4_2gpu" ]; then
 elif [ "$GPU_MODE" = "t4_1gpu" ]; then
     BSZ=4; GA=8; EVAL_BSZ=16
 elif [ "$GPU_MODE" = "p100" ]; then
-    BSZ=8; GA=4; EVAL_BSZ=4
+    BSZ=16; GA=2; EVAL_BSZ=32
 else
     BSZ=16; GA=2; EVAL_BSZ=128
 fi

improve_gainlora/src/cl_trainer_specroute.py

@@ -12,7 +12,6 @@ Key differences from GainLoRA_InfLoRA_Trainer:
 
 import gc
 import torch
-import math as _math
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 from transformers import GenerationConfig
@@ -141,30 +140,9 @@ class SpecRoute_Trainer(Seq2SeqTrainer):
                 self._precond_matrices[id(lora.lora_B)] = inv_sqrt.to(lora.lora_B.data.device)
         print(f"[C4] Precomputed {len(self._precond_matrices)} preconditioner matrices")
 
-    def _compute_spectral_entropy_loss(self):
-        """Compute spectral entropy regularization via efficient QR trick.
-        For each LoRA layer: QR(B) and QR(A^T) -> SVD of r×r matrix -> entropy."""
-        ent_loss = torch.tensor(0.0, device=self.args.device)
-        count = 0
-        for module in self.model.modules():
-            if not (hasattr(module, 'lora_q') and hasattr(module, 'lora_v')):
-                continue
-            for lora in [module.lora_q, module.lora_v]:
-                B = lora.lora_B.float()    # [d_out, r]
-                A = lora.lora_A.float()    # [r, d_in]
-                if B.norm() < 1e-8:
-                    continue
-                _, R_B = torch.linalg.qr(B)         # B: [d_out, r] -> R_B: [r, r]
-                _, R_A = torch.linalg.qr(A.T)       # A.T: [d_in, r] -> R_A: [r, r]
-                sigma_hat = torch.linalg.svdvals(R_B @ R_A.T)  # [r, r] @ [r, r] -> [r, r] -> [r]
-                sigma_hat = sigma_hat / (sigma_hat.sum() + 1e-12)
-                ent = -(sigma_hat * torch.log(sigma_hat + 1e-12)).sum()
-                max_ent = _math.log(sigma_hat.size(0))
-                ent_loss = ent_loss + (max_ent - ent)
-                count += 1
-        if count > 0:
-            ent_loss = ent_loss / count
-        return ent_loss
+    # NOTE: _compute_spectral_entropy_loss() removed in V7.
+    # Entropy regularization conflicts with C5 data-informed init philosophy.
+    # lambda_entropy is hardcoded to 0.0 and this method was dead code.
 
     def _apply_preconditioning(self):
         """Apply (AA^T + eps*I)^{-1/2} preconditioner to lora_B gradients after backward."""
@@ -412,7 +390,7 @@ class SpecRoute_Trainer(Seq2SeqTrainer):
                 else:
                     labels = None
                 outputs = self.model(**inputs)
-                if step > 200:  # 200 batches sufficient for stable SVD (reduced from 1000 for speed)
+                if step > 100:  # 100 batches sufficient for stable SVD
                     break
         print('end get representation')
 

improve_gainlora/src/run_t5.py

@@ -638,21 +638,15 @@ def main():
         previous_lora_list_sig = model_args.previous_lora_path.split(',')
         previous_lora_list_sig.reverse()
         spectral_sigs = []
-        task_protos = []
         for path in previous_lora_list_sig:
             sig_path = os.path.join(path, "spectral_signatures.pt")
             if os.path.exists(sig_path):
                 sig = torch.load(sig_path, map_location=device)
                 spectral_sigs.append(sig)
-            proto_path = os.path.join(path, "task_prototype.pt")
-            if os.path.exists(proto_path):
-                task_protos.append(torch.load(proto_path, map_location='cpu'))
         model.encoder.spectral_signatures = spectral_sigs
-        model.encoder.task_prototypes = task_protos
         del previous_lora_list_sig
         gc.collect()
-        print(f"----------Loaded {len(spectral_sigs)} spectral signatures, "
-              f"{len(task_protos)} task prototypes----------")
+        print(f"----------Loaded {len(spectral_sigs)} spectral signatures----------")
 
     for name, param in model.named_parameters():
         if  training_args.model_name in ['gainlora_olora', 'olora']:
@@ -1034,12 +1028,6 @@ def main():
             from t5_specroute import compute_spectral_signatures
             signatures = compute_spectral_signatures(trainer.model, config)
             torch.save(signatures, os.path.join(save_path, 'spectral_signatures.pt'))
-            # Finalize and save task prototype (V5)
-            trainer.model.encoder.finalize_prototype()
-            if trainer.model.encoder._current_task_prototype is not None:
-                torch.save(trainer.model.encoder._current_task_prototype,
-                           os.path.join(save_path, 'task_prototype.pt'))
-                print("----------Saved task prototype----------")
             print("----------Saved spectral signatures----------")
         # Only save tokenizer for non-specroute (specroute never reloads it)
         if training_args.model_name != 'specroute':

improve_gainlora/src/t5_specroute.py

@@ -178,28 +178,12 @@ class T5Stack(T5PreTrainedModel):
 
         self.prompt_config = prompt_config
 
-        if not self.is_decoder:
-            # Prototype-based inference routing (V5)
-            # Task centroids in input embedding space — immune to GPM constraints.
-            # Must be initialized for ALL tasks (including first/run_single) since
-            # _update_prototype() is called during every training forward pass.
-            self.task_prototypes = []          # List[Tensor(d,)] — normalized μ_k, old tasks
-            self._current_prototype_sum = None  # Running sum for current task (CPU float)
-            self._current_prototype_count = 0
-            self._current_task_prototype = None # Finalized normalized prototype
-
         if not self.is_decoder and not prompt_config["run_single"]:
             # ===== Spectral routing: NO learned parameters for routing =====
             # Spectral signatures loaded from previous tasks' saved weights
             self.spectral_signatures = []  # List[dict] — one dict per old task
             self.routing_temperature = prompt_config.get('attn_temperature', 1.0)
 
-            # Prototype routing temperature (V5): cosine similarities are in a much
-            # narrower range (~0.85-0.95) than spectral fits (~0-0.5). A separate,
-            # lower temperature ensures discriminative softmax for prototype routing.
-            # τ_proto = 0.01 → for gap=0.05, ratio = e^5 ≈ 148 (semi-hard routing).
-            self._prototype_temperature = 0.01
-
             # Adaptive training bias: β = T·ln(α·n_old/(1−α))
             # Ensures current task gets consistent routing weight ~α regardless
             # of total number of tasks (fixes softmax dilution with constant bias).
@@ -251,8 +235,6 @@ class T5Stack(T5PreTrainedModel):
 
         fits = []
 
-        _use_proto = False
-
         if self.training:
             # === TRAINING MODE: A-row fit + adaptive bias for current task ===
             # A rows give stable non-zero signal even when B=0 (cold-start).
@@ -276,83 +258,40 @@ class T5Stack(T5PreTrainedModel):
                 current_fit = current_fit + beta
             fits.append(current_fit)
         else:
-            # === INFERENCE MODE ===
-            # Try prototype routing (V5): cosine similarity in embedding space.
-            # Immune to GPM constraints — handles same-domain tasks correctly.
-            _n_expected = 1 + len(self.spectral_signatures)
-            _protos = []
-            if self._current_task_prototype is not None:
-                _protos.append(self._current_task_prototype)
-            elif self._current_prototype_sum is not None and self._current_prototype_count > 0:
-                # During training eval: use running mean as approximate prototype
-                _tmp = self._current_prototype_sum / self._current_prototype_count
-                _protos.append(_tmp / (_tmp.norm() + 1e-8))
-            _protos.extend(self.task_prototypes)
-
-            if len(_protos) == _n_expected:
-                # Prototype routing for ALL tasks
-                _use_proto = True
-                h_flat = h.squeeze(1)  # (B, d_model)
-                h_norm = h_flat.norm(dim=-1, keepdim=True).clamp(min=1e-8)  # (B, 1)
-
-                # One-time diagnostic: pairwise cosine similarity between prototypes
-                if self.is_inference and not getattr(self, '_proto_diag_done', False):
-                    self._proto_diag_done = True
-                    _pstack = torch.stack([p.to(h_flat.device, dtype=h_flat.dtype) for p in _protos])
-                    _pair = torch.matmul(_pstack, _pstack.T)  # already unit-norm → cosine sim
-                    _mask = ~torch.eye(_pair.size(0), dtype=torch.bool, device=_pair.device)
-                    _offdiag = _pair[_mask]
-                    print(f"[SpecRoute] Prototype cosine matrix (n={_pair.size(0)}): "
-                          f"off-diag min={_offdiag.min():.4f}, max={_offdiag.max():.4f}, "
-                          f"mean={_offdiag.mean():.4f}, std={_offdiag.std():.4f}")
-                    if _offdiag.max() > 0.99:
-                        print(f"[SpecRoute] WARNING: some prototypes are near-identical "
-                              f"(max cos={_offdiag.max():.4f}). Consider mean-centering.")
-
-                for proto in _protos:
-                    p = proto.to(h_flat.device, dtype=h_flat.dtype)  # (d,) normalized
-                    cos_sim = torch.matmul(h_flat, p.unsqueeze(-1)) / h_norm  # (B, 1)
-                    fits.append(cos_sim)
-            else:
-                # Fallback: SVD-based spectral routing (V3)
-                if self._current_task_svd is not None:
-                    cur_fits = []
-                    for key, sig_data in self._current_task_svd.items():
-                        if not key.startswith('enc.'):
-                            continue
-                        V = sig_data['V'].to(h.device, dtype=h.dtype)
-                        sigma = sig_data['sigma'].to(h.device, dtype=h.dtype)
-                        proj = torch.matmul(h, V.T)
-                        sigma_sq = sigma ** 2
-                        sigma_sq_sum = sigma_sq.sum() + 1e-8
-                        weighted_proj = (proj ** 2 * sigma_sq.unsqueeze(0).unsqueeze(0)).sum(dim=-1)
-                        fit = weighted_proj / (sigma_sq_sum * h_norm_sq)
-                        cur_fits.append(fit)
-                    if cur_fits:
-                        current_fit = torch.stack(cur_fits).mean(dim=0)
-                    else:
-                        current_fit = torch.zeros(h.shape[0], 1, device=h.device, dtype=h.dtype)
+            # === INFERENCE MODE: SVD-based spectral routing (symmetric) ===
+            if self._current_task_svd is not None:
+                cur_fits = []
+                for key, sig_data in self._current_task_svd.items():
+                    if not key.startswith('enc.'):
+                        continue
+                    V = sig_data['V'].to(h.device, dtype=h.dtype)
+                    sigma = sig_data['sigma'].to(h.device, dtype=h.dtype)
+                    proj = torch.matmul(h, V.T)
+                    sigma_sq = sigma ** 2
+                    sigma_sq_sum = sigma_sq.sum() + 1e-8
+                    weighted_proj = (proj ** 2 * sigma_sq.unsqueeze(0).unsqueeze(0)).sum(dim=-1)
+                    fit = weighted_proj / (sigma_sq_sum * h_norm_sq)
+                    cur_fits.append(fit)
+                if cur_fits:
+                    current_fit = torch.stack(cur_fits).mean(dim=0)
                 else:
-                    current_fits_layers = []
-                    for block in self.block:
-                        attn = block.layer[0].SelfAttention
-                        for lora in [attn.lora_q, attn.lora_v]:
-                            A = lora.lora_A.data.float()
-                            r = lora.r
-                            A_h = A.to(h.device, dtype=h.dtype)
-                            proj = torch.matmul(h, A_h.T)
-                            fit = (proj ** 2).sum(dim=-1) / (r * h_norm_sq)
-                            current_fits_layers.append(fit)
-                    current_fit = torch.stack(current_fits_layers).mean(dim=0)
-                fits.append(current_fit)
-
-        # Early return for prototype routing (uses separate temperature)
-        if _use_proto:
-            fit_scores = torch.cat(fits, dim=1)  # (B, n_tasks)
-            weights = torch.softmax(fit_scores / self._prototype_temperature, dim=1)
-            return weights.unsqueeze(2)  # (B, n_tasks, 1)
-
-        # Old tasks: spectral signatures (training or fallback inference)
+                    current_fit = torch.zeros(h.shape[0], 1, device=h.device, dtype=h.dtype)
+            else:
+                # Fallback: A-row fit (before prepare_inference_routing() is called)
+                current_fits_layers = []
+                for block in self.block:
+                    attn = block.layer[0].SelfAttention
+                    for lora in [attn.lora_q, attn.lora_v]:
+                        A = lora.lora_A.data.float()
+                        r = lora.r
+                        A_h = A.to(h.device, dtype=h.dtype)
+                        proj = torch.matmul(h, A_h.T)
+                        fit = (proj ** 2).sum(dim=-1) / (r * h_norm_sq)
+                        current_fits_layers.append(fit)
+                current_fit = torch.stack(current_fits_layers).mean(dim=0)
+            fits.append(current_fit)
+
+        # Old tasks: σ²-weighted spectral affinity (Rayleigh quotient)
         for sig_dict in self.spectral_signatures:
             task_fits = []
             for key, sig_data in sig_dict.items():
@@ -398,40 +337,13 @@ class T5Stack(T5PreTrainedModel):
                     B = lora.lora_B.data.float()
                     r = lora.r
                     S, Vt = _thin_svd_low_rank(B, A)
-                    sigs[f'enc.{j}.self.{name}'] = {
+                    sigs[f'enc.{j}.{name}'] = {
                         'V': Vt[:r].cpu(),
                         'sigma': S[:r].cpu()
                     }
         self._current_task_svd = sigs
         print(f"[SpecRoute] Prepared inference routing: {len(sigs)} encoder SVD signatures for current task")
 
-    def _update_prototype(self, h_batch):
-        """Accumulate running mean for current task prototype.
-        Called during training forward passes.
-        Args:
-            h_batch: (B, d_model) — attention-masked mean of input embeddings
-        """
-        with torch.no_grad():
-            batch_sum = h_batch.detach().float().sum(dim=0).cpu()
-            count = h_batch.size(0)
-            if self._current_prototype_sum is None:
-                self._current_prototype_sum = batch_sum
-                self._current_prototype_count = count
-            else:
-                self._current_prototype_sum = self._current_prototype_sum + batch_sum
-                self._current_prototype_count += count
-
-    def finalize_prototype(self):
-        """Finalize current task prototype from accumulated running mean."""
-        if self._current_prototype_sum is not None and self._current_prototype_count > 0:
-            proto = self._current_prototype_sum / self._current_prototype_count
-            proto = proto / (proto.norm() + 1e-8)
-            self._current_task_prototype = proto
-            print(f"[SpecRoute] Finalized task prototype "
-                  f"({self._current_prototype_count} samples, d={proto.shape[0]})")
-        self._current_prototype_sum = None
-        self._current_prototype_count = 0
-
     def parallelize(self, device_map=None):
         self.device_map = (
             get_device_map(len(self.block), range(torch.cuda.device_count()))
@@ -512,19 +424,15 @@ class T5Stack(T5PreTrainedModel):
 
         batch_size, seq_length = input_shape
 
-        # ============ SPECTRAL ROUTING + PROTOTYPE (V5) ============
+        # ============ SPECTRAL ROUTING ============
         self.key_attention_weights = None
         if not self.is_decoder:
-            # Properly masked mean of input embeddings (V5 fix)
+            # Properly masked mean of input embeddings
             _mask_count = attention_mask.sum(dim=1, keepdim=True).clamp(min=1).unsqueeze(-1)  # (B,1,1)
             avg_inputs_embeds = (attention_mask.unsqueeze(-1) * inputs_embeds).sum(dim=1, keepdim=True) / _mask_count
 
-            # Accumulate prototype during training (all tasks including first)
-            if self.training:
-                self._update_prototype(avg_inputs_embeds.squeeze(1))
-
         if not self.is_decoder and not self.prompt_config["run_single"]:
-            if len(self.spectral_signatures) > 0 or len(self.task_prototypes) > 0:
+            if len(self.spectral_signatures) > 0:
                 # Multi-task: compute routing
                 key_attention_weights = self.compute_spectral_routing(avg_inputs_embeds)
                 # Detach: routing weights are shared across all gradient-checkpointed