Attribution-guided Pruning for Compression, Circuit Discovery, and Targeted Correction in LLMs

2025-06-16

AI Summary

This paper addresses the dual challenges of deploying large language models (LLMs)—namely, high memory and computational overhead, and safety risks such as toxic outputs. To this end, we propose an attribution-guided unstructured pruning framework grounded in Layer-wise Relevance Propagation (LRP). This work is the first to extend LRP to fine-grained attribution and pruning in LLMs, unifying three objectives: (1) efficient model compression—achieving >50% parameter pruning on Llama and OPT with negligible performance degradation; (2) automatic identification of task-critical subcircuits (e.g., indirect object recognition pathways); and (3) targeted mitigation of harmful behaviors. Unlike conventional parameter-level pruning, our approach enables functional-level intervention, enhancing both interpretability and safety. By bridging explainability and structured model editing, it establishes a novel paradigm for lightweight, trustworthy LLM deployment.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Large Language Models (LLMs) are central to many contemporary AI applications, yet their extensive parameter counts pose significant challenges for deployment in memory- and compute-constrained environments. Recent works in eXplainable AI (XAI), particularly on attribution methods, suggest that interpretability can also enable model compression by identifying and removing components irrelevant to inference. In this paper, we leverage Layer-wise Relevance Propagation (LRP) to perform attribution-guided pruning of LLMs. While LRP has shown promise in structured pruning for vision models, we extend it to unstructured pruning in LLMs and demonstrate that it can substantially reduce model size with minimal performance loss. Our method is especially effective in extracting task-relevant subgraphs -- so-called ``circuits'' -- which can represent core functions (e.g., indirect object identification). Building on this, we introduce a technique for model correction, by selectively removing circuits responsible for spurious behaviors (e.g., toxic outputs). All in all, we gather these techniques as a uniform holistic framework and showcase its effectiveness and limitations through extensive experiments for compression, circuit discovery and model correction on Llama and OPT models, highlighting its potential for improving both model efficiency and safety. Our code is publicly available at https://github.com/erfanhatefi/SparC3.

Problem

• Compress LLMs for memory-constrained environments efficiently
• Discover task-relevant subgraphs (circuits) in LLMs
• Correct spurious behaviors by removing toxic output circuits

Innovation

• Attribution-guided pruning using Layer-wise Relevance Propagation
• Extracting task-relevant subgraphs for core functions
• Selectively removing circuits for model correction