Irrational Complex Rotations Empower Low-bit Optimizers

Zhen Tian, Xin Zhao, Ji-Rong Wen

Keywords: Irrational Complex Rotation, Training-oriented Quantization, Memory-efficient Optimization

Abstract:

In this paper, we propose a novel optimizer state compression algorithm, namely \textbf{$\pi$-Quant}, which leverages the properties of irrational numbers (\eg $\pi$) for memory-efficient training.
The core idea is based on our mathematical findings, which show that a pair of parameters can be represented by a single rotation angle using the complex rotation scheme.
Building on this insight, we map the parameters into a complex space and perform quantization using the corresponding rotation angles.
To efficiently integrate it into optimization process, we develop an efficient system of geometric equations that computes the precise rotation angles with linear complexity.
We evaluate $\pi$-Quant on a wide range of tasks. Our experiments show that it can reduce the bit-width of parameters to 3.32-bit, achieving a 41.8\% decrease in GPU memory usage, all while maintaining full accuracy. \textcolor{blue}{We have submitted the code in supplementary materials}.

Primary Area: Deep learning (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 9082