FlatQuant

Summary

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

Methodology

• Use Kronecker product instead of full-size transformation matrix to reduce computation.
• 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.