Short-circuiting Shortcuts: Mechanistic Investigation of Shortcuts in Text Classification

TL;DR

This work investigates how spurious shortcuts are processed inside LLMs by combining mechanistic interpretability with a controlled ActorCorr sentiment dataset. It identifies a shortcut circuit where early MLPs enrich shortcut features and later attention heads (Label Heads) extract and inject label-specific information, leading to premature predictions. The authors introduce Head-based Token Attribution (HTA), a targeted feature-attribution method that attributes token-level influence to specific heads, outperforming IG and LIME in detecting shortcuts and enabling mitigation via selective head ablation. The findings advance interpretability by revealing when and where internal decisions are made and offer practical mitigation strategies to curb shortcut effects in real-world models.

Abstract

Reliance on spurious correlations (shortcuts) has been shown to underlie many of the successes of language models. Previous work focused on identifying the input elements that impact prediction. We investigate how shortcuts are actually processed within the model's decision-making mechanism. We use actor names in movie reviews as controllable shortcuts with known impact on the outcome. We use mechanistic interpretability methods and identify specific attention heads that focus on shortcuts. These heads gear the model towards a label before processing the complete input, effectively making premature decisions that bypass contextual analysis. Based on these findings, we introduce Head-based Token Attribution (HTA), which traces intermediate decisions back to input tokens. We show that HTA is effective in detecting shortcuts in LLMs and enables targeted mitigation by selectively deactivating shortcut-related attention heads.

Figures (17)

• Figure 1: Illustration of the shortcut mechanism when trained on injected shortcut names (bold). Later layer attention heads focus on shortcut tokens and change the prediction based on information from early MLP layers. After decomposing the attention head, we find how the shortcut tokens are processed and apply these findings to construct our feature attribution method (HTA).
• Figure 2: Schematic of transformer architecture, illustrating the activations per component and decomposition of the MHSA, based on elhage2021mathematical.
• Figure 3: Effect of shortcuts on correlated and anti-correlated classes. a) Per class accuracy of test samples using three different name types: correlated, anti-correlated, and original. b&c) Effect of anti-correlated shortcuts (quantified by the ACAC metric of Equation \ref{['eq:acac_equation']}) when changing shortcut frequency (b) and purity ratio (c).
• Figure 4: Path Patching results on ActorCorr trained model for Bad actor in positive reviews. (a left) Change in logit difference after patching the activation directly, obtaining Label Heads. (a right) Change in logit difference after patching via Label Heads. (b) Cumulative attention on name tokens against the logit difference for Label Head 11.2
• Figure 5: a,b) Shortcut classification evaluated via distribution separation metrics for the three feature attribution methods HTA, LIME and IG, using the two aggregation functions (max, sum). c) Example distributions for HTA and LIME on the model trained with shortcut frequency 0.003.