Identifying Interactions at Scale for LLMs

🔑 Enhanced Key Takeaways

• •ProxySPEX fits gradient boosted trees (GBTs) as a proxy model to predict masked LLM outputs, then extracts interactions from the GBTs to exploit hierarchical structure in feature interactions.[1][2][3]
• •ProxySPEX applied to data attribution identifies interactions among CIFAR-10 training samples influencing test predictions, and to mechanistic interpretability uncovers attention head interactions within and across layers on question-answering tasks.[1][2][3]
• •ProxySPEX provides a scalable approximation of Shapley values by accounting for interactions and outperforms LASSO in faithfulness with limited inferences, while achieving higher test accuracy in attention head pruning tasks.[1][4]

🛠️ Technical Deep Dive

• •Fits gradient boosted trees (GBTs) to masked LLM outputs as a proxy model, leveraging the observation that LLM feature interactions are hierarchical (higher-order interactions accompanied by lower-order subsets).[1][2][3]
• •Extracts interactions from fitted GBTs, converting to Fourier representation for attribution definitions including Shapley-based ones.[1][5]
• •Evaluated on four high-dimensional datasets with hundreds of features; outperforms marginal attributions by 15-25% in faithfulness (R²) and uses 10× fewer inferences than SPEX to match performance.[1][4]
• •In attention head pruning, ProxySPEX identifies heads for removal across layer ranges (initial 1-3, middle 14-16, final 30-32), yielding higher test accuracies than LASSO baselines at various sparsity levels.[4]

🔮 Future ImplicationsAI analysis grounded in cited sources

⏳ Timeline