3BASiL: An Algorithmic Framework for Sparse plus Low-Rank Compression of LLMs

This paper introduces 3BASiL-TM, an efficient one-shot post-training framework that combines a novel 3-Block ADMM algorithm with a transformer-matching refinement step to significantly improve the accuracy and speed of sparse plus low-rank compression for Large Language Models.

The Gist

Imagine you have a giant, incredibly detailed encyclopedia (a Large Language Model, or LLM) that knows almost everything. It's brilliant, but it's also massive. It takes up so much space that it won't fit in your backpack (your phone or laptop), and it's too heavy to carry around quickly. You need to shrink it down without losing its ability to tell good stories or solve math problems.

This paper introduces a new, super-smart way to shrink these giant models called 3BASiL.

Here is how it works, using some everyday analogies:

1. The Problem: The "Heavy Suitcase"

Think of the original AI model as a suitcase packed with thousands of heavy bricks.

• Old methods tried to shrink it by either:
  • Throwing away bricks: Removing many bricks (Pruning/Sparsity). This makes the suitcase lighter, but if you throw away the wrong ones, the suitcase falls apart.
  • Flattening the bricks: Compressing them into thin sheets (Low-Rank). This saves space, but you lose some of the 3D detail.
• The Issue: Previous attempts to do both at the same time were like trying to juggle while blindfolded. They would remove some bricks, then flatten some, then remove more, often messing up the structure and making the AI "forget" things (losing accuracy).

2. The Solution: 3BASiL (The "Master Organizer")

The authors created a new algorithm called 3BASiL. Think of it as a super-organized packing robot that doesn't just throw things away; it rearranges the whole suitcase perfectly.

It uses a mathematical trick called ADMM (which sounds fancy, but think of it as a "Three-Step Dance"):

1. Step 1 (The Sparse Step): The robot looks at the suitcase and says, "Okay, let's remove the bricks that aren't doing much work." It creates a "sparse" version (lots of empty space).
2. Step 2 (The Low-Rank Step): Then it says, "For the bricks we kept, let's flatten the ones that are redundant." It creates a "low-rank" version (compressed details).
3. Step 3 (The Harmony Step): Instead of doing these steps one after another and hoping for the best, 3BASiL does them simultaneously. It constantly checks: "If I remove this brick, does the flattened part need to change to compensate?"

The Result: The suitcase is now half the size, but the contents are arranged so perfectly that the AI still works almost as well as the giant original.

3. The Secret Sauce: "Transformer Matching" (The "Soundcheck")

Even with a great packing job, sometimes the suitcase feels a little "off" when you try to walk with it. The authors added a second step called Transformer Matching (TM).

• The Analogy: Imagine you've packed a band's instruments into a van. You think you did a good job. But before you hit the road, you do a soundcheck. You play a few notes and listen to how the whole band sounds together.
• What it does: Instead of just checking if individual instruments (layers) are packed right, this step checks if the whole band (the whole transformer block) sounds right. It makes tiny adjustments to the packing so that the final output matches the original giant model perfectly.
• Why it's cool: This step is like a "universal adapter." You can use it with any compression method, not just 3BASiL, to make it work better.

4. The Payoff: Fast, Light, and Smart

The paper shows that this new method is a winner in three ways:

• Smarter: It shrinks the model (specifically the Llama-3-8B model) so much that it loses very little intelligence. In fact, it reduced the "confusion" (perplexity) of the AI by 30% compared to other methods. It's like shrinking a backpack but keeping the map inside perfectly legible.
• Faster: The packing process itself is 2.5 times faster than the current best methods. It's like going from hand-packing a suitcase to using a vacuum-seal machine.
• Ready for the Future: The compressed model is set up perfectly for "LoRA" (a technique to teach the AI new tricks). It's like packing the suitcase so that when you get to your destination, you can instantly swap in a new set of clothes without unpacking everything.

Summary

3BASiL is a new, efficient way to shrink giant AI brains. Instead of clumsily cutting and pasting parts of the brain, it uses a smart, three-step dance to reorganize the information, followed by a "soundcheck" to ensure everything still works perfectly. The result is an AI that fits in your pocket, runs fast on your phone, and still knows how to write poetry, code, and solve math problems.

Technical Summary

1. Problem Statement

Large Language Models (LLMs) face significant deployment challenges due to their massive computational and memory requirements. While model compression techniques like pruning and quantization exist, they often lead to substantial performance degradation. A promising direction is Sparse plus Low-Rank (S + LR) decomposition, where pre-trained weights $W$ are approximated as $W \approx S + L$, with $S$ being a sparse matrix and $L$ being a low-rank matrix.

However, existing S + LR methods suffer from two main limitations:

1. Suboptimal Joint Optimization: Most current approaches use alternating minimization (separately optimizing $S$ and $L$), which lacks convergence guarantees and often fails to effectively optimize the interaction between the sparse and low-rank components.
2. Layer-wise Limitations: Existing methods typically optimize layers independently. This ignores the cumulative error propagation across transformer layers and fails to align the compressed model's output with the dense model at the transformer block level.

2. Methodology

The authors propose 3BASiL-TM, a two-stage framework designed to address these gaps.

A. 3BASiL: 3-Block ADMM for Layer-wise Decomposition

The core of the method is a novel 3-Block Alternating Direction Method of Multipliers (ADMM) algorithm tailored for the S + LR decomposition problem.

• Problem Formulation: The goal is to minimize the reconstruction error between the original layer output and the decomposed output, subject to sparsity and rank constraints:
  $$\min_{S, L} \frac{1}{2}\|XW - X(S+L)\|_F^2 + \frac{\lambda}{2}\|W - (S+L)\|_F^2$$
  subject to $S \in \mathcal{S}$ (sparsity set) and $\text{rank}(L) \leq r$.

• 3-Block ADMM Strategy: Unlike standard 2-block ADMM, the authors introduce an auxiliary variable $D$ (a copy of $S$) to reformulate the problem into three blocks: $S$ (sparse), $L$ (low-rank), and $D$ (constrained sparse).

  • S-update: Solves a quadratic subproblem with a closed-form solution involving matrix inversion.
  • L-update: Solves a rank-constrained subproblem using a closed-form solution based on Reduced-Rank Regression (via SVD), avoiding iterative gradient descent on low-rank factors.
  • D-update: Projects the result onto the sparsity constraint set (e.g., magnitude-based pruning).
  • Dual Update: Updates the Lagrange multiplier.

• Theoretical Guarantee: The paper provides a rigorous convergence proof (Theorem 1), showing that the sequence of decompositions converges to a limit matrix provided the penalty parameter $\rho_t$ increases sufficiently fast. This is a significant theoretical advancement over prior alternating minimization methods which lack such guarantees.

• Efficiency: The authors employ computational tricks, such as pre-computing eigenvalue decompositions and using randomized SVD, to achieve an overall time complexity of $O(N^3)$ per iteration, making it scalable for large LLMs.

B. Transformer Matching (TM): Global Refinement

After layer-wise decomposition, the authors introduce a Transformer Matching (TM) step to refine the components globally.

• Objective: Instead of optimizing layer-by-layer, TM optimizes the entire transformer block to minimize the discrepancy between the output of the compressed block and the original dense block:
  $$\min \left\| T_i(X_i; \{W\}) - T_i(X_i; \{S+L\}) \right\|_F^2$$
• Mechanism: It uses gradient-based optimization (Adam) to jointly update the weights of the sparse and low-rank components across all layers within a transformer block.
• Advantages:
  • Memory Efficient: It operates on small calibration data chunks and only requires forward/backward passes within a single transformer block.
  • Universal: It can be applied as a post-processing step to any S + LR decomposition method (including pure sparsity methods) to improve performance.
  • Smart Initialization: The refined low-rank components serve as a superior initialization for subsequent LoRA fine-tuning.

3. Key Contributions

1. 3BASiL Algorithm: A novel 3-Block ADMM framework that explicitly models the interaction between sparse and low-rank components under a unified objective, offering theoretical convergence guarantees and closed-form updates for the low-rank component.
2. Transformer Matching (TM): A universal, memory-efficient refinement procedure that aligns transformer-level outputs, significantly improving the quality of the sparse components and reducing accumulated errors.
3. State-of-the-Art Performance: The combined 3BASiL-TM method achieves superior results compared to existing S + LR baselines (like OATS and HASSLE-free) in terms of perplexity and zero-shot task accuracy.
4. Computational Efficiency: The method is significantly faster than SOTA baselines (e.g., >2.5x faster on A100 GPUs) due to the closed-form low-rank updates and efficient implementation.

4. Experimental Results

The authors evaluated 3BASiL-TM on Llama-3 (1B, 3B, 8B) and OPT-30B models using various configurations (e.g., 2:4 or 3:8 structured sparsity + 64 low-rank).

• Perplexity Reduction:
  • For Llama-8B under a (2:4 Sparse + 64 LR) configuration, 3BASiL-TM reduced the WikiText2 perplexity gap relative to the dense model by over 30% compared to prior methods.
  • It consistently outperformed baselines (OATS, Hf-SparseGPT, Hf-ALPS) across all datasets (WikiText2, Penn Treebank, C4).
• Zero-Shot Tasks: The method achieved higher average scores on eight zero-shot tasks (PIQA, ARC, HellaSwag, etc.) compared to competitors.
• LoRA Fine-Tuning: When used as a smart initialization for LoRA, 3BASiL-TM maintained a significant advantage over other methods even after fine-tuning, achieving ~8% lower perplexity than the next best competitor in aggressive compression regimes.
• Runtime: On an A100 GPU, 3BASiL-TM achieved >2.5x speedup in compression runtime compared to the HASSLE-free-ALPS baseline. On an L40 GPU, it was >3x faster.

5. Significance

This work represents a major step forward in LLM compression by bridging the gap between theoretical optimization and practical performance.

• Theoretical Rigor: It moves beyond heuristic alternating minimization to a provably convergent optimization framework for S + LR decomposition.
• Practical Utility: The "one-shot" nature of the method (no expensive retraining) combined with the TM refinement makes it highly viable for deploying LLMs on resource-constrained devices.
• Universality: The Transformer Matching procedure is a generalizable tool that can enhance any existing pruning or decomposition method, suggesting a new paradigm for post-training compression where layer-wise reconstruction is followed by block-level alignment.

In summary, 3BASiL-TM offers a faster, more accurate, and theoretically grounded approach to compressing Large Language Models, effectively balancing the trade-off between model size, inference speed, and task performance.