Investigating Linear Probe Robustness to Linguistic Register, Medical Specialty, and Corpus Shifts in Medical QA

AI summary

Overview

Research area: Natural language processing / LLM interpretability, specifically linear probing of hidden states for factual-error detection in medical question answering.

Technical level: Intermediate. The paper assumes familiarity with transformer hidden states, linear classifiers, and AUROC, but its experimental logic is straightforwardly described.

Scope (one sentence): The paper isolates three kinds of input change — writing style, medical specialty, and source dataset — to determine which ones break a linear "truth direction" probe trained on a medical LLM's internal representations.

What This Paper Is About

Recent work proposes that a simple linear classifier reading a frozen LLM's hidden states can detect factual errors from a single forward pass, implying that true and false statements separate along a stable geometric direction. Existing studies disagree about whether this signal survives changes in input, but their cross-dataset tests change several things at once, making failures impossible to attribute. This paper builds a controlled medical QA benchmark that varies writing style, specialty, and corpus one at a time, then measures how much each shift degrades probe performance.

Key Contributions

1. A three-axis decomposition of probe transfer. The authors separate register (textbook, patient-facing, clinical note, colloquial), specialty (15 clinical sub-domains via S-MedQA), and corpus (MedQA, MedMCQA, MMLU-medical, MedRedQA), holding two axes fixed while varying the third. This makes cross-dataset failures attributable rather than confounded.

2. A controlled rewriting benchmark. 500 MedQA facts, each rewritten into four registers with both correct and incorrect answers preserved, yielding 4,000 matched variants. The benchmark is regenerated with a second, independent LLM rewriter (Gemini 3 Flash) and validated against human-written patient questions from MedRedQA. Benchmark, code, prompts, and rubric are publicly released.

3. A quantified robustness profile. The truth direction survives register shifts (mean ≈0.095 AUROC drop) and specialty shifts (≈0.031), but degrades unevenly across corpora: 0.12 on MMLU-medical and 0.21 on MedMCQA. Question format is explicitly ruled out as the cause of the MedMCQA break.

4. A characterization of the signal itself. A parameter-free difference-of-means probe matches the learned logistic probe, and a nonlinear MLP adds almost nothing (+0.037 AUROC, inside confidence intervals in 15 of 16 conditions), confirming the signal is dominantly linear. The paper also shows the probe beats entropy and self-consistency baselines but only ties a self-evaluation baseline on larger models, and that raw probe scores are severely miscalibrated (mean ECE 0.341).

Main Findings

• Register is largely harmless. Training on textbook-style questions and testing on patient-facing, clinical-note, and colloquial rewrites of the same held-out facts costs only 0.095 AUROC on average. Clinical-note shorthand is the hardest register in all four LLMs; patient-facing prose is the easiest in three of four.

• The register effect replicates. A different LLM rewriter produces a mean gap of 0.079 versus 0.095 for the original, with the same per-register ordering. Dropping low-fidelity wrong-answer rewrites moves the gap by at most 0.0023 AUROC.

• Partial fix via training coverage. Training on a balanced mix of all four registers recovers an average of +0.061 AUROC on the non-textbook registers, concentrated on the hardest clinical-note register (+0.122 on Gemma-3-4B, +0.127 on Llama-3-8B), at a textbook cost of ≤0.03.

• Specialty barely matters, but per-specialty difficulty varies widely. Holding out 8 of 15 specialties costs only 0.031 AUROC on average, yet absolute AUROC spans 15 points across specialties: 0.710 for Rheumatology to 0.859 for Nephrology. Specialties with distinctive diagnostic vocabulary perform best; those sharing overlapping symptom vocabulary perform worst.

• Rare diseases are easier, not harder. Rare-disease AUROC meets or exceeds common-disease AUROC in all 16 (LLM, register) conditions, with four cells showing strictly disjoint confidence intervals favoring rare diseases. The authors suggest distinctive low-frequency vocabulary occupies more separable regions of representation space.

• Corpus shift is the real failure mode. The same frozen probe loses 0.21 AUROC transferring MedQA→MedMCQA and 0.12 on MMLU-medical. The failure is mutual: a MedMCQA-trained probe reaches only 0.607 AUROC transferring back to MedQA.

• Question format is exonerated. Reformatting 100 MedMCQA items into MedQA-style vignettes with answers held verbatim leaves AUROC essentially unchanged (0.559 vs 0.561). A worst-case additive combination of specialty and register effects (0.095–0.149) still falls short of the 0.21 MedMCQA drop. By elimination the authors attribute the break to how MedMCQA items are constructed, most plausibly distractor authorship and the adversarial pressure those distractors exert.

• The signal is linear, not a learned artifact. A label-permutation probe collapses to 0.488 AUROC (chance), confirming correctness is being read. The unregularized difference-of-means probe matches logistic regression overall and is uniformly better on the hardest clinical-note register (+0.070 AUROC), suggesting L2 regularization overweights the textbook training distribution.

• The probe beats output-only baselines, but not decisively on larger models. It outperforms self-consistency by 6–11 AUROC points and token entropy by 13–31 points (entropy is near or below chance, dropping to 0.46 for Llama-3-8B). Against P(True) — reading the model's own "Yes" token probability — the probe wins by +0.05 on Gemma-2-2B and +0.04 on Gemma-3-4B but only ties on Llama-3-8B and Qwen2.5-7B. P(True) also has a tighter register profile (0.063 vs 0.095). The probe's remaining advantage is inference-side: it reads the question position before any answer token is generated.

• Raw probe scores cannot be read as probabilities. Mean ECE is 0.341, with no condition below 0.19. Platt scaling more than halves it to 0.135; isotonic regression does worse (0.234), likely from overfitting at this sample size. All but one condition remain above the conventional 0.05 well-calibrated threshold even after Platt scaling. The Youden-optimal cutoff swings between 0.04 and 1.00 across registers for one model.

• Negation remains a confirmed weakness. The probe systematically rejects correct sentences containing explicit negation, replicating prior brittleness findings on that specific perturbation class.

Methodology in Plain English

The researchers take 500 MedQA exam questions, each paired with its correct answer and one incorrect distractor. Every pair is rewritten by Claude Sonnet 4.5 into four writing styles: textbook prose (the style medical benchmarks already use), patient-facing lay language, telegraphic clinical-note shorthand, and casual internet colloquialisms. Both the correct and incorrect versions are rewritten, yielding 4,000 matched items where the underlying medical fact is identical and only the surface form differs.

Each item is presented to the LLM as a yes/no prompt: "Is this answer medically correct? Respond with Yes or No." The researchers capture the model's internal hidden-state vector at the final question token, at every second layer, and train a small logistic regression classifier to separate correct from incorrect answers. Crucially, splits are done at the fact level: a fact's correct and incorrect variants always land on the same side, so the probe is never tested on a fact it saw in training.

To isolate specialty, they use existing S-MedQA annotations to train on seven clinical specialties and test on eight held-out ones, with register fixed at textbook in both halves. To isolate corpus, they apply the frozen MedQA-trained probe directly to three other datasets without retraining.

They run four open-weight models spanning 2B to 8B parameters and three architecture families: Gemma-2-2B-it, Gemma-3-4B-it, Qwen2.5-7B-Instruct, and Llama-3-8B-Instruct. They compare against a nonlinear MLP probe, a parameter-free difference-of-means probe, a label-permutation sanity check, and output-only baselines (token entropy, self-consistency, and the model's own yes-token probability). All AUROC numbers carry 95% bootstrap confidence intervals computed by resampling at the fact level.

Why This Matters

Impact on research. The probing literature has been split between groups reporting that truth directions are stable geometric properties of LLM representations and groups reporting that probes are brittle artifacts. This paper shows both are partly right, and explains why the disagreement persisted: prior cross-dataset evaluations confounded register, topic, format, and corpus construction. Decomposing those axes converts a confusing literature into a specific, actionable finding — the break is tied to how a dataset was built, not to medical knowledge itself. The three-axis method is applicable well beyond multiple-choice medical QA.

Real-world applications:

• Clinical decision support. A probe monitoring a medical LLM's internal states could flag potentially wrong dosages or contraindications before they reach a clinician — provided the tool is trained on the deployment corpus and never ported zero-shot.
• Patient-facing health apps. The probe retained usable signal on genuine human-written patient questions (0.695 AUROC on Qwen2.5-7B, 0.680 on Llama-3-8B), supporting lay-language safety monitoring, though weaker on smaller models.
• Electronic health record integrations. Clinical-note shorthand is the hardest register to transfer to, and mixed-register training partially closes that specific gap, making it the most relevant practical fix.
• Regulatory and assurance pipelines. The paper's warning that raw probe scores are badly calibrated and that per-specialty AUROC varies by 15 points bears directly on how such detectors would need to be documented and validated.

Industry relevance. For LLM vendors and health-AI companies, the findings set clear deployment boundaries: linear probes are viable within-distribution safety tools but not standalone cross-corpus factuality checks. Post-hoc calibration is presented as a mandatory safety layer, not an optional refinement. Any organization reporting a single dataset-level AUROC for a factuality monitor would be overstating its reliability.

Future Directions

1. Extend the three-axis decomposition beyond multiple-choice QA to long-form generation and real clinical text, where answer boundaries, negation, and hedging are far more complex.
2. Isolate exactly which construction property of MedMCQA breaks the probe. The authors point to distractor authorship and adversarial pressure, but the mechanism is inferred by elimination rather than identified directly.
3. Build calibration into probe design rather than applying it post hoc, given that Platt scaling leaves almost every condition above the conventional well-calibrated threshold. The authors cite PING as a reference approach that trains jointly for prediction and calibration.
4. Test how the picture changes with scale and architecture. The 2B–8B models show no clean scaling pattern in register gaps, and the weaker results on smaller models for human-written patient text leave open whether larger models recover more of the signal.

Target Audience

This paper is most useful to researchers working on LLM interpretability and probing methods, and to applied NLP practitioners building factuality or hallucination detectors for high-stakes domains. Medical informatics researchers evaluating LLM safety will find the three-axis experimental design directly reusable. ML engineers and product teams deploying medical QA systems should read it for the deployment boundaries it establishes — specifically the prohibition on zero-shot cross-corpus probe transfer and the necessity of post-hoc calibration. Regulatory and clinical validation teams will benefit from the per-specialty and calibration analyses, which show why single aggregate performance numbers are misleading for safety-critical medical tools.

Authors’ abstract

Linear classifiers trained on hidden states of a large language model (LLM), linear probes, can flag factual errors from a single forward pass. Geometrically, that implies that true and false statements separate along a stable direction in hidden state space, i.e., the truth direction. Prior work disagrees on whether this generalises across input shifts, but the disagreement is hard to interpret because cross-dataset probe transfer experiments confound several kinds of input change at once. We isolate three such variables in medical question-answering (QA): writing style (register), domain (medical specialty), and corpus (dataset). We build a benchmark using 500 MedQA entries, each rewritten into four styles (textbook, patient, clinical note, colloquial), annotated with clinical specialty, and grouped with two other exam corpora, MedMCQA and MMLU-medical, for cross-dataset evaluation. Probing four open-weight LLMs (2--8B), we find that the truth direction is largely robust to writing style (mean $Δ_\text{register} \approx 0.10$ AUROC on held-out facts) and to medical specialty ($Δ_\text{specialty} \approx 0.03$), but degrades unevenly across corpora: by $0.12$ AUROC on MMLU-medical and by $0.21$ on MedMCQA, roughly twice the register gap. The register result replicates with a second generator and carries over to human-written patient questions. The truth direction is therefore largely stable within the medical domain but breaks under some corpus shifts, and question format does not explain the break, which suggests that the signal a linear probe recovers is partly bound to dataset structure rather than to medical knowledge alone.