PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Partially Defer for Sequences

Sahana Rayan, Ambuj Tewari

arXiv:2502.01459v2stat.MEcs.LGstat.ML

TL;DR

Existing L2D rejectors defer whole predictions, which can make deferral costly for long sequences. This paper proposes post-hoc rejectors for partial sequence deferral, and reports better cost-accuracy tradeoffs than whole deferrals across three tasks.

  • Problem

    Existing L2D methods support only complete prediction deferral, so costs can scale with long sequence outputs even when only parts are inaccurate.

    Source: §1 · p. 1

  • Method

    The paper proposes token-level and one-time model-based post-hoc rejectors that interleave machine and expert predictions at granularities suited to expert capabilities.

    Source: §1 · p. 2

  • Results

    Granular model-based rejectors achieve better cost-accuracy tradeoffs than whole-sequence rejectors on Traveling salesman solvers, News summarization, and Weather prediction.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Partial deferral offers a way to improve the cost-effectiveness of learning-to-defer systems, with the demonstrated setting using one expert.

    Source: §7 · p. 9

    The proposed partial-deferral setting uses one expert; extensions to multiple experts and larger outputs remain future work.

    Source: §7 · pp. 9–10

Abstract

from arXiv · show

In the Learning to Defer (L2D) framework, a prediction model can either make a prediction or defer it to an expert, as determined by a rejector. Current L2D methods train the rejector to decide whether to reject the {\em entire prediction}, which is not desirable when the model predicts long sequences. We present an L2D setting for sequence outputs where the system can defer \textit{specific outputs} of the whole model prediction to an expert in an effort to interleave the expert and machine throughout the prediction. We propose two types of model-based post-hoc rejectors for pre-trained predictors: a token-level rejector, which defers specific token predictions to experts with next token prediction capabilities, and a one-time rejector for experts without such abilities, which defers the remaining sequence from a specific point onward. In the experiments, we also empirically demonstrate that such granular deferrals achieve better cost-accuracy tradeoffs than whole deferrals on Traveling salesman solvers, News summarization, and Weather prediction.

1 Introduction

The introduction extends Learning to Defer to sequence outputs by allowing partial deferrals, with rejectors that interleave machine and expert predictions. The paper reports better cost-accuracy tradeoffs than whole-sequence deferral across three tasks.

1 Introduction

  • Existing L2D methods support complete deferral, whose costs scale with sequence output size even when only parts of a prediction are inaccurate.Source: §1 · p. 1
  • Partial deferral can reduce expert effort by limiting work to uncertain sequence parts, and model predictions may provide context for completing deferred parts.Source: §1 · p. 1
  • Token-level deferral suits next-token experts, while one-time deferral lets an expert complete a sequence from a selected point using a partial model-generated subsequence.Source: §1 · p. 2
  • The proposed L2D framework interweaves expert and machine predictions through token-level or one-time deferrals, accommodating different expert capabilities.Source: §1 · p. 2
  • The paper proposes convex surrogate losses with Bayes consistency guarantees for implementable model-based post-hoc rejectors on pre-trained predictors.Source: §1 · p. 2
  • Granular model-based rejectors improve cost-accuracy tradeoffs over whole-sequence rejectors for Traveling salesman solvers, News summarization, and Weather prediction.Source: §1 · p. 2

2 Background

The background distinguishes confidence-based from model-based deferral and frames L2D as minimizing prediction or expert-query cost. It also motivates convex, Bayes-consistent surrogates because direct loss minimization is challenging.

2.1 Learning to Defer

  • Confidence-based deferral rejects predictions below a cost-dependent confidence threshold, whereas model-based deferral trains the rejector jointly with or after the predictor.Source: §2.1 · p. 2
  • L2D’s loss selects either the model’s prediction loss or the cost of deferring to an expert; the related multi-expert formulation requires choosing which expert to query.Source: §2.1 · p. 2

2.2 Consistent Losses

  • Direct minimization of the deferral losses is challenging because of non-convexity and discontinuities, motivating convex surrogate optimization.Source: §2.2 · p. 2
  • The paper introduces surrogates for token-level and one-time deferral and proves they satisfy the stated inequality, establishing Bayes consistency.Source: §2.2 · p. 3

3 Token-level partial deferral

Token-level deferral lets a next-token expert replace selected predictions in an autoregressive sequence, with expert outputs available as context for later predictions. The section develops a convex surrogate and gives consistency and generalization results under stated assumptions.

3 Token-level partial deferral

  • A next-token expert can replace selected predictions, and its supplied tokens can inform subsequent predictions, addressing uncertainty that may propagate from earlier autoregressive steps.The rejector acts on the input and the mixed predictor–expert prefix; training analysis assumes teacher-forced contexts, while experiments are not limited to teacher forcing.Source: §3 · p. 3

3.1 Surrogate Loss

  • A convex surrogate upper-bounds token-level deferral loss up to a scaling factor and is Bayes consistent when expert costs and predictor losses are bounded and the surrogate is classification calibrated.The paper gives logistic and square losses as examples of convex binary surrogates.Source: §3.1 · p. 4

3.2 Generalization Bounds

  • Finite-sample excess-risk bounds relate the empirical surrogate-risk minimizer to token-level loss, under bounded losses and costs and a Lipschitz, classification-calibrated surrogate.The bound depends on the rejector hypothesis class’s Rademacher complexity; Bayes consistency alone is an infinite-data guarantee.Source: §3.2 · pp. 4–5

4 One-time partial deferral

One-time deferral hands the suffix of a sequence to an expert at a selected position, allowing partial deferral even when the expert cannot predict token by token. The section also considers restricting the candidate handoff positions and analyzes a surrogate loss.

4 One-time partial deferral

  • The rejector selects a handoff position: the predictor supplies the prefix and the expert completes the suffix, with the loss accounting for both prefix error and expert-completion cost.Choosing position 1 defers the whole sequence; choosing L + 1 means no deferral.Source: §4 · p. 5

4.1 Influence of large L

  • Allowing a selected set J of handoff positions limits the available choices; every position may be over-resolved because adjacent positions can yield similar accuracy.The paper connects this design to diminishing accuracy gains as the number of experts increases [Verma et al., 2023, Hemmer et al., 2023].Source: §4.1 · p. 5

4.2 Surrogate loss

  • The one-time surrogate is convex, upper-bounds the target loss up to a scale, and is Bayes consistent under bounded costs and losses with a nonzero cost difference for at least one pair of positions.Cross entropy and mean absolute error are given as surrogate examples.Source: §4.2 · pp. 5–6

4.3 Generalization Bounds

  • A generalization bound for the empirical surrogate-risk minimizer holds with bounded costs and losses and a bounded, Lipschitz surrogate, for a fixed predictor.The bound is stated for a finite rejector hypothesis class.Source: §4.3 · p. 6

5 Related Works

The paper distinguishes partial deferral from cascades, query routing, cost-sensitive learning, and selective prediction. It emphasizes finer-grained collaboration and surrogate losses with consistency guarantees.

Related work

  • Unlike cascades that assign each prediction to one model, partial deferral lets two models collaborate on a prediction; its learned token-level rule avoids the repeated expert-confidence queries used by Narasimhan et al. [2024].The paper also describes cascades as progressively escalating harder cases.Source: §5 · p. 6
  • The paper contrasts its more granular routing decisions and stronger statistical guarantees with query-routing losses that lack consistency guarantees.The cited query-routing losses are from Shnitzer et al. [2023], Lu et al. [2023], and Ding et al. [2024].Source: §5 · pp. 6–7
  • Unlike selective classification and regression, which treat abstention as a terminal action with a fixed penalty, L2D accounts for downstream expert-prediction accuracy.The cited selective-classification and regression work includes El-Yaniv and Wiener [2010], Geifman and El-Yaniv [2017, 2019], Shah et al. [2022], Jiang et al. [2020], and Zaoui et al. [2020].Source: §5 · p. 7

6 Experiments

Experiments compare partial and whole-sequence deferral on TSP, XSUM, and multi-step weather prediction, using model-based and confidence-based rejectors. Model-based partial methods outperform whole-sequence baselines, while the effect of candidate handoff-set size varies.

6.1 Tasks

  • The experiments use PointerNet with Gurobi for 50-city TSP, T5-small with T5-base for XSUM, and an LSTM with a simulated human expert for weather prediction.Gurobi completes a predictor-provided partial tour; the weather task predicts six 10-minute observations from two hours of temperature history.Source: §6.1 · p. 7

6.2 Baselines

  • Whole-sequence baselines defer based on aggregate token confidence, while partial-deferral comparisons include token-level and one-time confidence rules and a model-based whole-sequence rejector.The confidence baselines include ChowSum, ChowMean, and ChowQuantile; the model-based baseline predicts whether the expert is more accurate.Source: §6.2 · pp. 7–8

6.4 Comparison with other deferral methods

  • Model-based token-level and one-time deferral outperform whole-sequence baselines, and token-level deferral attains the best AUDC where applicable; one-time gains are modest on TSP but markedly better in XSUM and MWP.The authors attribute modest TSP gains to its non-autoregressive expert, which cannot act on partial subsequences.Source: Table 1 · p. 8, §6.4 · p. 8

6.3 Evaluation

  • AUDC summarizes the loss-versus-deferred-token tradeoff, and reported results average five runs of 100 test instances each.The study also reports percentage improvement over a random rejector; one-time curves use a policy based on the no-deferral score and the best handoff position.Source: §6.3 · p. 8

6.5 Size of J for One time Deferral

  • The best choice is not always to allow every handoff position: using all L + 1 positions is suboptimal, especially in XSUM, while one-time rejectors generally outperform confidence-based counterparts across candidate-set sizes.Figure 2 plots candidate-set size against AUDC improvement over the random baseline for TSP and XSUM.Source: §6.5 · p. 9, Figure 2 · p. 9

7 Discussion

The discussion frames partial rejection as a way to improve L2D cost-effectiveness and identifies extensions beyond the proposed one-expert setting. The accompanying proofs establish surrogate properties for token-level and one-time deferral.

7 Discussion

  • The framework currently covers one expert; the authors identify multiple experts, intermediate deferral granularities, and larger outputs such as protein structures as directions for further study.A rejector could select contiguous sequence parts, extending beyond token-level and one-time deferrals.Source: §7 · pp. 9–10

A.1 Proof of Theorem A.1

  • For token-level deferral, a convex binary-calibrated surrogate loss upper-bounds the target loss up to a scale factor.The proposition’s proof ties convexity to the surrogate and calibration to the scaled upper bound.Source: §A.1 · p. 15

A.2 Proof of Theorem 3.1

  • The proof of Theorem 3.1 assumes bounded positive deferral costs and bounded prediction loss, and relates surrogate and target risks through a non-decreasing concave function.The supplied lemma states these conditions, while the theorem passage introduces the risk inequality.Source: §A.2 · p. 15

A.3 Proof of Theorem 3.2

These appendix proofs support token-level and one-time deferral with calibrated surrogate losses, including consistency and finite-sample risk bounds under stated assumptions. They also establish a convex upper bound for one-time deferral.

A.3 Proof of Theorem 3.2

  • Under the stated bounded-cost and bounded-loss assumptions, a binary classification-calibrated surrogate is Bayes consistent for token-level deferral.The proof invokes binary classification calibration to relate surrogate risk to target risk.Source: §A.3 · p. 17

A.4 Proof of Theorem 3.3

  • For a fixed predictor, the empirical minimizer’s estimation error is bounded with probability 1 − δ when the surrogate is ρ-Lipschitz.The bound is stated under bounded deferral costs and prediction loss.Source: §A.4 · p. 17

A.5 Proof of Theorem 3.4

  • With probability 1 − δ, the empirical surrogate-risk minimizer has an upper bound on excess token-level risk under the stated boundedness and Lipschitz assumptions.The bound includes estimation and approximation terms.Source: §A.5 · p. 18

B.1 Proof of Theorem B.1

  • For one-time deferral, the cross-entropy and mean-absolute-error surrogates are convex upper bounds up to a scale when cmax > 1/(|J| − 1).The proposition states the condition for both surrogate choices.Source: §B.1 · p. 19

B.2 Proof of Theorem B.3

  • Under the stated condition on prediction losses, a multiclass surrogate-risk inequality is established for one-time deferral.The theorem passage states that the inequality holds for any measurable g and predictor h.Source: §B.2 · p. 20

B.3 Proof of Theorem 4.1

The appendix establishes consistency and finite-sample risk bounds for one-time deferral under stated assumptions. It also describes confidence scores and rejector scoring functions used in the baselines.

B.3 Proof of Theorem 4.1

  • For fixed h, cross-entropy and mean-absolute-error surrogates are consistent for one-time deferral under the stated bounded-cost and bounded-loss conditions.The result also requires the stated positive-gap condition.Source: §B.3 · pp. 21–22

B.4 Proof of Theorem 4.2

  • With probability 1 − δ, the empirical minimizer’s surrogate risk has an upper bound under bounded losses, bounded costs, and the stated function-class conditions.The theorem assumes a ρ-Lipschitz surrogate and bounded surrogate values.Source: §B.4 · p. 22

B.5 Proof of Theorem B.6

  • For cross-entropy or mean absolute error, the empirical minimizer also receives a stated high-probability excess-risk bound for one-time deferral.The corollary assumes bounded costs and losses and the function-class conditions inherited from Theorem B.5.Source: §B.5 · p. 24

C Baseline Details

  • The baseline details use token predictive uncertainty scores, including negative log probability and predictive entropy for XSUM and TSP.The text also describes Monte Carlo variance as a confidence measure when the predictor has dropout layers.Source: C Baseline Details · pp. 24–25

D Training Details and Hyperparameters

The paper reports training details and examines how deferral costs affect cost-accuracy tradeoffs. The best reported αL settings differ between TSP and XSUM.

D Training Details and Hyperparameters

  • The experiments use early stopping, and Table 2 contains the hyperparameter settings; the TSP One-Time rejector takes at most 10 minutes to train.TSP graphs have 50 nodes, generated by sampling coordinates from a standard normal distribution.Source: D Training Details and Hyperparameters · p. 25, Table 2 · p. 25

F Effects of αj

  • For TSP, results are best when αL = 0.01; for XSUM, they are best at αL = 0.01 or αL = 0.015.Extreme cost values often produce poor cost-accuracy tradeoffs.Source: F Effects of αj · p. 26

G AUDC and Percent Improvement for Size of J experiment

Tables 3 and 4 report AUDC and percentage improvement over a random baseline across sizes of J for TSP and XSUM. The supplied passages identify the measures and scope but do not provide the table values.

G AUDC and Percent Improvement for Size of J experiment

H Token-level Rejector Ablation

The token-level rejector ablation compares MLP and LSTM architectures and three training strategies on XSUM; the selected model is an LSTM trained without teacher forcing.

H Token-level Rejector Ablation

  • The token-level rejector follows a classifier-chain approach, with each rejector using earlier rejection predictions as inputs.The chain mirrors multi-label classification, where each classifier receives the feature vector and preceding labels.Source: H Token-level Rejector Ablation · p. 28
  • Teacher forcing uses correct rejection decisions during training, whereas scheduled sampling sometimes uses predicted decisions to reduce dependence on ground-truth rules.The supplied passages describe teacher forcing as potentially causing train-test distribution shift and error propagation through the chain.Source: H Token-level Rejector Ablation · p. 28
  • On XSUM, LSTMs outperform MLPs, teacher forcing adversely affects performance, and the selected rejector is an LSTM trained without teacher forcing.The reported AUDC scores are in Table 5, but their values are not included in the supplied passages.Source: H Token-level Rejector Ablation · p. 28, Figure 5 · p. 29, Table 5 · p. 29
  • Training compares teacher forcing, no teacher forcing, and scheduled sampling for the token-level rejector on XSUM.Scheduled sampling exponentially decays the teacher-forcing probability by 0.95 per epoch after five warmup epochs, stopping at 0.5.Source: H Token-level Rejector Ablation · p. 28

Related papers

Loading 2502.01459v2…