Second Response to Reviewer d65g
Thank you again for your comments. Regrading the questions:
**Results on generative modeling task.**
We trained a StyleGAN2-ADA with both 32-bit Adam and our 4-bit Adam on the conditional CIFAR-10 dataset.
We used the same hyperparameters as the official implementation, except for a larger batch size that facilitated quicker training. The following table shows the best achieved FID score throughout the training process.
During the training dynamics, several quantites, including generator loss, descriminator loss, fake image scores, and real image score closely align for both optimizers.
When evaluated in terms of FID, the two models trained with different optimizers exhibit a slight gap in performance.
Furthermore,
in the earlier conducted diffusion model task, we observed a close alignment in the loss dynamics between the 32-bit Adam and our 8-bit Adam optimizers.
| Method | FID |
|-------------------|-------------:|
| 32-bit Adam | 2.40|
| 4-bit Adam (ours)| 2.89|
**Theory about block size.**
In our analysis, the quantization variance only impacts the magnitude of the stationary point, i.e., the case when $T \to \infty$, and it does not relate to the convergence speed.
This is similar to the variance from stochastic gradient.
Then, we aim to characterize the influence of block-wise normalization on quantization variance.
When employing a linear quantizer with an interval of $\delta$ (without block-wise normalization), the quantization variance is explicitly given by
$$
\sigma_m^2 = \mathbb{E}\left[||Q_{\delta}(x) - x||^2\right]
\le \frac{\delta^2 d}{4}, \forall x\in \mathbb R^d.
$$
The value of $\delta$ is determined by the maximum absolute value (aka absmax) within the vector and the number of representable values, i.e., the number of bits used.
Upon incorporating block-wise normalization with a block size of $B$, assuming there are $N$ blocks in total, and denoting the interval within each block as $\delta_i$, the quantization variance is given by
$$
\sigma_m^{2} = \mathbb{E}\left[||Q_{\delta}(x) - x||^2\right]
\le \sum_{i=1}^N \frac{\delta_i^2 B}{4}, \forall x\in \mathbb R^d.
$$
However, the extent of improvement in quantization variance through blocking heavily relies on the tensor's structure. Here, we qualitatively analyze the impact of block size in various scenarios:
- In an extreme case, if outliers comparable to, or same as, the absmax value occur within each block, block-wise quantization fails to yield any improvement.
This instance finds empirical support in our investigations, where we have empirically observed regular distribution of outliers within moment tensors along rows and/or columns. This pattern suggests the advantages of employing smaller block sizes.
- When the first moments are i.i.d. from $N(0, 1)$, we have $\delta_i \propto \sqrt{\log(B)}$ thanks to the nice property of the Gaussian distribution.
Consequently, if two distinct block sizes, $B_1$ and $B_2$, are employed in quantization, the quantization variance ratio is given by
$$
\frac{\sigma_{m, B_1}^{2}}{\sigma_{m, B_2}^{2}}
= \frac{\log(B_1)}{\log(B_2)}.
$$
For neural network training, considering $B_1 = 128$ and $B_2$ as the size of a single tensor, significant enhancements in quantization variance are achieved. Furthermore, with $B_1 = 128$ and $B_2 = 2048$, the quantization variance improves by 7/11.
- In practical scenarios, the distribution of moments is difficult to characterize analytically, thus making it difficult to determine the impact of block size on quantization variance. We leave the theoretical results for future work and instead present empirical analyses regarding block size and quantization variance. We utilize block-wise normalization with different block sizes to quantize the first moment tensors in a GPT-2 medium model.
We report both the mean and maximum relative quantization errors across tensors. The relative quantization error is defined as $||Q(x) - x||/||x||$.
In this specific setting, the results show a rough log-correlation between relative error and block size.
| Block size | 128 | 256 | 512 | 1024 | 2048|
|---------------|----:|----:|----:|-----:|-----:|
| first moment(mean) | 0.158|0.170|0.183 | 0.194|0.205 |
| first moment(max) | 0.173|0.190|0.208 | 0.222|0.235 |
However, it is important to note that the actual performance metric (e.g., accuracy) exhibits a complex relationship with quantization variance, and even training loss, exceeding the scope of our block size analysis.