Probing Large Language Model Hidden States for Adverse Drug Reaction Knowledge
摘要
Large language models (LLMs) integrate knowledge from diverse sources into a single set of internal weights. However, these representations are difficult to interpret, complicating our understanding of the models’ learning capabilities. Sparse autoencoders (SAEs) linearize LLM embeddings, creating monosemantic features that both provide insight into the model’s comprehension and simplify downstream machine learning tasks. These features are especially important in biomedical applications where explainability is critical. Here, we evaluate the use of Gemma Scope SAEs to identify how LLMs store known facts involving adverse drug reactions (ADRs). We transform hidden state embeddings of drug names from Gemma2-9b-it into monosemantic features, then train a linear classifier on these to classify ADR likelihood, evaluating against an established benchmark. These embeddings provide strong predictive performance, giving AUC-ROC of 0.957 for identifying acute kidney injury, 0.902 for acute liver injury, 0.954 for acute myocardial infarction, and 0.963 for gastrointestinal bleeds. Notably, there are no significant differences (p > 0.05) in performance between the simple linear classifiers built on SAE outputs and neural networks trained on the raw embeddings, suggesting that the information lost in reconstruction is minimal. Our primary contribution is a systematic demonstration that LLM hidden states contain robust biomedical signals relevant to ADRs, and that applying SAEs preserves these signals in an interpretable format. We believe this approach can help optimally apply the biomedical knowledge that world models learn in training to be used for downstream applications, such as improving and expanding ADR reference sets.