Pyramid MoA: A Probabilistic Framework for Cost-Optimized Anytime Inference

Arindam Khaled

cs.CLFeb 23, 2026

Pyramid MoA: A Probabilistic Framework for Cost-Optimized Anytime Inference

Authors: Arindam Khaled

Organizations: Independent Researcher

Abstract

We observe that LLM cascading and routing implicitly solves an anytime computation problem -- a class of algorithms, well-studied in classical AI, that improve solutions as additional computation is allocated. We formalize this connection and propose Pyramid MoA, a hierarchical Mixture-of-Agents architecture governed by a decision-theoretic router that escalates queries only when necessary. We establish a Probabilistic Anytime Property with provable monotonicity guarantees and derive a generalized escalation rule from Value of Computation theory that accounts for imperfect oracles, extending the Hansen-Zilberstein monitoring framework to stochastic LLM inference. On MBPP, the router intercepts 81.6% of bugs; on GSM8K/MMLU, the system nearly matches the 68.1% Oracle baseline while achieving up to 42.9% compute savings. The router transfers zero-shot to unseen benchmarks: matching Oracle accuracy on HumanEval (81.1%) and MATH 500 (58.0%) with significant cost reductions. We further discover a context-conditioned anchoring effect across four benchmarks: passing correct SLM reasoning improves Oracle accuracy by up to +19.2pp, while incorrect reasoning degrades it by up to -18.0pp, revealing a fundamental tension in hierarchical MoA architectures.

Figures & tables

Anytime Computation	Pyramid MoA
Initial solution	Layer 1 (SLM ensemble) output
Extended computation	Escalation to higher-capability layers
Computation time axis	Model capacity / inference cost
Performance profile	Accuracy as a function of cost
Monitoring problem	Router’s escalation decision at each layer
Quality guarantee (deterministic)	Quality guarantee (probabilistic)

Table 1: Mapping from classical anytime computation to Pyramid MoA.

Panel of Figure 1 — Figure 1: Pyramid MoA Architecture: The system extracts ensemble-wide features from the Layer 1 models to estimate $P_{fail}$ . The router solves the anytime monitoring problem—deciding whether to allocate additional computation via the Oracle.

Feature	Consensus Router	Anytime Router
Algorithm	XGBoost	XGBoost
Training Domain	MBPP (Code)	GSM8K/MMLU (Math)
Primary Signal	Semantic Agreement	Token Log-Probabilities
Routing Characteristic	Error-catching (high recall)	Selective escalation (high precision)
Transfer Evaluation	HumanEval	MATH 500

Table 2: Router configurations by domain.

Panel of Figure 2 — Figure 2: Consensus Mechanism: Evaluation on MBPP showing that peer-agreement signals significantly outperform intrinsic model confidence for error detection.

Figure 3: Zero-Shot Transfer to HumanEval: The MBPP-trained Consensus Router transfers effectively, achieving the Oracle baseline (81.1%) and enabling up to 62.7% cost savings in economy mode.

Panel of Figure 4 — Figure 4: Math Router Analysis: The XGBoost router leverages candidate model correctness and token-level uncertainty signals for escalation decisions on convergent tasks.

Figure 5: Anytime Performance Profile (GSM8K/MMLU Holdout): The dual-axis plot shows accuracy (red) and compute savings (green) as a function of router threshold. The concave accuracy profile confirms efficient allocation of Oracle computation. The monotonically decreasing accuracy curve as the threshold increases (i.e., as less computation is allocated) empirically demonstrates the probabilistic anytime property: expected solution quality is non-decreasing with computational depth.

Figure 6: Zero-Shot Transfer to MATH 500: The GSM8K/MMLU-trained router transfers to out-of-distribution problems, preserving the Oracle ceiling (58.0%) and enabling efficiency gains at higher thresholds.

Benchmark	$p_{R}$	$\alpha_{L1}(R)$	$\alpha_{L2}(R)$	Gap	$\Delta$ Acc	Satisfied?
MBPP (Code)	0.558	0.379	0.690	+0.311	+17.3%	✓
HumanEval (Code, OOD)	0.561	0.696	0.804	+0.109	+6.1%	✓
GSM8K/MMLU (Math)	0.271	0.632	0.768	+0.137	+3.7%	✓
MATH 500 (Math, OOD)	0.110	0.145	0.582	+0.436	+4.8%	✓

Table 3: Verification of the monotonicity condition (Theorem 1 ) at threshold

t=0.4

\alpha_{L1}(R)

is computed using majority vote over the Layer 1 ensemble. In all cases,

\alpha_{L2}(R)>\alpha_{L1}(R)

, confirming strict probabilistic anytime improvement.

Benchmark	SLM Majority	$N$	Baseline	With Context	$\Delta$
MATH 500	Correct	239	70.3%	89.5%	+19.2 pp
	Wrong	261	59.0%	41.0%	$-$ 18.0 pp
MMLU	Correct	301	89.4%	96.0%	+6.6 pp
	Wrong	251	59.8%	43.8%	$-$ 16.1 pp
GSM8K	Correct	1079	99.0%	99.4%	+0.5 pp
	Wrong	121	88.4%	73.6%	$-$ 14.9 pp

Table 4: Context-aware escalation across four benchmarks. “Baseline” denotes the Oracle solving independently; “With Context” denotes the Oracle receiving Layer 1 outputs. Results are segmented by whether the SLM majority answer was correct. The anchoring effect—context helps when SLMs are correct, hurts when SLMs are wrong—is consistent across both mathematical reasoning and code generation tasks.

LLM Routing

Large language model routing selects models, tools, experts, or inference strategies to balance answer quality, cost, latency, and deployment constraints.

See all 345 papers

ReM-MoA: Reasoning Memory Sustains Mixture-of-Agents Scaling
Jun 23, 2026Heng Ping, Arijit Bhattacharjee, Peiyu Zhang +5Multi-Agent System Optimization LLM Agent Memory
University of Southern California · Iowa State University · Cisco AI Research
Signed Rescue Routing: Harm-Aware Cascades for Efficient LLM Inference
Sep 7, 2026Zheyuan Wang, Siyu Li, Peiqiao Song +3LLM Inference Efficiency LLM Routing
School of Artificial Intelligence, Beijing Normal University · School of Information Science and Engineering, Chongqing Jiaotong University · School of Computer Science and Technology, Dalian University of Technology +1
RLCascadeRouter: Quality-Estimator-Free Cascade Routing via Reinforcement Learning
Aug 16, 2026Shihong Huang, Shengjie Wang, Hong Ma +1LLM Routing Language Model Cascading
Polytechnic Institute, Zhejiang University, Hangzhou, China · Department of Logistics and Maritime Studies, Faculty of Business, The Hong Kong Polytechnic University, Hong Kong, China

Jun 23, 2026J/K move · Enter open · S save

ReM-MoA: Reasoning Memory Sustains Mixture-of-Agents Scaling

Heng Ping, Arijit Bhattacharjee, Peiyu Zhang, Shixuan Li+4

University of Southern California · Iowa State University · Cisco AI Research

Multi-Agent System Optimization LLM Agent Memory Multi-Agent Systems Multi-Agent Reasoning Memory-Augmented Language Model Agents Mixture-of-Agents

Open paper AskPDF arXiv

Mixture-of-Agents (MoA) architectures improve inference-time scaling by organizing multiple LLM agents into layered reasoning pipelines. However, existing MoA variants fail to sustain gains as depth increases, exhibiting degradation, early plateauing, or saturation. We propose ReM-MoA, a memory-augmented MoA framework that sustains scaling through two mechanisms: (1) a Ranked Reasoning Memory that persistently stores and ranks reasoning traces from all layers using a comparative Reviewer Agent, and (2) a Curated Diversified Memory Routing scheme that exposes different agents to distinct combinations of successful and failed traces, preserving exploration diversity while propagating high-quality reasoning. We further introduce an optional multi-domain Reviewer distillation pipeline that improves ranking quality through frontier-model supervision. Across five reasoning benchmarks spanning math, formal logic, code, knowledge, and commonsense, ReM-MoA consistently outperforms prior MoA variants across both depth and width scaling, and its advantage widens with depth, establishing structured cross-layer reasoning memory as a key missing mechanism for scalable multi-agent inference.

Abstract

LLM Routing

Explore similar work

ReM-MoA: Reasoning Memory Sustains Mixture-of-Agents Scaling

Signed Rescue Routing: Harm-Aware Cascades for Efficient LLM Inference

RLCascadeRouter: Quality-Estimator-Free Cascade Routing via Reinforcement Learning