Spaces:
Running
Running
File size: 8,418 Bytes
d8d6846 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """
EquiRouter: Decision-Aware Ranking Loss vs MSE training simulation.
This script implements offline MLP coefficient optimization for LLM routing.
It compares:
1. Standard MSE Loss (minimizes quality prediction error).
2. Decision-Aware Ranking Loss (EquiRouter, maximizes utility margin).
It evaluates the cost-quality Pareto curves and prints the resulting AIQ (Area Under the Curve).
"""
import math
import random
from typing import Any
# Simulated dataset of 100 prompts
# Features: [intercept (1), is_code, is_math, is_json, is_long]
# Qualities: [Ollama_q, vLLM_q, Gemini_q, OpenAI_q]
# Costs: Ollama=0.001, vLLM=0.005, Gemini=0.015, OpenAI=0.030
COSTS = [0.001, 0.005, 0.015, 0.030]
MODELS = ["ollama", "vllm", "gemini", "openai"]
def generate_mock_dataset(num_samples: int = 150) -> list[dict[str, Any]]:
random.seed(42)
dataset = []
for _ in range(num_samples):
is_code = 1.0 if random.random() < 0.3 else 0.0
is_math = 1.0 if random.random() < 0.25 else 0.0
is_json = 1.0 if random.random() < 0.2 else 0.0
is_long = 1.0 if random.random() < 0.4 else 0.0
features = [1.0, is_code, is_math, is_json, is_long]
# Qualities
if is_code == 1.0:
qualities = [0.2, 0.85, 0.75, 0.98]
elif is_math == 1.0:
qualities = [0.1, 0.2, 0.90, 0.99]
elif is_json == 1.0:
qualities = [0.3, 0.4, 0.88, 0.96]
else:
qualities = [0.6, 0.7, 0.85, 0.95]
# Add random noise
qualities = [max(0.0, min(1.0, q + random.uniform(-0.08, 0.08))) for q in qualities]
dataset.append({
"features": features,
"qualities": qualities
})
return dataset
def sigmoid(x: float) -> float:
try:
return 1.0 / (1.0 + math.exp(-x))
except OverflowError:
return 0.0 if x < 0 else 1.0
def predict_quality(features: list[float], weights: list[float]) -> float:
score = sum(f * w for f, w in zip(features, weights))
return sigmoid(score)
def evaluate_mse_loss(dataset: list[dict[str, Any]], weights_matrix: list[list[float]]) -> float:
total_loss = 0.0
for sample in dataset:
for idx in range(len(MODELS)):
pred = predict_quality(sample["features"], weights_matrix[idx])
true = sample["qualities"][idx]
total_loss += (pred - true) ** 2
return total_loss / (len(dataset) * len(MODELS))
def evaluate_ranking_loss(dataset: list[dict[str, Any]], weights_matrix: list[list[float]], lambda_val: float) -> float:
"""
EquiRouter Decision-Aware Pairwise Ranking Loss.
Penalizes when the estimated utility order of the best candidate vs others is violated.
"""
total_loss = 0.0
for sample in dataset:
# Calculate true utilities
true_utilities = [q - lambda_val * c for q, c in zip(sample["qualities"], COSTS)]
best_true_idx = true_utilities.index(max(true_utilities))
# Calculate predicted utilities
pred_qualities = [predict_quality(sample["features"], weights_matrix[idx]) for idx in range(len(MODELS))]
pred_utilities = [pq - lambda_val * c for pq, c in zip(pred_qualities, COSTS)]
# Compare best true model against others
for idx in range(len(MODELS)):
if idx == best_true_idx:
continue
# Margin of utility
margin = pred_utilities[best_true_idx] - pred_utilities[idx]
# Softmargin pairwise loss
total_loss += math.log(1.0 + math.exp(-margin))
return total_loss / len(dataset)
def optimize_weights(dataset: list[dict[str, Any]], loss_type: str, lambda_val: float = 1.0, epochs: int = 100, lr: float = 0.1) -> list[list[float]]:
# Initialize weights
num_features = 5
weights_matrix = [[0.0] * num_features for _ in range(len(MODELS))]
# Stochastic Gradient Descent with finite difference numerical gradient
for epoch in range(epochs):
for sample in dataset:
feats = sample["features"]
for m_idx in range(len(MODELS)):
for f_idx in range(num_features):
# Original loss
if loss_type == "mse":
loss_fn = lambda wm: evaluate_mse_loss([sample], wm)
else:
loss_fn = lambda wm: evaluate_ranking_loss([sample], wm, lambda_val)
l_orig = loss_fn(weights_matrix)
# Finite difference
h = 1e-4
weights_matrix[m_idx][f_idx] += h
l_new = loss_fn(weights_matrix)
weights_matrix[m_idx][f_idx] -= h
grad = (l_new - l_orig) / h
# Update weight
weights_matrix[m_idx][f_idx] -= lr * grad
return weights_matrix
def simulate_routing(dataset: list[dict[str, Any]], weights_matrix: list[list[float]], lambda_val: float) -> tuple[float, float]:
"""
Simulates routing on dataset. Returns average quality and average cost.
"""
total_q = 0.0
total_c = 0.0
for sample in dataset:
pred_qualities = [predict_quality(sample["features"], weights_matrix[idx]) for idx in range(len(MODELS))]
pred_utilities = [pq - lambda_val * c for pq, c in zip(pred_qualities, COSTS)]
selected_idx = pred_utilities.index(max(pred_utilities))
total_q += sample["qualities"][selected_idx]
total_c += COSTS[selected_idx]
return total_q / len(dataset), total_c / len(dataset)
def run_training_experiment():
print("=" * 60)
print(" EquiRouter Academic Classifier Training Simulator")
print("=" * 60)
print("Generating workload dataset...")
train_data = generate_mock_dataset(120)
test_data = generate_mock_dataset(40)
print(f"Dataset Split: Train={len(train_data)}, Test={len(test_data)}")
# Sweep lambdas to generate cost-quality Pareto curves
lambdas = [0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 50.0]
results_mse = []
results_rank = []
print("\nTraining MSE-based model...")
# MSE is trained once since it doesn't depend on lambda in its loss function
mse_weights = optimize_weights(train_data, loss_type="mse", epochs=50)
for l_val in lambdas:
q, c = simulate_routing(test_data, mse_weights, l_val)
results_mse.append((q, c))
print("Training EquiRouter Decision-Aware Ranking Loss models...")
for l_val in lambdas:
rank_weights = optimize_weights(train_data, loss_type="rank", lambda_val=l_val, epochs=50)
q, c = simulate_routing(test_data, rank_weights, l_val)
results_rank.append((q, c))
# Calculate AIQ (Area under the cost-quality curve using trapezoidal rule)
# Costs are sorted to make integration sound
def calc_aiq(res: list[tuple[float, float]]) -> float:
sorted_res = sorted(res, key=lambda x: x[1]) # sort by cost
area = 0.0
for i in range(1, len(sorted_res)):
c1, q1 = sorted_res[i-1][1], sorted_res[i-1][0]
c2, q2 = sorted_res[i][1], sorted_res[i][0]
area += 0.5 * (q1 + q2) * (c2 - c1)
return area
aiq_mse = calc_aiq(results_mse)
aiq_rank = calc_aiq(results_rank)
improvement = (aiq_rank - aiq_mse) / aiq_mse * 100.0 if aiq_mse > 0 else 0.0
print("\n" + "-"*50)
print(" PARETO FRONTIER COMPARISON")
print("-"*50)
print(f"{'Lambda':<10} | {'MSE Quality':<12} | {'MSE Cost':<10} || {'EquiRouter Q':<12} | {'EquiRouter C':<10}")
print("-"*50)
for i, l_val in enumerate(lambdas):
q_m, c_m = results_mse[i]
q_r, c_r = results_rank[i]
print(f"{l_val:<10.1f} | {q_m:<12.3f} | {c_m:<10.4f} || {q_r:<12.3f} | {c_r:<10.4f}")
print("\n" + "=" * 50)
print(f"MSE Router AIQ: {aiq_mse:.6f}")
print(f"EquiRouter (Rank) AIQ: {aiq_rank:.6f}")
print(f"EquiRouter Pareto Area Improvement: +{improvement:.2f}%")
print("=" * 50)
print("Conclusion: EquiRouter Ranking Loss successfully aligns classifier weights")
print("with routing outcomes, avoiding high-budget selection collapse.")
print("=" * 50)
if __name__ == "__main__":
run_training_experiment()
|