Optimal Model Activation Policies for Inference Networks of Large Language Models

Introduction

• Inference networks are directed graphs of LLMs with conditional activations; this work focuses on serial chains where a cheaper model’s confidence determines whether to invoke a costlier model.The framework’s design problem is to choose a topology that best addresses the cost-performance trade-off.Source: §1 · p. 2
• Cost falls by up to 97% in some instances while performance stays within a defined budget constraint, and the policies outperform other state-of-the-art strategies.Experiments use off-the-shelf open-source LLMs on text classification and generation datasets.Source: §1 · p. 2
• Under mild assumptions, optimal activation uses the lowest-cost LLM first and escalates when confidence falls below a threshold; classification uses class-specific thresholds.The policies rely on models’ self-confidence rather than an external routing model requiring additional training and data collection.Source: §1 · p. 2

• Natural-language tasks can be expressed as query-answer tuples, with an LLM mapping each query to an answer.Source: §2 · p. 2
• Discriminative tasks restrict outputs to predefined classes, whereas generative tasks sample responses from a predicted distribution over the vocabulary [Holtzman et al., 2020, Fan et al., 2018, Ficler and Goldberg, 2017].Source: §2 · p. 2
• Confidence estimation scores answer reliability; this work derives a [0, 1] score from predicted probabilities, for a full sequence in generative tasks or a predicted class in discriminative tasks.The paper leaves more sophisticated confidence estimation approaches for future work.Source: §2 · p. 2

3 SYSTEM MODEL

• The two-model policy accepts the student LLM’s answer when confidence meets a threshold and otherwise defers; the objective minimizes expected inference cost subject to a target performance constraint.The paper considers analytic or numerical solutions when distributions are known and empirical solutions from finite samples when they are not.Source: §3 · pp. 2–3

3.1 GENERATIVE TASKS - SINGLE-THRESHOLD POLICY

• For generative tasks, the optimal threshold is the smallest feasible one: oracle error decreases as cost rises with the threshold, while non-oracle error may be nonmonotone.In both settings, monotone cost makes the leftmost feasible threshold optimal, including when the non-oracle feasible set is disjoint.Source: §3.1 · p. 4

3.2 DISCRIMINATIVE TASKS - MULTI-THRESHOLDS POLICY

• For classification, class-specific thresholds equalize boundary error across interior classes in the oracle case and excess error (student minus master) in the non-oracle case.Classes may instead always accept with the student or always defer to the master when their errors fall entirely below or above the common level.Source: §3.2 · pp. 5–6

3.3 SUMMARY OF THEORETICAL RESULTS

• Across the four settings, cost monotonicity selects the least costly feasible policy, while multi-threshold optima equalize boundary error or excess boundary error across interior classes.The relevant equalized quantity is student error in the oracle case and student-minus-master error in the non-oracle case.Source: §3.3 · p. 6

4 EXPERIMENTAL RESULTS

• Experiments use an 80% optimization and 20% test split, with a computationally restricted subset of 8,000 samples per dataset.Source: §4 · p. 7, §4.1 · p. 7

4.1 GENERATIVE TASKS

• For generation, the single-threshold policy uses average predicted-token log probability as confidence and cosine similarity against a threshold to determine correctness.The similarity is computed between predicted and true answer embeddings.Source: §4.1 · p. 7
• On SQUAD, the policy maintains mLLM performance with a 26.66% inference-cost reduction; on WMT, it preserves best-possible performance with maximal cost savings.Results compare GEMMA3-1B as sLLM with GEMMA3-4B as mLLM in the non-oracle setting.Source: Table 1 · p. 7, §4.1 · p. 7

4.2 DISCRIMINATIVE TASKS

• For zero-shot classification, the multi-threshold policy optimizes a separate threshold per category using maximum softmax probability, with the model constrained to output one class-label token.Source: §4.2 · p. 7
• On SST-2, the policy matches full-deferral accuracy with ∼82% cost savings; AGNEWS and EMOTION reach full-deferral performance with ∼16% and ∼68% savings.On FAKENEWS, performance degrades ∼2% from full-deferral while inference cost falls ∼45%.Source: §4.2 · p. 8

5 CONCLUSION

• The proposed policies reduce inference costs while meeting a specified performance budget, using class-specific thresholds for discriminative tasks and one threshold for generative tasks.The authors identify learning thresholds with reinforcement learning, optimizing richer network topologies, and improving generative confidence estimation as future directions.Source: §5 · p. 8

A ADDITIONAL RELATED WORK

• Model-chaining work escalates queries from smaller to larger models using confidence or quality criteria, while predictive classification can estimate answer quality without querying the next model.Examples include multi-level inference, FrugalGPT [Chen et al., 2024], and classification deferral [Chen et al., 2020].Source: A ADDITIONAL RELATED WORK · p. 15
• Model routing assigns each query to one model before inference, using learned difficulty, reward, predicted cost and performance, or benchmark-driven classifiers.The passage contrasts these approaches with hybrid classification systems that train the router and both models from scratch, a costly approach for LLMs.Source: A ADDITIONAL RELATED WORK · p. 15

B PROOFS

• Appendix B provides extended versions of the proofs presented in the main text.Source: B PROOFS · p. 16

B.1 SINGLE-THRESHOLD ORACLE CASE

• In the oracle case, the optimal threshold is the left endpoint of the feasible interval, where the performance constraint is tight and the threshold is unique if the constraint function strictly decreases.Ties can occur when the constraint function is flat at the budget level; if the error rate is zero near threshold zero, multiple thresholds may have the same cost.Source: §B.1 · p. 16

B.2 SINGLE-THRESHOLD NON-ORACLE CASE

• In the non-oracle case, the smallest feasible threshold exists and minimizes cost; if threshold zero is infeasible, the performance constraint is tight at the optimum.If threshold zero is feasible, it is optimal.Source: §B.2 · p. 17

B.3 MULTI-THRESHOLDS ORACLE CASE

• At an oracle optimum, all interior classes share a common boundary error level, with the error constraint tight unless zero thresholds are already feasible.The tight constraint uniquely determines the common level except when coordinates saturate at 0 or 1.Source: §B.3 · p. 17
• The oracle common threshold can be computed by solving the tight-budget equation g(t) = b, numerically or analytically when assumptions permit.Bisection or Brent’s method can solve the equation; the Lagrangian gives the equivalent stationarity rule εk(θk*) = cm/λ*.Source: §B.3 · pp. 17–18
• Oracle thresholds need not be unique when zero-density plateaus preserve cost and error, or when multiple classes have flat segments at the optimal error level.Both cases permit nontrivial threshold changes without changing the relevant cost and error totals.Source: §B.3 · p. 17

B.4 MULTI-THRESHOLDS NON-ORACLE CASE

• In the non-oracle case, interior classes with positive marginal risk share a common level; classes whose marginal risk never reaches it can saturate at the lowest-cost boundary.The passage identifies that boundary as θ*=0 for classes whose marginal risk is everywhere at most the common level.Source: §B.4 · p. 19

B.5 MULTI-THRESHOLDS POLICY JUSTIFICATION

• In an illustrative balanced binary-class example, two thresholds achieve higher average accuracy than a single threshold.The example compares the average across equally sized classes; Figure 2 illustrates the accuracy–cost trade-off.Source: §B.5 · p. 19, Figure 2 · p. 20

C OPTIMALITY OF MULTI-THRESHOLDS POLICY OVER GENERALIZED POLICIES

• Under a mild monotone-risk assumption on the small model’s confidence, Bayes-optimal deferral among arbitrary policies reduces to per-class single thresholds on calibrated uncertainty.The section gives an alternative proof using Lagrangian relaxation and Bayes-optimal decision theory.Source: C OPTIMALITY OF MULTI-THRESHOLDS… · p. 19

C.1 THE RISK–COST REGION AND CONVEXITY

• Allowing randomization makes the attainable risk–cost region the convex hull of deterministic policies, and Pareto-optimal policies lie on its lower-left boundary.Randomizing between two policies yields convex combinations of their error and cost.Source: §C.1 · p. 21

C.2 LAGRANGIAN FORMULATION

• For every Pareto-optimal risk–cost point, some nonnegative Lagrange multiplier makes its policy a minimizer of the Lagrangian objective.Thus, characterizing the constrained problem’s solutions reduces to finding minimizers for each multiplier.Source: §C.2 · p. 22

C.3 BAYES-OPTIMAL POLICY

• For fixed λ>0, accept the small model when its conditional risk is no greater than deferral risk, and defer otherwise; this policy globally minimizes Jλ.The risks are rs(k, β) = cs + λεs,k(β) and rm = cs + cm + λ¯εm,k.Source: §C.3 · pp. 22–23
• The rule accepts whenever the small model’s excess error over the master is at most the cost-normalized margin cm/λ.The excess error is ∆k(β) = εs,k(β) −¯εm,k; otherwise, the policy defers and pays the additional cost.Source: §C.3 · p. 23

C.4 FROM LAGRANGIAN MINIMIZERS BACK TO THE CONSTRAINED PROBLEM

• Every Pareto-optimal risk–cost vector corresponds to a Bayes rule for some λ≥0, accepting when excess error is at most cm/λ.Source: §C.4 · p. 24
• Any constrained optimum on the error–cost frontier must accept the small model exactly on the corresponding excess-error acceptance set.The set is determined by some λ≥0.Source: §C.4 · p. 24

C.5 MONOTONE PER-CLASS ERROR CURVES

• Under the monotone-risk assumption, the excess-error acceptance region can be represented by a single confidence threshold per class.Source: §C.4 · p. 24
• The monotone-risk property requires each class’s error probability to be essentially non-increasing as confidence β increases.Equivalently, within each class, the probability that the prediction is correct is non-decreasing in β.Source: §C.5 · p. 24

C.6 FROM ERROR THRESHOLD TO CONFIDENCE-THRESHOLD

• With non-increasing risk, each sublevel set is empty, all of [0, 1], or an interval extending from a left endpoint to 1.The endpoint may or may not belong to the interval.Source: §C.6 · p. 24
• Applying this structure to the Bayes acceptance condition yields an optimal per-class single-threshold policy in confidence.The threshold’s endpoint decision does not affect Bayes risk; with right-continuity or continuity there, the acceptance set is exactly [θk(λ), 1].Source: §C.6 · p. 25

C.7 LAGRANGIAN IN TERMS OF θ

• At an interior optimum, all classes equalize the small-model excess error at their thresholds, and per-class threshold policies are optimal among measurable score-gating policies.In the oracle case, this reduces to equalizing boundary errors at cm/λ.Source: §C.7 · p. 25

C.8 COMPUTATION STRATEGY VIA LAGRANGE

• For fixed λ, invert each class’s error curve at the common excess-error level, then find λ* by solving the scalar error-budget equation with one-dimensional root-finding.Under mild regularity, error is monotone in λ; thresholds can collapse to always accept or always defer when the equation has no interior solution.Source: §C.8 · p. 26

D OPTIMALITY OF SINGLE-THRESHOLD POLICY OVER GENERALIZED POLICIES

• With arbitrary measurable score-based policies allowed, a mild monotone-risk assumption still makes the Bayes-optimal deferral policy a single threshold.Source: D OPTIMALITY OF SINGLE-THRESHOLD… · p. 26

D.1 LAGRANGIAN FORMULATION

• The Lagrangian minimizes expected loss, charging small-model cost plus λ-weighted error on acceptance and both model costs plus λ-weighted master-model error on deferral.Source: §D.1 · p. 27

D.2 BAYES-OPTIMAL DEFERRAL RULE IN THE BINARY CASE

• The Bayes-optimal rule chooses action 0 when the small-model risk is no greater than the other action’s risk, and action 1 when it is greater.Source: §D.2 · p. 27
• The rule’s excess conditional loss is nonnegative for every score, and is strictly positive when the actions differ and their conditional risks are unequal.Source: §D.2 · p. 28
• In error terms, the rule selects action 0 when εs(β) − ¯εm ≤ cm/λ and action 1 when it exceeds cm/λ.Source: §D.2 · p. 28

D.3 MONOTONE-RISK ASSUMPTION AND THRESHOLD STRUCTURE IN CONFIDENCE SCORE

• The threshold result assumes small-model error is non-increasing in confidence score, equivalently making correctness non-decreasing and excess error non-increasing.Source: §D.3 · p. 28
• Under this assumption, each nontrivial sublevel set of excess error is an interval [θ(t), 1], rather than an arbitrary set of scores.Source: §D.3 · p. 28
• Consequently, for each λ, the Bayes-optimal policy is a single-threshold policy in confidence score β.Source: §D.3 · p. 28

E CONTINUITY OF THE CONSTRAINT FUNCTION

• For single-threshold policies, the oracle error function g is absolutely continuous on [0, 1].Source: E CONTINUITY OF THE CONSTRAINT FUNCTION · p. 29
• The continuity proof uses that a bounded measurable error function times an integrable density is integrable, so its tail integral is absolutely continuous.Source: E CONTINUITY OF THE CONSTRAINT FUNCTION · p. 29
• The single-threshold non-oracle error function gNO is absolutely continuous because it adds a constant-weighted density CDF to g.Source: E CONTINUITY OF THE CONSTRAINT FUNCTION · p. 29
• For K-class policies with normalized class weights, each class-wise error contribution is absolutely continuous and the total oracle and non-oracle maps are continuous on [0, 1]^K.Source: E CONTINUITY OF THE CONSTRAINT FUNCTION · pp. 29–30
• The multi-threshold error maps have coordinate-wise partial derivatives almost everywhere.Source: E CONTINUITY OF THE CONSTRAINT FUNCTION · p. 30

F DETAILS ABOUT THE EXPERIMENTAL SETUP

• Section F provides additional information about the datasets, algorithms, and supplementary results.Source: F DETAILS ABOUT THE EXPERIMENTAL SETUP · p. 30

F.1 DATASETS

• Five datasets cover binary and multi-class classification plus generative question answering: SST-2, FAKENEWS, AGNEWS, EMOTION, and SQUAD.The four classification datasets cover two binary and two multi-class tasks.Source: §F.1 · p. 30
• SST-2 uses movie-review sentences with human-annotated positive or negative sentiment labels for sentence-level binary classification.Source: §F.1 · p. 30
• FAKENEWS pairs English-language news articles with veracity labels distinguishing misleading from factual content.Source: §F.1 · p. 30
• EMOTION classifies tweets into six emotion categories, while AGNEWS assigns news articles one of four topic labels.The EMOTION categories are anger, fear, joy, love, sadness, and surprise; AGNEWS categories are World, Sports, Business, and Science/Technology.Source: §F.1 · p. 30
• SQUAD pairs questions with Wikipedia context passages, but the experiments use free-form answers rather than restricting responses to spans or fixed labels.Source: §F.1 · p. 30

F.2 FRUGALGPT AND HYBRIDLLM DETAILS

• The auxiliary router uses DistilBERT as a binary classifier, with training data constructed following FrugalGPT [Chen et al., 2024] and Hybrid-LLM [Ding et al., 2024].Source: §F.2 · pp. 30–31
• For HybridLLM, the router labels an sLLM answer as 1 when its BARTScore against the ground truth exceeds the mLLM answer’s score, and 0 otherwise.Source: §F.2 · p. 31
• For FrugalGPT, the router labels an sLLM answer 1 when it matches the ground truth and 0 otherwise; both routers train up to 20 epochs with patience 3 and η = 2 × 10−5.Source: §F.2 · p. 31

F.4 MONTE CARLO SAMPLING

• Monte Carlo evaluation approximates theoretical cost and error expectations with sample averages of per-example indicator functions.The samples are observed tuples drawn from the available examples.Source: §F.4 · p. 31

F.4.1 Single-threshold policy

• The single-threshold policy accepts the small model’s prediction when its confidence β_i ≥ θ and otherwise defers to the larger model.Source: §F.4.1 · p. 31
• For generative tasks, exact-match error indicators can be difficult to use; continuous indicators based on a performance metric are proposed as future work.The metric could be cosine similarity between answer embeddings.Source: §F.4.1 · p. 31
• The single-threshold cost model charges a constant small-model cost and adds the larger-model cost when deferring.Empirical deferral rate determines the additional cost, while coverage is one minus that rate.Source: §F.4.1 · p. 32

F.4.2 Multi-thresholds policy

• The multi-threshold policy assigns each predicted class k its own threshold θ_k to decide whether to accept or defer.Source: §F.4.2 · p. 32
• The multi-threshold policy uses the same cost-expression form as the single-threshold policy.Source: §F.4.2 · p. 33

F.4.3 Finite-sample effects and empirical feasibility

• Finite-sample estimates produce stepwise error functions, so the best feasible threshold may leave slack; Monte Carlo evaluation weights predicted classes by their empirical frequencies.This discretization does not contradict population-level boundary optimality when the constraint is active and the error curve is continuous.Source: §F.4.3 · p. 33

F.5 DETAILS ABOUT IMPLEMENTED ALGORITHMS

• On finite data, error and deferral change only at observed confidence levels, reducing constrained cost minimization to a finite set of sample-induced candidates.The implemented algorithms minimize expected per-sample cost, equivalently the deferral rate, subject to a global error budget.Source: §F.5 · p. 34

F.5.1 Single-threshold policy

• For a single threshold, the minimum-cost feasible choice is the smallest threshold meeting the error budget; sorting enables an O(N log N) scan, while oracle selection admits O(N) worst-case time.The linear-time guarantee relies on oracle feasibility; an exact worst-case O(N) algorithm is generally unavailable in the non-oracle case without additional assumptions.Source: §F.5.1 · pp. 34–35

F.5.2 Multi-thresholds policy

• For multi-threshold policies, each class offers options that defer a low-confidence prefix, and the optimization minimizes total deferrals under an error-count budget of B = ⌊bN⌋.Tie-aware per-class sorting constructs at most U ≤ N + K options overall.Source: §F.5.2 · p. 35
• The multi-threshold problem is a multi-choice knapsack instance with an exact dynamic-programming solution over the integer error budget, though its worst-case runtime can be heavy for large datasets.In practice, the implementation keeps only Pareto-optimal states.Source: §F.5.2 · pp. 35, 36, 38
• A Lagrangian sweep selects each class’s option independently for each multiplier and retains the feasible candidate with fewest deferrals, but it is not guaranteed globally optimal.Its runtime is O(N log N) + O(LU) and memory is O(U).Source: §F.5.2 · p. 39

F.6 ALTERNATIVE APPROXIMATION METHODS

• Alternative nonparametric estimation uses KDE or histograms for the confidence density and binning or kernel regression for conditional error, then numerically integrates the estimated functions.The passage identifies KDE and kernel regression as examples.Source: §F.6 · pp. 39–40
• For moderate-to-large samples with recognizable marginal shape, a parametric density can simplify computation when goodness-of-fit is satisfactory; simple analytic distribution and conditional functions may yield closed-form integrals, while complex or nonlinear forms may call for numerical integration.Source: §F.6 · p. 40

G SUFFICIENCY OF THE CONFIDENCE SCORE

• A score β is sufficient relative to gating information J when the expected risk difference Δ(j) depends only on β(j), making equal-score queries have equal deferral trade-offs.Source: G SUFFICIENCY OF THE CONFIDENCE SCORE · p. 42
• If β is sufficient and Δ is non-increasing in β, the optimal policy accepts above a cutoff—or always or never—and no policy using all of J can do better.The acceptance set is either empty, all of R, or a half-line [β*, ∞).Source: G SUFFICIENCY OF THE CONFIDENCE SCORE · p. 42
• In preliminary experiments, adding input embeddings or logits statistics to β did not significantly improve prediction of which LLM was more accurate, so the authors used β alone for deferral.They compared cross-fitted Gradient Boosting Regressors using mean squared error and a permutation test; Table 18 shows results for two model pairs.Source: G SUFFICIENCY OF THE CONFIDENCE SCORE · p. 43, Table 18 · p. 45