Function train_and_evaluate

scripts/train_predictor.py:400–676

Train expert predictor and evaluate prediction accuracy. Accepts pre-loaded numpy arrays (from load_multiple_routing_files). Supports loading a checkpoint for finetuning via load_model_path. The loss combines standard BCE with a hinge penalty on false positives above `threshold` (d

(layers, hiddens, experts, K, n_tokens,
                       hidden_size=256, epochs=20, lr=1e-3, K_pred=4,
                       model_type='mlp',
                       hidden_dim=HIDDEN_DIM,
                       num_experts=NUM_EXPERTS,
                       num_layers=NUM_LAYERS,
                       load_model_path=None,
                       save_model_path=None,
                       threshold=0.75,
                       fp_penalty_weight=5.0,
                       skip_layers=None,
                       dropout=0.1)

def train_and_evaluate(layers, hiddens, experts, K, n_tokens,
                       hidden_size=256, epochs=20, lr=1e-3, K_pred=4,
                       model_type='mlp',
                       hidden_dim=HIDDEN_DIM,
                       num_experts=NUM_EXPERTS,
                       num_layers=NUM_LAYERS,
                       load_model_path=None,
                       save_model_path=None,
                       threshold=0.75,
                       fp_penalty_weight=5.0,
                       skip_layers=None,
                       dropout=0.1):
    """Train expert predictor and evaluate prediction accuracy.

    Accepts pre-loaded numpy arrays (from load_multiple_routing_files).
    Supports loading a checkpoint for finetuning via load_model_path.

    The loss combines standard BCE with a hinge penalty on false positives above
    `threshold` (default 0.75). Any expert the model predicts with probability
    >= threshold must be in the actual top-4; `fp_penalty_weight` (beta) controls
    how hard the model is pushed to suppress confident-but-wrong predictions.
    """
    try:
        import torch
        from torch.utils.data import DataLoader, TensorDataset
    except ImportError:
        print("ERROR: pip install torch")
        sys.exit(1)

    cross_layer = model_type.endswith('-cross')

    print(f"  {len(layers)} samples, K={K}, layers 0-{layers.max()}")
    print(f"  Hidden state shape: {hiddens.shape}")
    print(f"  Hidden RMS: {np.sqrt(np.mean(hiddens**2)):.4f}")
    print(f"  Model type: {model_type}")

    # Temporal locality baseline
    print("\n=== Temporal Locality Baseline ===")
    prev_experts = {}
    temporal_hits = 0
    temporal_total = 0
    for i in range(len(layers)):
        li = int(layers[i])
        ei = set(int(x) for x in experts[i] if x >= 0)  # drop -1 padding
        if li in prev_experts:
            temporal_hits += len(ei & prev_experts[li])
            temporal_total += len(ei)  # real experts only, not padded K
        prev_experts[li] = ei
    if temporal_total > 0:
        print(f"  Temporal hit rate: {temporal_hits}/{temporal_total} = "
              f"{temporal_hits/temporal_total*100:.1f}%")

    # Cosine similarity analysis (Fate §4.2)
    print("\n=== Adjacent-Layer Cosine Similarity (Fate §4.2) ===")
    layer_sims = analyze_cosine_similarity(layers, hiddens, num_layers)
    avg_sim = layer_sims.mean()
    print(f"  Mean cos(h[l], h[l+1]): {avg_sim:.4f}  Min: {layer_sims.min():.4f}")
    print("  Per-layer: " +