LLM Pruning using ADMM (ADMM-Iter)

This repository implements and extends ADMM-Iter, an optimization-based pruning framework for Large Language Models (LLMs), as proposed in
“Fast and Optimal Weight Update for Pruned Large Language Models.”

The project focuses on reproducing the original results and exploring structured and adaptive pruning strategies for efficient LLM deployment.

---

Overview

Large Language Models achieve strong performance but are expensive to deploy due to their size and compute requirements.
This work applies ADMM-based layer-wise pruning to reduce model parameters while preserving accuracy and convergence properties.

Key goals:

• Reproduce ADMM-Iter pruning across different model scales
• Analyze pruning tolerance and layer sensitivity
• Extend the framework to structured sparsity (hardware-friendly pruning)

---

Key Contributions

• Reimplementation of ADMM-Iter with layer-wise weight reconstruction
• Validation on models ranging from small MLPs to billion-parameter LLMs
• Extension to Structured ADMM-Iter Pruning (e.g., 2:4, 4:8 sparsity)
• Empirical analysis of:
  • Convergence behavior
  • Perplexity vs sparsity trade-offs
  • Attention vs MLP layer sensitivity

---

Models Evaluated

• MNIST MLP

  • Baseline sanity check for ADMM convergence

• DistilBERT (SST-2)

  • 66M parameters
  • Binary sentiment classification

• GPT-2

  • Small (124M)
  • Medium (355M)
  • Large (774M)

• TinyLlama-1.1B

  • Decoder-only LLaMA-style architecture
  • Evaluated on WikiText-2 (language modeling)

---

Pruning Methods

• Unstructured ADMM Pruning
• Structured ADMM Pruning
  • 2:4 sparsity
  • 4:8 sparsity

ADMM enables efficient pruning by reusing a single matrix inversion per layer, reducing computational complexity from
O(m³n) → O(m²n).

---

Key Observations

• ADMM converges within 5–10 iterations for most layers
• Attention layers (especially q_proj) are more sensitive than MLP layers
• Unstructured pruning achieves the best perplexity at high sparsity
• Structured pruning incurs higher overhead but enables hardware acceleration
• GPT-2 Large pruning closely matches results reported in the original paper

---

Metrics

• Accuracy (classification tasks)
• Perplexity (PPL) for language modeling
• Convergence error
• Pruning time vs sparsity

---

Motivation for Ongoing Work

Current research directions include:

• Adaptive pruning
• Model-guided pruning (e.g., letting the LLM decide what and how much to prune)
• Importance-aware and data-driven pruning criteria
• Combining pruning with quantization

---

References

• Fast and Optimal Weight Update for Pruned Large Language Models
• WikiText-2 Language Modeling Benchmark
• Hugging Face Transformers

---

This repository is intended as a research prototype and experimental framework rather than a production-ready library.