OI-MAS: Confidence-Aware Routing for Multi-Agent Multi-Scale Collaboration — L2 #

§1 TL;DR #

OI-MAS introduces a hierarchical conductor that jointly routes agent roles and model scales per reasoning step via confidence-aware RL. Token log-prob confidence modulates cost penalty: high confidence → penalize expensive models, low confidence → allow escalation. Result: +7.68% avg accuracy over best baseline at up to 79.78% cost reduction.

§2 痛点 · 方法 · 结果 #

Q1 痛点 #

Multi-agent systems (MAS) achieve strong reasoning performance but deploy the same large LLM backbone uniformly across all agent roles and all reasoning steps. This ignores that many subtasks (drafting, aggregation, simple decomposition) are easily handled by small models, while only critical steps (core reasoning generation, complex refinement) require large-capacity backbones. The result is massive computational waste: every agent call pays the full cost of the largest model regardless of the step's actual difficulty.

Existing cost-reduction strategies attack the wrong level — either optimizing agent organization (who talks to whom) or performing query-level routing (choose one model before inference starts). Neither enables state-dependent, per-step model allocation that adapts as the reasoning trajectory unfolds.

Q2 方法 #

OI-MAS decomposes the routing problem into two hierarchical stages per reasoning turn:

1. Role Router $\mathcal{F}_\phi$: encodes query $q$, current context $c_t$, and candidate roles via a pretrained text encoder (MiniLM). A learnable projection network computes softmax-normalized similarity scores. Roles are activated by accumulating probability mass until threshold $\theta$ is reached, selecting a variable-size subset $R_t$ per turn. An explicit EarlyStop role enables adaptive termination.
1. Model Router $\mathcal{G}_\psi$: for each selected role $r \in R_t$, evaluates the suitability of every backbone $m_j \in \mathcal{M}$ via a second learnable network conditioned on the triplet $(q, c_t, r)$. The highest-scoring model is assigned, yielding role-specific capacity allocation.

The two routers are jointly optimized via a confidence-aware RL objective:

$$\min_{\phi, \psi} \mathbb{E}_{(q,a) \sim \mathcal{D}} \left[ -r(q,a;\phi,\psi) + \sum_t \lambda \cdot \text{Conf}_{\text{adj}}(\tilde{s_t}) \cdot C(r_t, m_t) \right]$$

where $\text{Conf}_{\text{adj}}(\tilde{s_t}) \in [0,1]$ is a calibrated confidence score derived from average token log-probability. The key mechanism: higher confidence amplifies cost penalty (discouraging unnecessary use of large models when a small one suffices), while lower confidence relaxes the penalty (permitting escalation when the system detects difficulty).

核心技术壁垒: The confidence-modulated cost penalty is the central insight that makes the performance-cost trade-off learnable without explicit difficulty labels. The system learns to estimate per-state complexity from output token log-probabilities, then uses this signal to weight the cost term in the RL objective. This creates an implicit curriculum: easy states are routed to cheap models with strong penalty pressure, while hard states see reduced penalty and gain access to large backbones. Replicating this requires (1) stable confidence calibration across heterogeneous model families (handled by the percentile normalization + cold-start interpolation scheme in Appendix B), and (2) avoiding mode collapse where the policy converges to always-large or always-small routing.

Q3 结果 #

Metric	Value
Avg accuracy (5 benchmarks)	78.23% (+7.68% over MasRouter, +4.76% over MaAS-Large)
Best single-benchmark gain	MBPP 91.59% (+13.35 pp over MasRouter)
Cost reduction vs baselines	17.05%–78.47%
Latency (GPQA)	23.12s/query vs 36.82s (MasRouter), 192.62s (MaAS)
OOD generalization (MBPP→HumanEval)	91.46% Pass@1 at lowest cost (\$0.097/query)

Ablation reveals confidence is accuracy-critical (removing it: −2.52% MedQA, −4.20% MBPP), while the model router is cost-critical (removing it: +76–84% cost with marginal accuracy gain).

§3 架构 / 方法図 #



Paper Figure 1. Three paradigms contrasted: (a) static MAS with fixed roles and shared LLM, (b) dynamic agent routing but still a single backbone, (c) OI-MAS — dynamic role routing + per-role model selection from a heterogeneous pool (Qwen2.5-3B/7B, Llama3.1-8B/70B). The key transition from (b) to (c) is decoupling "who acts" from "how much capacity they get."



Paper Figure 2. Top: the State-Dependent Role-Model Router. Role Router $\mathcal{F}_\phi$ selects a subset of roles from the pool (Generator, Decomposer, Verifier, etc.) via embedding similarity. Model Router $\mathcal{G}_\psi$ then independently assigns each selected role a backbone from the LLM pool, producing role-model pairs like (Generator, Llama70B) and (Refiner, Qwen7B). Bottom: the film-strip shows how the agent team composition evolves across 4 reasoning turns — different roles activated with different-sized models at each step.

stateDiagram-v2 [*] --> Observe: input query q Observe --> RoleRoute: state s_t = (q, c_t) RoleRoute --> ModelRoute: selected roles R_t ModelRoute --> Execute: role-model pairs {(r, m)} Execute --> UpdateContext: agent outputs → c_{t+1} UpdateContext --> ConfidenceEval: compute Conf_adj(s̃_t) ConfidenceEval --> CheckStop: EarlyStop ∈ R_t? CheckStop --> RoleRoute: no, t < L CheckStop --> [*]: yes or t = L note right of ConfidenceEval High conf → amplify cost penalty Low conf → relax, allow large model end note

The system operates as a discrete-time loop over at most $L=4$ turns. At each turn, the reasoning state $s_t = (q, c_t)$ drives role selection (which functional operations are needed), then model assignment (what capacity each role requires). The confidence signal from the executed agents feeds back into the RL reward, training the routers to learn the mapping from reasoning-state complexity to optimal role-model configurations.

§4 作者証明 #

無形式化作者証明 — 仅实证

No formal convergence guarantee or success-rate bound is provided. The paper relies entirely on empirical validation across 5 benchmarks + 1 OOD transfer. The following could have been bounded but were not: (a) regret of the routing policy relative to oracle per-step model assignment, (b) convergence rate of the RL optimization, (c) calibration error of the confidence estimator.

記号表 #

Symbol	Definition	Introduced
$\mathcal{R}$, $\mathcal{M}$	Role set, model (backbone) set	§3.1
$r_i$, $m_j$	A specific role, a specific model	§3.1
$s_t = (q, c_t)$	Reasoning state at turn $t$: query + context	§3.1
$\mathcal{F}_\phi$	Learnable role routing network	§3.2
$\mathcal{G}_\psi$	Learnable model routing network	§3.2
$p_t^{(r)}(r_i \mid q, c_t)$	Role activation probability	§3.2, Eq. 2
$p_t^{(m)}(m_j \mid q, c_t, r)$	Model assignment probability conditioned on role	§3.2, Eq. 3
$\theta$	Role selection threshold (cumulative probability mass)	§3.2
$\text{Conf}_{\text{base}}(\tilde{s_t})$	Raw confidence: average token log-prob of output	§3.3, Eq. 4
$\text{Conf}_{\text{adj}}(\tilde{s_t})$	Calibrated confidence ∈ [0,1]	§3.3 / App. B
$\lambda$	Cost penalty coefficient	§3.3, Eq. 5
$C(r_t, m_t)$	Computational cost at step $t$	§3.3, Eq. 5
$L$	Maximum reasoning turns	§3.1
$\alpha$	Cost-scaling exponent for pricing model	App. C, Eq. 6

方程物理意義 #

Eq. 2 — $p_t^{(r)}(r_i \mid q, c_t) = \mathcal{F}_\phi(q, c_t, r_i)$: role routing via embedding-space similarity. The pretrained encoder projects query, context, and role descriptions into a shared space; the learnable network computes projected cosine similarity with softmax normalization. The threshold-based selection ($\theta = 0.3$) means the system activates 1–3 roles per turn depending on how concentrated the probability mass is.

Eq. 3 — $p_t^{(m)}(m_j \mid q, c_t, r) = \mathcal{G}_\psi(q, c_t, r, m_j)$: model routing conditioned on the already-selected role. Structurally identical to role routing but operates in a role-augmented embedding space. Argmax selection at inference.

Eq. 4 — $\text{Conf}_{\text{base}}(\tilde{s_t}) = \frac{1}{T} \sum_{k=1}^{T} \log P(y_k \mid \tilde{s_t}, y_{

Eq. 5 — the RL objective: $-r(q,a) + \sum_t \lambda \cdot \text{Conf}_{\text{adj}} \cdot C(r_t, m_t)$. The confidence-weighted cost penalty creates an asymmetric pressure: confident-and-cheap is strongly rewarded, confident-but-expensive is penalized, uncertain-and-expensive is tolerated. This is the mechanism that prevents both failure modes (always-large = wasteful; always-small = inaccurate).

Eq. 6 (App. C) — $C(m) = C_{\text{base}} \cdot (P(m)/P_{\text{base}})^\alpha$: power-law cost model. $\alpha = 0.73$ calibrated from Llama 8B/70B API pricing. Used to estimate Qwen2.5-3B cost from Qwen2.5-7B pricing.

6 检查项 #

#	Check	Status
1	Confidence metric well-defined and computable from standard LLM outputs?	✓ Average token log-prob (Eq. 4) — universally available from any autoregressive LLM
2	Confidence calibration handles cross-model scale differences?	✓ Appendix B: percentile-based normalization per model + cold-start fallback via geometric mean token probability, with smooth interpolation
3	Cost model explicit and reproducible?	✓ Appendix C: API pricing as proxy with power-law interpolation; full price table provided (Table 4)
4	Reward signal well-specified for RL?	✓ Binary sparse reward $r(q,a) \in \{0,1\}$ for exact-match correctness; no reward shaping
5	Role set and threshold justified?	Partial — 9 roles chosen without ablation on role set composition; $\theta = 0.3$ supported by sensitivity analysis (§5.4) but not validated across datasets
6	EarlyStop correctness validated?	Empirical only — case study (Appendix D) shows plausible termination; no analysis of premature vs. late stopping rates

Agent-specific checks #

• Success-rate model: empirical only. Figure 5 shows model selection shifts with MATH difficulty levels (1–5), confirming monotonic increase in large-model usage with difficulty. No formal sweep matrix over (difficulty × planning depth × tool set × backbone).
• Latency budget: 23.12s/query on GPQA, dominated by LLM inference (routing network overhead is negligible — MiniLM embeddings). Interactive latency not explicitly claimed.
• Failure mode classification: not systematically addressed. The paper does not categorize failure types (wrong routing, premature stopping, confidence miscalibration). The dominant failure mode — inferred from ablation — is confidence miscalibration leading to under-allocation (removing confidence drops accuracy 2–4%).

§5 実験与数据 #

Main results: accuracy across 5 benchmarks #



Paper Table 1. OI-MAS achieves the highest average accuracy (78.23%) using the heterogeneous LLM pool, outperforming all baselines including MasRouter (70.55%, also uses LLM pool) and MaAS-Large (73.47%, uses only Llama3.1-70B). The gains are largest on MATH (+7.56 pp over MasRouter) and MBPP (+14.28 pp over MasRouter), benchmarks where task complexity varies most within a dataset.

Key observations from Table 1:

• OI-MAS ties MasRouter on GSM8K (95.80%) — the easiest benchmark where model routing provides little lift.
• The gap widens on harder benchmarks: MATH 79.83 vs 72.27, MedQA 78.99 vs 71.43, GPQA 44.94 vs 35.96.
• Against all-large baselines (LLM-Debate Large, MaAS Large), OI-MAS matches or exceeds accuracy while using a mix of small/medium/large models.

Latency comparison #



Paper Figure 4. OI-MAS achieves the lowest per-query latency (23.12s) — 37% faster than MasRouter (36.82s), 41% faster than GPTSwarm (39.31s), and 8.3× faster than MaAS (192.62s). The latency advantage comes from routing most steps to lightweight models and enabling early termination via the EarlyStop role.

Model selection adapts to task difficulty #



Paper Figure 5. Stacked area chart of model proportion vs. MATH difficulty level (1=easiest, 5=hardest). At level 1, Qwen2.5-3B dominates (~35% of assignments). As difficulty increases, Llama3.1-70B share grows monotonically from ~15% to ~40%. This confirms the routing policy learns a meaningful difficulty-to-capacity mapping without explicit difficulty labels — the confidence signal provides sufficient supervision.

Role-level model allocation patterns #



Paper Figure 6. Heatmap of role × model selection frequency. Generator has the highest overall activation (0.119 for Llama3.1-70B) and the strongest preference for large models — it constructs core reasoning trajectories. Programmer and Ensembler concentrate on medium models (Llama3.1-8B, Qwen2.5-7B) for structured operations. Refiner shows the bimodal pattern noted in L1: high usage of both Llama3.1-70B (0.045) and Qwen2.5-3B (0.037) — easy refinements use the smallest model, complex ones escalate to the largest.

Hyperparameter sensitivity #



Paper Figure 7. Left: cost penalty coefficient $\lambda$. Accuracy peaks at $\lambda=200$, drops sharply at $\lambda=400$ (over-penalization forces too-small models). Cost decreases monotonically with $\lambda$. Right: max turns $L$. Accuracy peaks at $L=4$; $L=6{-}8$ degrades accuracy (accumulated noise from redundant interactions) while increasing cost. The sweet spot at $L=4$ balances reasoning depth against noise accumulation.

Ablation results #

Variant	MedQA Acc (%)	MedQA Cost	MBPP Pass@1 (%)	MBPP Cost
OI-MAS (full)	78.99	1.79	91.59	1.67
w/o $\mathcal{G}_\psi$ (no model router)	81.51 (+2.52)	3.16 (+76%)	93.28 (+1.69)	3.07 (+84%)
w/o $C(\cdot)$ (no cost term)	79.83 (+0.84)	2.14 (+20%)	92.43 (+0.84)	1.96 (+17%)
w/o $\text{Conf}(\cdot)$ (no confidence)	76.47 (−2.52)	1.68 (−6%)	87.39 (−4.20)	1.53 (−8%)

The ablation reveals a clear decomposition of component contributions:

• Model router ($\mathcal{G}_\psi$) is the cost-efficiency mechanism: removing it forces all-large routing, achieving marginally higher accuracy but nearly doubling cost.
• Confidence is the accuracy-critical component: removing it causes the largest accuracy drops (−2.52% MedQA, −4.20% MBPP) because the system loses its complexity-sensing ability and misallocates models.
• Cost term prevents drift toward expensive configurations without significantly affecting accuracy.

OOD generalization (MBPP → HumanEval) #

Method	Pass@1 (%)	Cost ($10^{-1}$ \$)
MaAS	89.63	2.22
LLM-Debate	78.05	2.97
AFlow	78.66	1.25
MasRouter	74.39	1.17
OI-MAS	91.46	0.97

Routing policy trained on MBPP transfers to HumanEval with zero retraining, achieving highest accuracy at lowest cost. This suggests the policy captures generalizable patterns about code-task complexity rather than dataset-specific shortcuts.

Case study: GPQA chemistry problem #



Paper Figure 8. A GPQA neutralization-enthalpy problem routed across 3 turns. Turn 1: Decomposer (Qwen2.5-3B — simple categorization) + Generator (Llama3.1-70B — error-sensitive quantitative setup). Turn 2: Refiner (Qwen2.5-7B) + Verifier (Llama3.1-8B) — medium models for structured checking. Turn 3: Ensembler (Qwen2.5-7B) consolidates, EarlyStop triggers. Demonstrates the adaptive allocation principle: lightweight for decomposition, heavy for critical generation, medium for verification, early termination once resolved.

§6 論証鎖 #

Step	Claim	Evidence	Strength
1	Uniform LLM deployment in MAS wastes compute because subtask complexity varies widely	Motivating observation (§1): drafting/aggregation are trivially handled by small models; only core reasoning needs large ones. Supported by Figure 5 showing difficulty-dependent model allocation emerges naturally.	Moderate — intuitive argument without formal characterization of subtask complexity distribution
2	Hierarchical role-model routing decouples "what to do" from "how much capacity to use," enabling per-step adaptation	Architecture design (§3.2): Role Router selects functional roles independently of model scale, then Model Router assigns capacity. Table 1 shows +7.68% avg accuracy over MasRouter (query-level routing).	Strong — the accuracy gap over static routing (MasRouter) demonstrates per-step adaptation adds value
3	Token log-probability confidence is a sufficient proxy for task complexity, enabling learned cost modulation without difficulty labels	Confidence-aware objective (§3.3, Eq. 5) + ablation (§5.3): removing confidence causes the largest accuracy drops (−2.52% to −4.20%). Figure 5 shows the learned routing aligns with ground-truth difficulty levels.	Moderate-strong — the ablation is convincing, but only average log-prob is tested (no comparison with entropy, consistency, or other confidence measures)
4	The combined system achieves Pareto-superior accuracy-cost trade-off	Table 1 (accuracy), Figure 4 (latency), cost comparison (§4.2): OI-MAS dominates all baselines on the accuracy-cost frontier. OOD transfer (Table 2) shows generalizability.	Strong — consistent across 5 benchmarks + 1 OOD test, with large margins on cost (up to 79.78% reduction)
5	The routing policy generalizes across datasets	OOD evaluation (§5.2): MBPP-trained policy achieves 91.46% Pass@1 on HumanEval at lowest cost, outperforming all baselines including those trained/tuned on HumanEval.	Moderate — single transfer pair (MBPP → HumanEval), both are code generation; cross-domain transfer (e.g., code → math) not tested

§7 実装 cross-reference #

[実現未公開]

No public code repository is referenced in the paper. Implementation details that are specified:

• Agent roles: Generator, GeneratorCoT, Decomposer, Critique, Ensembler, Verifier, Refiner, Programmer, EarlyStop (9 roles total)
• Text encoder: MiniLM for embedding $q$, $c_t$, and role/model descriptions
• LLM serving: vLLM on NVIDIA A100 80G GPUs, temperature = 0
• Routing networks ($\mathcal{F}_\phi$, $\mathcal{G}_\psi$): learnable projection networks trained with RL, learning rate $\alpha = 0.01$
• Hyperparameters: $L = 4$, $\lambda = 200$, $\theta = 0.3$

核心技術壁壘 #

The hardest-to-replicate aspect is the confidence calibration pipeline (Appendix B). Raw token log-probabilities are semantically inverted and incomparable across model families. The solution — percentile-based normalization with running statistics per model, a cold-start geometric-mean fallback, and smooth interpolation as observations accumulate — is a multi-part calibration scheme whose hyperparameters (percentile bins, interpolation schedule, normalization bounds) are not fully specified. Getting this wrong leads to either collapsed routing (all-large or all-small) or oscillatory instability.

関鍵実装細節 #

1. Confidence calibration cold-start: during early training when per-model statistics are sparse, the system falls back to a model-agnostic transformation using $\exp(\text{Conf}_{\text{base}})$ (geometric mean token probability). The interpolation from fallback to model-specific normalization is gradual — without this, early-stage routing would be poorly calibrated and the RL policy could converge to a degenerate solution.
1. Cost accounting via cross-family API pricing: all models are served locally but costs are computed using Together AI API token pricing as a proxy. The Qwen2.5-3B price (not available via API) is estimated by power-law extrapolation from Qwen2.5-7B pricing using exponent $\alpha = 0.73$ derived from the Llama 8B/70B pricing ratio. This cross-family extrapolation means the cost metric is an approximation — actual compute costs on shared hardware would differ due to batching effects, memory bandwidth bottlenecks, and GPU utilization differences between model sizes.