PaperlayerView exact revision on arXiv

Source-linked AI summary

Making Clinical Language Models Auditable: Concept-Guided Fine-Tuning for Robust Prediction

Jin Mu, Guanhua Chen

arXiv:2608.27397v1cs.CLcs.AI

TL;DR

Clinical language models may rely on clinically meaningless artifacts, making their internal reasoning difficult to understand and audit. CAST uses SAE-derived concepts to steer fine-tuning away from spurious artifacts, while maintaining strong predictive performance, improving calibration, and exposing concept-level evidence.

  • Problem

    Clinical models can exploit dataset-specific artifacts rather than patient physiology, while input-level explanations may poorly reflect the internal representations driving predictions.

  • Method

    CAST uses Sparse Autoencoders to identify human-understandable clinical concepts and uses them as training-time steering signals to suppress spurious artifacts during fine-tuning.

  • Results

    CAST maintains strong predictive performance while improving calibration and exposing concept-level evidence for model decisions on MIMIC-IV 30-day mortality prediction.

  • Takeaways & Limitations

    The framework provides a practical path toward auditable and reliable clinical NLP systems by directly suppressing spurious artifacts and exposing model evidence.

  • Takeaways & Limitations

    Evaluation covers one prediction task and discharge notes from MIMIC-III and MIMIC-IV at the same institution, so external-institution and broader clinical evaluation remains necessary.

Abstract

from arXiv · show

Clinical language models can achieve strong in-hospital accuracy yet fail under deployment shifts because they exploit note-specific artifacts (e.g., templates, separators, boilerplate) that do not reflect patient state. We propose CAST (Concept-guided Artifact Suppression Tuning), an SAE-based framework for auditable clinical text classification. CAST uses Sparse Autoencoders to expose sparse, human-auditable features from intermediate Transformer activations, labels SAE latents with an LLM-assisted interpretation pipeline and ICD-10 retrieval constraints, suppresses verified artifact latents via residual subtraction during fine-tuning, and provides post-hoc per-concept attributions for auditing model decisions. On MIMIC-IV discharge-note mortality prediction, CAST improves over its corresponding fine-tuned encoder baselines and remains competitive with strong LLM baselines, while producing a feature-level audit trail of the clinical concepts that support each prediction and the artifact concepts suppressed during training.

1 Introduction

Clinical language models can exploit documentation artifacts that undermine trustworthy ICU risk prediction. CAST addresses this gap by using SAE-derived clinical concepts to suppress spurious features during fine-tuning and expose evidence supporting predictions.

  • Motivation: Artifact-based signals such as templates, headers, boilerplate, separators, and formatting can predict outcomes without reflecting patient physiology.These shortcuts may be unstable under deployment shift.
  • Motivation: Opaque shortcut-driven alerts are difficult to justify in high-stakes ICU decisions such as escalating life support or initiating palliative care.Clinical interpretability must expose internal concepts influencing predictions, not only salient text spans.
  • Limitations of Existing Interpretability: Token-level attribution methods identify salient spans but do not reveal the underlying concepts or functional logic governing model processing.Their reliability is also contested because attention weights can be poorly correlated with model outputs.
  • Research Gap: Sparse Autoencoders decompose internal activations into distinct human-understandable features, but medical applications of explicit artifact suppression during task-specific fine-tuning remain underexplored.Prior medical work demonstrated dictionary-based explanations and feature steering without integrating these features as suppression targets during adaptation.
  • CAST: CAST uses SAE features to identify clinically meaningful concepts, suppress spurious shortcuts, and steer clinical classification during training.The framework is presented as transforming mechanistic interpretability from passive analysis into active model steering.
  • Evaluation: On MIMIC-IV ICU discharge notes, CAST targets 30-day out-of-hospital mortality prediction while providing supporting clinical concepts and ICD-10 codes.The reported contribution combines predictive performance with a transparent feature-level audit trail.

2 Related Work

Related work develops SAE architectures and uses SAE features for interpretation, steering, and model adaptation. CAST extends this direction by making clinically interpreted artifact concepts explicit suppression targets during task-specific fine-tuning.

  • Sparse Autoencoder Architectures: Vanilla, TopK, BatchTopK, and Matryoshka SAEs use different sparsity or hierarchical design choices to improve feature quality and stability.These include ℓ1 regularization, hard activation budgets, batch-level selection, and nested dictionaries.
  • SAE-Guided Control and Adaptation: Inference-time SAE interventions amplify or ablate selected latents to steer generation and analyze information flow or in-context learning behavior.The cited work applies SAE features to causal intervention beyond passive interpretation.
  • SAE-Guided Control and Adaptation: Recent adaptation methods use SAE-derived features to suppress unintended generalizations or regularize LLM-based classifiers.Examples include concept ablation for gender bias and sparse-autoencoder-based regularization.
  • Mechanistic Interpretability in Clinical Classification: General-domain mechanistic classification methods integrate task-relevant neurons or SAE features as compact representations that can transfer across models and modalities.These approaches focus on classifier representations rather than clinically interpreted artifact suppression during fine-tuning.
  • CAST’s Position: CAST updates clinical representations during training to suppress spurious documentation artifacts using SAE-derived concepts as explicit targets.This distinguishes the framework from prior approaches that do not combine clinical artifact interpretation with task-specific suppression.

3 Method

CAST turns SAE-discovered clinical concepts into an intervention and audit interface for fine-tuning clinical classifiers. It extracts and interprets sparse latents, suppresses verified artifact contributions while preserving residual information, and attributes predictions to individual concepts.

  • 3.1 Concept Extraction via SAEs: CAST maps Transformer hidden states into sparse, interpretable concept representations using an encoder-decoder SAE module.The SAE is trained to keep latent activations sparse while reconstructing the original hidden states faithfully.
  • 3.2 Automated Concept Interpretation: An LLM interprets each latent from its highest-activating clinical contexts, classifying task relevance or artifacts and generating medical keywords.The keywords are constrained by ICD-10 retrieval so assigned codes are selected from retrieved candidates.
  • 3.3 Fine-Tuning with Concept Steering: The framework steers fine-tuning by rewriting layer-K representations after a frozen prefix, then passing the steered states through a trainable suffix.Long discharge notes are processed as overlapping chunks, pooled into a document embedding, and classified with a linear head.
  • 3.3 Fine-Tuning with Concept Steering: Residual correction subtracts only decoder contributions of suppressed artifact latents, preserving the original hidden state and unreconstructed clinical signal.This avoids replacing the hidden state with an SAE reconstruction and leaves the SAE residual unchanged.
  • 3.3 Fine-Tuning with Concept Steering: CAST uniformly suppresses all latents labeled artifactual, including low-magnitude directions that may encode institution- or template-specific documentation patterns.The residual intervention makes this conservative by removing only the explicit selected decoder contributions.
  • 3.4 Per-Concept Attribution for Auditability: Post-hoc attribution scores each SAE latent's contribution without changing trained parameters, providing per-document evidence for concepts that influence predictions.The estimator uses a gradient-based approximation to counterfactual logit changes; on 30 held-out documents, it achieved Spearman ρ = 0.976 and Pearson r = 0.935 against exact effects.

4 Results and Discussion

CAST uses SAE-derived clinical concepts to steer fine-tuning away from verified documentation artifacts while preserving clinically meaningful evidence. Across mortality prediction experiments, it improves or maintains predictive performance, calibration, and auditability relative to alternative encoder-based approaches.

  • Experimental setup: The MIMIC-IV task predicts 30-day out-of-hospital mortality from ICU discharge notes using ClinicalBERT and Clinical-Longformer encoders.The dataset contains 49,832 admissions/notes from 39,705 patients and has approximately 27:1 class imbalance.
  • Experimental setup: CAST compares against matched fine-tuning, zero-shot LLMs, input removal, and SAE-based auxiliary-feature or regularization baselines.The SAE baselines do not use latents as explicit artifact-suppression interventions inside fine-tuning.
  • Feature interpretation: SAE latents are interpreted with an LLM-assisted pipeline, and strict three-run consensus identifies formatting or de-identification concepts that are not mortality-related for suppression.Agreement ranges from 0.707–0.749 for mortality labels and 0.749–0.806 for artifact labels, with unanimous agreement of 82.6–87.7% and 87.9–93.3%, respectively.
  • Feature interpretation: Many latents capture mortality-related concepts, roughly half can be grounded to ICD-10 codes, and layer-11 SAEs capture richer clinical semantics than layer-8 SAEs.ClinicalBERT yields a higher clinical-concept fraction than Longformer.
  • Auditing: Artifact suppression corrected a representative confident false positive while preserving activations corresponding to genuine clinical evidence.The case study demonstrates CAST’s intended steering behavior at the individual-note level.
  • Performance analysis: CAST provides a stronger balance of discrimination, calibration, and concept-level interpretability than standard fine-tuning, frozen SAE-Probe features, and Self-Regul regularization.Standard fine-tuning remains a strong but opaque discriminative baseline.
  • Performance analysis: Input removal is competitive on selected metrics but inconsistent across backbones and layers, and it lacks CAST’s concept-level suppression and audit trail.The strongest selected input-removal behavior occurs for ClinicalBERT at layer 11.
  • Performance analysis: CAST’s gains are not tied to a single encoder or SAE variant, with improved operating-point performance and substantially lower Brier score, NLL, and ECE for Clinical-Longformer layer-8 Matryoshka.ClinicalBERT TopK at layer 11 gives the strongest discrimination, while Matryoshka at layer 8 yields the most reliable risk estimates.

5 Conclusion

CAST is an SAE-guided fine-tuning framework for interpretable and controllable clinical prediction. On MIMIC-IV mortality prediction, it maintains strong predictive performance while improving calibration and exposing concept-level evidence for decisions.

  • Conclusion: CAST turns sparse SAE concepts into training-time steering signals that suppress spurious artifacts and provide concept-level evidence for clinical predictions.The framework is presented as a practical path toward auditable and reliable clinical NLP systems.

Limitations

The evaluation is limited to one mortality-prediction task using discharge notes from MIMIC-III and MIMIC-IV, which originate from the same institution. Clinical deployment also requires prospective, clinician-validated, and subgroup-specific assessment.

  • Scope: The study evaluates a single prediction task on discharge notes from MIMIC-III and MIMIC-IV, both from the same institution.External institutions, additional note types, and other clinical tasks remain future evaluation directions.
  • Deployment: CAST is a research-stage auditing and steering framework rather than a deployable clinical model because the task has modest absolute F1 scores.Clinically useful operating points require prospective evaluation and clinician-in-the-loop assessment of the audit trail.
  • Validation: Concept labeling relies on an LLM judge and ICD-10-CM retrieval, so clinician-annotated validation is needed to assess labeling reliability more strongly.The strict three-run consensus rule is designed to reduce erroneous suppression but does not replace clinician validation.
  • Safety and cost: Artifact and workflow-related features may carry valid clinical or demographically correlated information, requiring expert review of suppression sets and subgroup-specific evaluation before clinical use.CAST also introduces additional offline cost for SAE pretraining and latent interpretation, although the LLM is not used during test-time prediction.

A.1 Prompt Template

Figure 5 specifies the system and user prompts used to interpret SAE latents from Bio_ClinicalBERT-based clinical discharge-note models.

  • Prompt components: The system prompt frames the model as a clinical NLP expert analyzing SAE features for 30-day out-of-hospital mortality prediction.The prompt targets features from an SAE trained on Bio_ClinicalBERT.
  • Prompt components: The user prompt supplies each latent’s identifier, top 10 maximally activating tokens, activation strengths, and surrounding clinical-note contexts.These examples are used to interpret the latent concept.
  • Prompt components: Figure 5 presents the complete system-and-user prompt template for concept interpretation.

A.2 Per-Concept Attribution: Linear Limit, Validation, and Use Cases

CAST estimates how individual SAE concepts influence predictions, with a closed-form linear limit and a validated approximation for nonlinear classifiers that supports per-concept auditing.

  • Attribution mechanism: The linear case reduces Equation (2) to an analytic ablation expression evaluable without an additional network forward pass.This applies when the suffix and prediction head collapse to a single linear map over a document-level SAE aggregate.
  • Attribution mechanism: The gradient-based estimator in Equation (2) extends the closed-form calculation to the deployed nonlinear classifier.
  • Validation: ρ = 0.976 mean Spearman correlation and r = 0.935 mean Pearson correlation were obtained against exact counterfactual effects on 30 held-out documents.The validation used approximately 180 concepts per document, combining highly activated and highly attributive latents.
  • Use cases: The attribution module provides per-prediction evidence trails, global behavioral audits, and cost-efficient interpretation.Top-ranked active concepts and ICD-10 anchors support review alongside risk scores, while cohort aggregation exposes dominant concepts and residual mislabeled latents.

A.3 LLM Interpretation Details

CAST interprets SAE latents with repeated LLM labeling constrained by ICD-10 retrieval, then suppresses only artifacts meeting a strict three-run consensus criterion.

  • Interpretation pipeline: Each latent is shown to an LLM through 10 maximally activating contexts, which are summarized into a concept category and mortality-relatedness judgment.Diagnostic claims are grounded against retrieved terminology.
  • Consensus filtering: A latent enters suppression set T only when all three independent runs identify it as formatting or de-identification artifact and not mortality-related.
  • Consensus filtering: 44.3–51.2% of latents flagged as artifacts in at least one run satisfy the strict criterion, reflecting a precision-oriented design.
  • Agreement reporting: Table 3 reports Fleiss’ κ, unanimous 3-of-3 labeling fractions, and the strict-to-any artifact-set fraction for mortality and artifact decisions.

A.4 Per-Configuration Interpretation Summary

Across configurations and evaluation settings, CAST uses strictly interpreted artifact concepts to steer fine-tuning, with stable performance, improved calibration, and interpretable clinical concept summaries.

  • Configuration comparisons: ClinicalBERT yields a higher clinical-concept fraction than Clinical-Longformer at matched layers and SAE types.The paper plausibly attributes this to full versus local attention mechanisms.
  • Configuration comparisons: Layer 11 SAEs produce richer clinical semantics than layer 8 SAEs across backbone–SAE combinations.
  • Configuration comparisons: Matryoshka SAEs match or slightly exceed TopK SAEs in clinical-concept fraction under comparable settings.
  • Artifact selection: The strict-consensus artifact set T is a smaller nonmortality subset requiring agreement across three independent LLM runs.
  • Sensitivity analyses: Early-layer performance depends on SAE architecture: TopK performs best for ClinicalBERT, while Matryoshka performs best for Clinical-Longformer.
  • Sensitivity analyses: Performance remains stable across tested SAE dictionary sizes and sparsity levels for ClinicalBERT layer 11 CAST.
  • MIMIC-III evaluation: On MIMIC-III, every best-performing CAST variant exceeds its matched fine-tuning baseline in F1, with gains up to +0.11 F1 for Clinical-Longformer layer 8 with Matryoshka.CAST also gives the best calibration metrics in every backbone–layer cell.
  • MIMIC-III evaluation: Calibration and F1 gains transfer from MIMIC-IV to MIMIC-III, while SAE-family preferences differ between ClinicalBERT and Clinical-Longformer.
Loading 2608.27397v1…