PaperlayerView exact revision on arXiv

Source-linked AI summary

Trajectory-Guided Steering: Mitigating Spurious Feature Bias in Clinical Classification

Fardin Ahsan Sakib, Ziwei Zhu, Ozlem Uzuner

arXiv:2610.22618v1cs.CL

TL;DR

Clinical predictions can rely on spurious features, and existing inference-time corrections may require feature-specific counterfactual data. TGS uses layer-wise prediction trajectories to select whether and how to correct each document; on i2b2 it fixes errors while breaking few correct predictions and is net-positive compared with baselines.

  • Problem

    Clinical predictions can change with irrelevant input features, while existing inference-time corrections may require counterfactual data or knowledge of the feature driving errors.

    Source: §1 · p. 1

  • Method

    TGS uses answer-token preferences across layers to decide from each model forward pass whether and in which direction to correct, without counterfactual data or feature knowledge.

    Source: §1 · pp. 1–2

  • Results

    On i2b2, TGS fixes 17.7% of prediction errors while breaking only 1.3% of correct predictions, with a larger net gain than any baseline.

    Source: §1 · p. 2, §5.1 · p. 6

  • Takeaways & Limitations

    TGS is net-positive across three model families and four clinical and biomedical datasets, and its corrections make predictions more consistent under gender swaps.

    Source: §7 · p. 8

    Experiments use 7–8B parameter instruction-tuned models in English on biomedical datasets, so generalization to larger models, other languages, or non-biomedical domains remains untested.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large language models rely on spurious features in clinical classification: swapping a single gender token in clinical notes alters gender-irrelevant diagnoses in up to 53.5% of cases. Existing inference-time corrections either need counterfactual data built around the specific spurious feature or behave unpredictably. We observe that how a model's preference between answers evolves across layers, its prediction trajectory, indicates for each document whether a correction is likely to help. We propose Trajectory-Guided Steering (TGS), which decides from the trajectory alone whether to correct each prediction and in which direction, with no counterfactual data and no knowledge of the spurious feature. On 16 clinical classification tasks from i2b2, TGS fixes 17.7% of errors while breaking only 1.3% of correct predictions and stays net-positive across all tested hyperparameter settings; all baselines are net-negative or near zero. The same corrections reduce reliance on the spurious feature: on the most affected diseases, gender-swap inconsistency falls sharply (hypertension 53.5% -> 38.0%, diabetes 35.1% -> 21.5%). Applying the correction to every document instead raises breaks on correct positive predictions from 0.1% to 6.4%, so per-document selection is essential. TGS also transfers to three model families and three other clinical and biomedical datasets.

1 Introduction

Clinical predictions can change when a gender token changes despite unchanged clinical evidence, while existing corrections may require retraining, counterfactual data, or feature knowledge. TGS uses each document’s prediction trajectory to choose whether and how to correct, with strong i2b2 results and reduced gender-swap inconsistency.

1 Introduction

  • Gender-token swaps change predictions for gender-irrelevant diseases in 9.7% to 53.5% of cases, despite identical clinical evidence.Such changes are errors because the clinical evidence alone determines the correct answer in this evaluation.Source: §1 · p. 1
  • Existing approaches may require retraining, labels, counterfactual training pairs, or advance knowledge of which feature drives errors, while prompt changes can have unpredictable effects.Clinical records may be restricted and counterfactual data requires expert annotation; changing a correct prediction is itself a harm.Source: §1 · p. 1
  • TGS uses answer preferences across layers to choose whether to correct each prediction and in which direction, without counterfactual data or knowledge of the spurious feature.It derives this decision from the model’s own forward pass and adjusts answer-token logits when the trajectory supports a correction.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • On i2b2, TGS fixes 17.7% of prediction errors while breaking only 1.3% of correct predictions, outperforming baselines and remaining net-positive across tested model families and datasets.The introduction also reports that steering and prompting baselines are net-negative or near zero, and that per-document selection is essential for preserving correct predictions.Source: §1 · p. 2
  • TGS reduces gender-swap inconsistency most on diseases where gender sways the model most, including hypertension, where inconsistency falls from 53.5% to 38.0%.The introduction also reports transfer across three model families and four clinical and biomedical datasets.Source: §1 · p. 2

2 Prediction Trajectories Reveal Three Trajectory Shapes

The section investigates whether layer-by-layer answer preferences can guide selective inference-time correction. It describes how prediction trajectories are measured, distinguishes three trajectory shapes, and maps each shape to a correction action.

2.1 Prediction Trajectories

  • A prediction trajectory tracks the answer-token logit margin across layers, with its direction indicating the preferred answer and its size measuring conviction.The analysis uses the logit lens to project each layer’s hidden state through the unembedding matrix; trajectories are oriented so the model’s preferred answer points upward.Source: §2.1 · pp. 2–3

2.3 From Shapes to a Correction Rule

  • On i2b2, trajectories fall into growing conviction, static conviction, and internal conflict; each shape implies reinforcing, opposing, or withholding correction, respectively.This mapping is applied per document and does not require distinguishing correct from incorrect predictions.Source: §2.2 · p. 3, §2.3 · p. 3

2.2 Three Trajectory Shapes

  • Growing-conviction trajectories overlap for correct and incorrect predictions, yet reinforcing their shared preferred direction can recover errors while most correct predictions keep their answer.In this shape, early and late layers agree and conviction grows substantially.Source: §2.2 · p. 3
  • With static conviction, correct predictions continue gaining conviction while incorrect predictions do not, making the unstrengthened preference a cue for errors that could be safely opposed.Early and late layers agree, but conviction growth is limited overall.Source: §2.2 · p. 3
  • When early and late layers conflict, correct and incorrect trajectories overlap and no consistent correction direction exists, so the prediction is left unchanged.The same qualitative shapes and outcome overlaps appear in Qwen3 8B, suggesting they are not an artifact of a single architecture.Source: §2.2 · p. 3

3 Method

TGS uses each document’s layerwise answer-preference trajectory to decide whether to correct a prediction and in which direction. Its two-pass design extracts the trajectory, routes the document to an action, and selectively biases answer logits.

3 Method

  • TGS formalizes per-document inference-time correction without counterfactual data or prior knowledge of the spurious feature, while aiming to leave correct predictions intact.The binary formulation uses positive and negative answer tokens and generalizes to more than two classes.Source: §3 · p. 3

3.1 Trajectory Extraction

  • A first forward pass projects each layer’s final-token hidden state to answer logits, whose positive-minus-negative margin tracks the model’s preference across layers.The logit-lens trajectory uses an adaptive steering layer selected at the first significant binary-entropy drop, together with the final-layer margin.Source: §3.1 · pp. 3–4

3.2 Trajectory-Guided Routing

  • Trajectory shape routes documents to reinforce a growing preference, oppose static positive conviction, or leave predictions unchanged when early and late layers conflict.The conservative routing signal encodes whether to correct and in which direction; static negative early margins also remain unchanged.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 Steered Generation

  • For documents selected for correction, a second pass adds opposing signed biases to the positive and negative answer-token logits; all other logits and unselected generations remain unchanged.The correction magnitude is controlled by α.Source: §3.3 · pp. 4–5

3.4 Hyperparameters and Runtime

  • The correction strength α is set per model and dataset, while the conviction-growth threshold τ is robust across its tested range; held-out selection gives the same conclusions.Trajectory extraction adds one forward pass per document, and only documents selected for correction receive a modified generation pass; the formulation extends to K > 2 classes.Source: §3.4 · p. 5

4 Experimental Setup

The experiments evaluate TGS across clinical and biomedical classification settings, multiple models, and comparisons with prompt-based and activation-steering baselines. Fix, break, and net-gain measures assess correction outcomes, with dataset-specific counting and exclusions.

4.1 Datasets

  • The evaluation spans i2b2 disease classification, n2c2 trial eligibility, SHAC drug-use status, and PubMedQA biomedical questions.i2b2 covers 16 diseases; n2c2 covers 13 eligibility criteria, and both use predictions under two gender conditions.Source: §4.1 · p. 5

4.2 Models

  • Experiments use 7–8B instruction-tuned models, with LLaMA 3.1 8B Instruct as the primary model and Mistral 7B and Qwen3 8B for cross-model generalization.Source: §4.2 · p. 5

4.3 Baselines

  • Baselines include three training-data-free prompt modifications and two activation-steering methods that require counterfactual training pairs.CAA and FairSteer use 30% of documents for training, and their best universal configurations are selected by sweeping test-split hyperparameters.Source: §4.3 · p. 5

4.4 Evaluation

  • Fix rate is fixes divided by errors, break rate is breaks divided by correct predictions, and net gain is total fixes minus total breaks.For i2b2 and n2c2, gender-condition predictions count independently; original non-answers are excluded, steered non-answers count as incorrect, and TGS is also evaluated on the baselines’ approximately 70% test subset.Source: §4.4 · p. 6

5 Results

TGS improves i2b2 predictions while limiting damage to correct ones, reduces gender-swap inconsistency, and transfers across models and datasets. Its gains remain positive across tested hyperparameters and prompt variations, with per-document routing protecting correct predictions.

5.1 TGS Outperforms All Baselines

  • TGS fixes 17.7% of i2b2 prediction errors while breaking 1.3% of correct predictions, outperforming baselines that are net-negative or only marginally positive.The comparison covers the full test set and a shared approximately 70% subset.Source: §5.1 · p. 6, Table 1 · p. 6
  • CoT reaches 35.7% fix rate but 47.5% break rate; on the shared test subset, CAA and FairSteer achieve net +3 and +19 versus TGS’s +146.Prompt strategies are net-negative, while Figure 3 shows TGS’s curve remaining in the low-break region as correction strength rises.Source: §5.1 · p. 6, Figure 3 · p. 7
  • Gender-swap inconsistency drops from 25.1% to 21.5% across 16 diseases, with hypertension falling from 53.5% to 38.0% and diabetes from 35.1% to 21.5%.TGS fixes 19.8% of errors that flip with gender, versus 4.9% of errors that persist under both genders.Source: §5.1 · p. 6

5.2 TGS Transfers Across Models and Datasets

  • TGS transfers to n2c2 2018, SHAC, and PubMedQA; on SHAC it tests substance-use terminology and remains net-positive on both binary and three-class tasks.SHAC inconsistency falls from 31.7% to 16.8%; PubMedQA achieves net +35.Source: §5.2 · pp. 6–7
  • Across models, TGS is net-positive on 41 of 48 i2b2 model-disease combinations, and the same τ transfers without retuning while α adjusts to model logit scale.The reported cross-model results include three model families; no disease is consistently negative across models.Source: §5.2 · pp. 6–7, Table 2 · p. 7

5.3 TGS Is Robust to Its Hyperparameters and Prompt Wording

  • Every tested i2b2 (α, τ) configuration is net-positive, with net gain peaking at α = 0.5 and remaining between +184 and +209 across τ settings.TGS remains net-positive under a reworded prompt, unlike prompt baselines whose outcomes vary widely.Source: §5.3 · p. 7
  • Held-out tuning preserves the conclusions, and a paraphrased i2b2 prompt raises net gain from +201 to +234 while reducing gender-swap inconsistency from 21.2% to 16.1%.On SHAC and PubMedQA, held-out tuning is net-positive in all 200 random splits.Source: §5.3 · pp. 7–8
  • Removing trajectory-based routing raises true-positive breaks from 0.1% to 6.4%, while the routed system leaves 40% of i2b2 documents unchanged.A constant logit bias is net-negative at every tested strength and in both directions.Source: §5.3 · p. 8

6 Related Work

Related work addresses spurious correlations through training-time mitigation, supervised activation steering, and layer-wise prediction methods. TGS instead derives a per-document correction decision from the model’s own trajectory.

Spurious correlations and shortcut learning

  • Training-time approaches to spurious correlations often require retraining and labels for the feature or distinct environments, whereas TGS leaves the model unchanged and corrects at inference time.Clinical models may rely on demographic or substance-related surface cues rather than clinical evidence.Source: §6 · p. 8

Activation steering

  • Activation-steering methods use external supervision such as contrastive pairs, labeled examples, or trained classifiers; TGS derives its correction signal from the current input’s forward pass.FairSteer conditionally applies steering when a trained bias classifier fires.Source: §6 · p. 8

Logit lens and prediction trajectories

  • Layer-wise methods have used intermediate predictions to alter outputs, but those methods act on every input while TGS uses trajectory shape to decide whether and how to correct each document.TGS’s trajectories span layers of an autoregressive model, unlike diffusion-model trajectories over decoding steps.Source: §6 · p. 8

Logit-level correction

  • Unlike contextual calibration’s same correction for every input, TGS adjusts answer-token logits per document and chooses whether and which direction to correct using the model’s trajectory.Source: §6 · p. 8

7 Conclusion

The paper concludes that logit-lens prediction trajectories can guide selective inference-time correction without counterfactual data, knowledge of the spurious feature, or retraining. TGS limits breaks on i2b2 and is net-positive across the tested model and dataset settings.

7 Conclusion

  • TGS uses prediction trajectories read with the logit lens to decide whether and in which direction to correct each clinical classification prediction.The correction uses no counterfactual data, no knowledge of the spurious feature, and no retraining.Source: §7 · p. 8
  • On i2b2, TGS fixes 17.7% of errors while breaking 1.3% of correct predictions, and its corrections concentrate on errors driven by the spurious feature.The paper reports net-positive results across three model families and four clinical and biomedical datasets.Source: §7 · p. 8
  • Larger multi-class tasks and non-medical domains are identified as natural next settings for evaluation.Source: §7 · p. 8

Limitations

The method’s scope and generalization remain bounded by its answer format, tested models and datasets, and assumptions about intermediate predictions.

Limitations

  • TGS requires a fixed set of candidate answer tokens, covering classification and QA but not open-ended generation.Source: Limitations · p. 8
  • The hyperparameters α and τ require a small grid search for each model-dataset combination.Source: Limitations · p. 9
  • Experiments cover English biomedical datasets and 7–8B instruction-tuned models; generalization to larger models, other languages, and non-biomedical domains remains untested.The trajectory assumption has been validated for transformer-based LLMs, not other architectures.Source: Limitations · p. 9

Ethics Statement

The paper identifies limits on interpreting its gender-swap probe and recommends further validation before clinical deployment. It selects a steering layer adaptively, with tested threshold and offset changes producing only minor differences.

Ethics Statement

  • Gender swapping is only a diagnostic probe, TGS does not eliminate spurious-feature errors, and further validation is recommended before clinical deployment.The clinical experiments use de-identified patient data and were conducted under data-use agreements and IRB protocols.Source: Ethics Statement · p. 9

A Steering Layer Selection

  • The steering layer is selected per document by finding an early layer where binary entropy drops substantially, then choosing a layer slightly before that transition.The search uses layers [10, L−4], and selects ℓ∗ = max(10, ℓ′−2) after the entropy drop exceeds 50% of the maximum observed drop.Source: A Steering Layer Selection · p. 11
  • Thresholds from 30% to 70% and offsets from 1 to 3 produce only minor changes, supporting robustness to these selection choices.Across i2b2, median selected layers were 14 for LLaMA, 14 for Mistral, and 10 for Qwen3.Source: A Steering Layer Selection · p. 11

B Prediction Inconsistency

TGS reduces gender-swap inconsistency across most i2b2 diseases, and also lowers inconsistency associated with alcohol and smoking mentions on SHAC. Its correction success differs between gender-dependent and persistent errors.

B Prediction Inconsistency

  • Gender-swap inconsistency measures the share of documents whose prediction changes when only the gender token is swapped; Table 3 reports rates across 16 diseases.The table is sorted by inconsistency rate.Source: Table 3 · p. 11, B Prediction Inconsistency · p. 11
  • 25.1% to 21.5%: TGS reduces overall gender-swap inconsistency, lowering it in 14 of 16 i2b2 diseases.The two diseases with increased inconsistency are also the two most net-negative in the fix/break evaluation.Source: Table 4 · p. 11, B Prediction Inconsistency · p. 11
  • 19.8% vs. 4.9%: TGS fixes this share of gender-dependent errors and errors persisting under both genders, respectively.A uniform push toward the correct answer fixes both error types at nearly the same rates, 13.7% vs. 14.8%.Source: B Prediction Inconsistency · p. 11
  • 31.7% to 16.8%: On SHAC, inconsistency associated with alcohol and smoking mentions falls after TGS.Here inconsistency measures how often removing those mentions from a note changes the prediction.Source: B Prediction Inconsistency · p. 11

C Per-Disease Results

On i2b2, TGS is net-positive across most model-disease combinations, while per-disease outcomes vary with the prevalence of errors in the original predictions.

C Per-Disease Results

  • 41 of 48 model-disease combinations are net-positive with TGS.Diseases with more original prediction errors, such as Diabetes and Hypertension, show the largest net gains; diseases with few errors yield smaller or negative outcomes.Source: C Per-Disease Results · pp. 11–12
  • Table 5 reports per-disease i2b2 results for LLaMA at α=0.5 and τ=0.5, with Fix% computed over errors and Break% over correct predictions.The disease-level pattern varies with original error prevalence.Source: C Per-Disease Results · p. 11, Table 5 · p. 12

D Hyperparameter Grids

The appendices report hyperparameter and dataset extensions, output-level controls, and implementation details. Results include positive outcomes on n2c2 and a three-class SHAC task, while fixed nonzero logit biases are net-negative.

D Hyperparameter Grids

  • On PubMedQA, all 16 tested configurations are net-positive; on SHAC, configurations with α ≥0.75 are net-positive (+21 to +42), while those with α ≤0.50 are net-negative at every τ.The i2b2 sensitivity analysis is reported in Table 6, and SHAC and PubMedQA grids in Tables 7 and 8.Source: D Hyperparameter Grids · p. 12, Table 6 · p. 12, Table 7 · p. 12

E n2c2 2018 Cohort Selection

  • On n2c2 cohort selection, TGS is net-positive (+51; 6.6% fix, 1.2% break) and reduces gender-swap inconsistency on 10 of 13 criteria (37.2% →33.9%).The corpus contains 288 patients annotated for 13 trial-eligibility criteria.Source: E n2c2 2018 Cohort Selection · p. 12

F Output-Level Controls

  • Every tested nonzero constant logit bias is net-negative on i2b2; choosing its strength on the full test set or a 30% split selects no bias.Constrained decoding leaves existing Yes/No predictions unchanged and changes only outputs that were neither Yes nor No.Source: F Output-Level Controls · pp. 12–13

G Multi-Class Extension

  • In a three-class SHAC extension, TGS remains net-positive (net +29; 27.8% fix, 5.8% break), while broader multi-class study remains future work.The positive class is split into current and past, alongside negative.Source: G Multi-Class Extension · p. 13

H Extended Related Work

  • The related-work discussion contrasts paired-prompt activation steering, learned bias-signature detection, layer-wise prediction methods, and fixed logit calibration with TGS’s approach.It notes that the constant-bias control is net-negative at every tested strength.Source: H Extended Related Work · p. 13

I Prompt Templates

  • For the prompt used in the reported result, TGS fixes 35.0% of errors and breaks 3.2% of correct predictions; these are the two reported outcome rates.Source: I Prompt Templates · p. 14

I.1 Prompt-Based Baselines

  • The prompt-based baselines alter system instructions or the final-answer instruction, while retaining the user message unless otherwise noted.The variants include clinical-evidence and objective-evidence instructions, plus a step-by-step reasoning instruction.Source: §I.1 · p. 14

J Computational Resources

  • Experiments use three instruction-tuned models and add one forward pass per document for the trajectory; documents receiving a nonzero correction also undergo a modified generation pass.The reported experiments required approximately 100 GPU hours.Source: J Computational Resources · p. 14

Related papers

Loading 2610.22618v1…