Weight Clustering

Weight clustering operates by analyzing the distribution of a model's weight values and grouping them into a predefined number (k) of clusters. The centroid of each cluster represents a shared weight value. The original high-precision weights are then replaced with integer indices pointing to their assigned centroid. This process transforms a large matrix of unique 32-bit floating-point numbers into a much smaller matrix of integers (e.g., 8-bit) and a tiny codebook of centroid values, drastically reducing the model's size.

Unlike quantization-aware training, weight clustering is typically applied after a model has been fully trained. This makes it a straightforward, low-cost compression method. The process involves:

• Running a clustering algorithm (like k-means) on the trained weights.
• Replacing weights with cluster indices.
• Optionally performing a short fine-tuning step where the centroids (not the indices) are adjusted to recover accuracy lost during the replacement. This post-training nature allows for rapid compression of pre-existing models without access to the original training pipeline.

The primary benefit of weight clustering is a significant reduction in model storage. However, it introduces a runtime trade-off. During inference, the integer indices must be de-referenced through a lookup table to fetch the actual centroid weight value for computation. This adds a small overhead compared to direct weight access. The technique is therefore ideal for deployment scenarios where storage (e.g., flash memory on a microcontroller) is the primary constraint, and the latency of occasional table lookups is acceptable.

The k-means algorithm is commonly used for clustering. The most critical hyperparameter is k, the number of clusters. A smaller k yields higher compression (fewer bits needed for indices) but risks greater accuracy loss. A larger k preserves accuracy better but offers less compression. For example, k=256 allows indices to be stored in 8 bits (1 byte). The choice of k is a direct balance between the target compression ratio and the acceptable accuracy drop for a given application.

Weight clustering is highly complementary to other TinyML compression methods and is often used in a pipeline:

• Pruning First: Applying pruning to remove insignificant weights creates a sparse model. Clustering is then applied to the remaining non-zero weights for additional compression.
• Quantization After: The centroid values in the codebook, which are typically stored in full precision after fine-tuning, can themselves be quantized (e.g., to INT8) for final deployment, squeezing out further memory savings. This layered approach is key to achieving extreme compression for microcontroller deployment.

For efficient execution on microcontrollers, the inference engine must support the clustered weight format. This requires:

• A mechanism to store the codebook (centroid values) in fast memory.
• Kernel operations that efficiently perform the index lookup and multiplication during the convolution or matrix multiplication step.
• Careful management of memory bandwidth, as fetching weights now involves an indirect read. Frameworks like TensorFlow Lite Micro provide operators and conversion tools that handle this format, abstracting the complexity from the developer.

This method approximates a large weight matrix (e.g., in a fully connected or convolutional layer) as the product of two or more smaller matrices. It exploits the idea that weight matrices are often low-rank—meaning they contain redundant information.

• Mechanism: A weight matrix W of size [m x n] is factorized into U ([m x r]) and V ([r x n]), where the rank r is much smaller than m and n.
• Compression Gain: Reduces parameters from m*n to r*(m+n).
• Trade-off: Introduces an extra sequential computation step during inference.

Sparsity is a property of a model where a significant portion of its weights are zero. It is a direct outcome of pruning. The key challenge is leveraging sparsity for actual speedup, which depends on hardware support.

• Structured Sparsity: Enables immediate speedups on standard hardware (CPU/GPU) because it removes whole blocks of computation.
• N:M Fine-Grained Sparsity: A pattern like 2:4 sparsity (2 non-zero values in every block of 4) is natively accelerated on modern NVIDIA Ampere GPU Tensor Cores.
• Software Libraries: Frameworks like TensorFlow Lite and PyTorch provide APIs to induce and leverage sparsity.