Making Clinical Language Models Auditable: Concept-Guided Fine-Tuning for Robust Prediction
Organizations: University of Wisconsin–Madison
Abstract
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.
Figures & tables
| SAE | Latent | Concept (ICD-10-CM) | |
| Push prediction mortality | |||
| BatchTopK | 1108 | palliative care (Z51.5) | |
| TopK | 6730 | COPD exacerbation (J44.1) | |
| Matryoshka | 101 | metastatic disease / secondary malignancy (C79.9) | |
| TopK | 915 | acute kidney failure (N17.9) | |
| Push prediction mortality | |||
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Backbone | SAE / Layer | 3/3 Mort. | 3/3 Art. | Strict/any | ||
| ClinicalBERT | BatchTopK / 11 | .749 | .796 | 87.7% | 93.0% | 47.5% |
| Matryoshka / 11 | .748 | .793 | 87.2% | 92.8% | 48.7% | |
| TopK / 11 | .746 | .806 | 87.5% | 93.3% | 50.2% | |
| BatchTopK / 8 | .745 | .804 | 87.0% | 93.1% | 51.2% | |
| Matryoshka / 8 | .722 | .790 | 85.2% | 91.9% | 47.6% | |
| TopK / 8 | .733 | .801 | 86.3% | 92.9% | 48.0% |
| Backbone | SAE type / Layer | ICD-10 (%) | Mortality (%) |
| ClinicalBERT | TopK / 8 | ||
| TopK / 11 | |||
| BatchTopK / 8 | |||
| BatchTopK / 11 | |||
| Matryoshka / 8 | |||
| Matryoshka / 11 |
| Backbone | Method | F1 | AUROC | PR-AUC |
| ClinicalBERT | Fine-tuning | .2935 | .8415 | .2429 |
| TopK CAST | .3061 | .8511 | .2574 | |
| Matryoshka CAST | .2306 | .8538 | .2342 | |
| Clinical- Longformer | Fine-tuning | .2746 | .8464 | .2378 |
| TopK CAST | .2432 | .8570 | .2485 | |
| Matryoshka CAST | .3025 | .8484 | .2526 |
| F1 | AUROC | PR-AUC | ||
| 4,096 | 64 | .2703 | .8511 | .2531 |
| 8,192 | 32 | .2569 | .8525 | .2530 |
| 8,192 | 64 | .2961 | .8579 | .2460 |
| 8,192 | 128 | .2801 | .8502 | .2344 |
| 16,384 | 64 | .2881 | .8538 | .2504 |
| Backbone | Layer | SAE | F1 | F1 | AUROC | PR-AUC | Brier | NLL | ECE |
| ClinicalBERT | 11 | TopK | |||||||
| 11 | Matryoshka | ||||||||
| ClinicalBERT | 8 | TopK | |||||||
| 8 | Matryoshka | ||||||||
| Clinical- Longformer | 11 | TopK | |||||||
| 11 | Matryoshka |
| SAE | L# | ICD-10 | Concept (LLM description) | Representative activating context | |
| Push prediction mortality | |||||
| TopK | 4415 | R19.1 | Physical exam finding of a “soft” abdomen, often with descriptors of normal bowel sounds. | / abd : soft , nt , nd , normal | |
| BatchTopK | 1108 | Z51.5 | Palliative-care services, often in the context of serious illness or cancer treatment. | second / round of palliative chemo | |
| TopK | 6730 | G70.01, J44.1 | Worsening of existing symptoms or emergence of new symptoms. | temp > 101.5, worsening pain , / drainage | |
| BatchTopK | 7197 | I27.23, J60–J70 | Mentions of the lung, particularly lung cancer, lung masses, and other pulmonary pathology. | lung sounds very diminished l lung fie... | |
| TopK | 5454 | F41.1, F06.4 | Anxiety, often in the context of psychiatric conditions such as depression. | 20mg daily for depression / anxiety | |
| Sub-category | L# | LLM concept description | Representative activating context |
| formatting | 7488 | Line breaks and formatting elements, especially the <cr> token in lists. | please take all pills ) / / followup instructions |
| formatting | 4310 | Structural transition from response to plan sections in clinical notes. | cr > response : / plan : / chronic o |
| formatting | 1625 | Parenthetical confirmation notations, e.g., (x) yes or ( ) no. | services / considered ? ( x ) yes - ( ) no |
| formatting | 8123 | List-item or code-snippet delimiters, especially the semicolon character. | care , asking app questions . a ; / loving p ; c |
| formatting | 2201 | Phrases indicating that specific details or instructions are provided elsewhere. | of systems is unchanged from admission except as n |
| de-id | 2195 | De-identified first- or last-name tokens of the form [FirstName5(NamePattern1) X] . | needed . [ * * first name5 ( namepattern1 ) 1875 * |
Language Model Robustness
Language model robustness studies how reliably language models preserve correct, safe behavior under input changes, distribution shifts, and adversarial pressure.
See all 493 papers