Enhancing deep neural networks with feature attribution : Scaling from edge models to foundation models

Abstract

Feature attribution is an effective technique in interpreting deep neural networks across a wide range of scales, from compressed convolutional neural networks on edge devices to cloud-based large language models. While feature attribution methods are traditionally used for explanation, this dissertation explores the potential of feature attribution for identifying and addressing model vulnerabilities, ultimately improving model performance and reliability. The study progressively expands its model scale, deriving appropriate techniques for each scale. At the smallest scale, this dissertation identifies that while prediction accuracy is preserved during model compression, attribution map integrity is significantly compromised. To address this, we propose an attribution map matching loss function that simultaneously enhances both model explainability and accuracy. For mid-scale applications, we develop a novel tabular data oversampling framework using transformer-based language models. This approach leverages column-wise self-attention attribution to identify class-representative features, enabling targeted column imputation for generating high-quality synthetic samples. The dissertation then expands to large-scale language models, exploring both input and output token-level attribution. From the input perspective, we propose an attribution-guided key-value cache compression methodology that identifies and differentially preserves crucial information, optimizing memory usage while maintaining model performance. From the output perspective, we develop a token-level importance-based verifier that enhances mathematical reasoning capabilities in large language models through fine-grained supervision signals.