Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix

TL;DR

This paper presents a new non-linear pruning method for attention matrices in LLMs that improves efficiency and maintains performance, addressing the limitations of linear approximation methods.

Contribution

It introduces a gradient descent-based pruning approach that directly optimizes attention matrix approximation, with theoretical guarantees and superior empirical results.

Findings

Significant reduction in computational costs compared to state-of-the-art methods

Maintains model performance after pruning

Provides theoretical convergence guarantees for the pruning method

Abstract

Large Language Models (LLMs) have shown immense potential in enhancing various aspects of our daily lives, from conversational AI to search and AI assistants. However, their growing capabilities come at the cost of extremely large model sizes, making deployment on edge devices challenging due to memory and computational constraints. This paper introduces a novel approach to LLM weight pruning that directly optimizes for approximating the attention matrix, a core component of transformer architectures. Unlike existing methods that focus on linear approximations, our approach accounts for the non-linear nature of the Softmax attention mechanism. We provide theoretical guarantees for the convergence of our Gradient Descent-based optimization method to a near-optimal pruning mask solution. Our empirical results demonstrate the effectiveness of our non-linear pruning approach in maintaining model performance while significantly reducing computational costs, which is beyond the current state-of-the-art methods, i.e., SparseGPT and Wanda, by a large margin. This work establishes a new theoretical foundation for pruning algorithm design in LLMs, potentially paving the way for more efficient LLM inference on resource-constrained devices.