README.md

---
license: mit
library_name: sklearn
pretty_name: "Cost-Aware Protocol Routing: Metadata-Only Protocol Router"
# A real checkpoint ships here: router_metadata_only/model.joblib, with its
# fitted feature builders, an explicit label map, and a runnable example.
# `inference: false` because this is a local scikit-learn artifact -- there is
# no hosted inference endpoint for it.
inference: false
datasets:
  - AgentsSci/EMNLP_Cost-Aware-Protocol-Routing
tags:
  - protocol-routing
  - llm-routing
  - multi-agent-systems
  - cost-aware-inference
  - protocol-selection
  - calibration
  - reasoning
  - reproducibility
  - arxiv:2608.14927
  - emnlp2026
language:
  - en
---

# Cost-Aware Protocol Routing — router model card

**EMNLP 2026 Accepted Paper**
*LLMs Can Predict Failure Risk, But Struggle to Predict Which Collaboration Protocol Pays Off: Cost-Aware Protocol Routing Across Reasoning Tasks*

[Paper (arXiv)](https://arxiv.org/abs/2608.14927) ·
[PDF](https://arxiv.org/pdf/2608.14927) ·
[Project page](https://chihhsuan-yang.github.io/EMNLP_Cost-Aware-Protocol-Routing/) ·
[Code](https://github.com/ChihHsuan-Yang/EMNLP_Cost-Aware-Protocol-Routing) ·
[Dataset](https://huggingface.co/datasets/AgentsSci/EMNLP_Cost-Aware-Protocol-Routing)

---

## Read this first: what this model is, and is not

**This is not a language model.** It contains no foundation-model weights. The
paper evaluates two third-party solver families — `openai/gpt-oss-120b` and
`google/gemma-4-31B-it` — which we did **not** train, do not own, and do not
redistribute. Their licenses and terms are their own.

**What this is:** the paper's lightweight *router* — a small scikit-learn
classifier that picks which collaboration protocol to run for a given problem.
A trained checkpoint ships here, along with its fitted feature builders, an
explicit label map, held-out predictions, and a runnable example.

> **Using the "Use this model" snippet?** It loads `sklearn_model.joblib` from
> the repository root, which is a byte-identical copy of the checkpoint provided
> so that snippet works. It gives you the bare estimator and nothing else — **no
> feature builders and no label map**, so its integer output is not
> interpretable on its own and features built by hand will not match the fitted
> pipeline. Use `router_metadata_only/predict_example.py` instead; it loads the
> same model together with everything needed to use it correctly.

### What ships here

**A trained router checkpoint**, in `router_metadata_only/`:

| File | What it is |
|---|---|
| `model.joblib` | The fitted multinomial logistic regression (220 features, 5 actions) |
| `feature_builders.joblib` | The fitted `OneHotEncoder`, `MultiLabelBinarizer` and `StandardScaler` |
| `label_mapping.json` | **Read this before decoding predictions.** See the warning below |
| `predict_example.py` | A runnable, offline end-to-end example |
| `model_metadata.json` | Feature order, split, hyperparameters, checksums, environment |
| `metrics.csv`, `class_recalls.csv`, `test_predictions.csv` | Held-out results and per-problem predictions |
| `selected_config.json` | The dev-only hyperparameter search and what it chose |

A byte-identical copy of `model.joblib` also sits at the repository root as
`sklearn_model.joblib`, purely so Hugging Face's auto-generated snippet resolves.
`router_metadata_only/model_metadata.json` records the SHA-256 of both, so the
two cannot drift apart unnoticed: if they ever disagree, the one under
`router_metadata_only/` is canonical.

```bash
pip install scikit-learn pandas scipy joblib
python router_metadata_only/predict_example.py
```

### Two things that will bite you

**1. The class indices are not alphabetical.** The estimator stores integer
classes `[0, 1, 2, 3, 4]` with no embedded label map. The order is the paper's
**fixed cost order**, not `sorted()`:

| index | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| **correct** | `baseline_llm` | `single_agent` | `per` | `broadcast` | `none` |
| ~~alphabetical~~ | ~~baseline_llm~~ | ~~broadcast~~ | ~~none~~ | ~~per~~ | ~~single_agent~~ |

This matters because the wrong decoding looks plausible. Decoding the released
test set alphabetically agrees with the published predictions on **352 of 423**
rows — it silently mislabels 71 and still produces a sensible-looking
distribution. Use `label_mapping.json`, which we verified reproduces
**423 of 423**.

**2. The pickle is version-fragile.** It was fitted under scikit-learn 1.8.0 and
emits `InconsistentVersionWarning` on other versions. In our test environment
(1.6.1) `predict()` works but `predict_proba()` raises `AttributeError`, because
an internal attribute moved between versions. `predict_example.py` catches this
and falls back to a softmax over `decision_function`. If you need exact
probabilities, either match 1.8.0 or retrain with the released code — training
takes minutes on CPU.

A checkpoint is a snapshot of a library version as much as of a model. The
training code is the durable artifact; this file is a convenience.

### What this checkpoint actually achieves

Held-out results for **this** artifact, read from `router_metadata_only/metrics.csv`:

| Split | n | Oracle-action accuracy | Macro-F1 | Realized solve rate | Avg tokens |
|---|--:|--:|--:|--:|--:|
| dev | 416 | 0.5697 | 0.2679 | 0.6274 | 51.7K |
| test | 423 | 0.5721 | 0.2295 | 0.6076 | 44.3K |

Read those numbers honestly. A macro-F1 of **0.23** on five actions means the
router is good at the majority action and weak on the rare, expensive ones —
which is exactly the paper's finding, not a defect in this file. On the same
split, always-Baseline solves 56.8% and the retrospective oracle reaches 92.4%.
This router lands at 60.8%: better than the fixed policy, far short of the
ceiling, and it is the *metadata-only ablation*, deliberately the weakest of the
three variants.

Do not deploy it as-is. It is a reproducible baseline to beat.

### Which router is this?

The **metadata-only** variant, trained on the primary 423-problem split. It is
the paper's ablation that isolates what metadata alone can do, and it is the one
variant with serialized weights. The main **text+metadata** router and the
six-setting evaluations have no saved estimator; their per-problem predictions
and full hyperparameter searches are released under `artifacts/`, and the
training code reproduces them.

## What the router predicts

For each problem, **before any protocol is executed**, the router picks one of
**five actions**:

| Action | Meaning | Oracle rank |
|---|---|:--:|
| `baseline_llm` | Direct one-shot solving | 1 |
| `single_agent` | Iterative single-agent self-correction | 2 |
| `per` | Planner–Executor–Reviewer collaboration | 3 |
| `broadcast` | Multi-agent deliberation | 4 |
| `none` | Do not escalate: no observed protocol succeeded | 5 |

`none` is a **router action**, not a fifth protocol. It is the correct choice
when every protocol would have failed, so spending on escalation is wasted.

The target label is the **fixed-order oracle**: the first protocol that actually
succeeded, in the order above. That label is retrospective — it is what the
router is trained to *predict*, not something available at inference.

## Feature inputs, and the leakage boundary

Routers see **only** problem-level information available before execution:

- problem text (text+metadata variant only), as TF-IDF features
- difficulty tier, source, and domain metadata
- benchmark and condition identifiers

They **never** see: gold answers, Baseline correctness, any protocol outcome,
the oracle label, or any downstream collaboration result. This is enforced in
code, not merely by convention — `protocol_routing.features` maintains an
explicit forbidden-column set and raises `LeakageError` on contact, including
for case variants. The repository's test suite proves the guard refuses
deliberately constructed violations.

## Router variants

| Variant | Features | Notes |
|---|---|---|
| **Text + metadata** | TF-IDF over problem text, plus metadata | The paper's main learned router |
| **Metadata only** | Metadata, no text | Ablation; isolates the contribution of text |
| **Tier-majority** | None (train-split lookup) | Reference policy: majority oracle label within each difficulty tier, with the train-split global majority as fallback |

All are multinomial logistic regression over the five actions.

## Training and selection procedure

**Six-setting evaluation** (the main router results): stratified 70/15/15 split
by oracle label, seed `20260712`. Hyperparameters selected on **dev only**, then
refit on train+dev, then evaluated once on untouched test problem identifiers.
All routers are scored on **identical held-out ids**, which we verified.

**Primary split** (Table 1): stratified 80/10/10, seed `42`, test n=423.

Confidence intervals throughout are 2,000-resample **problem-level percentile
bootstraps**. They quantify uncertainty over which problems are in the
benchmark — *not* run-to-run variability, since each problem-protocol pair has
exactly one realized outcome.

## Supported settings

Router evaluations cover **six** settings only: both solvers on OmniMath,
LAB-Bench strict, and LAB-Bench text-no-tool. The matched four-protocol outcome
data covers ten settings, but the router analysis does not. Please do not read
these results as holding across all ten.

## Results

Held-out router performance (from `results/aggregate/heldout_router_evaluation.csv`):

| Setting | n | Baseline | Router | Oracle | Oracle gap |
|---|--:|--:|--:|--:|--:|
| Gemma-4-31B-it OmniMath | 628 | 0.6943 | 0.7659 | 0.9570 | 0.1911 |
| gpt-oss-120b OmniMath | 628 | 0.5685 | 0.6672 | 0.9268 | 0.2596 |
| Gemma-4-31B-it LAB-Bench strict | 112 | 0.4464 | 0.7054 | 0.9107 | 0.2054 |
| gpt-oss-120b LAB-Bench strict | 112 | 0.1875 | 0.5625 | 0.7679 | 0.2054 |
| Gemma-4-31B-it LAB-Bench text-no-tool | 232 | 0.4181 | 0.7759 | 0.9612 | 0.1853 |
| gpt-oss-120b LAB-Bench text-no-tool | 232 | 0.3060 | 0.5733 | 0.8621 | 0.2888 |

The router beats Baseline everywhere. **It also leaves 18.5–28.9 points on the
table against the retrospective oracle**, and oracle-label macro-F1 stays in the
0.27–0.50 range. That gap is the paper's point, not a footnote.

## Why this does not solve collaboration-value prediction

The paper's central negative result: a model's own confidence ranks *Baseline
failures* well (0.8847 AUROC) but is much weaker at identifying *which* protocol
pays off (0.1674 AUPRC for PER-specific value, 0.1041 for Broadcast-specific).
Learned routers improve on fixed policies but do not close the oracle gap.

**Cost-aware selection among Single, PER, Broadcast, and None remains
unresolved.** These artifacts are published so that others can attack that
problem on matched data, not because the problem is solved.

## Intended use and limitations

**Intended:** reproducing the paper; research on cost-aware routing and
escalation policies; a baseline to beat.

**Not intended:** production routing without re-validation on your own workload.
These routers are fit to specific benchmarks and two specific solver families.
Protocol value varies substantially by task — the paper shows Broadcast's
advantage over PER ranging from about 10 to 45 points across settings — so a
router fit here should not be assumed to transfer.

**Further limitations**

- One realized execution per problem-protocol pair; no run-to-run variance.
- The fixed-order oracle is retrospective and is not a deployable policy.
- Cost is measured in logged tokens only — not latency, price, energy, or
  parallelism.
- No model snapshot was version-pinned at run time, so exact historical
  execution cannot be recreated. We claim functional, not bit-for-bit,
  reproducibility.
- Routing errors are asymmetric in a way these metrics do not capture:
  under-escalating loses a solvable problem, over-escalating only wastes tokens.

## Reproducing

```bash
git clone https://github.com/ChihHsuan-Yang/EMNLP_Cost-Aware-Protocol-Routing.git
cd EMNLP_Cost-Aware-Protocol-Routing
make setup
make test
python scripts/train_router.py --help
python scripts/evaluate_router.py --help
```

See `docs/REPRODUCE.md` in the repository for the full guide.

## Citation

```bibtex
@article{yang2026protocolrouting,
  title  = {{LLMs Can Predict Failure Risk, But Struggle to Predict Which
            Collaboration Protocol Pays Off: Cost-Aware Protocol Routing
            Across Reasoning Tasks}},
  author = {Yang, Chih-Hsuan and Jiang, Jingyan and Yang, Cheng-Hau and
            Vasudevan, Vikram and Zheng, Huihuo and Vishwanath, Venkatram and
            Thakur, Rajeev},
  journal = {arXiv preprint arXiv:2608.14927},
  year    = {2026},
  note    = {To appear at EMNLP 2026},
  url     = {https://arxiv.org/abs/2608.14927}
}
```

## Acknowledgments

This research used resources of the Argonne Leadership Computing Facility, a
U.S. Department of Energy (DOE) Office of Science user facility at Argonne
National Laboratory (ANL) operated under Contract No. DE-AC02-06CH11357.

## Contact

Chih-Hsuan (Bella) Yang — bellayang@anl.gov