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 note models may rely on dataset-specific artifacts rather than patient physiology, and integrating interpretable internal features into artifact suppression remains underexplored. CAST uses SAE-derived concepts to steer fine-tuning; on MIMIC-IV 30-day mortality prediction, it maintains strong predictive performance while improving calibration and exposing concept-level evidence.

  • Problem

    Although EHR-note models can exploit clinically meaningless artifacts, integrating interpretable internal features as artifact-suppression interventions during task-specific fine-tuning remains underexplored.

    Source: §1 · pp. 1–2

  • Method

    CAST uses SAE-derived human-understandable concepts to steer fine-tuning by suppressing features associated with spurious shortcuts.

    Source: §1 · p. 2, §3.3 · p. 4

  • Results

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

    Source: §5 · p. 9

  • Takeaways & Limitations

    CAST offers a practical path toward auditable and reliable clinical NLP systems.

    Source: §5 · p. 9

    Evaluation covers one prediction task and same-institution datasets; CAST remains research-stage, requiring prospective evaluation and clinician-in-the-loop assessment before clinical use.

    Source: Limitations · p. 9

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

CAST uses Sparse Autoencoders to identify clinical concepts and artifact features in model representations, then guides fine-tuning to suppress shortcuts and produce feature-level prediction audits. On MIMIC-IV ICU discharge notes, the paper reports improved predictive performance for 30-day out-of-hospital mortality prediction and an audit trail with clinical concepts and ICD-10 codes.

1 Introduction

  • Clinical note models may exploit dataset-specific templates, headers, boilerplate, separators, or documentation patterns that predict outcomes but are clinically meaningless or unstable under deployment shift.In ICU settings, predictions driven by shortcuts rather than deterioration evidence can have limited clinical value for time-sensitive decisions.Source: §1 · p. 1
  • Token-level attribution methods identify salient spans but do not reveal underlying concepts or model logic, and attribution reliability is contested (Mullenbach et al., 2018; Vu et al., 2020; Jain and Wallace, 2019).The introduction argues that clinical interpretability must expose internal concepts influencing predictions and assess whether they reflect medically meaningful evidence (Rudin, 2019).Source: §1 · pp. 1–2
  • Although SAEs decompose activations into human-understandable features, using such features as explicit artifact-suppression interventions during medical task fine-tuning remains underexplored (Bricken et al., 2023; Gao et al., 2025; Wu et al., 2024).Prior medical-domain work used dictionary features for mechanistic explanations and feature-based steering, while general-domain SAE work disentangled polysemantic neurons into monosemantic features.Source: §1 · p. 2
  • CAST uses SAE-identified concepts to suppress shortcut features and encourage clinically meaningful representations during fine-tuning, improving prediction and providing an audit trail with clinical concepts and ICD-10 codes.The evaluation uses MIMIC-IV ICU discharge notes for 30-day out-of-hospital mortality prediction.Source: §1 · p. 2

2 Related Work

Related work advances SAE feature extraction and uses sparse features for model intervention and classification. Existing approaches do not explicitly suppress clinically interpreted documentation artifacts during task-specific clinical fine-tuning.

Feature Extraction via Sparse Autoencoders

  • Vanilla SAEs use ℓ1-sparse latents, while TopK, BatchTopK, and Matryoshka variants address shrinkage, dead neurons and stability, and hierarchical feature organization (Gao et al., 2025; Bussmann et al., 2024; Bussmann et al., 2025).These developments improve the quality, stability, and semantic organization of SAE-derived features.Source: Feature Extraction via Sparse… · pp. 2–3
  • SAE features support inference-time steering and ablation, reasoning-faithfulness probing, and information-flow tracking for in-context learning (Bayat et al., 2025; Chen et al. (2026); Cho and Hockenmaier (2025)).Adaptation work also uses SAE-based concept ablation to suppress unintended generalizations or sparse-feature regularization for interpretable classification (Casademunt et al. (2025); Wu et al., 2025).Source: Feature Extraction via Sparse… · p. 3
  • CAST adapts SAE-guided learning to clinical prediction by integrating SAE-derived concepts into training to suppress spurious documentation artifacts.The related methods primarily target general-domain generation, reasoning, behavioral steering, or controllable classification rather than clinical risk prediction.Source: Feature Extraction via Sparse… · p. 3
  • SPIN integrates task-relevant internal neurons into compact interpretable classifiers, while SAE features have served as classifier representations that transfer across models and modalities (Jiao et al., 2024; Gallifant et al. (2025)).These approaches do not use clinically interpreted artifact concepts as explicit suppression targets during task-specific fine-tuning.Source: Feature Extraction via Sparse… · p. 3

3 Method

CAST uses sparse autoencoders to expose interpretable clinical concepts, uses an LLM-assisted pipeline with ICD-10 retrieval to interpret them, and suppresses artifact concepts during fine-tuning. It also provides post-hoc per-concept attributions to audit prediction evidence.

3.1 Concept Extraction via SAEs

  • SAEs encode intermediate token representations as sparse concept activations and decode them into faithful reconstructions, creating an interpretable concept interface.The sparsity mechanism can be instantiated in different SAE variants.Source: §3.1 · p. 3

3.2 Automated Concept Interpretation

  • For each latent, an LLM interprets top-activating token contexts, labels task relevance versus artifact status, and generates keywords for ICD-10 code retrieval and verification (National Center for Health Statistics, 2024).The LLM selects verified clinical codes exclusively from the retrieved candidate set.Source: §3.2 · pp. 3–4

3.3 Fine-Tuning with Concept Steering

  • CAST subtracts decoder contributions for every latent labeled artifactual while preserving the SAE residual, so unmodeled hidden-state signal passes through unchanged.Uniform suppression includes low-magnitude artifacts; the frozen prefix and SAE steer representations before the trainable suffix processes them.Source: §3.3 · pp. 4–5

3.4 Per-Concept Attribution for Auditability

  • On 30 held-out documents, per-document attribution correlated with exact counterfactual effects at Spearman ρ = 0.976 and Pearson r = 0.935.The post-hoc signed first-order attribution estimates each latent’s contribution to the positive-class logit; computing it uses one forward and one backward pass per batch instead of |D| separate forward passes.Source: §3.4 · p. 5

4 Results and Discussion

On imbalanced MIMIC-IV discharge-note mortality prediction, CAST uses SAE concepts for training-time artifact suppression and post-hoc auditing. The results examine its performance, the clinical meaning of learned concepts, and the effects of concept-level interventions.

4.1 Dataset

  • The 30-day task contains approximately 1,830 positive and 48,002 negative admissions among 49,832 notes, with patient-level splits and exclusions for in-hospital death and hospice disposition.The notes come from 39,705 patients in MIMIC-IV v2.2, and labels link encounters to an external death registry.Source: §4.1 · p. 5

4.3 Training setup and hyperparameters

  • For clinical note encoding, ClinicalBERT takes inputs up to 512 tokens and Clinical-Longformer handles notes up to 4,096 tokens; both were continually pre-trained on MIMIC.SAEs are frozen during all downstream fine-tuning. Models are fine-tuned for five epochs using class-weighted focal loss for the approximately 27:1 class imbalance, with F1 reported at a fixed threshold of 0.5 and a validation-selected threshold.Source: §4.2 · p. 5, §4.3 · p. 6

4.2 Models and Baselines

  • CAST is compared with fine-tuned clinical encoders, zero-shot GPT-4 and Llama-3-8B, input removal, Self-Regul, and SAE-Probe.Self-Regul uses SAE-derived sparse features for regularization, while SAE-Probe uses frozen SAE activations as document-level features for a lightweight classifier.Source: §4.2 · pp. 5–6

4.4 Feature Extraction and Interpretation

  • A three-run unanimous-consensus rule includes a latent for suppression only when all runs label it a non-mortality-related formatting or de-identification artifact.Fleiss’ κ ranges from 0.707–0.749 for mortality-related labels and 0.749–0.806 for artifact labels; aggregate findings indicate richer clinical semantics in layer-11 than layer-8 SAEs and a higher clinical-concept fraction for ClinicalBERT than Longformer. In a representative case, suppressing artifact latents corrects a confident false positive while preserving clinical evidence.Source: §4.4 · pp. 6–7, Figure 3 · p. 6

4.5 Performance Analysis

  • CAST provides a stronger balance of discrimination, calibration, and concept-level interpretability than encoder-based alternatives; GPT-4 is a strong reference baseline, while Llama3-8B performs poorly.The reported comparisons concern 30-day out-of-hospital mortality prediction from long discharge notes; the supplied passages do not provide the table’s cell values.Source: Table 1 · p. 7, §4.5 · pp. 7–8

4.6 Feature Importance

  • Positive attributions include palliative care, COPD exacerbation, metastatic disease, and acute kidney failure, while negative attributions include mobility and independent ambulation across all three SAE variants.The analysis ranks signed mean per-concept attributions over test documents where each latent is active.Source: §4.6 · p. 8, Table 2 · p. 8

4.7 Ablation and Sensitivity Analysis

  • In matched-size ClinicalBERT layer-11 ablations, blocking artifact-related concepts improves PR-AUC and F1, whereas blocking mortality-related concepts degrades performance.CAST remains competitive across the tested earlier intervention layer, SAE dictionary sizes, and TopK sparsity settings.Source: §4.7 · pp. 8–9, Figure 4 · p. 9

5 Conclusion

CAST uses SAE concepts as training-time steering signals to suppress spurious artifacts in clinical prediction. On MIMIC-IV 30-day mortality prediction, it maintains strong predictive performance while improving calibration and exposing concept-level evidence for model decisions.

5 Conclusion

  • CAST turns sparse SAE concepts into training-time steering signals, moving beyond post-hoc interpretation to directly suppress spurious artifacts.The framework is presented as an interpretable and controllable approach to clinical prediction.Source: §5 · p. 9
  • On MIMIC-IV 30-day mortality prediction, CAST maintains strong predictive performance while improving calibration and exposing concept-level evidence for model decisions.These results support the framework's stated aim of auditable clinical prediction.Source: §5 · p. 9
  • The framework offers a practical path toward auditable and reliable clinical NLP systems.Source: §5 · p. 9

Limitations

CAST’s evaluation is limited to one prediction task on MIMIC-III and MIMIC-IV discharge notes from the same institution, while its modest absolute F1 leaves external validation and clinically useful operating points unestablished. Its audit trail depends on LLM- and ICD-10-CM-assisted interpretation and suppression choices that need clinician review, and SAE pretraining and latent interpretation add offline cost.

Limitations

  • Evaluation covers one task and discharge notes from MIMIC-III and MIMIC-IV, both from the same institution, leaving external institutions, note types, and tasks for future work.The authors identify these as important directions for future evaluation.Source: Limitations · p. 9
  • Modest absolute F1 on the highly imbalanced task positions CAST as a research-stage auditing and steering framework, not a deployable model, pending prospective evaluation and clinician-in-the-loop assessment.Clinically useful operating points remain to be established.Source: Limitations · p. 9
  • LLM-judge and ICD-10-CM-assisted concept interpretation uses three-run consensus, but clinician-annotated validation and expert review of suppression sets remain needed to assess reliability and protect valid or demographically correlated information.Subgroup-specific evaluation is also needed before clinical use.Source: Limitations · p. 9
  • SAE pretraining and latent interpretation add offline cost, although the LLM is not used during test-time prediction.Source: Limitations · p. 9

A Appendix · A.1 Prompt Template

The appendix includes a prompt template showing the system and user prompts used for concept interpretation.

A.1 Prompt Template

  • Figure 5 presents the system and user prompts used for concept interpretation.Source: Figure 5 · p. 12

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

CAST’s per-concept attribution uses a closed-form ablation in the linear case and a gradient-based estimator for the nonlinear classifier. Validation shows strong agreement with exact counterfactual effects, while attributions support human-readable explanations, behavioral audits, and focused expert review.

Closed-form linear case

  • The linear case reduces to an analytic ablation expression that requires no additional network forward pass; the gradient-based estimator extends it to the nonlinear deployed classifier.This closed form applies when the suffix and prediction head collapse to a single linear map over a document-level aggregate.Source: §A.2 · p. 12

Validation protocol

  • ρ = 0.976 (median 0.978; Q1 = 0.970; Q3 = 0.986), r = 0.935, and mean relative L1 error = 0.38 validate attribution against exact effects across approximately 180 concepts in 30 documents.The residual error comes from higher-order Taylor terms and does not substantially affect ranking or sign.Source: §A.2 · p. 12

Use cases

  • Attributions do not feed back into training and support per-document evidence trails, cohort-level behavioral audits, and |Aj(x)|-ranked expert review of influential latents.Per-document explanations pair top-ranked active concepts with ICD-10 anchors; ranking focuses review on a few hundred influential latents rather than the full dictionary.Source: §A.2 · p. 12

A.3 LLM Interpretation Details

CAST interprets alive SAE latents with an LLM and suppresses only latents unanimously labeled as artifacts and not mortality-related across three runs. ICD-10-CM retrieval grounds diagnostic claims.

A.3 LLM Interpretation Details

  • Gemini 2.5 Flash Lite interprets each alive latent from its 10 maximally activating contexts, assigning a concept summary, one of seven categories, and a mortality-related label.Diagnostic claims are checked against an ICD-10-CM retrieval lookup to reduce code hallucinations.Source: §A.3 · p. 12
  • Only 44.3–51.2% of latents flagged as artifacts in at least one run entered the strict-consensus suppression set.Entry required all three runs to label a latent as a formatting or de-identification artifact and not mortality-related.Source: §A.3 · pp. 12–13, Table 3 · p. 13

A.4 Per-Configuration Interpretation Summary … A.13 Artifacts and licensing

The appendices examine CAST’s interpretability, sensitivity, cross-corpus results, and implementation, and illustrate both the concepts used for attribution and the artifacts suppressed during training. They also specify research-use and licensing constraints.

A.4 Per-Configuration Interpretation Summary

  • ClinicalBERT and layer 11 yield richer clinical concepts than Longformer and layer 8, while Matryoshka matches or slightly exceeds TopK in comparable settings.The authors attribute the backbone difference plausibly to attention mechanisms and the layer difference to deeper task-specific representations. The strict-consensus artifact set T is a much smaller subset of nonmortality latents, requiring 3-of-3 agreement across LLM runs.Source: §A.4 · p. 13

A.5 Training Details

  • SAEs are pretrained on 200,000 MIMIC notes, and each is optimized for 100,000 steps, or approximately 410M tokens.The corpus comprises 120,000 MIMIC-III narrative records and 80,000 MIMIC-IV discharge summaries; downstream runs use AdamW, effective batch size 256, and random seed 42.Source: §A.5 · p. 13

A.6 Layer and Hyperparameter Sensitivity

  • At layer 4, the best SAE architecture depends on the backbone, while CAST performance remains stable across the tested dictionary sizes and TopK sparsities.TopK performs best for ClinicalBERT and Matryoshka for Clinical-Longformer in the layer-4 comparison.Source: Table 5 · p. 14, §A.6 · p. 14, Table 6 · p. 14

A.7 Baseline implementations.; A.8 Paired Bootstrap Analysis

  • The SAE baselines are adapted to long notes, and CAST and fine-tuning baselines train upper backbone layers while Self-Regul and SAE-Probe keep encoder and SAE frozen.The paired bootstrap draws 10,000 shared-test-set resamples and tests CAST against matched fine-tuning baselines in the direction of improvement for each metric.Source: §A.7 · p. 14, §A.8 · p. 14

A.9 Additional BatchTopK SAE Results on MIMIC-IV

  • BatchTopK’s MIMIC-IV test-set trends mirror those in the main results table and closely track TopK SAE behavior.The appendix reports BatchTopK across four backbone–layer configurations alongside fine-tuning, Self-Regul, and SAE-Probe baselines.Source: §A.9 · p. 14, Table 8 · p. 16

A.11 Top Clinical Concepts by Attribution

  • Table 10 presents top signed-attribution clinical concepts for ClinicalBERT layer 11 across three SAE variants, with LLM-generated descriptions and representative activating contexts.The contexts are sampled from MIMIC-IV discharge notes.Source: §A.11 · p. 15, Table 10 · p. 17

A.12 Examples of Suppressed Artifacts; A.13 Artifacts and licensing

  • Representative suppressed artifacts include formatting and de-identification patterns that may correlate with outcomes but do not directly reflect patient physiology.Their decoder contributions are subtracted during fine-tuning; MIMIC data are used under credentialed access for retrospective clinical NLP research, and the work is not intended for clinical deployment without further validation and review.Source: §A.12 · p. 15, §A.13 · p. 15, Table 11 · p. 18

A.10 Evaluation on MIMIC-III

  • On MIMIC-III, every backbone–layer setting’s best-performing CAST variant beats its fine-tuning baseline in F1, with up to +0.11 F1 for Clinical-Longformer layer 8 Matryoshka; CAST also has the best calibration in every setting.Models are trained on MIMIC-IV and evaluated on MIMIC-III without further tuning; the authors report reductions in Brier, NLL, and ECE across all four settings.Source: §A.10 · p. 15, Table 9 · p. 16

Related papers

Loading 2608.27397v1…