CR^2: Cost-Aware Risk-Controlled Routing for Wireless Device-Edge LLM Inference

Technical Summary: CR2 – Cost-Aware Risk-Controlled Routing for Wireless Device-Edge LLM Inference

Problem Statement
As Large Language Models (LLMs) migrate from centralized clouds to mobile-edge environments, efficient serving requires balancing latency, energy consumption, and accuracy under constrained device-edge resources. While query-level routing between lightweight on-device models and stronger edge models offers a flexible mechanism to navigate this trade-off, existing routing methods are primarily designed for centralized cloud settings. These prior approaches optimize token-level costs or static quality-cost trade-offs, failing to capture the dynamic latency and energy overheads inherent in wireless edge deployments.

A critical gap exists in the two-stage decision structure required by device-edge systems:

1. Local Decision (UE): The user equipment (UE) must decide whether to execute a query locally or defer it to the edge using only locally available information (query embedding and cost preference), without access to edge model scores or real-time edge state.
2. Edge Decision: Only deferred queries are routed among available edge models based on full system state.

Existing methods often assume global knowledge of edge states or treat routing as a single-shot selection, ignoring the asymmetric decision risk: unnecessary deferral incurs communication latency and energy, while false local acceptance irreversibly forgoes better edge models, degrading user experience. Furthermore, standard routers do not explicitly control the risk of false local acceptance under partial information.

Methodology: The CR2 Framework
The authors propose CR2, a two-stage device-edge routing framework that decomposes the routing problem into a deployable local acceptance gate and a state-aware edge selector.

1. System and Cost Modeling:

   • The system models end-to-end latency and energy consumption, incorporating stochastic wireless conditions (uplink/downlink channel gains), computation workloads (prefill and decoding stages), and device/edge power profiles.
   • A deployment cost is defined as a weighted sum of normalized latency and energy, relative to a reference model.
   • The routing objective is formulated as maximizing accuracy subject to a deployment cost budget, parameterized by a scalarized utility function $u_m = y_m - \lambda c_m$, where $\lambda$ is a cost weight.

2. Two-Stage Architecture:

   • Stage 1: UE Local Acceptance (Margin Gate): A lightweight on-device module computes a "utility margin" to estimate if local execution is competitive with the best edge alternative.
     • It uses a frozen query encoder and a lightweight scoring head.
     • It takes the query embedding and the operator-specified cost weight $\lambda$ as input.
     • It does not access edge model utilities or runtime state $\xi$ before making the decision.
   • Stage 2: Edge Utility Selector: For deferred queries, the edge server selects the optimal model from the edge pool based on the full utility estimate (incorporating runtime state $\xi$).

3. Training and Calibration:

   • Teacher-Student Distillation: A "teacher" selector is trained on full-information data (query + state + correctness labels) to estimate model correctness and utility margins. The on-device gate is trained to distill the teacher-induced utility margin ($\Delta$) between local and edge execution.
   • Loss Functions: The gate is optimized using a combination of:
     • Sign Loss: Aligning the predicted margin's sign with the teacher's local-vs-edge preference.
     • Margin Loss: Regressing the predicted margin toward the continuous teacher margin (using Huber loss).
     • Monotonicity Regularizer: Enforcing that the predicted margin is non-decreasing with respect to the cost weight $\lambda$, reflecting the structural property that higher cost weights favor local execution.
   • Conformal Risk Control (CRC): To address the asymmetric risk of false local acceptance, the authors introduce a CRC calibration procedure.
     • For a target risk level $\alpha$, CRC maps the learned score to a calibrated acceptance threshold $\tau^*(\lambda; \alpha)$.
     • This ensures that the marginal false-acceptance risk (the probability of accepting locally when an edge model is strictly better) is bounded by $\alpha$ under the assumption of exchangeability between calibration and test data.

Key Contributions

1. Deployment-Constrained Framework: The paper establishes a two-stage routing framework that explicitly respects the information constraints of mobile-edge deployment, separating UE local acceptance from edge model selection.
2. CR2 Algorithm: The authors develop CR2, which features:
   • A UE-observable margin gate that distills full-information utility margins into a deployable on-device score.
   • A monotonicity regularizer to ensure coherent decisions across different cost-weight operating points.
   • A per-$\lambda$ CRC calibration rule that provides explicit control over the marginal false-acceptance risk without requiring edge model scores at the UE.
3. Empirical Performance: Evaluation on a benchmark-derived routing task (using Qwen models and profiled deployment costs) demonstrates that CR2 achieves the best deployable accuracy–cost Pareto frontier. It reduces normalized deployment cost by up to 16.9% at matched accuracy compared to strong query-level baselines (KNN, MLP, EmbedLLM, LLM-Rank).

Results

• Accuracy-Cost Trade-off: CR2 consistently outperforms baselines in the accuracy-cost plane. In zoomed operating regions, it forms the strongest frontier among deployable methods.
• Cost Reduction: At fixed accuracy levels (e.g., 0.81 and 0.82), CR2 reduces normalized deployment cost by 16.9% and 16.2% respectively compared to KNN.
• Risk Control: The CRC-calibrated thresholds successfully control the empirical marginal false-acceptance rate, keeping it within the specified bounds (e.g., $\alpha=0.01$) while maintaining competitive accuracy.
• Efficiency: The on-device margin gate is computationally lightweight (approx. 105k parameters, 207k FLOPs), significantly smaller than full-information references or other parametric routers like LLM-Rank.
• Error Decomposition: Analysis shows that the remaining performance gap between the deployable CR2 and the full-information oracle is primarily driven by "false deferrals" (over-escalation) at lower costs, while false local acceptance remains low (<1.5%).

Significance and Claims
The paper claims that CR2 bridges the gap between query-level routing and edge LLM serving by formulating model selection under a two-stage device-edge deployment. Its significance lies in:

• Decoupling: Successfully decoupling the local decision (which must be made with partial information) from the edge decision (which has full information).
• Risk Control: Providing a rigorous, distribution-free mechanism (via CRC) to explicitly regulate the risk of false local acceptance, a critical requirement for reliable edge inference where communication overhead is costly.
• Practicality: Demonstrating that a lightweight, frozen-encoder-based gate can closely approximate a full-information oracle, enabling efficient, cost-aware routing in real-world wireless edge scenarios without requiring dense cross-tier coordination or global state knowledge at the device.

The authors conclude that CR2 offers a robust solution for mobile-edge collaborative inference, balancing the trade-offs between latency, energy, and accuracy while adhering to the strict information constraints of wireless edge deployments.