Parameter-Efficient Token Embedding Editing for Clinical Class-Level Unlearning
Authors: Iyad Ait Hou, Shrenik Borad, Harsh Sharma, Pooja Srinivasan, Rebecca Hwa, Aya Zirikly
Organizations: Department of Computer Science, The George Washington University, Washington, DC · Center for Speech and Language Processing, Johns Hopkins University, Baltimore, MD
Machine unlearning is increasingly important for clinical language models, where privacy regulations and institutional policies may require removing sensitive information from deployed systems without retraining from scratch. In practice, deletion requests must balance effective forgetting of targeted information with preservation of model utility and minimal parameter modification. We introduce Sparse Token Embedding Unlearning (STEU), a parameter-efficient method for behavioral class-level unlearning that updates only PMI-selected token embeddings together with a small classifier head while keeping all encoder layers frozen. Across experiments on MIMIC-IV, MIMIC-III, and eICU using BioClinicalBERT, BERT-base, and DistilBERT, STEU consistently suppresses the target class while largely preserving retained task performance. In the primary MIMIC-IV setting, STEU achieves near-complete forgetting (forget F1 = 0.0004) while maintaining competitive retained utility (retain avg F1 = 0.4766) after modifying only 0.19% of model parameters. These results suggest that targeted behavioral unlearning can be achieved through sparse embedding edits without modifying deeper encoder representations.
Figures & tables
Figure 1: Conceptual overview of Sparse Token Embedding Unlearning (STEU). While existing machine unlearning methods modify large portions of encoder parameters, STEU performs localized edits to selected token embeddings, enabling class-level behavioral unlearning with minimal parameter updates.
Figure 2: STEU workflow figure combining token selection and constrained unlearning. The pipeline first identifies class-discriminative tokens via frequency-weighted PMI, then updates only the selected embedding rows and classifier head with gradient masking while all encoder layers remain frozen.
Dataset
Source
Total Admissions
Target Classes
MIMIC-IV
Discharge summaries 9 ; 11
148,399
Selected 6 CCSR groups
MIMIC-III
Discharge summaries 10
52,726
Selected 6 CCSR groups
eICU
ICU clinical notes 13
200,859
Selected 6 CCSR groups
Table 1: Dataset summary. Experiments use 50,000/20,000 train/val subsamples with patient-disjoint splits from each corpus.
Encoder
Total Params
Layers
Hidden Dim
Val Macro F1
Bio_ClinicalBERT 2
108.3M
12
768
0.499
BERT-base 6
109.5M
12
768
0.484
DistilBERT 15
66.4M
6
768
0.242
Table 2: Pre-trained encoder configurations. Val Macro F1 is reported on the MIMIC-IV validation split.
Method
Forget F1
Retain F1
Params Updated
% Model
Gradient ascent
0.0000
0.3127
21.3M
19.6%
Direct suppression
0.0000
0.4912
21.3M
19.6%
Influence-weighted
0.0000
0.3228
21.3M
19.6%
STEU (ours)
0.0004
0.4766
201K
0.19%
Table 3: Comparison of forgetting effectiveness, retained utility, and parameter budget.
Dataset
Encoder
Forget Base
Forget Final
Retain Base
Retain Final
Δ Utility Loss
MIMIC-IV
BioClinicalBERT
0.4685
0.0000
0.5064
0.5038
-0.0026
MIMIC-IV
BERT-base
0.4188
0.0000
0.4975
0.4921
-0.0054
MIMIC-IV
DistilBERT
0.3521
0.0000
0.4413
0.4350
-0.0063
MIMIC-III
BioClinicalBERT
0.4210
0.0000
0.6296
0.6261
-0.0035
MIMIC-III
BERT-base
0.0393
0.0000
0.2750
0.2698
-0.0052
MIMIC-III
DistilBERT
0.0285
0.0000
0.2147
0.2091
-0.0056
Table 4: Cross-dataset and cross-encoder evaluation of STEU. The final column reports retained utility loss ( Δ=Final−Baseline ).
Figure 3: Utility–parameter tradeoff. Each point shows retained F1 versus number of updated parameters (log scale), with forget F1 annotated. STEU occupies the high-utility, low-parameter region.
Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.