Year: 2025
Type: preprint
Abstract: Large Language Models (LLMs) have achieved remarkable performance by capturing complex interactions between input features. To identify these interactions, most existing approaches require enumerating all possible combinations of features up to a given order, causing them to scale poorly with the number of inputs $n$. Recently, Kang et al. (2025) proposed SPEX, an information-theoretic approach that uses interaction sparsity to scale to $n \approx 10^3$ features. SPEX greatly improves upon prior methods but requires tens of thousands of model inferences, which can be prohibitive for large models. In this paper, we observe that LLM feature interactions are often hierarchical -- higher-order interactions are accompanied by their lower-order subsets -- which enables more efficient discovery. To exploit this hierarchy, we propose ProxySPEX, an interaction attribution algorithm that first fits gradient boosted trees to masked LLM outputs and then extracts the important interactions. Experiments across four challenging high-dimensional datasets show that ProxySPEX more faithfully reconstructs LLM outputs by 20% over marginal attribution approaches while using $10\times$ fewer inferences than SPEX. By accounting for interactions, ProxySPEX efficiently identifies the most influential features, providing a scalable approximation of their Shapley values. Further, we apply ProxySPEX to two interpretability tasks. Data attribution, where we identify interactions among CIFAR-10 training samples that influence test predictions, and mechanistic interpretability, where we uncover interactions between attention heads, both within and across layers, on a question-answering task.
Source: arXiv (Cornell University)
Authors: Butler, Landon, Agarwal, Abhineet, Kang, Justin Singh, Erginbas, Yigit Efe, Yu, Bin, Ramchandran, Kannan
Institution: University of California, Berkeley
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: G1353381076, Theory for the Practice of Deep Learning: Insights into Autoencoder, LLM Fine-Tuning, and Transfer Learning, G1453181110, G1503004319, Track 4: Advanced CI Coordination Ecosystem: Monitoring and Measurement Services, G2647087871, G298330452, G34258105, G3562386289, Collaboration on the Theoretical Foundations of Deep Learning, Track 1: ACCESS Resource Allocations Marketplace and Platform Services (RAMPS), G4010157413, G4275568664, G4403015908, Advancing Theory and Methodology for Tree-Based Algorithms in High Dimensions, ACO: An Open CI Ecosystem to Advance Scientific Discovery (OpenCI), Track 3: COre National Ecosystem for CyberinfrasTructure (CONECT), G5156331529, G5367127336, G5464919783, Graduate Research Fellowship Program (GRFP), Foundations of Data Science Institute, G7526614437, G8108599699, Track 2: Customized Multi-tier Assistance, Training, and Computational Help (MATCH) for End User ACCESS to CI, G8499740153, G8751405378
OpenAlex ID: W6979339545
Cites:
Cited by:
Related to: