01. Risk-Constrained, Distribution-Shift-Aware Adaptive Routing Between LLM/Compact Vision Models Under Fixed Compute Budgets

Executive Summary Working Title: Shift-Aware Conformal Cascades Under Dual Risk and Compute Constraints

Deploying a flagship Vision Transformer or $14\text{B}+$ Large Language Model on every incoming request is too slow and expensive for most real-world applications, yet relying solely on a compact model ($f_L$) sacrifices accuracy on hard or out-of-distribution inputs. Model cascades solve this on clean benchmarks by running $f_L$ first and escalating uncertain inputs to a heavy model $f_H$. However, in my experiments with cross-domain classification, I repeatedly hit a structural failure mode: existing deferral rules trust $f_L$'s softmax margin even when the input distribution shifts. Under covariate shift, $f_L$'s predictive entropy and actual error probability decouple—either $f_L$ silently commits confident mistakes, or it escalates almost every sample and violates the system's latency/compute cap $B_{\max}$. This research topic develops a risk-constrained, distribution-shift-aware routing controller that fuses Predictive Entropy with representation-level OOD energy scores and solves a constrained primal-dual allocation problem to bound selective risk under a hard compute budget.

Tags: #Adaptive Routing #Model Cascades #Distribution Shift #Conformal Risk Control #Fixed Compute Budget

Research Motivation & Feasibility

Best overall balance of reliable ML, adaptive inference, single-GPU feasibility, and industry relevance. Novelty comes from our specific shift-aware primal-dual routing rule and proving how existing confidence-based cascades break under covariate drift.

Core Identified Research Gap

Existing two-tier model cascades (e.g., Jitkrittum et al., 2023; Dekoninck et al., 2024) assume the deferral threshold calibrated on an i.i.d. validation split remains valid at test time. When covariate shift hits, the compact Stage-1 model becomes overconfident on corrupted inputs—failing to defer right when the heavy Stage-2 model is needed most—or else triggers a cascade storm that blows past the fixed compute budget.

Basic Information

Field	Details
Core Focus	Top Thesis & Research Direction
Final Research Title	Risk-Constrained, Distribution-Shift-Aware Adaptive Routing Between LLM/Compact Vision Models Under Fixed Compute Budgets
Working Title	Shift-Aware Conformal Cascades Under Dual Risk and Compute Constraints
Research Area	Reliable Machine Learning & Adaptive Inference Systems
Subfield	Model Cascades, Conformal Risk Control & Distribution Shift
Target Academic Level	MSc Thesis / PhD Foundation (Single-GPU Feasible)
Current Status	Validated & Ready for Full Proposal

Idea Summary

2.1 One-Sentence Core Thesis

Design a shift-aware routing gate $g(x) \in \{\text{Accept } f_L,\, \text{Escalate to } f_H,\, \text{Abstain } \perp\}$ that couples predictive uncertainty with feature-space drift scores and updates its thresholds online via primal-dual optimization so that target error constraints and average compute budgets are respected simultaneously under distribution shift.

2.2 Research Motivation & Strategic Fit

1
Direct Continuity With My Core Research Identity: It unifies my three primary research notes—Understanding Predictive Uncertainty, Distribution Shift: A Practical Introduction, and Risk-Constrained AI Decision-Making—into one cohesive, testable problem.
2
100% Single-GPU Feasible: Once forward-pass logits and penultimate embeddings $(z_L(x), h_L(x), z_H(x))$ are extracted and cached on disk for CIFAR-100-C, ImageNet-C/R, WILDS, or 1.5B/14B LLMs, running 500 routing ablations and online primal-dual simulations takes minutes on a single workstation GPU.
3
High Industry & Academic Value: Every company serving LLMs or edge-cloud vision models needs cost-bounded routing that does not break silently when user traffic drifts.

Problem Exploration & Identified Research Gap

3.1 Where Existing Cascades Break Down

Standard two-stage cascades (Wang et al., 2022; Jitkrittum et al., 2023) define a static deferral rule $g_\tau(x) = \mathbf{1}\{u_L(x) > \tau\}$, where $u_L(x) = 1 - \max_y \hat{p}_L(y\mid x)$ is calibrated on an in-distribution validation set $\mathcal{D}_{\text{cal}} \sim p_S(x, y)$. When the test stream shifts to $p_T(x, y) \neq p_S(x, y)$, two failure modes occur:

Failure Mode A (Silent Under-Routing / Risk Violation): If $f_L$ suffers from feature collapse or miscalibration under corruption (Ovadia et al., 2019), $u_L(x)$ remains artificially small ($\hat{p}_L \approx 0.92$) while true accuracy drops from $90\% \to 58\%$. The router keeps accepting $f_L$ and breaches the maximum allowable risk $r^\star$.
Failure Mode B (Budget Explosion / Cascade Storm): Conversely, if a mild domain shift increases predictive entropy across all samples uniformly, $u_L(x) > \tau$ triggers on $85\%$ of traffic instead of the budgeted $25\%$, saturating $f_H$ and violating the hardware FLOPs/latency budget $B_{\max}$.

3.2 Explicit Research Gap

Current literature treats selective risk control (Angelopoulos et al., 2022) and compute-budgeted model routing (FrugalGPT, Chen et al., 2023; RouteLLM, Ong et al., 2024) as separate problems—and almost all existing routers assume stationarity between calibration and test distributions. No existing framework jointly enforces a finite-sample risk ceiling $r^\star$ and a hard average compute budget $B_{\max}$ under unlabeled test-time distribution shift by explicitly separating cases where $f_H$ can recover the answer from cases where both $f_L$ and $f_H$ fail (requiring abstention $\perp$).

Proposed Research Direction

4.1 Core Hypothesis

Hypothesis: A composite routing score that estimates the expected accuracy gain of escalating from $f_L$ to $f_H$ conditioned on both predictive entropy $\mathcal{H}(\hat{p}_L)$ and feature-space shift energy $E_L(x)$, paired with an online primal-dual threshold controller, will reduce budget violations by $>80\%$ and lower selective risk by $15\text{–}30\%$ under covariate shift compared to static softmax/entropy cascades.

4.2 Testable Research Questions

RQ1 (Failure Characterization): Under what severity of covariate and domain shift do standard confidence-based routers route the wrong subset of examples to $f_H$ (i.e., escalating unrecoverable noise while accepting overconfident $f_L$ errors)?
RQ2 (Shift-Aware Routing Score): Can a lightweight verifier $\hat{\Delta}(x) = \widehat{\mathbb{P}}(f_H(x)=y \mid h_L(x)) - \widehat{\mathbb{P}}(f_L(x)=y \mid h_L(x))$ incorporating Mahalanobis/Energy drift scores distinguish inputs where $f_H$ succeeds from inputs where both models should abstain?
RQ3 (Online Primal-Dual Budget & Risk Control): How effectively can an online dual variable $\lambda_t$ (for compute budget $B_{\max}$) and Adaptive Conformal Inference quantile $\tau_t$ (for risk $r^\star$) track non-stationary shift streams without ground-truth test labels?

4.3 Mathematical Formulation

Let $c_L < c_H$ denote the normalized compute costs of the compact model $f_L$ and heavy model $f_H$. For each input $x_t$, the router chooses an action $a_t \in \{0\text{ (use }f_L\text{)},\, 1\text{ (escalate to }f_H\text{)},\, \perp\text{ (abstain)}\}$. We formulate a shift-corrected escalation utility score:

$$ S_{\text{route}}(x) = \underbrace{\tilde{\mathcal{H}}\big(\hat{p}_L(y\mid x)\big)}_{\text{Local Decision Ambiguity}} + \beta_1 \underbrace{\phi_{\text{shift}}\big(h_L(x)\big)}_{\text{Feature Drift Score}} - \beta_2 \underbrace{\widehat{\text{Risk}}_H\big(h_L(x)\big)}_{\text{Predicted } f_H \text{ Failure}} $$

where $\phi_{\text{shift}}(h_L(x))$ is the normalized Mahalanobis or Helmholtz free-energy distance of $f_L$'s penultimate embedding $h_L(x)$. To enforce the average compute budget $\frac{1}{T}\sum_{t=1}^{T} \big(c_L + a_t c_H\big) \le B_{\max}$, we maintain an online dual price of compute $\lambda_t \ge 0$ updated after each request:

$$ \lambda_{t+1} = \Big[\lambda_t + \eta_\lambda \big( (c_L + a_t c_H) - B_{\max} \big)\Big]_+, \qquad a_t = \mathbf{1}\Big\{ S_{\text{route}}(x_t) > \tau_{\text{risk}} + \lambda_t c_H \Big\} $$

When the shift score $\phi_{\text{shift}}(h_L(x_t))$ exceeds the support of $f_H$ or when $\lambda_t$ prevents escalation on a high-risk sample, the router routes to Abstention ($\perp$) rather than emitting an uncalibrated $f_L$ guess.

Interactive Diagram SpecificationMermaid Flow / Schedule

flowchart LR
    X["Test Input x_t"] --> FL["Compact Model f_L(x_t)<br/>Cost = c_L"]
    FL --> FEAT["Extract Entropy H(p̂_L)<br/>& Shift Energy φ_shift(h_L)"]
    FEAT --> GATE{"Primal-Dual Router<br/>S_route(x_t) vs (τ_risk, λ_t)"}
    GATE -- "Low Risk & In-Dist" --> OUT_L["Output ŷ_L<br/>(Fast Path: Cost c_L)"]
    GATE -- "Recoverable by f_H & Budget OK" --> FH["Heavy Model f_H(x_t)<br/>(Escalated Path: Cost c_L + c_H)"]
    FH --> OUT_H["Output ŷ_H"]
    GATE -- "Severe Shift or Budget Exhausted on High-Risk" --> ABST["Abstain (⊥)<br/>Preserve Risk Ceiling r*"]

Preliminary Technical Plan

Model Pairs to Benchmark:

1
Vision: MobileNetV3-Small ($0.06\text{ GFLOPs}$) $\to$ ConvNeXt-Base / ViT-B/16 ($15.4\text{–}17.6\text{ GFLOPs}$) on CIFAR-100-C, ImageNet-C, ImageNet-R, and Camelyon17-WILDS.
2
Language / Reasoning: Qwen2.5-1.5B-Instruct $\to$ Qwen2.5-14B-Instruct (4-bit AWQ) on MMLU-Pro and GSM8K with synthetic domain/prompt perturbations.

Baselines:

1
Static Softmax Margin & Predictive Entropy Cascades.
2
Chow's Rule + Post-Hoc Temperature Scaling (Guo et al., 2017).
3
Learned MLP Router (RouteLLM / FrugalGPT style) trained on source validation split.
4
Random Routing at matched FLOPs budget $B_{\max}$.

Evaluation Metrics:
Selective Risk at Fixed Compute Budget ($B_{\max} \in \{0.20, 0.35, 0.50\} \times c_H$).
Area Under the Risk-Compute-Coverage Surface ($\text{AURCC}$).
Budget Violation Rate ($\text{BVR}$) and Risk Ceiling Violation ($\max_t [R_t - r^\star]_+$) across 15 corruption severities.

Expected Contributions

1
Empirical Failure Proof: A systematic demonstration showing that standard confidence-based cascades suffer from inverse routing under moderate-to-severe covariate shift.
2
Shift-Aware Primal-Dual Routing Algorithm: A lightweight, zero-retraining routing rule combining epistemic/shift scores with online dual budget control and conformal abstention.
3
Open-Source Benchmark Suite: A reproducible logit-and-embedding routing benchmark (shift-cascade-bench) enabling rapid evaluation on a single consumer GPU.

Feasibility, Scope, and Honest Risks

Why Scope Is Manageable: Pre-extracting logits and penultimate embeddings for 50,000 test images across 5 models takes $< 3$ hours on one GPU; all subsequent router training, threshold optimization, and online stream simulations run on cached tensors in seconds.
Main Risk & Mitigation: If $f_L$'s penultimate features $h_L(x)$ also collapse on certain corruptions, $\phi_{\text{shift}}(h_L(x))$ may under-detect shift. Fallback: Use lightweight input-space/early-layer stem statistics (first-layer Gram/norm discrepancy) alongside penultimate energy.

Initial Literature Notes

[1]
W. Jitkrittum et al., "When Does Confidence-Based Cascade Deferral Suffice?" NeurIPS, 2023 — Proves theoretical conditions where confidence deferral fails even in-distribution when $f_L$ and $f_H$ have mismatched error regions.
[2]
L. Chen, M. Zaharia, and J. Zou, "FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance," TMLR, 2024 — Establishes LLM cascade cost formulation under stationary queries.
[3]
A. N. Angelopoulos et al., "Conformal Risk Control," ICLR, 2024 — Provides finite-sample monotonic risk bounds that we extend to two-stage compute-constrained cascades.
[4]
I. Gibbs and E. Candès, "Adaptive Conformal Inference Under Distribution Shift," NeurIPS, 2021 — Online quantile tracking under non-exchangeable distribution shifts.

Idea Validation Checklist

Specific limitation of existing cascades under distribution shift identified mathematicallyVerified
Single-GPU experimental protocol designed via cached embedding/logit bankVerified
Clear differentiation between escalation ($f_H$) and abstention ($\perp$)Verified
Both vision (ImageNet-C/WILDS) and LLM (1.5B $\to$ 14B) tracks scopedVerified
Run initial inverse-routing diagnostic plot on CIFAR-100-C (ResNet-18 $\to$ ViT-B/16)Pending

Current Assessment

Criterion	Score	My Assessment Notes
Problem clarity	10 / 10	Crisp mathematical trade-off between Risk ($r^\star$), Coverage ($\phi$), and Compute ($B_{\max}$).
Novelty potential	9 / 10	Directly targets the gap between stationary cascades and shift-robust selective prediction.
Technical feasibility	9 / 10	Primal-dual + ACI updates are closed-form and convex per step.
Resource feasibility	10 / 10	Runs easily on $1\times$ 24GB GPU.
Dataset availability	10 / 10	CIFAR-100-C, ImageNet-C/R, WILDS, and MMLU-Pro are public.
Publication potential	9 / 10	Strong fit for NeurIPS, ICML, ICLR, or TMLR.
Personal interest	10 / 10	My #1 pick combining reliable ML, uncertainty, and systems efficiency.

Next Actions

1
Extract and save penultimate features + logits for ResNet-18 and ViT-B/16 across all 15 corruptions of CIFAR-100-C.
2
Plot the empirical Oracle Routing Gain vs. Softmax Confidence Routing under severity 1–5 to visually demonstrate the Research Gap in Figure 1 of the paper.

Open Questions and Working Notes

Can we bound the cumulative regret of the online dual variable $\lambda_t$ when the test stream alternates between clean bursts and corrupted bursts?
How much does pre-calibrating $f_L$ with Test-Time Adaptation (AdaBN) reduce unnecessary escalations to $f_H$?