UniQL: Unified Quantization and Low-rank Compression for Adaptive Edge LLMs

UniQL is a unified post-training framework that combines quantization and low-rank compression with on-device configurable pruning to efficiently deploy diverse LLM architectures on resource-constrained edge devices, achieving significant memory and throughput gains while maintaining high accuracy.

The Gist

Imagine you have a brilliant, world-class chef (a Large Language Model, or LLM) who can write poems, solve math problems, and code software. This chef is amazing, but they are also huge. They need a massive, state-of-the-art kitchen (a powerful server) with unlimited ingredients (memory) and a team of sous-chefs (computational power) to work.

Now, imagine you want to take this chef to a tiny food truck (your smartphone or a small edge device) to cook for people on the go. The problem? The food truck has a tiny fridge, a small stove, and the chef is too big to fit inside. If the truck is already carrying heavy boxes (other apps running), there's even less room.

UniQL is the revolutionary new "kitchen remodeling kit" that allows you to shrink this giant chef down to fit into the food truck without losing their cooking skills, and even lets you adjust the size of the chef on the fly depending on how crowded the truck is.

Here is how UniQL works, broken down into simple concepts:

1. The Problem: The "Fixed-Size" Trap

Usually, if you want to fit a big model onto a small device, you have to pick a version beforehand: a "Small Chef" or a "Medium Chef."

• The Issue: If you pick the "Medium Chef" and the food truck gets super crowded (high workload), the Medium Chef still won't fit, and the truck crashes.
• The Old Way: You'd have to stop cooking, go back to the big kitchen, shrink the chef further, and bring them back. This takes hours and is too slow.

2. The Solution: UniQL (The "Smart Shrink Ray")

UniQL is a one-time process done in the cloud that prepares the model to be elastic. Think of it like a Lego set that is pre-sorted.

• Sorting the Bricks (Weight Sorting):
  Imagine the chef's brain is made of millions of Lego bricks. Some bricks are critical (the ones that hold the structure together), and some are just decorative. UniQL sorts these bricks by importance. It puts the most important bricks at the front of the line and the least important ones at the back.

  • The Magic: It does this sorting incredibly fast (20x faster than old methods) by using a clever trick that avoids doing heavy math calculations (like "pseudo-inverses") that usually slow things down.

• The "Smart" Cut (Quantization & Low-Rank Compression):
  Once sorted, UniQL does two things:

  1. Shrinks the bricks: It turns the giant, heavy bricks into tiny, lightweight ones (Quantization) without breaking them.
  2. Removes the extras: It knows exactly which decorative bricks can be removed without the structure collapsing (Low-rank compression).

3. The Best Part: "Adaptive" Cooking

This is where UniQL shines. Because the bricks are pre-sorted, the food truck doesn't need to go back to the big kitchen to shrink the chef.

• The Scenario: You are driving. Suddenly, the truck gets very crowded (your phone is running many apps).
• The Action: The truck simply says, "Okay, we need 35% less space." It instantly pushes the "least important" bricks to the back and removes them. The chef is now smaller, fits perfectly, and keeps cooking.
• The Reversal: When the traffic clears, the truck says, "We have space again!" and the chef grows back to full size.
• Result: The model adapts to the device's current memory availability in real-time, something previous methods couldn't do.

4. Special Tricks for Different Kitchens

The paper mentions different types of "chefs" (Transformers, SSMs, and Hybrids). UniQL has specific tools for each:

• For the "Rotary" Chefs (RoPE): Some chefs use a special spinning technique to remember context. UniQL built a custom "fused engine" that handles this spinning even after the chef has been shrunk, saving time and energy.
• For the "State" Chefs (SSMs/Mamba): These chefs rely on a specific "state" to remember things. UniQL realized that cutting the wrong part of this state ruins the memory, so it developed a "state-aware" strategy to only cut the safe parts.

5. The Results: Faster, Smaller, Smarter

• Size: The models are 4 to 5.7 times smaller. A model that used to take 16GB of space now fits in 3GB.
• Speed: Because the chef is smaller and the kitchen is organized, they cook 2.7 to 3.4 times faster.
• Accuracy: Even after shrinking, the chef is still 95% as good as the original giant version.

Summary Analogy

Think of UniQL as a smart wardrobe organizer for your clothes (the AI model).

• Old Way: You have to buy three different wardrobes (Small, Medium, Large) and hope you pick the right one for the day. If you pick the wrong one, you can't fit your clothes.
• UniQL Way: You have one giant wardrobe where every shirt is pre-folded and sorted by importance. If you need to fit into a small suitcase (low memory), you just zip up the bottom half of the wardrobe, and the least important shirts disappear. If you need more space, you unzip it, and they reappear. You get the perfect fit instantly, every time, without repacking.

In short: UniQL makes giant AI models small enough to run on your phone, fast enough to be useful, and flexible enough to adapt to whatever your phone is doing at that exact moment.

Technical Summary

1. Problem Statement

Deploying Large Language Models (LLMs) on edge devices (e.g., mobile phones, VR/AR glasses) faces critical bottlenecks due to limited and dynamic memory resources.

• Resource Uncertainty: Edge device memory is shared and dynamically managed by the OS. A model compressed to a fixed size (e.g., 4GB) may fail to load if the current system workload consumes significant memory.
• Limitations of Existing Solutions:
  • Fixed Compression: Pre-compressing models into specific sizes requires storing multiple replicas, which is storage-inefficient.
  • Elastic Training: Methods like Flextron require extensive retraining on curated datasets with GPU resources, limiting applicability to specific model families.
  • Re-compression: Re-quantizing or re-pruning a model on-device to fit available memory is computationally prohibitive (taking hours on cloud GPUs).
• The Gap: There is a lack of a post-training framework that supports adaptive, on-device pruning rates across diverse architectures (Transformers, State Space Models, and Hybrids) without requiring retraining or multiple model replicas.

2. Methodology: The UniQL Framework

UniQL is a unified post-training framework that combines quantization and structured low-rank pruning. It operates in a "one-shot" fashion on a single server GPU, producing a model capable of adaptive pruning on the edge.

Core Pipeline

1. Structured Weight Sorting: Weights are reordered based on importance scores to facilitate structured pruning.
2. Masked Fine-Tuning: A LoRA-based recovery process where global pruning rates are randomly sampled during training to make the model robust to various pruning levels.
3. Quantization: The refined model is quantized to low bit-widths (e.g., INT4).
4. On-Device Adaptation: The deployed model can dynamically prune channels based on real-time memory availability (e.g., pruning 35% if memory is low, 15% if high).

Key Technical Innovations

A. Efficient Structured Weight Sorting
Instead of using time-consuming pseudo-inverses ($O(n^3)$ complexity), UniQL uses Ridge Leverage Scores derived from channel correlation matrices.

• MLP Layers: Computes correlation of intermediate activations to sort columns of $W_{up}$, $W_{gate}$, and rows of $W_{down}$. This avoids pseudo-inverse, speeding up sorting by 20×.
• MHSA (Self-Attention):
  • Query/Key: Uses symmetric sorting to handle Rotary Positional Embeddings (RoPE). The sorting indices are split to allow efficient gathering of sin/cos embeddings.
  • Value/Output: Employs Quantization-Aware SVD. Instead of standard truncation, the diagonal matrix $\Sigma$ (eigenvalues) is fused into the $U$ matrix ($W = (U\Sigma)V$). This ensures the scaling factors are preserved within the quantization group, minimizing quantization error in INT4.
• Mamba (SSM) & Hybrids:
  • Input Masks ($W_B, W_C$): Uses state-aware sorting considering the input-dependent step size $\Delta$.
  • State/Output ($W_Z, W_X, W_O$): Uses State-Aware Weight Sorting based on the SSM state matrix correlations, which are critical for SSM performance.

B. Masked LoRA Fine-Tuning

• The model is fine-tuned using LoRA on the sorted (un-pruned) weights.
• During training, a global pruning rate $P_t$ is randomly sampled, and the least important channels are masked out.
• This trains the model to be robust to any pruning rate within a range, enabling a single model to support multiple compression levels.

C. Fused Rotary Embedding (RoPE) Kernel

• Since structured sorting breaks the original positional indices, UniQL introduces a fused kernel that gathers the correct sin/cos embeddings based on the new sorted indices.
• This minimizes memory access overhead, providing a 10% latency reduction on edge devices.

3. Key Contributions

1. Unified Framework: The first post-training framework to systematically combine quantization and structured pruning for Transformers, SSMs (Mamba), and Hybrid models in a single pass.
2. Adaptive Edge Deployment: Enables configurable on-device pruning rates (up to 35%) without retraining or storing multiple model replicas. The model adapts to dynamic system workloads.
3. Algorithmic Efficiency:
   • Pseudo-inverse-free sorting: Achieves a 20× speedup in the compression pipeline compared to prior art (MoDeGPT).
   • Quantization-aware SVD: Reduces quantization errors by fusing eigenvalues into the weight matrix.
   • State-aware sorting: Specifically addresses the sensitivity of SSM state matrices.
4. Fused Kernels: A custom fused RoPE kernel that handles index gathering efficiently, reducing latency.

4. Experimental Results

The framework was evaluated on Llama-3.1-8B, Qwen-2.5-7B, Mamba-2-8B, Nemotron-H-8B, and Bamba-v2-9B.

• Accuracy Retention:
  • At 15% pruning, UniQL maintains accuracy within ~5% of the original FP16 model across all architectures.
  • At 35% pruning, it still delivers reasonable performance (e.g., 62.7% on Llama-3.1-8B vs. 74.0% baseline).
  • Outperforms baselines like MoDeGPT and SVD-LLM, which suffer significant accuracy drops or require extensive retraining.
• Memory Reduction:
  • Achieves 4× to 5.7× memory reduction compared to FP16.
  • With 35% pruning, models fit into 2.7GB–2.9GB, enabling deployment on devices with tight memory constraints.
• Latency & Throughput:
  • 2.7× to 3.4× improvement in token throughput (TPOT) on edge devices (NVIDIA Orin Nano).
  • On the Nano 8G, UniQL is 1.7× faster than TAO-HQQ (HQQ quantization) and 2.1× faster when combined with 35% pruning.
  • Energy Efficiency: Reduces energy consumption per request by ~40% compared to non-pruned quantized models.
• Compression Speed:
  • The entire pipeline (sorting, fine-tuning, quantization) takes ~7 hours on a single A6000 GPU, significantly faster than SVD-LLM (16+ hours) and MoDeGPT (7+ hours without FT, but slower with it).

5. Significance

• Practical Edge AI: UniQL solves the "fixed-size" problem in edge AI. It allows a single model to dynamically shrink to fit available memory, making LLMs viable for real-world scenarios where device load fluctuates.
• Architecture Agnostic: Unlike many compression methods tailored only to Transformers, UniQL successfully adapts to the emerging State Space Models (SSMs) and Hybrid architectures, which are crucial for efficient long-context modeling.
• Efficiency: By eliminating the need for pseudo-inverses and enabling one-shot compression, it lowers the barrier for deploying compressed models.
• Open Source: The authors release code and quantized models, fostering further research in efficient edge LLM deployment.

In summary, UniQL represents a significant step forward in making large language models truly "adaptive" for resource-constrained environments, balancing high compression ratios with minimal accuracy loss and enabling real-time responsiveness to device constraints.