Measuring and Mitigating Shortcut Reliance in Language Models with Probe-Based Representation Entanglement
Year: 2026
Type: conference-paper
Abstract: Shortcut learning occurs when models solve tasks by exploiting superficial correlations in training data like formatting patterns or lexical heuristics instead of learning the intended underlying task.These shortcuts often produce high in-distribution accuracy but fail under distribution shift.We study whether shortcut reliance can be diagnosed and mitigated in small instruction-tuned language models using a simple representation-level quantity.We fine-tune Gemma 3 1B Instruct and Llama 3.2 1B on two synthetic sentiment shortcuts in SST-2 and one natural shortcut in MNLI based on lexical overlap.During training, we fit linear probes for the task label and the shortcut attribute at every layer and define Correlation-Dependent Representational Entanglement (CDRE) as the absolute cosine similarity between the two probe directions.Across settings, increasing shortcut prevalence produces a sharp rise in the robustness gap between shortcut-aligned and shortcutfree test sets, and higher deep-layer CDRE tracks this degradation.At a 99% shortcut ratio, Llama's clean accuracy on capitalizationbiased SST-2 drops from 93.2% at 0% bias to 49.0%, while Gemma drops from 91.8% to 60.2%.A CDRE-regularized objective substantially improves robustness for capitalization and lexical-overlap shortcuts, but offers little benefit for a speaker-prefix shortcut whose learned directions are already nearly orthogonal.These results show that probe-derived representation entanglement provides a reliable signal of harmful shortcut reliance and offers a practical criterion for determining when shortcut mitigation is likely to be effective.
Author: Divyajot Singh
Institution: Birla Institute of Technology and Science, Pilani
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: hybrid
OpenAlex ID: W7166835376