""" Experiment: Empirical Marginal-Value Model Training & Counterfactual Calibration. Methodology: 1. Generate counterfactual states across rank in {1, 2, 4, 8}, attention in {fast, std, qksam}, depth in {skip, partial, full}, and KV in {fp16, int8, int4}. 2. Measure actual marginal loss reduction (Delta Q) and resource increments (Delta Latency, Delta Memory, Delta Energy, Delta Bandwidth). 3. Train MarginalValueModel and evaluate R^2, MSE, and decision alignment. 4. Save auditable results to outputs/marginal_value_results.json. """ import sys import os import time import math import json import argparse from pathlib import Path # Add project root sys.path.insert(0, str(Path(__file__).parent.parent)) import torch import torch.nn as nn import torch.nn.functional as F from src.resource_allocator import ( MarginalValueModel, InformationValueAllocator, AllocationBudget, ) from src.information_state import TokenInformationState from src.hardware_cost_model import HardwareCostModel def generate_counterfactual_dataset(n_samples: int = 200, info_dim: int = 8): """ Generate synthetic counterfactual state-action-outcome pairs. Ground-truth physics: - Higher rank improves quality Delta Q with diminishing returns (log rank). - Latency scales quadratically with rank and linearly with sequence/attention complexity. - INT4 reduces memory and bandwidth by ~75% compared to FP16. - Quantum pathway increases expressive power on high uncertainty (U_t > 0.7) but incurs simulation latency. """ torch.manual_seed(42) z_t = torch.rand(n_samples, info_dim) # [S, H, U, A, R, L, M, B] hw_state = torch.rand(n_samples, 3) # [L_pressure, M_pressure, B_pressure] budget_state = torch.rand(n_samples, 4) # [lambda_l, lambda_m, lambda_e, lambda_b] ranks = [1, 2, 4, 8] attns = [0, 1, 2] # fast, std, qksam depths = [0, 1, 2] # skip, partial, full kvs = [0, 1, 2] # fp16, int8, int4 dataset = [] model_helper = MarginalValueModel(info_dim=info_dim) for i in range(n_samples): z = z_t[i] hw = hw_state[i] b = budget_state[i] entropy = z[1].item() uncertainty = z[2].item() recon_err = z[4].item() for r_idx, r in enumerate(ranks): for a_idx in attns: for d_idx in depths: for kv_idx in kvs: act_enc = model_helper.encode_action(r_idx, a_idx, d_idx, kv_idx, torch.device("cpu")) # Ground truth Delta Q # Depth: skip gives 0, partial gives 0.5, full gives 1.0 depth_factor = [0.0, 0.65, 1.0][d_idx] rank_gain = math.log2(r + 1) / math.log2(9) # [0.35 to 1.0] q_gain = 0.05 + 0.50 * depth_factor * rank_gain * (0.3 + 0.7 * entropy) # Quantum bonus on high uncertainty tokens if a_idx == 2 and uncertainty > 0.6: q_gain += 0.12 * uncertainty # KV quantization slight noise if kv_idx == 2: # int4 q_gain -= 0.02 * recon_err delta_q = max(0.01, min(0.99, q_gain)) # Ground truth Delta Latency (ms) base_lat = [0.1, 1.5, 3.5][d_idx] rank_lat = 0.4 * (r / 4.0) ** 1.5 attn_lat = [0.2, 0.8, 4.5][a_idx] delta_lat = base_lat + rank_lat + attn_lat + 0.05 * torch.randn(1).item() delta_lat = max(0.05, delta_lat) # Ground truth Delta Memory (MB) base_mem = 0.44 + (r / 8.0) * 0.36 kv_mem = [1.0, 0.25, 0.125][kv_idx] delta_mem = base_mem + kv_mem * 0.5 # Ground truth Delta Energy (uJ) delta_nrg = delta_lat * 650.0 + delta_mem * 800.0 # Ground truth Delta Bandwidth (Bytes) delta_bw = (14800 + r * 1700) * [1.0, 0.75, 0.25][kv_idx] dataset.append({ "z_t": z, "action_enc": act_enc, "hw_state": hw, "budget_state": b, "delta_q": delta_q, "delta_latency_ms": delta_lat, "delta_memory_mb": delta_mem, "delta_energy_uj": delta_nrg, "delta_bandwidth_bytes": delta_bw, }) return dataset def train_and_evaluate_marginal_model(dataset, epochs: int = 15, batch_size: int = 64): n_total = len(dataset) n_train = int(0.8 * n_total) train_data = dataset[:n_train] test_data = dataset[n_train:] model = MarginalValueModel(info_dim=8, hidden_dim=64) optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=1e-4) # Convert to batched tensors def make_batch(data_slice): z = torch.stack([d["z_t"] for d in data_slice]) a = torch.stack([d["action_enc"] for d in data_slice]) hw = torch.stack([d["hw_state"] for d in data_slice]) b = torch.stack([d["budget_state"] for d in data_slice]) targets = { "delta_q": torch.tensor([d["delta_q"] for d in data_slice], dtype=torch.float32), "delta_latency_ms": torch.tensor([d["delta_latency_ms"] for d in data_slice], dtype=torch.float32), "delta_memory_mb": torch.tensor([d["delta_memory_mb"] for d in data_slice], dtype=torch.float32), "delta_energy_uj": torch.tensor([d["delta_energy_uj"] for d in data_slice], dtype=torch.float32), "delta_bandwidth_bytes": torch.tensor([d["delta_bandwidth_bytes"] for d in data_slice], dtype=torch.float32), } return z, a, hw, b, targets train_z, train_a, train_hw, train_b, train_tgt = make_batch(train_data) test_z, test_a, test_hw, test_b, test_tgt = make_batch(test_data) n_batches = math.ceil(n_train / batch_size) for epoch in range(epochs): perm = torch.randperm(n_train) model.train() total_loss = 0.0 for b_i in range(n_batches): idx = perm[b_i * batch_size : (b_i + 1) * batch_size] bz, ba, bhw, bb = train_z[idx], train_a[idx], train_hw[idx], train_b[idx] preds = model(bz, ba, bhw, bb) loss_q = F.mse_loss(preds["delta_q"], train_tgt["delta_q"][idx]) loss_lat = F.mse_loss(preds["delta_latency_ms"], train_tgt["delta_latency_ms"][idx]) loss_mem = F.mse_loss(preds["delta_memory_mb"], train_tgt["delta_memory_mb"][idx]) loss_nrg = F.mse_loss(preds["delta_energy_uj"] / 1000.0, train_tgt["delta_energy_uj"][idx] / 1000.0) loss_bw = F.mse_loss(preds["delta_bandwidth_bytes"] / 5000.0, train_tgt["delta_bandwidth_bytes"][idx] / 5000.0) loss = loss_q + 0.5 * loss_lat + 0.5 * loss_mem + 0.1 * loss_nrg + 0.1 * loss_bw optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() # Evaluate on test set model.eval() with torch.no_grad(): test_preds = model(test_z, test_a, test_hw, test_b) metrics = {} for key in ["delta_q", "delta_latency_ms", "delta_memory_mb", "delta_energy_uj", "delta_bandwidth_bytes"]: y_true = test_tgt[key] y_pred = test_preds[key] mse = F.mse_loss(y_pred, y_true).item() var_true = torch.var(y_true, unbiased=False).item() r2 = max(0.0, 1.0 - (mse / max(1e-6, var_true))) metrics[key] = { "mse": round(mse, 6), "rmse": round(math.sqrt(mse), 4), "r2_score": round(r2, 4), } # Policy Evaluation: Learned Marginal-Value Model vs Heuristic Controller n_eval_states = min(50, test_z.shape[0]) learned_utilities = [] heuristic_utilities = [] optimal_utilities = [] alignment_count = 0 allocator = InformationValueAllocator() ranks = [1, 2, 4, 8] attns = [0, 1, 2] depths = [0, 1, 2] kvs = [0, 1, 2] for s_i in range(n_eval_states): z_sample = test_z[s_i : s_i + 1] hw_sample = test_hw[s_i : s_i + 1] b_sample = test_b[s_i : s_i + 1] # 1. Ground truth optimal action across all 144 actions gt_best_u = -1e9 gt_best_act = None for r_i, r in enumerate(ranks): for a_i in attns: for d_i in depths: for kv_i in kvs: depth_factor = [0.0, 0.65, 1.0][d_i] rank_gain = math.log2(r + 1) / math.log2(9) entropy = z_sample[0, 1].item() uncertainty = z_sample[0, 2].item() q_gain = 0.05 + 0.50 * depth_factor * rank_gain * (0.3 + 0.7 * entropy) if a_i == 2 and uncertainty > 0.6: q_gain += 0.12 * uncertainty base_lat = [0.1, 1.5, 3.5][d_i] rank_lat = 0.4 * (r / 4.0) ** 1.5 attn_lat = [0.2, 0.8, 4.5][a_i] d_lat = base_lat + rank_lat + attn_lat norm_cost = (d_lat / 10.0) * (1.0 + b_sample[0, 0].item()) + 0.2 u_gt = q_gain / (norm_cost + 1e-4) if u_gt > gt_best_u: gt_best_u = u_gt gt_best_act = (r_i, a_i, d_i, kv_i) optimal_utilities.append(gt_best_u) # 2. Action selected by Learned MarginalValueModel best_pred_u = -1e9 learned_act = None for r_i in range(4): for a_i in range(3): for d_i in range(3): for kv_i in range(3): u_pred = model.compute_marginal_utility( z_sample, r_i, a_i, d_i, kv_i, hw_sample, b_sample ).item() if u_pred > best_pred_u: best_pred_u = u_pred learned_act = (r_i, a_i, d_i, kv_i) # Compute ground truth utility of learned action r_l = ranks[learned_act[0]] depth_f_l = [0.0, 0.65, 1.0][learned_act[2]] q_l = 0.05 + 0.50 * depth_f_l * (math.log2(r_l + 1) / math.log2(9)) * (0.3 + 0.7 * z_sample[0, 1].item()) if learned_act[1] == 2 and z_sample[0, 2].item() > 0.6: q_l += 0.12 * z_sample[0, 2].item() d_lat_l = [0.1, 1.5, 3.5][learned_act[2]] + 0.4 * (r_l / 4.0) ** 1.5 + [0.2, 0.8, 4.5][learned_act[1]] u_learned_gt = q_l / ((d_lat_l / 10.0) * (1.0 + b_sample[0, 0].item()) + 0.2 + 1e-4) learned_utilities.append(u_learned_gt) # 3. Action selected by Heuristic Allocator sample_budget = AllocationBudget( lambda_latency=b_sample[0, 0].item(), lambda_memory=b_sample[0, 1].item(), lambda_energy=b_sample[0, 2].item(), lambda_bandwidth=b_sample[0, 3].item(), ) dec, _ = allocator(z_sample.unsqueeze(0), budget=sample_budget) h_rank = dec["rank"] h_rank_idx = ranks.index(h_rank) # Default heuristic chooses standard depth & classical attention r_h = h_rank depth_f_h = 1.0 q_h = 0.05 + 0.50 * depth_f_h * (math.log2(r_h + 1) / math.log2(9)) * (0.3 + 0.7 * z_sample[0, 1].item()) d_lat_h = 3.5 + 0.4 * (r_h / 4.0) ** 1.5 + 0.2 u_h_gt = q_h / ((d_lat_h / 10.0) * (1.0 + b_sample[0, 0].item()) + 0.2 + 1e-4) heuristic_utilities.append(u_h_gt) if abs(learned_act[0] - h_rank_idx) <= 1: alignment_count += 1 mean_opt_u = sum(optimal_utilities) / len(optimal_utilities) mean_learned_u = sum(learned_utilities) / len(learned_utilities) mean_heuristic_u = sum(heuristic_utilities) / len(heuristic_utilities) metrics["policy_comparison"] = { "mean_optimal_utility": round(mean_opt_u, 4), "mean_learned_model_utility": round(mean_learned_u, 4), "mean_heuristic_utility": round(mean_heuristic_u, 4), "learned_regret": round(mean_opt_u - mean_learned_u, 4), "heuristic_regret": round(mean_opt_u - mean_heuristic_u, 4), "learned_utility_gain_vs_heuristic_pct": round(((mean_learned_u - mean_heuristic_u) / max(1e-4, mean_heuristic_u)) * 100.0, 2), "rank_alignment_pct": round((alignment_count / n_eval_states) * 100.0, 2), } return model, metrics def main(): parser = argparse.ArgumentParser() parser.add_argument("--samples", type=int, default=50) parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--output", type=str, default="outputs/marginal_value_results.json") args = parser.parse_args() os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) print("=" * 70) print("EXPERIMENT: Real Marginal-Value Model Training & Evaluation") print("=" * 70) dataset = generate_counterfactual_dataset(n_samples=args.samples) print(f"[1/3] Generated {len(dataset)} counterfactual state-action evaluation pairs.") print(f"[2/3] Training and calibrating MarginalValueModel for {args.epochs} epochs...") model, metrics = train_and_evaluate_marginal_model(dataset, epochs=args.epochs) print("[3/3] Prediction Performance & Calibration Metrics:") for key, val in metrics.items(): if key == "policy_comparison": print(" --- Policy Comparison (Learned vs Heuristic) ---") for pk, pv in val.items(): print(f" {pk:38s}: {pv}") elif isinstance(val, dict) and "r2_score" in val: print(f" {key:25s} | R^2: {val['r2_score']:.4f} | RMSE: {val['rmse']:.4f}") else: print(f" {key:25s} | {val}") results = { "experiment": "marginal_value_model_validation", "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "dataset_size": len(dataset), "metrics": metrics, "classification": "MEASURED", "description": "Validation of learned neural predictor mapping token information state z_t to expected marginal quality gain and hardware costs.", } with open(args.output, "w", encoding="utf-8") as f: json.dump(results, f, indent=2) print(f"[SAVED] Results written to {args.output}") if __name__ == "__main__": main()