AgentsSci/EMNLP_Cost-Aware-Protocol-Routing

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) · PDF · Project page · Code · Dataset

---

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.

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

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

@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