PaperlayerView exact revision on arXiv

Source-linked AI summary

Route, Don't Fix: Regime-Dependent Decoding Correction and a Trajectory-Gated Router for Reliable Clinical LLM Answer Selection

Zeyu Dong, Benjamin Wang, Joyee W. Jin

arXiv:2609.14825v1cs.CLcs.AIcs.LG

TL;DR

Clinical question answering needs hallucination mitigation, yet fixed decoding correction helps on a truthfulness stress test while yielding negligible gains on clinical multiple-choice benchmarks. ALTAS uses terminal entropy and late-layer linearity to route questions between greedy decoding and trajectory correction, retaining TruthfulQA gains while keeping evaluated clinical benchmarks within a one-point do-no-harm band.

  • Problem

    Fixed inference-time correction can help on TruthfulQA yet be indistinguishable from greedy on MedQA and PubMedQA, while clinical mitigations such as retrieval and verification can add infrastructure, latency, or model calls.

    Source: §I · p. 1

  • Method

    ALTAS uses terminal entropy and late-layer logit-trajectory linearity from one forward pass to choose per question between greedy decoding and late-layer trajectory correction, without training a probe, classifier, or head.

    Source: §I · p. 1, A. Trajectory correction · p. 2

  • Results

    ALTAS gains 9.5 and 8.3 points on TruthfulQA, while its deltas on the three clinical benchmarks are non-significant and within a one-point do-no-harm band at every scale.

    Source: B. Main results · p. 3

  • Takeaways & Limitations

    On questions it does not flag, ALTAS stands down and leaves the output identical to the base model’s.

    Source: §VI · p. 5

    R2-cutoff transfer is demonstrated across two sizes of one model family, while transfer to a different architecture fails; calibration probes use only 15 prompts.

    Source: §VII · p. 5

Abstract

from arXiv · show

Large language models (LLMs) are often deemed unsafe for clinical question answering because of their tendency to hallucinate. Retrieval augmentation, fine-tuning, and external verifiers require new infrastructure that clinical governance must approve and may add latency or extra model calls. Inference-time correction uses the model's internal logit signals, but a fixed transformation need not suit every question. A corrector that improves accuracy by about ten percentage points on a truthfulness stress test yields negligible gains on clinical multiple-choice benchmarks, where instruction tuning concentrates output probability on one answer and leaves low terminal entropy. We introduce ALTAS, which reads terminal entropy and late-layer linearity ($R^2$) from one forward pass to choose per question between greedy decoding and late-layer trajectory correction. No classifier, probe, or head is trained; the router operates on candidate-answer logits and adds 6.5% latency overhead. Applied to every question, the correction improves TruthfulQA over greedy at 3B and 8B by 11.4 and 10.0 percentage points, respectively ($p<10^{-10}$). Gated per question, ALTAS retains gains of 8.3 to 9.5 percentage points while keeping MedQA, PubMedQA, and MedHallu within a one-percentage-point do-no-harm band, with no statistically significant differences from greedy. The method passes verification sweeps over frozen thresholds, the scoring rule, and the domain label.

I. INTRODUCTION

Clinical QA needs hallucination mitigation that avoids added infrastructure, but fixed decoding corrections do not suit every question. ALTAS routes per question, extending decoding interventions from token-level transformations to choices conditioned on the model’s trajectory.

I. INTRODUCTION

  • A confident unsupported clinical answer can be dangerous, while retrieval, self-verification, and supervised probes add infrastructure, latency, calls, or governance burdens,,,.Source: §I · p. 1
  • Fixed correction helps on TruthfulQA but is indistinguishable from greedy on MedQA and PubMedQA, where instruction tuning concentrates probability on one answer.Source: §I · p. 1
  • ALTAS reads terminal entropy and late-layer trajectory linearity in one forward pass to route each question between greedy decoding and correction without training a probe, classifier, or head.Source: §I · p. 1

II. RELATED WORK

  • DoLa contrasts mature and premature layers, while DeLTa extrapolates the trajectory toward a converged distribution; both are token-level maps without question-level conditioning.Source: §II · p. 2
  • Sampling-based verification uses agreement or self-checking,, while ITI steers hidden states using labeled probe data; these approaches add generations or require supervision and per-model fitting.Source: §II · p. 2
  • ALTAS routes among decoding strategies within one fixed model from its late-layer logit geometry, unlike coarser model-selection routers such as RouteLLM.The paper characterizes per-question, intra-model routing as unexplored.Source: §II · p. 2

III. METHOD

ALTAS combines a late-layer trajectory correction with an entropy-based control so correction strength falls as the model’s terminal distribution concentrates. The design also fixes the trajectory anchor, fitting window, and weights.

III. METHOD

  • ALTAS reads terminal entropy, model-level late-layer linearity, and a clinical/general tag; its specified thresholds are 0.50 for R2, 0.5 nats for entropy, and 0.65 for cosine acceptance.The first two thresholds were calibrated on 15 prompts; correction weights αc and αe are both 0.3.Source: §III · p. 2

A. Trajectory correction

  • The correction combines a DoLa-style contrast against layer 7 with a one-virtual-layer extrapolation of per-token least-squares fits over layers 14:N for the 200 highest-logit candidates.Source: A. Trajectory correction · p. 2
  • The entropy gate g_t reduces correction strength as terminal entropy falls, and at zero the operator reduces to greedy scoring so a resolved answer is not overwritten.Source: A. Trajectory correction · p. 2

B. Gate readings

A single prompt pass supplies ALTAS’s entropy and late-layer linearity readings; scalar thresholds and a domain tag determine whether it applies correction or falls back to greedy decoding.

B. Gate readings

  • From half depth onward, ALTAS fits the 50 highest-logit candidates’ final-position logits against layer depth and averages the fits into late-layer R2, while also reading terminal entropy.Source: B. Gate readings · p. 2

C. Decision rule

  • The R2 cutoff of 0.50 and entropy threshold of 0.5 nats were calibrated on 15 TruthfulQA prompts without labels or gold-answer information.A separate cosine threshold of 0.65 governs scoring.Source: C. Decision rule · p. 2
  • ALTAS uses scalar comparisons rather than external training, learned probes, or classifiers, and reuses layer logits already computed during decoding.Source: C. Decision rule · pp. 2–3
  • The correction is used only when model-level R2 clears its cutoff, the item is tagged general, and terminal entropy exceeds its threshold; otherwise decoding is greedy and matches the base model.Reported runs use benchmark-supplied clinical labels, with a lexical keyword test as fallback when labels are absent.Source: C. Decision rule · p. 3

D. One R2 cutoff per model family

The model-level R2 cutoff transfers across the two Llama scales but does not establish correction viability across architectures.

D. One R2 cutoff per model family

  • A 15-prompt R2 estimate supports one cutoff across Llama scales, but Qwen2.5-14B’s 0.53 clears it despite the correction being nonviable, limiting the test’s demonstrated transfer to a model family.Source: D. One R2 cutoff per model family · p. 3

E. Cost

ALTAS adds one prefill before generation, with a feature pass costing 6.5% overhead on the greedy branch. The scoring rule uses cosine similarity for free-form answers, while option and polarity matching resolve single-token clinical answers first.

E. Cost

  • A single feature-pass prefill adds 6.5% overhead to the greedy branch; corrected decoding takes 9.78 s, or 1.68× greedy.On one RTX 5060 laptop GPU, the feature-pass median was 377 ms against a greedy median of 5.82 s; peak memory was unchanged within 0.03 GB.Source: E. Cost · p. 3

F. Cosine answer scoring

  • Option and polarity matching resolve single-token clinical answers before cosine similarity is consulted.For these answers, cosine agreement coincides with exact match to the gold answer.Source: F. Cosine answer scoring · p. 3
  • Free-form answers receive graded cosine scores, and unscorable generations count as incorrect rather than being dropped.Accuracy is reported over the full benchmark, with comparisons paired at the question level.Source: F. Cosine answer scoring · p. 3

IV. EXPERIMENTS

Across the evaluated models and benchmarks, ALTAS retains significant TruthfulQA gains at the two Llama scales while remaining within a one-point do-no-harm band on clinical benchmarks. At 14B, correction behavior differs across free-form and single-token scoring, with the router defaulting to greedy on clinical questions.

A. Setup

  • The evaluation covers three instruction-tuned models, four benchmarks, and greedy, DoLa+DeLTa ablation, always-on correction, and ALTAS protocols.Correction and ALTAS are compared with greedy using paired bootstrap and exact McNemar tests at 95%.Source: A. Setup · p. 3

B. Main results

  • ALTAS gains 9.5 and 8.3 points on TruthfulQA at the two Llama scales, while every clinical-benchmark delta is non-significant and within one point at every scale.Always-on correction gains 11.4 and 10.0 points on TruthfulQA at the two Llama scales; all four TruthfulQA gains are highly significant.Source: B. Main results · p. 3
  • At 14B, corrected free-form MedHallu gains of 5.9 and 4.4 coexist with a −0.5 TruthfulQA delta, while ALTAS yields −0.8 on MedHallu, all non-significant.On clinical single-token benchmarks, correction collapses MedQA and PubMedQA to 0.9% and 36.6%; ALTAS sidesteps this by defaulting to greedy.Source: B. Main results · p. 3

C. Regime dependence

At 8B, always-on correction improves TruthfulQA and MedHallu but is inert on clinical multiple-choice, while ALTAS retains TruthfulQA gains and stays within the clinical do-no-harm band. The oracle gap indicates that action selection remains a bottleneck for the gate.

C. Regime dependence

  • At 8B, always-on correction gains 10.0 on TruthfulQA and 11.8 on MedHallu; ALTAS retains an 8.3 TruthfulQA gain and routes every MedHallu item to greedy, retaining 1.0 of the gain.The ALTAS MedHallu delta is non-significant.Source: C. Regime dependence · pp. 3–4
  • Correction is inert on MedQA and PubMedQA, with deltas of 0.6 and 1.0 (ns), while ALTAS stays inside the do-no-harm band on clinical multiple-choice..Source: C. Regime dependence · pp. 3–4
  • The late-layer trajectory separates candidate support in layers 15–28: Georgia keeps gaining while California and South Carolina plateau.The late-window linearity shown is 0.62, the feature read by the gate.Source: Figure 1 · p. 4

D. Oracle gap

  • ALTAS captures 55% of the oracle’s available TruthfulQA gain at 3B and 47% at 8B, and only 15% at 14B despite a 70.3 oracle score.On clinical multiple-choice, ALTAS captures almost none of the per-question oracle gain.Source: D. Oracle gap · p. 4

V. SENSITIVITY ANALYSIS

Sensitivity sweeps find that the entropy threshold lies on a broad plateau, while viability-estimator reliability depends on prompt sampling and the measurement window. Re-scoring and domain-label checks also probe whether the reported conclusions hold under alternative settings.

Sensitivity analysis

  • Any entropy threshold in [0, 1.0] leaves TruthfulQA gains above 5 at both Llama scales, with the frozen threshold on a plateau rather than a peak.Sweeping τH across [0, 3.0] moves routing from near-total coverage to near zero and traces the gain smoothly.Source: §V · p. 4
  • Across 20,000 resamples of 817 logged per-question values, 15-prompt probes at 8B ranged [0.557, 0.606] and never fell below the cutoff; at 3B, [0.454, 0.517], below in 82% of draws.A deployment should use far more than 15 prompts.Source: §V · p. 4
  • Using the top 30% of layers lowers the 14B viability statistic below cutoff: 0.4436 over 50 prompts and 0.4393 across logged items, versus 0.5303 at half depth.The measurement window determines whether the viability test detects the Qwen2.5 collapse.Source: §V · p. 4
  • Cosine rescoring thresholds from 0.50 to 0.90 preserve every sign; ALTAS TruthfulQA gains stay positive and peak near 0.75 at 10.0.The fallback lexical test has precision 0.99 and recall 0.82 overall, but MedHallu recall is 57.4%.Source: §V · p. 4

VI. DISCUSSION

The results vary by regime: ALTAS improves TruthfulQA at 3B and 8B but not significantly at 14B, while it leaves clinical benchmarks close to greedy. The discussion relates this pattern to entropy and verification limits, and identifies routing as an important remaining bottleneck.

VI. DISCUSSION

  • Clinical multiple-choice answers often have concentrated probability on one option, leaving correction little to change and risking overwrite of a correct answer.TruthfulQA’s misconceptions are associated with higher terminal entropy, which gives correction room to work; self-checking is also limited in clinical content.Source: §VI · pp. 4–5
  • TruthfulQA gains about ten points at 3B and 8B, but ALTAS gains only 1.2 points (ns) at 14B; clinical benchmarks stay within one point of greedy.Fig. 2 reports accuracy by model size and ALTAS changes relative to greedy at 8B.Source: Figure 2 · p. 5
  • Qwen2.5-14B clears the linearity cutoff despite correction failure; candidate explanations include architecture-dependent logit-lens decodability and absolute layer indices.Across 48 layers, the correction window opens at 29% depth, versus 50% for Llama-3.2-3B.Source: §VI · p. 5
  • ALTAS routes every MedHallu item to greedy, retaining 1.0 of the always-on correction’s 11.8-point gain (ns).The discussion argues that a better gate, rather than a stronger corrector, is the next advance.Source: §VI · p. 5

VII. LIMITATIONS

The reported boundaries include limited transfer evidence for the linearity cutoff, small calibration probes, and benchmark and scoring constraints. The oracle comparison also leaves substantial available gain uncaptured, while public benchmarks do not substitute for clinical workflow validation.

VII. LIMITATIONS

  • Clinical do-no-harm relies on a benchmark-supplied clinical tag; the lexical fallback recalls 82% of clinical items, and the worst-case bound applies only to these three benchmarks.The reported worst-case change is −0.6 points.Source: §VII · p. 5
  • Calibration probes are small (15 prompts); the consequence is asymmetric, harmless at 8B and one draw from flipping the verdict at 3B.Source: §VII · p. 5
  • Free-form answers are graded with cosine scoring, and the residual risk concerns score magnitude.This limits how directly graded free-form results characterize answer quality.Source: §VII · p. 5
  • A 0.50 linearity cutoff transfers between two Llama sizes but also admits Qwen2.5-14B, where correction fails; ALTAS captures 15% of available TruthfulQA gain there.The figure reports 55% and 47% captured at 3B and 8B, respectively.Source: Figure 3 · p. 6
  • Public benchmarks remain proxies for clinical workflows; ALTAS is not yet a decision-support system and requires external validation.Against a per-question oracle, the gated decoder recovers at most 55% of available gain, placing the open problem in routing.Source: §VII · p. 6

VIII. CONCLUSION

ALTAS routes decoding correction by question because a fixed corrector helps in one regime but can be inert or harmful in others. It retains TruthfulQA gains without significant accuracy reduction on the evaluated clinical benchmarks, though it does not close the per-question oracle gap.

VIII. CONCLUSION

  • ALTAS delivers the TruthfulQA gain without significant accuracy reduction on evaluated clinical benchmarks, and this result holds at every tested threshold.The entropy-and-R2 gate routes per question between greedy decoding and late-layer trajectory correction.Source: §VIII · p. 6
  • A fixed corrector helps on a truthfulness stress test but is inert on clinical multiple-choice and destructive on one architecture.The calibration transfers within a model family.Source: §VIII · p. 6
  • ALTAS does not close the per-question oracle gap, so the next advance would belong to the gate rather than the corrector.Source: §VIII · p. 6
Loading 2609.14825v1…