Oscillation-Reduced MXFP4 Training for Vision Transformers

We thank the reviewer for valuable comments and the acknowledgment of our contributions. Below we respond to the questions.

Question 1: … Are unbiased gradients necessary to get oscillation reduction to work, or is oscillation reduction necessary to get unbiased gradients to have an effect on the training curve?

Thanks for the valuable questions. In short, the methods for forward (oscillation-reduction) and for backward (unbiased gradients) are orthogonal. We provide empirical evidence as follows.

For the first question, we conduct an additional experiment to show that oscillation reduction is not limited to a certain method of gradient calculation (Table 1). We can see the Q-EMA also works on the baseline Microscaling (MS), which does not unbiasedly estimate gradients.

For the second question, we have done experiments in our work without oscillation reduction techniques. In Table 2 (part of Table 2 in our paper), our method TetraJet(TJ) with quantization improvement on unbiased estimation (do not add oscillation reduction technique) outperforms baseline Microscaling (MS).

Table 1. Oscillation Reduction on DeiT-S 60-epoch pre-training.

Table 2. Accuracy improvement without additional oscillation reduction techniques (90-epoch pre-train)

Question 2: Is double quantization actually necessary?

Thanks for the valuable question. There are mainly two core issues that necessitate double quantization：

• (1) The correctness of optimization goal. The forward pass of our optimized network is $Y= Q_D^{(1)}(X) \times Q_D^{(2)}(W^\top)$. So the gradient for weight is $\nabla_{W}\mathcal{L}\overset{\text{STE}}{\approx} \nabla_{Q_D^{(2)}(W)} \mathcal L = (\nabla_Y\mathcal L)^\top \times Q_D^{(1)}(X)$. To unbiasedly calculate this gradient using MXFP4, we need to estimate $Q_D^{(1)}(X)$ rather than $X$, to align with the forward pass.
• (2) The non-square group shape. To enable efficient MXFP4 Matrix Multiplication, the group shapes should be (1x32) x (32x1) for each multiplication. Because $Q_D^{(1)}(X)$ is in group shape 1x32 in forward, we need another unbiased quantizer for $Q_D^{(1)}(X)$ to achieve group shape 32x1 during backward, so we get the double-quantized $Q_S^{(6)}\left(Q_D^{(1)}(X)\right)$ (see Eq.5 of our paper).

In conclusion, the correctness of optimization goal & the non-square group shape (1x32 / 32x1) make double quantization necessary in our MXFP4 training.

Question 3: Is there a reason you chose to evaluate on small vision transformers? This method does not seem specific to vision transformers and could be applied to language models as well.

Thanks for the valuable comment. We chose to begin our exploration of FP4 training with Vision Transformers (ViTs) primarily due to computational resource constraints, as pre-training large language models (LLMs) in FP4 is significantly more resource-intensive. Importantly, developing efficient FP4 training algorithms for vision tasks is itself a meaningful and underexplored direction, with many prior works also focusing on ViTs (e.g., [1,2,3]).

Our ultimate goal is to enable FP4 training for LLMs, and we actually observe that our proposed method, TetraJet, generalizes well to LLMs and outperforms FP4 baselines (see Table 3). We also find that oscillation problem still exists in LLM pertaining task. However, fully resolving convergence issues and matching BF16 performance in large-scale LLMs still requires further systematic investigation, particularly due to their large scale and dynamics. We believe that our work on ViTs Training with FP4 formats would be inspiring and valuable for further exploration of LLMs training.

Table 3. OLMo-2 150M pre-training with 20 Billion tokens

Question 4: What is the cost of storing the EMA components for oscillation reduction?

During training, we only need to store an extra EMA weight for each linear layer in transformer blocks. For DeiT-B with 85M linear parameters in Transformer blocks, this adds about 340MB of storage, which is relatively small compared to the overall training memory footprint (around 20GB per GPU in a 4-GPU setup).

If we use FP4-trained models for inference, we can calculate FP4 values for each parameter in advance, so we don’t need the EMA component anymore during inference, which means no additional cost.

[1] Y. Li et al.,"Q-vit: Accurate and fully quantized low-bit vision transformer," NeurIPS, 2022.
[2] Y. Liu et al.,"NoisyQuant: Noisy bias-enhanced post-training activation quantization for vision transformers," CVPR, 2023.
[3] Z. Wang et al.,"Quantformer: Learning extremely low-precision vision transformers," IEEE TPAMI,2022.

We thank the reviewer for valuable comments and the acknowledgment of our contributions. Below we respond to the questions.

Broader Scientific Literature:
The authors might enhance their context by explicitly citing some recent low-precision training surveys...

Thanks for the valuable advice. We have found some surveys on low-precision training [1,2,3]. We will include them in our revised version.

Other Comments Or Suggestions:
I think the sentence "Therefore, quantizer …. group shape." can be better phrased. It's hard to follow the sentence and what quantizer(*) refer to.

Thanks for the valuable advice. We will clarify this by pointing out that they are Quantizers $Q^{(*)}$ in Eq. 3-5.

Weakness 1:
The analysis primarily focuses on ViT architectures; evaluating more diverse architectures (e.g., CNNs, LLMs) could further validate general applicability.

Thanks for the valuable comments. According to NVIDIA's design[4], block scaling formats (e.g., MXFP4) are mostly designed for matrix multiplications rather than convolutions. Therefore, our work focuses on designing methods tailored to Transformers rather than CNNs.

Due to space limitations, we left the detailed response to this question in the rebuttal to “Question 3” of Reviewer TZFB. In short, we include results demonstrating that our method generalizes to LLMs, suggesting that our approach can be a promising direction for FP4 training of LLMs as well.

Weakness 2:
Discussion of potential overhead or computational costs associated with the oscillation tracking in Q-Ramping could be expanded.

Thanks for the constructive comment. In Q-Ramping method, we do not track oscillations all the time. For example, in ImageNet-1K pre-training, an epoch consists of ~2500 iterations. We only track the first 30 iterations to find the oscillating weights. As a result, this only adds 1.64% to the total training wall time compared to w/o Q-Ramping in our DeiT-B training on RTX4090.

Question 1:
Could you clarify the computational overhead introduced by Q-EMA and Q-Ramping methods in actual training?

Thanks for the valuable questions. In conclusion, both methods only introduce little overhead into the training time. For Q-EMA method, we only need to use EMA for weight quantization during the forward pass (other quantizers don’t have additional cost), and we update the EMA-weight for quantized layers when parameters are updated. In our DeiT-B training on RTX4090, we found Q-EMA only adds 0.35% to the total training wall time compared to w/o Q-EMA. For Q-Ramping method, we have discussed above in response to Weakness 2.

Question 2:
Are the proposed methods sensitive to network architecture or hyperparameters beyond those tested? Specifically, how robust are they for larger models beyond the evaluated ViTs?

Thanks for the valuable questions. We have evaluated the robustness of hyperparameters for DeiT-B. As Tables 1&2 show, Q-EMA’s decay factor β and Q-Ramping’s update frequency factors $k_1$, $k_2$ are not sensitive within certain intervals.

We conducted additional fine-tuning experiments for ViT-base model [6] using the same default hyperparameters from our paper, and found our methods are not sensitive to network architectures or specific tasks (Table 3). Furthermore, our methods significantly outperform the baseline in LLMs pre-training using the larger OLMo-2-150M (Table 4), showing generalization beyond ViT models. These results demonstrate that our method can be adopted to different kinds of architectures and tasks.

Table 1: Insensitivity to hyperparameters (Q-EMA) on DeiT-B.

Table 2: Insensitivity to hyperparameters (Q-Ramping) on DeiT-B.

Table 3. Results of 50epoch MAE ViT-base Fine-tuning (MS: Microscaling; TJ: TetraJet)

Table 4. OLMo-2 [6] 150M Pre-training with 20 Billion tokens

[1] Wei L et al., "Advances in the Neural Network Quantization: A Comprehensive Review". Applied Sciences. 2024.
[2] Chitsaz K et al., "Exploring Quantization for Efficient Pre-Training of Transformer Language Models," arXiv:2407.11722, 2024.
[3] Kumar T et al., "Scaling Laws for Precision," arXiv:2411.04330, 2024.
[4] https://docs.nvidia.com/cuda/pdf/ptx_isa_8.7.pdf
[5] OLMo Team et al., "2 OLMo 2 Furious," arXiv:2501.00656, 2024.
[6] https://github.com/facebookresearch/mae

We thank the reviewer for valuable comments. Below, we respond to the questions.

Claims And Evidence & Theoretical Claims:
Lack formal theoretical validation & Leaving some theoretical gaps

Thank you for the insightful comment. We agree that theoretical analysis is important. While our current work primarily focuses on empirical evaluation, we acknowledge the need for a stronger theoretical foundation to support our findings on oscillation reduction. However, theoretically analyzing the convergence properties of low-precision Transformer training remains an open and challenging problem, primarily due to the non-differentiability of quantization functions. We view this as a valuable direction for future research and are actively exploring ways to address it.

Other Weaknesses:
The study focuses only on Vision Transformers and does not evaluate other architectures (e.g., CNNs, NLP models like LLMs).

Thanks for the valuable comment. According to NVIDIA's design[1], block scaling formats (e.g., MXFP4) are mostly designed for matrix multiplications rather than convolutions. Therefore, our work focuses on designing methods tailored to Transformers rather than CNNs.

Due to space limitations, we left the detailed response to this question to the rebuttal to “Question 3” of Reviewer TZFB. In short, we include results demonstrating that our method generalizes to LLMs, suggesting that our approach can be a promising direction for FP4 training of LLMs as well.

Experimental Designs Or Analyses:
the study lacks statistical significance testing (e.g., confidence intervals or variance analysis), which would further strengthen the validity of the reported improvements.

Thanks for the valuable advice and questions. We report results with standard deviation based on three runs using different random seeds in our additional fine-tuning experiments (see Table 1). Results demonstrating that our method performs consistently better than baselines.

Question 1:
Have you tested MXFP4 fine-tuning on pre-trained models? Does the oscillation issue persist in fine-tuning settings?

Our research is mainly focused on ViTs pre-training, but can also be generalized to fine-tuning. Our further conducted experiments show that the oscillation problem still exists, which aligns with previous literature on oscillation problem about low-precision fine-tuning [2].

We finetune MAE-ViT-base [3] for 50 epochs based on the pre-trained model, and we report the average accuracy and standard deviation over 3 seeds. Our method TetraJet still outperforms the baseline, and Q-EMA/Q-Ramping provides additional enhancement by alleviating the oscillation problem.

Table 1. Results of ViT-base Fine-tuning (MS: Microscaling; TJ: TetraJet)

Question 2:
How sensitive are Q-EMA’s decay factor (β) and Q-Ramping’s update frequency adjustments to different datasets and architectures? Did you find any optimal ranges?

Thanks for the valuable questions. The choices of these hyperparameters are not sensitive through our experiments. As Table 2&3 show, Q-EMA’s decay factor β and Q-Ramping’s update frequency factors $k_1$, $k_2$ are not sensitive within certain intervals. We find [0.997, 0.999] is an optimal range for Q-EMA’s decay factor (β); a good choice for oscillation detection factor $k_1$ is 16, and good weight update frequency factor $k_2$ can be from 3 to 5.

We used default settings ($\beta=0.998, k_1=16,k_2=5$) for our reported pre-training and fine-tuning experiments, and found additional enhancement of Q-EMA and Q-Ramping across different architectures and tasks (DeiT/Swin pre-training & ViT fine-tuning). These results further confirmed the robustness and the ability of generalization of our oscillation-reduction methods.

Table 2: Insensitivity to hyperparameters (TetraJet + Q-EMA) on DeiT-B.

Table 3: Insensitivity to hyperparameters (TetraJet + Q-Ramping) on DeiT-B.

[1] https://docs.nvidia.com/cuda/pdf/ptx_isa_8.7.pdf
[2] S.-Y. Liu, Z. Liu, and K.-T. Cheng, "Oscillation-free quantization for low-bit vision transformers," ICML, PMLR, 2023.
[3] https://github.com/facebookresearch/mae