TLDR: STaMP (Sequence Transformation and Mixed Precision) is a novel quantization method for generative AI models that improves accuracy at low bit-widths (e.g., 4-bit activations). It achieves this by applying linear transformations along the sequence dimension to exploit local data correlations, then uses a mixed-precision strategy to allocate more bits to important “energy-concentrated” tokens. This approach is complementary to existing feature-based quantization methods, offers significant accuracy gains for both large language and vision models, and introduces minimal computational overhead.
In the rapidly evolving world of artificial intelligence, generative models like large language models (LLMs) and large vision models (LVMs) are achieving remarkable feats. However, their immense computational and memory demands pose significant challenges for efficient deployment, especially on devices with limited resources. A key technique to address this is quantization, which reduces the precision of model weights and activations to lower bit-widths, thereby cutting down latency, power consumption, and memory footprint.
While quantization is crucial, pushing activation bit-widths below eight bits often leads to a sharp decline in model accuracy. Previous research has explored using invertible linear transformations, such as rotations, to reparameterize feature channels and weights, helping to mitigate this accuracy degradation. These methods primarily operate along the ‘feature dimension’ of the data.
A new research paper titled “STaMP: Sequence Transformation and Mixed Precision for Low-Precision Activation Quantization” introduces a novel strategy called STaMP, which stands for Sequence Transformation and Mixed Precision. This approach takes a different route by applying linear transformations along the ‘sequence dimension’ of the data. This is particularly insightful because language and visual data exhibit strong local correlations—think of adjacent words in a sentence or neighboring pixels in an image. STaMP leverages this inherent structure to improve quantization efficiency.
The core idea behind STaMP, developed by Marco Federici, Riccardo Del Chiaro, Boris van Breugel, Paul Whatmough, and Markus Nagel from Qualcomm AI Research, is to concentrate the ‘energy’ or importance of activations into a small number of tokens within a sequence. Once this energy is concentrated, a mixed-precision strategy is employed: these high-energy tokens are kept at a higher precision (e.g., 8 bits), while the majority of other tokens are quantized to a much lower precision (e.g., 4 bits). This clever allocation allows the model to maintain overall accuracy at a significantly lower average activation bit-width.
The researchers found that while the Karhunen-Loève Transform (KLT) is theoretically optimal for energy concentration, it is too computationally intensive for practical use. Instead, they identified that the autocorrelation matrix of typical LVM and LLM activations has a structured, Toeplitz-like form, which can be efficiently approximated by transforms like the Discrete Cosine Transform (DCT) or the Discrete Wavelet Transform (DWT). The DWT was chosen for its computational efficiency, reducing complexity significantly while still effectively concentrating energy.
STaMP is not designed to replace existing quantization methods but rather to complement them. Unlike feature transformations that affect weights, sequence transformations do not, making them orthogonal to advanced weight quantization techniques. The paper demonstrates that STaMP can be combined with popular feature transformation and weight quantization methods, leading to even greater improvements in model accuracy.
Experimental results on recent LVM architectures like PixArt-Σ and SANA, and LLM architectures including Llama 3 8B and Qwen 2.5 3B instruct, consistently show that STaMP significantly enhances low bit-width activation quantization. For instance, when combined with other methods, STaMP leads to visually more accurate image generations and improved perplexity scores for language models, especially in scenarios where baselines struggle at 4-bit quantization.
Furthermore, the computational overhead of STaMP with DWT is minimal, comparable to other efficient transforms, accounting for less than 5% of the total runtime in a denoising step. This suggests that STaMP offers a practical, training-free solution for deploying high-performance generative models in resource-constrained environments.
Also Read:
- Direct Semantic Learning from Compressed Files with TEMPEST
- Optimizing LLM Memory for Extended Text Processing
By drawing inspiration from traditional signal processing techniques like those used in JPEG and MP3, STaMP opens new avenues for optimizing generative AI models. It highlights the potential of exploiting sequence-level correlations to push the boundaries of low-precision quantization, making powerful AI models more accessible and efficient. You can read the full research paper here.
