汤不热吧详解神经网络权重的聚类压缩算法：如何利用 Codebook 降低移动端内存带宽压力
在移动端和边缘设备上部署深度学习模型时，模型体积和推理时的内存带宽往往是最大的性能瓶颈。传统的量化（如INT8）可以压缩数据，但聚类压缩提供了一种更为灵活且通常更接近原始精度的方法。
本文将聚焦于权重聚类压缩技术，详细解释如何通过K-Means算法生成Codebook（码本），从而将模型参数从浮点数（Float32）转化为紧凑的索引（Index），达到大幅降低存储和内存带宽的目的。
1. 为什么需要聚类压缩？
对于一个典型的移动端模型，其权重通常以Float32存储（4字节/参数）。在推理过程中，这些权重需要不断地从内存中加载到计算单元（如DSP或NPU）。如果模型参数量巨大，频繁地加载4字节数据会严重占用内存带宽，导致推理速度受限。
权重聚类压缩的核心思想是：找到权重值中的少量代表性数值（质心），然后用这些质心的索引来替换原始的权重值。
2. K-Means聚类与Codebook生成
我们使用K-Means聚类算法来寻找最优的质心（Centroids），这些质心构成了我们的Codebook。
假设我们希望将权重压缩到$B$位（例如$B=8$），那么Codebook中将包含 $K=2^B$ 个质心。
步骤概述：
1. 展平模型权重。(W \rightarrow w_{flat})
2. 对 $w_{flat}$ 执行K-Means聚类，得到 $K$ 个质心 $C$ (Codebook)。
3. 将 $w_{flat}$ 中的每个原始权重替换为其最近质心的索引 $I$。
最终存储的数据是：Codebook $C$ (少量Float32) + 索引矩阵 $I$ (大量Uint8)。
3. 代码实操：利用Codebook进行权重压缩
我们将使用Python和scikit-learn来模拟对一个全连接层权重进行聚类压缩的过程。
环境准备：
pip install numpy scikit-learn
Python 代码示例：
import numpy as np
from sklearn.cluster import KMeans
# 1. 模拟一个全连接层的权重矩阵 (Float32)
INPUT_DIM = 256
OUTPUT_DIM = 512
weights_float32 = np.random.randn(INPUT_DIM, OUTPUT_DIM).astype(np.float32)
TOTAL_ELEMENTS = weights_float32.size
print(f"原始权重形状: {weights_float32.shape}")
print(f"总元素数量: {TOTAL_ELEMENTS}")
# 设定目标压缩位数 (例如：8位，Codebook大小 K=256)
K = 256
# 2. 展平权重并执行 K-Means 聚类
# K-Means 在这里用于找到 K 个最佳的代表值（质心）
weights_flat = weights_float32.reshape(-1, 1)
print(f"开始进行 K-Means 聚类 (K={K})...")
kmeans = KMeans(n_clusters=K, random_state=42, n_init=10, max_iter=100)
kmeans.fit(weights_flat)
# Codebook: 质心是压缩后的浮点值
codebook = kmeans.cluster_centers_.astype(np.float32).flatten()
# 索引矩阵: 原始权重属于哪个质心
indices = kmeans.labels_.astype(np.uint8) # 8位索引足够存储 0-255
indices_matrix = indices.reshape(weights_float32.shape)
# 3. 结果分析与压缩率计算
# 原始存储大小 (Bytes)
size_original = TOTAL_ELEMENTS * 4
# 压缩后存储大小 (Bytes)
# 索引矩阵: 1 byte * 总元素数量
size_indices = indices.size * 1
# Codebook大小: 4 bytes * K 个质心
size_codebook = K * 4
size_compressed = size_indices + size_codebook
compression_ratio = size_original / size_compressed
print("\n--- 压缩结果 ---")
print(f"Codebook大小 (K={K}): {codebook.size} 个 Float32 值")
print(f"Codebook存储大小: {size_codebook / 1024:.2f} KB")
print(f"索引矩阵存储大小: {size_indices / 1024:.2f} KB")
print(f"原始模型权重大小: {size_original / 1024:.2f} KB")
print(f"压缩后模型总大小: {size_compressed / 1024:.2f} KB")
print(f"存储压缩比率: {compression_ratio:.2f}:1")
# 4. 模型重建（用于验证精度损失）
# 重建方法：使用索引去查Codebook
weights_reconstructed = codebook[indices].reshape(weights_float32.shape)
# 计算重建误差
error_mse = np.mean((weights_float32 - weights_reconstructed)**2)
print(f"重建均方误差 (MSE): {error_mse:.6f}")
4. Codebook 如何降低内存带宽压力？
这个压缩方案带来的核心益处在于推理阶段：
- 降低存储体积： 上述代码示例展示了存储压缩比接近 4:1（因为我们将4字节的Float32替换成了1字节的索引，且Codebook体积占比极小）。
- 降低内存带宽： 在实际运行推理时，计算单元需要读取权重。在聚类压缩的模型中，计算单元不再需要读取完整的4字节浮点数，而是读取1字节的索引。
虽然在计算过程中，计算单元仍然需要查找Codebook（即：output = input * Codebook[Index]），但由于Codebook非常小（例如256个浮点数），它可以很容易地被缓存到芯片上的SRAM或寄存器中，基本不需要消耗外部DRAM的带宽。
因此，在矩阵乘法操作中，大部分时间消耗在从主存中读取权重数据，此时带宽需求从 $4N$ 降低到了 $N$ (其中 $N$ 是参数量)，从而释放了大量内存带宽，显著提升了移动端推理速度。