FlatQuant

Summary

• Identify the optimal affine transformation for each linear layer.
• Block-wise training strategy.
• Kronecker product.

Methodology

• Optimizes ${\mathbf{P}}^{\ast }=\arg {\min }_{\mathbf{P}}{\Vert \mathbf{Y}-\mathcal{Q}(\mathbf{XP})\mathcal{Q}({\mathbf{P}}^{-1}{\mathbf{W}}^{\top })\Vert }_F^2$
• Use Kronecker product instead of full-size transformation matrix to reduce computation and memory.
  • $\text{vec}(\mathbf{V})({\mathbf{P}}_1\otimes {\mathbf{P}}_2)=\text{vec}({\mathbf{P}}_1^{\top }\mathbf{V}{\mathbf{P}}_2)$
• Learnable per-channel scaling merged into weight.
• Learnable clipping thresholds $\in (0,1)$.
• Decoder-block-level training using MSE loss.
• Self-attention:
  • ${\mathbf{P}}_a,{\mathbf{P}}_o$ are decomposed.
  • ${\mathbf{P}}_h,{\mathbf{P}}_v$ are kept in original shape due to per-head quantization already facilitates cheap transformations.
  • ${\mathbf{P}}_o,{\mathbf{P}}_v$ are fused.
• MLP: ${\mathbf{P}}_{ug},{\mathbf{P}}_d$ are both decomposed.