# InMemQK: A Product Quantization Based MatMul Module for Compute-in-Memory Attention Macro
## 3 Optimized Product Quantization and InMemQK Based on Time-Domain MAC
### Hierarchical Feature Reconstruction and Reverse Layer-Wise Fine-Tuning  
Substitution of approximate matrix multiplication (AMM), based on the L2 norm distance metric PQ, typically results in decreased model performance. Experiments [25] show model accuracy diminishes as the number of AMM replacement layers increases. This is because the K-means clustering algorithm uses a greedy algorithm to minimize Euclidean distance between original data and centroids in the subvector space, which may lead to local optima.  
To address this issue, we propose a hierarchical feature reconstruction (HFR) algorithm. This algorithm employs mean squared error (MSE) to quantify the disparity between keys before and after quantization and utilizes gradient descent to optimize the loss function. The algorithm's pseudo-code is presented in Algorithm 1. Following hierarchical feature reconstruction, nearly every layer of the network model approaches an optimal solution. Nevertheless, this still diverges from the ultimate objective of Deep Neural Networks (DNNs), which aims to minimize the overall loss function. Furthermore, fine-tuning the model using soft-PQ techniques [26] fails to fully restore accuracy. This is attributed to the accumulated errors in the approximate attention module, which causes the distribution of quantized model parameters to deviate substantially from that of the pre-trained model, whereas model fine-tuning depends on well-initialized parameters [27].  
```
1:Input: Calibration dataset, pre-trained model
2:Output: Codebooks for key and value
3:Hierarchical feature reconstruction
4:for layer in [\(layer_{1}\), \(layer_{2}\), \(\dots\), \(layer_{12}\)] do
5:\(layer_{PQ}\leftarrow\) PQ(layer) [PQ means product quantizing this layer]
6:\(codebook_{key}\), \(codebook_{value}\gets layer_{PQ}\)(key, value)
7:\(loss_{key}\leftarrow\) MSE(\(key\), \(codebook_{key}\)) [MSE means mean squared error]
8:\(loss_{value}\leftarrow\) MSE(\(value\), \(codebook_{value}\))
9:\(loss\gets loss_{key}\) + \(loss_{value}\)
10: Minimize the loss using gradient descent
11:endfor
12:Reverse layer-wise fine-tuning
13:for i = N to 1 do
14: Load the model after fine-tuning the codebook at layer i + 1
15: Fine-tune the codebook at layer i
16: Save model
17:endfor
```  
**Algorithm 1** Optimized product quantization.  
To tackle this challenge, we introduced a reverse layer-wise fine-tuning (RL-FT) algorithm. This algorithm leverages the characteristic that the error magnitude increases with the number of quantization layers, fine-tuning the centroids in the codebooks layer by layer from back to front. To enhance fine-tuning efficiency, we initiate the RL-FT from the N-th last layer, where N (commonly set to 6) is the layer index where the PQ accuracy drops below 95% of the original accuracy. First, we load the codebooks from the N-th layer to the 12-th layer and freeze its parameters, then fine-tune the codebooks of the (N-1)-th layer, continuing in this manner by loading the codebooks from the (N-1)-th layer to the 12-th layer and fine-tuning the (N-2)-th layer.  
During this process, we incorporated the soft-PQ technique, employing softmax to approximate argmax to circumvent non-differentiability issues. Furthermore, we considered the variations and resolution of NVM devices, simulating the non-ideal behavior of real devices by quantizing the centroids and introducing Gaussian noise. Comprehensive results are presented in Section 4.2.