Detecting and Understanding Vulnerabilities in Language Models via Mechanistic Interpretability

TL;DR

The paper tackles the problem of locating and understanding vulnerabilities in large language models by leveraging Mechanistic Interpretability (MI). It introduces a framework that fuses circuit discovery via activation patching, gradient-based adversarial sample generation in embedding space, and logit attribution to identify and explain task-specific vulnerabilities within a model's circuit. The method is demonstrated on GPT-2 Small for the acronym third-letter prediction task, revealing key contributing attention heads (e.g., 10.10, 8.11, 9.9) and revealing vulnerabilities toward certain letters (notably A and S). This MI-guided approach provides a principled path to audit, understand, and potentially mitigate vulnerabilities without additional adversarial training, with broad applicability to other differentiable language models.

Abstract

Large Language Models (LLMs), characterized by being trained on broad amounts of data in a self-supervised manner, have shown impressive performance across a wide range of tasks. Indeed, their generative abilities have aroused interest on the application of LLMs across a wide range of contexts. However, neural networks in general, and LLMs in particular, are known to be vulnerable to adversarial attacks, where an imperceptible change to the input can mislead the output of the model. This is a serious concern that impedes the use of LLMs on high-stakes applications, such as healthcare, where a wrong prediction can imply serious consequences. Even though there are many efforts on making LLMs more robust to adversarial attacks, there are almost no works that study \emph{how} and \emph{where} these vulnerabilities that make LLMs prone to adversarial attacks happen. Motivated by these facts, we explore how to localize and understand vulnerabilities, and propose a method, based on Mechanistic Interpretability (MI) techniques, to guide this process. Specifically, this method enables us to detect vulnerabilities related to a concrete task by (i) obtaining the subset of the model that is responsible for that task, (ii) generating adversarial samples for that task, and (iii) using MI techniques together with the previous samples to discover and understand the possible vulnerabilities. We showcase our method on a pretrained GPT-2 Small model carrying out the task of predicting 3-letter acronyms to demonstrate its effectiveness on locating and understanding concrete vulnerabilities of the model.

Figures (5)

• Figure 1: Summary of the workflow of our proposal to detect and locate vulnerabilities in language models.
• Figure 2: Variation in logit difference when patching different heads on GPT-2 Small.
• Figure 3: Distribution of the words of the dataset that begin with each letter vs. the distribution of generated adversarial acronym in terms of the initial letter of the third word.
• Figure 4: Logit attribution for every attention head on adversarial samples with the letter A. This attribution is obtained by projecting into the logit difference direction.
• Figure 5: Logit attribution of head 10.10 on adversarial samples with the letter A. This attribution is obtained by projecting into the directions of the different capital letters.