Text Classification
router
model-routing
cma-es
benchgen
benchgen's picture
Upload folder using huggingface_hub
2e03a6a verified
|
Raw
History Blame Contribute Delete
5.03 kB
metadata
license: apache-2.0
base_model: Qwen/Qwen3-1.7B
tags:
  - router
  - model-routing
  - cma-es
  - benchgen
datasets:
  - benchgen/router-pilot
  - benchgen/router-pilot-tasks
pipeline_tag: text-classification

BenchGen Router Lite (head-Qwen3-1.7B-e99b449fde)

Repo: benchgen/benchgen-router-lite

A router head, not a fine-tuned LLM: a tiny (10,245-parameter) linear classifier that sits on top of a frozen Qwen/Qwen3-1.7B backbone and learns which model in a pool should answer a given query. The backbone never updates — only this head does. Trained with separable CMA-ES (a gradient-free evolutionary search) against a measured reward matrix, not gradient descent, and never calls any pool model during training.

Full write-up (dataset, training run, benchmark results): BenchGen Router Lite docs guide.

How it works

query --> frozen Qwen3-1.7B (mean-pooled embedding, 2048-d)
      --> this head (linear: logits = embedding @ W + b)
      --> argmax over 5 pool agents
      --> picked agent's REAL reply is what gets returned

The router itself never generates an answer — it only picks who should.

Pool (5 agents this head was trained to route between)

Slot Agent id (agent_order) Underlying model Role
Frontier frontier_a openai/gpt-oss-120b Reasoning
Frontier frontier_b deepseek/deepseek-v4-flash-0731 Reasoning
Frontier frontier_c google/gemma-3-27b-it Direct
Mid-tier open_mid mistralai/mistral-nemo Direct (deliberately the weakest agent)
Cheap reasoning open_cheap_reasoning inclusionai/ling-3.0-flash Reasoning

The underlying model slugs above are not baked into the head — they're the OpenRouter pool this specific run was trained against. A caller resolves picked_agent to a real model separately (BenchGen's gateway does this via its own pool config).

Files

File What it is
head_weights.npy Flat float32 parameter vector (10,245 values) — theta
manifest.json Architecture + training metadata (backbone id, agent order, dims, scores)
router_head.py Standalone reference implementation: embed a query, run the head, argmax

Usage

pip install torch transformers numpy
python router_head.py "What is the derivative of x^3 + 2x?"
from router_head import pick_agent

result = pick_agent("What is the derivative of x^3 + 2x?")
print(result["picked_agent"])   # e.g. "frontier_a"

manifest.json's head_hidden_dim is 0 for this run, so the head is a pure linear layer: logits = embedding @ W + b, W shape (2048, 5), b shape (5,), unflattened from head_weights.npy in that order. The embedding step (mean-pool the backbone's last hidden state over real, non-padding tokens) must match exactly — see router_head.py's embed_query().

Training data & held-out results

Trained on benchgen/router-pilot (reward matrix: 46 tasks x 5 agents x 3 repetitions, correctness measured, not assumed) joined with benchgen/router-pilot-tasks (prompt text) by task_id. Split 35 train / 11 held-out test rows.

Metric Score
Train reward (CMA-ES fit) 0.9429
Held-out test reward 0.8182
Held-out random-agent baseline 0.6364
Held-out best-single-fixed-agent baseline 0.7879
Held-out per-question oracle (upper bound) 0.8485
Beats best fixed agent? Yes

The router beats always-calling-the-best-single-agent on held-out data — the actual test of whether the extra routing step is worth it over a naive "always use one model" strategy.

Limitations

  • In-distribution only. The reward dataset is English math/knowledge/ reasoning tasks (MATH500, MMLU, MMLU-Pro, ARC-Challenge, GSM8K, AIME2025). Evaluated outside that distribution (a different language, or a very different domain), the classifier's embeddings fall outside anything it learned to discriminate and its pick becomes closer to arbitrary than a real routing decision.
  • Pool-specific. This head only knows how to choose between the exact 5 agents listed above, in that exact order. Retraining is required for a different pool.
  • Only 46 tasks currently carry reward labels (of 1,110 published prompts in router-pilot-tasks), so held-out numbers are directional, not a large-n benchmark result.

Citation / provenance

Trained on the BenchGen platform. Architecture and training approach are similar in spirit to Sakana AI's work on model routing, and the reward-matrix pool design mirrors the coordinator setup in the Trinity paper (arXiv:2512.04695).