PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nishant Mishra, Ameen Abu-Hanna, Iacer Calixto

arXiv:2609.01361v1cs.CLcs.LG

TL;DR

The paper asks whether linear-probe truth directions generalize across medical QA input shifts, a question complicated by confounded cross-dataset evaluations. It isolates register, specialty, and corpus changes using controlled benchmarks and cross-corpus tests, finding strong within-domain robustness but uneven corpus-transfer degradation. The results support truth probes as within-dataset monitoring tools while leaving their safety use outside exam-style QA unresolved.

  • Problem

    Prior work disagrees about whether truth directions are stable, while cross-dataset tests confound register, specialty, wrong-answer construction, and corpus difficulty.

  • Method

    The study probes four 2–8B open-weight LLMs using 500 MedQA facts rewritten into four registers, specialty-held-out splits, and transfer to other medical QA corpora.

  • Results

    Truth directions are largely robust to register and specialty but lose 0.12 AUROC on MMLU-medical and 0.21 on MedMCQA under corpus transfer.

  • Takeaways & Limitations

    The signal transfers across writing style and specialty within medicine but depends partly on corpus construction, making probes viable for within-dataset monitoring rather than standalone cross-corpus safety checks.

  • Takeaways & Limitations

    The evidence is limited to exam-style multiple-choice medical QA and does not test free-form clinical generation with partially correct, context-dependent, or long errors.

Abstract

from arXiv · show

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.

1 Introduction

The paper investigates whether medical QA truth directions generalize across register, specialty, and corpus shifts without conflating these changes. It finds robustness to register and specialty but uneven corpus-transfer degradation, supporting within-dataset monitoring rather than standalone cross-corpus safety checking.

  • Motivation: Linear probes use frozen LLM hidden states to separate true from false statements along a learned truth direction.The probe flags factual incorrectness from one forward pass by learning a separating hyperplane in hidden-state space.
  • Research gap: Prior transfer studies conflate writing style, topic, wrong-answer authorship, and dataset difficulty, obscuring why probes fail to generalize.The paper isolates these factors by measuring transfer along register, specialty, and corpus axes separately.
  • Approach: The benchmark rewrites 500 MedQA questions into four registers, annotates specialties, and evaluates transfer to other medical QA corpora.The controlled benchmark contains 4,000 variants, while specialty annotations and additional datasets support separate transfer tests.
  • Findings: ∆specialty ≈0.03 AUROC and ∆register ≈0.10 AUROC indicate robustness to medical specialty and substantial robustness to writing style.The register result replicates with a second generator and transfers to human-written patient questions.
  • Findings: Corpus transfer decreases AUROC by 0.12 on MMLU-medical and 0.21 on MedMCQA, with the larger drop not explained by question format.The findings suggest the recovered signal depends partly on corpus construction rather than medical content alone.
  • Implication: The probe is a viable within-dataset medical QA monitoring tool but not yet a standalone safety check that transfers across corpora.Raw probe scores are also severely miscalibrated.

2 Methodology

The methodology isolates register, specialty, and corpus shifts by varying one input axis while holding the others fixed. It uses matched prompts and hidden states from a fixed token position to train linear probes across medical QA conditions.

  • Controlled shifts: The study varies register, specialty, or corpus while holding the other two axes fixed to attribute changes to a single shift.Register changes surface form, specialty uses held-out clinical specialties, and corpus changes the medical exam dataset.
  • Register design: Textbook, patient-facing, clinical-note, and colloquial registers preserve meaning while differing in vocabulary, syntax, shorthand, and informality.Clinical notes use fragments and abbreviations, whereas patient-facing text uses lay vocabulary and full sentences.
  • Probing protocol: Each register uses the same chat-templated yes/no prompt, with the hidden state at the last question token supplied to an L2-regularised logistic probe.This fixed completion format reduces prompt-format confounding while probing every second layer.

3 Experimental Setup

The experiments use fact-level splits, controlled MedQA rewrites, specialty-held-out evaluation, cross-corpus transfer, and four open-weight LLMs. AUROC is the primary ranking metric, supplemented by calibration and probe-variant analyses.

  • Data: MedQA supplies USMLE-style multiple-choice items, while MedMCQA supplies Indian medical entrance questions with different distractor styles, specialty mixes, and surface conventions.These differences make MedMCQA a cross-corpus stability test rather than a register-only comparison.
  • Data: 500 MedQA facts become 4,000 controlled variants by preserving correct and incorrect polarities across four generated registers.The benchmark is generated with Claude Sonnet 4.5, and MedMCQA contributes 500 validation facts for transfer.
  • Evaluation sets: Evaluation-only sets include 500 MMLU-medical items, 100 reformatted MedMCQA items, and 84 human-written MedRedQA questions.The reformatted set changes question format while holding answers and distractors fixed, enabling a format comparison.
  • Splits: All experiments keep both polarities of each fact on the same side of fact-level splits, with textbook held-out data serving as the in-distribution baseline.This prevents paired variants from crossing train and test partitions.
  • Specialty transfer: 351 sampled MedQA facts map to 15 specialties, which are divided into seven training and eight held-out specialties across five random partitions.Specialty is fixed to textbook register in both halves, isolating the specialty axis.
  • Models: Four open-weight instruction-tuned LLMs spanning three architecture families range from 2B to 8B parameters.The models are Gemma-2-2B-it, Gemma-3-4B-it, Qwen2.5-7B-Instruct, and Llama-3-8B-Instruct.
  • Metrics: AUROC measures ranking quality, while ECE, accuracy, and F1 provide calibration and threshold-based evaluations.AUROC is interpreted as the probability that a correct example receives a higher score than an incorrect one.
  • Probe experiments: Register transfer uses an L2-regularised logistic probe, label-permutation controls, mixed-register training, a second generator, and low-fidelity rewrite filtering.Corpus transfer applies the MedQA-trained probe without retraining at the same selected layer, while specialty transfer evaluates held-out specialties.

4 Results and Discussion

Across controlled medical QA shifts, the truth direction transfers well across registers and specialties but degrades unevenly across corpora. Linear probes remain useful within datasets, outperform output-only uncertainty baselines, and require calibration because raw scores are poorly calibrated.

  • Register transfer: −0.095 AUROC is the mean drop when textbook-trained probes transfer to patient-facing, clinical-note, and colloquial registers.The gap remains below cross-corpus degradation despite substantial surface-form changes.
  • Register transfer: 0.079 AUROC is the replicated mean register gap with Gemini-generated variants, versus 0.095 for Sonnet.Filtering low-fidelity wrong-answer rewrites changes the mean gap by at most 0.0023 AUROC.
  • Specialty transfer: −0.031 AUROC is the mean specialty-transfer drop across five partitions, roughly seven times smaller than the cross-corpus drop.Absolute AUROC still varies from 0.710 in Rheumatology to 0.859 in Nephrology, separating transfer robustness from specialty difficulty.
  • Specialty transfer: Rare-disease AUROC matches or exceeds common-disease AUROC in all 16 LLM-register conditions.The largest reported gain is 0.880 rare versus 0.748 common for Qwen2.5-7B on the patient register.
  • Corpus transfer: −0.21 AUROC is the mean MedQA-to-MedMCQA drop, compared with −0.12 on MMLU-medical and roughly −0.095 within MedQA registers.Reformatting MedMCQA into MedQA-style vignettes leaves AUROC nearly unchanged, so question format does not explain the larger MedMCQA loss.
  • Probe structure: The nonlinear MLP adds only +0.037 AUROC on average, while the difference-of-means probe matches the logistic probe's mean register drop.These results support a predominantly linear signal in the hidden states.
  • Calibration: 0.341 is the mean raw ECE across 16 LLM-register conditions, falling to 0.135 after Platt scaling but remaining above 0.05 in all but one condition.High AUROC therefore does not make raw probe scores interpretable as correctness probabilities.

5 Related Work

Prior work provides evidence both for a stable, linearly recoverable truth direction and for brittleness under perturbations, task-family shifts, and dataset changes. This paper positions its medical QA study as a controlled separation of register, specialty, and corpus effects, alongside related medical probing and calibration work.

  • Linear truth: Marks and Tegmark and Bürger et al. report truth directions or subspaces that survive causal interventions and negation-related evaluations.These studies motivate treating truthfulness as a geometric property of hidden representations.
  • Linear truth: Levinstein and Herrmann and Haller et al. report that probes can rely on superficial training-data resemblance and break under paraphrasing or negation.The prior literature therefore disagrees about whether probe signals are stable representation-level properties.
  • Task specificity: Orgad et al. find transfer within task families but failure across task families, while this paper separates task-internal medical QA shifts into register, specialty, and corpus.The controlled decomposition addresses confounding among writing style, topic, wrong-answer authorship, and dataset difficulty.
  • Medical probing and benchmarks: PING and related medical probing studies recover clinical knowledge or adverse-drug-reaction information, with PING jointly training prediction and calibration.This work treats PING as a strong calibration reference for medical probe evaluation.

6 Conclusions and Future Work

The study finds that medical truth signals transfer across writing styles and specialties but degrade variably across corpora, with question format not explaining the MedMCQA drop. The findings support within-corpus monitoring while identifying corpus construction and calibration as deployment boundaries.

  • The truth signal transfers across human-written patient questions and variable corpora, while mixed-register training partially closes the register-transfer gap.
  • The predominantly linear signal matches a parameter-free difference-of-means probe, outperforms entropy and self-consistency, and matches self-evaluation.
  • The findings reconcile prior debate by supporting within-corpus geometric stability and across-corpus brittleness.
  • Linear probes are viable within-dataset monitoring tools but should not transfer zero-shot across corpora, and raw scores require post-hoc calibration.
  • The three-axis decomposition should be extended to long-form generation, real clinical text, larger models, item-construction properties, and truth-subspace alternatives.

7 Limitations

The evidence is limited to English exam-style multiple-choice medical QA and controlled or expert-authored distractors, while several aspects of real clinical generation and probe scaling remain untested.

  • All corpora are exam-style multiple-choice datasets, so the claims concern exam-format medical QA.
  • Free-form clinical generation is untested, including errors that are partially correct, context-dependent, or long.
  • The human-written patient set is small at n=84 and differs from MedQA on more than register, preventing isolated register-effect quantification.
  • The register-rewriting pipeline relies primarily on one instruction-tuned generator, although Gemini replication and fidelity filtering reduce but do not eliminate concern.
  • The four probed models span only 2–8B parameters, so linearity and brittleness may change at larger scales.
  • Probe performance on spontaneous deployed-LLM fabrications is not directly measured because wrong answers are designed multiple-choice distractors.

Ethics Statement

The paper uses public, research-licensed exam data and models, with AI systems serving as probed models, rewrite generators, and judging assistants. The benchmark and evaluation pipeline are released for research use.

  • The dataset uses public MedQA items and 84 publicly posted, de-identified MedRedQA questions, with no patient data involved beyond those questions.
  • The patient-facing and colloquial registers are model-generated rewrites rather than representations of real patient speech.
  • The probed LLMs, rewrite generators, and judging models occupy distinct roles in the study.
  • Linear probes are frozen-model diagnostic classifiers that output an estimate of whether a question–answer pair is correct without modifying the LLM.
  • The generator pilot selected Sonnet with a pooled selection score of 0.938, versus 0.925 for Gemini and 0.872 for GPT-4o-mini.
  • Figure 5 compares per-register AUROC for Sonnet-generated solid variants and Gemini-generated hatched variants, with the same directional result.

B.3 Robustness check: Gemini generator

The Gemini robustness check reproduces the main register finding: generator choice leaves the modest cross-register gap broadly stable. Additional controls show that the result is not explained by rewrite fidelity, layer selection, specialty partitioning, question format, or corpus-specific phrasing.

  • Generator replication: 0.079 AUROC was the mean register gap with Gemini, versus 0.095 with Sonnet, keeping the two-generator means within 0.02.The qualitative pattern also matched across generators, although Llama-3-8B was a low outlier on the 100-fact Gemini subset.
  • Label-noise sensitivity: The register result persists across fidelity filters, with the largest mean-gap change only 0.0023 AUROC.Removing 40% of wrong-answer variants changed the gap by 0.0008, while removing almost none changed it by 0.0004.
  • Stylistic distance: Perplexity ordering matches probe-drop ordering across registers, but stylistic distance remains much larger than the modest AUROC degradation.Mean perplexity rises from 8.6 for textbook to 62.0 for clinical note, while mean probe drops range from 0 to 0.144.
  • Protocol robustness: The textbook-blind layer rule increases the register gap to 0.160, yet the qualitative conclusion that register shifts are smaller than corpus shifts remains unchanged.Per-cell and best-mean rules yield mean gaps of 0.095 and 0.080, both below the 0.21 cross-corpus drop.
  • Specialty and position controls: Specialty-transfer estimates remain stable across five random partitions, with first-answer-token Δspecialty of 0.035 versus 0.031 at the last question token.The same per-model ordering is preserved across token positions, and split variability is small.
  • Corpus controls: Reformatting MedMCQA items into MedQA-style vignettes leaves AUROC nearly unchanged at 0.559 versus 0.561, while rewritten items reach 0.727.These controls indicate that question format alone does not recover the cross-corpus drop, whereas matching rewriting conventions produces much higher transfer without isolating one factor.

C.9 Training register: matrix, pairs, and mixed training

Register-trained probes transfer broadly across the four styles, and mixing registers yields little overall benefit. Mixed training mainly recovers performance on the hardest clinical-note register, with a small textbook cost.

  • Single-register matrix: The three prose registers transfer to one another at 0.70 to 0.79, while panel means range from 0.725 for clinical note to 0.755 for patient.A textbook-trained probe has panel mean 0.741, between clinical note and colloquial, with patient marginally highest.
  • Pair and mixed training: The best two-register pair reaches panel mean 0.732, compared with 0.729 for all four registers and 0.728 for patient alone.These results indicate that combining registers provides little overall improvement when evaluated on shared held-out facts.
  • Register-specific recovery: Mixed-register training recovers performance on all three non-textbook registers, with the largest gains on clinical note for Gemma-3-4B (+0.122) and Llama-3-8B (+0.127).The recovery comes at a textbook cost of at most 0.033 AUROC under the matched protocol.

D Human-written patient questions (MedRedQA)

The MedQA-trained probe was tested on human-written patient questions from MedRedQA, and larger models retained most register-transfer performance. Interpretation is limited because claims were generator-constructed and the evaluation also introduces a corpus and topic shift.

  • Construction: 84 of 100 MedRedQA items passed both medical-fidelity judges, yielding 168 evaluated claim variants.Both judges agreed on wrong-is-false labels at 0.99 agreement.
  • Results: Qwen2.5-7B scored 0.695 and Llama-3-8B scored 0.680 on human-written patient questions, within the corresponding Sonnet-register confidence intervals.The two Gemma models lost more, and Gemma-2-2B’s interval included chance.
  • Limitations: The human-question test does not isolate register because correct and wrong claims were generator-constructed and MedRedQA differs from MedQA in corpus and topic mix.With n=84, confidence intervals are wide.
  • Conclusion: The authors conclude that textbook-trained directions are recoverable on genuine patient writing for the more capable models.The MedRedQA text itself is not redistributed; code and item identifiers are released.

F Accuracy and threshold metrics

Raw probe scores rank correctness reasonably but are poorly calibrated and operationally unstable across registers. Register-specific threshold tuning improves apparent accuracy but does not provide a transferable deployment cutoff.

  • Accuracy: Accuracy ranges from 0.64 to 0.74 on textbook at a fixed 0.5 cutoff and from 0.54 to 0.68 on shifted registers.Because each test set contains one correct and one wrong variant per fact, balanced accuracy equals accuracy.
  • Thresholds: For Llama-3-8B, the Youden-optimal cutoff swings from 0.64 on textbook to 0.97 on patient, 1.00 on clinical note, and 0.04 on colloquial.Similar threshold swings occur for every LLM.
  • Operational interpretation: A threshold tuned on one register is not transferable to another even when AUROC transfers, because the ranking survives while the score scale changes.The Youden threshold evaluated on the same test set is an optimistic ceiling rather than a deployable estimate.
  • Calibration: Raw probe ECE exceeds 0.05 in all 16 model-register conditions.The worst cells are 0.395 for Gemma-3-4B on Rcolloquial and 0.376 for Qwen2.5-7B on Rnote.
  • Calibration: Platt scaling reduces mean ECE from 0.341 to 0.135, but only one of 16 conditions reaches the 0.05 calibration threshold.Isotonic regression performs less well at this sample size.

H Error analysis

Error analysis identifies three recurring failure modes among highly confident Llama-3-8B errors: contradictory confidence, plausible-but-wrong acceptance, and rejection of correct negations. Their prevalence varies by register.

  • Pattern 1: confident in both directions: Pattern 1 gives both variants of the same fact extreme confidence in opposite directions, preventing any single raw threshold from separating them.For one p-values item, the wrong and correct answers receive raw probabilities 1.000 and 0.003.
  • Pattern 2: plausible-but-wrong accepted: Pattern 2 accepts a medically plausible answer that is wrong for the specific fact, and it concentrates in patient and colloquial registers.A rheumatoid-arthritis explanation was accepted for an osteoarthritis case.
  • Pattern 3: correct negation rejected: Pattern 3 rejects correct negative statements, reflecting confusion between negated truths and wrong positive claims.The authors suggest balanced training with affirmative and negative correct answers could improve performance.
  • Register distribution: Pattern 1 dominates textbook errors, Pattern 2 dominates patient and colloquial errors, and Pattern 3 concentrates in textbook and clinical-note errors.The frequencies are summarized among the top-eight confidence-ranked errors per register.

I Reproducibility

The study releases its benchmark, evaluation sets, pipeline, configurations, and result tables, while documenting model, hardware, software, seed, and cost details. Re-extraction on different GPU generations may vary slightly.

  • Models and extraction: The four probed models are open-weight instruction-tuned chat models spanning Gemma, Qwen, and Llama families in the 2B–8B range.Activations were extracted every second transformer layer at the last question token.
  • Compute reproducibility: Re-extraction on different GPU generations reproduces tables only to within about 0.02 AUROC per cell because bf16 forward passes are not bit-reproducible.The full four-model extraction requires roughly 4–6 hours of A100 time per model.
  • Software and seeds: The software stack uses PyTorch 2.5.1 with CUDA 12.1, Transformers, and scikit-learn logistic regression with L2 regularization and C=1.0.Facts and wrong-answer selection use seed 42; mixed-register splits use seed 0.
  • Cost: The reported total cost is approximately $25 in API fees and 400 SBU A100 compute.The Sonnet run cost $20, the Gemini subset $0.20, and judge calls $3.
  • Released artefacts: The release includes the benchmark, evaluation-only sets, prompts, judge rubric, aggregate tables, full pipeline, scripts, commands, and configuration files.The MedRedQA release contains a build script and item identifiers rather than redistributed forum text.
  • Sanity checks: Word-count-only classification reaches 0.60–0.77 AUROC on rewritten registers but only 0.52 on native MedMCQA and 0.51 on minimal-edit MedRedQA pairs.This sanity check shows answer length can contribute above-chance signal in generated rewritten variants.
Loading 2609.01361v1…