Oscillation-Reduced MXFP4 Training for Vision Transformers
Authors: Yuxiang Chen
Haocheng Xi
Jun Zhu
Jianfei Chen
Venue: ICML 2025 poster (2025)
Keywords: Efficient Machine Learning
Low-Precision Training
Quantization
FP4
Microscaling
Primary area: deep_learning->algorithms

Abstract: Pre-training Transformers in FP4 precision is becoming a promising approach to gain substantial speedup, but it comes with a considerable loss of accuracy. Microscaling (MX) data format provides a fine-grained per-group quantization method to improve the representation ability of the FP4 format and is supported by the next-generation Blackwell GPU architecture. However, training with MXFP4 data format still results in significant degradation and there is a lack of systematic research on the reason.

In this work, we propose a novel training method TetraJet for a more accurate FP4 training. We comprehensively evaluate all of the quantizers involved in the training, and identify the weight oscillation problem in the forward pass as the main source of the degradation in MXFP4 training. Therefore, we introduce two novel methods, EMA Quantizer (Q-EMA) and Adaptive Ramping Optimizer (Q-Ramping), to resolve the oscillation problem. Extensive experiments on Vision Transformers demonstrate that TetraJet consistently outperforms the existing 4-bit training methods, and Q-EMA \& Q-Ramping can provide additional enhancement by effectively reducing oscillation. We decreased the accuracy degradation by more than 50% compared to the baseline, and can even achieve competitive performance compared to full precision training.

Oscillation-Reduced MXFP4 Training for Vision Transformers
Yuxiang Chen 12
Haocheng Xi
Jun Zhu'
Jianfei Chen
Abstract
1. Introduction
Low Low-precision training has emerged as a promising tech- nique for accelerating the training process of large-scale neural networks. By quantizing tensors in both the for- ward and backward passes to lower-precision formats, low- precision training leverages specialized compute units in modern hardware to enhance computational efficiency sig- nificantly. While BF16 and FP16 precision remain the most widely used formats for deep learning training (Narang et al., 2017; Kalamkar et al., 2019), FP8 training (Sun et al., 2019; Micikevicius al., 2022; NVIDIA, 2024c; Xi al., 2025) is becoming increasingly mature in these years, with suc- cessful application in training state-of-the-art large language models (Liu et al., 2024).
Pre-training Transformers in FP4 precision is be- coming a promising approach to gain substan- tial speedup, but comes with a considerable loss of accuracy. Microscaling (MX) data for- mat provides fine-grained per-group quantiza- tion method improve the representation ability of the FP4 format and is supported by the next- generation Blackwell GPU architecture. How- ever, training with MXFP4 data format still re- sults significant degradation and there a lack of systematic research on the reason. In this work, we propose novel training method TetraJet for a more accurate FP4 training. We comprehensively evaluate all of the quantizers involved in the training, and identify the weight oscillation problem in the forward pass as the main source of the degradation in MXFP4 train- ing. Therefore, we introduce two novel methods, EMA Quantizer (Q-EMA) and Adaptive Ramp- ing Optimizer (Q-Ramping), to resolve the oscil- lation problem. Extensive experiments on Vision Transformers demonstrate that Tetralet consis- tently outperforms the existing 4-bit training meth- ods, and Q-EMA & Q- O-Ramping can provide addi- tional enhancement by effectively reducing oscil- lation. We decreased the accuracy degradation by more than 50% compared to the baseline, and can even achieve competitive performance compared to full precision training. The codes are avail- able at https: htt.ttttthu.ooooccoo com/ hon/thu-ml/ Tet Tetrrauet-WFFP4Traii
There is a growing interest in pushing the training precision down to 4-bit. While earlier works attempt to train the net- work with FP4 (Sun et al., 2020), logarithm format (Chmiel et al., 2021), and INT4 (Xi et al., 2023), these works have rather large accuracy degradation (e.g.. 1-2%) even sim- ple tasks such as ResNet training, and are not practically favorable. Recently, a Microscaling (MX) data format has been proposed for accurate low-precision training and in- ference (Rouhani et al., 2023b;a), MX applics fine-grained per-group quantization, where each small group of 32 ele- ments shares scaling factor. This fine-grained quantization scheme significantly mitigates the impact of outliers, and thus reduces quantization error. Particularly, the MXFP4 format utilizes an E2M1 (Exponent / Mantissa) FP4 with an E8M0 scaling factor. MXFP4 is supported on the latest Nvidia Blackwell architecture and is 2 times faster than FP8/MXFP6 and 4 times faster than FP16/BF16 (NVIDIA, 2024a;b) when doing matrix multiplications. However, the low-precision training method proposed in the original Mi- croscaling paper uses MXFP6 activation/gradient which is as slow as FP8 training. The fast pure MXFP4 training still has major accuracy degradation as tested in our experiments, which makes it infeasible to use in practice.
1 Dept. of Comp. Sei. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua-Bosch Joint ML Center, Tsinghua University "Zhili College, Tsinghua University Uniiversity of California Berkeley. Correspondence to: Jianfei Chen <jian- tice@tsinhuu..cuu.cn>
In this work, we propose TetraJet, a novel training method for transformer (Vaswani, 2017) with MXFP4 computation in both forward and backward pass. All weight/activationn/ gradient tensors in linear layers are quantized to MXFP4 to fully unlock the acceleration potential of the hardware. We propose several techniques to improve the accuracy of
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

MXFP4 Training for Vision Transformers
MXFP4 training. First, we propose a truncation-free scaling method for quantizing full-precision values to MXFP4 to avoid information loss in truncation. We further propose a double quantization method to deal with the non-square quantization group of MXFP4. With these techniques, we prove that TetraJet can estimate the gradient unbiasedly.
the 4-bit MX format as MXFP4, and refer their train- ing method as Microscaling. We propose better training method TetraJet with accuracy improvement compared to Microscaling.
Oscillation Problem In low-precision training, weight oscillation has been proven to be serious problem that affects optimization. Nagel et al. (2022) revealed that the os- cillation of weight quantization does harm to Quantization- Aware Training (QAT) of CNNs. Besides, Liu et al. (2023) proved that oscillation was a key factor causing the degrada- tion of accuracy in QAT of Vision Transformers. However, they were both based on QAT, that is, they fine-tunes a low-precision model based a pre-trained full-precision network rather than pre-training from scratch. They both utilized pre-tensor Learned Step Size Quantization (I.QQ) to train the models. There is a lack of research on oscil- lation problems about pre-training and more fine-grained quantization methods (e.g., MX Format).
We then conduct comprehensive evaluation of the impact of individual quantizers on the final model performance, and find that activation and weight quantizers in the forward pass contribute the most to accuracy degradation, due to a weight oscillation problem: the master weight fluctuates around the quantization boundary, causing the model to be quantized into different values across iterations, which consequentially brings significant instability the optimization process. We propose two methods to alleviate the oscillation problem: the EMA quantizer (Q-EMA) conducts rounding based the moving average of historical weights rather than only depending on the current weight matrix; and the Adapting Ramping optimizer (Q-Ramping) adaptively identifies and reduces the update frequency of oscillating weights.
To reduce weight oscillation, Liu et al. (2023) proposed several methods, but the application is restricted to LSQ or QAT, while Nagel et al. (2022) proposed methods that can be generalized to reduce oscillation in MXFP4 pre- training: The method "Dampen" tried to encourage latent weights to be closer to the quantized value to avoid fluctuat- ing around the quantization boundary, by adding regulation term Ldmmpen = | Q(W) |/ in the loss function: The method "Freeze" tracks the oscillation frequency for each weight element, and freezes those frequently oscillating weights (f > fei) to running average value. The frozen weights would never be updated again in the whole training process, which may harm the optimization in pre-training. In this work, we propose two novel methods Q-EMA & Q- Ramping to better reduce oscillation in MXFP4 pre-training
Extensive experiments on Vision Transformers prove that TetraJet consistently outperforms Microscaling's original method (Rouhani et al., 2023b), and Q-EMA & Q-Ramping can provide additional improvement through oscillation re- duction. We decreased the accuracy degradation by more than 50% compared to the baseline, and even achieve com- petitive performance compared to full-precision training.
2. Related Work
Low-Precision Training Low-precision training has be- come prominent technique in modern deep learning to speed up the training process. FP16 and BF16 (half- precision) training (Narang et al., 2017; Kalamkar et al., 2019) is currently the most common low-precision method. FP8 and INT8 training (Sun al., 2019; Zhu et al., 2020: Micikevicius et al., 2022: Wortsman et al., 2023; NVIDIA, 2024c; Peng et al., 2023; Xi et al., 2024; 2025; Liu et al., 2024) further improves efficiency and uses more fine- grained per-tensor/ / per-row / per-block quantization. When comes down to bit training (Sun et al., 2020; Chmiel et al., 2021; Xi et al., 2023), more techniques are being applied (e.g. Hadamard transformation) to compromise the degradation caused by the low representation ability. Still, their accuracy degradation is not negligible.
3. Our TetraJet Training Method
In this section, we review and identify several drawbacks of the existing low-precision training method Microscal- ing (Rouhani et al., 2023b), and propose a more accurate training method TetraJet. The effectiveness of our method is shown in Section 7.
3.1. Preliminary
MXFP4 Format Floating points have three components: bign- exponent-bits, and mantissa bits. If format has exponent bits and y mantissa bits, we usually denote it as EaMy. We use Qp Qn. to represent the max positive value and the min negative value the format can represent. For E2MI, Qp = 6, Q. = -6.
For a more fine-grained quantization, the Microscaling (MX) format (Rouhani et al., 2023a;b) in the Blackwell architecture (NVIDIA, 2024a) offers a 1 x 32 per-group quantization and could potentially double the speed com- pared to FP8 training. Rouhani et al. (2023b) also propose low-precision training method with computation flow in MX formats in 4, 6, and 8 bits. In this paper, refer to
The MXFP4 (Microscaling Floating-Poin 4-bit) data for- mat (Rouhani et al., 2023a) follows per-group quantization
2

---

Oscillation MXFP4 Training for Vision Transformers
scheme where group of N 32 elements shares com- mon 8-bit exponential scaling factor 8. Each element X; in the group is represented by P. in E2MI format. The reconstruction of floating-pone point value X; from its MXFP4 representation follows the formula:
Deterministic & Stochastic Rounding of FP4 format
Now we diseuss the roundgr4()) operation in Eq. (1). With our scaling, all the values X:/S are in the range [Q.. Qp] Therefore, we can always find two consecutive FP4 value ạ1, (2(2) < (2) satisfying Q S X,/S & de for every X,.
A direct way of rounding X./S is to select the nearest FP4 value between q1 92, which we call deterministic rounding or round to nearest. Here we denote it as
X; P. x 2", i= 1,2, ...,, 32
Quantization To quantize a matrix to MXFP4, we need to split into blocks of size x 32 (or 32 x 1), and then quantize each block MXFP4. To quantize block of 32 first determine the ESMO scale factor S 2 with |s| S 127. Each value X; is then mapped to 4- bit FP4 representation P, such that:
[1,
|x./S qul XX./S q2l 92, otherwise
round (X:/S)
full-precision values {X; {
x= to MXFP4,
Microscaling always applies deterministic quantization to minimize the quantization enror. However, we find it subop- timal to apply deterministic quantization to gradients, since the gradient will no longer be unbiased. To this end, we apply stochastic rounding (Courbariaux et al., 2015) to gra- dients to maintain an unbiased gradient. Stochastic rounding generates random variable ~ Uniform( =2 for each value X, independently, and computes
X, S
P. roundEP1
P
(L)
The quantized representation is stored (1P)) S),
where S is an 8-bit exponent, and P; is an FP4 value.
{ Q1, X,/S+ 9
rounds(X,/5)) =
42. otherwise
3.2. Quantization with Fruncation-Free Scaling
Stochastic rounding is unbiased: E|rounds(X./5)) X,/S. We show the superiority of stochastic rounding in the ablation study in Sec. 7.3,
Computation of Scaling Factor the scale factor as follows
Microscaling computes
logg M]
Emax: 5 S= 2-,
3.3. TetraJet Linear Layer
where M- = maxisig32 [ is the largest absolute value of the block, Frnax represents the largest exponent in FP4 format'. A drawback of the approach is that the scaled value X may fall outside the range [Qn. Qpl. and exceeding values will be truneated. For instance, if M 31, the scaling factor will be S 2 = 24-2 4. Since M/S- 31/4- 7.75 exceeds the maximum representable value 2p = 6, the value M will be truncated to 6. Intuitively, large values carry more information, and such truncation will be harmful to retaining the precision of the network.
When training the transformer, linear layers usually take most of the computation. Following previous works on low-precision training (Xi et al., 2023), we mainly focus on accelerating the linear layer with MXFP4, whose forward and backward pass are defined as:
Y == =xxW"
(VxC), TWC=(VxX) where X E RNXD W E RCXD E RNCC is loss function, and VxC/VYL/VwL are the in- put/output/weight gradient matrices with the same size of X, Y. W.
TetraJet equips a truncation-free scaling method:
2M
To accelerate training, we need to compute all three matrix multiplications (MMs) in MXFP4. To achieve this, we need to quantize the six input matrices of the three MMs to MXFP4, which can be formulated as:
s =
logz Op - Qn
S= 2",
where M equals to M in most cases except when M = 0, we set M to a small number e 10-8 8 to avoid numerical issues. Compared to Eq. (2), we replace floor |.] with ceil- ing [-] to avoid truncation, and replace the numerical range from 2 +m to a more accurate range [Q.. Qpl- In this way, Qn S M/S S Qp always holds. For example, when M 31, the scaling factor will be S 23 8. so M/S 3.875 still lies in the representation range of FP4.
Y = X
(2) (WT)
(3)
VxL Q.00 (VrL)
x xQ= (o%
(4)
Vw Q ((VvL)
x Q (o) (x)
where QD/Qs refers to the deterministicsstocaastie round- ing quantizer. We explain the design of TetraJet linear layer as follows.
'or F2M1, Frnaxe 22-1 =
3

---

Oscillation- CesilattionReducedd MXFP4 Training for Vision Transformers
X
Matmul
Y
X
C [X]
Matmul
Microscallag's Scaling Factor
W
Quantted
P (WW)
W
Quandtor
4 (W"))
Taandat's Truncation Pree
Medel Weights
Duanile
C (X)
Feauring
4 (W)
Model Weights
2 (4 (X)
Kaaantian
Chenating
(WW(
Sealing Farter
Detorministle Bounding
C
VwC
Matmol
Quantte
Matmul
Quanettrr
Stochautic Bounding
C
V
VxC
V
Matmul
Quantize
Matmul
Quantian
(a) Microscaling's Linear Layer
(b) TetraJet's Linear Layer
Figure I: Visualization of MXFP4 Linear Layer
Block Format As fine-grann fine-grained format, doing MM with MXFP4 is more subtle than other grrsiergr formats such as per-tensor quantization. For hardware-accelleaated MM to be possible, MXFP4 requires quantization group shape to be x 32 for the first matrix and 32 1 for the sec- ond matrix. Therefore, quantizers Q(1, Q(3), Q ((5) should use 1 x 32 group shape, and quantizers Q Q(4) ((6)
Table I: Impact analysis on MXFP4 quantizers. We report the top-1 Acc.% % after 90-epoch pre- pre-training. Qi means we only activate the i-th quantizer Q())
DeïT-T
DeiT-S
Full Precision
63.73
73.33
Q1
61.50
71.66
Q2
62.77
72.45
should use 32 x 1 group shape. This means that weight W, activation X, and gradient V should be quantized along different axes in different quantizers. For example, the quantization block size of X should be 1 token x 32 channels in forward and 32 tokens x channel in backward.
Q3
63.46
72.97
Q4
63.37
72.79
Q5
63.81
73.25
Q6
63.78
73.13
All Quantizers
59.75
71.03
Double Quantization We propose double quantiza- tion strategy satisfy MXFP4's block format requirement. Specifically, @ (X) is a quantized activation with x 32 group size, which is used in the forward pass. We quantize the already quantized ! (X) again with a different 32 x group size to compute the gradient in Eq. (5). By doing so, we ensure the activation is quantized with the required group size for both forward and backward pass. Similarly, the weight is also doubly quantized.
Note that microscaling's gradient Eq. (6,7) does not equal to the correct gradient Eq. (8,9). Particularly, Q (W) +
D
(2) (w Microscaling is actually computing the gra- dient for another network with the forward pass Y ()T,, where both operands are quan- tized in the wrong direction.
Qa2x1 i
In contrast, TetraJet gives an unbiased estimation of Eq. (8,9). Take VxL in Eq. (4) as an example, since Q Q(4) are stochastic and truncation-free, the expecta- tion of our gradient is
In contrast, Microscaling takes different approach:
Vx -/0 (VY x (W)
Q (4)
(6)
B (s) (Vr = E
X 6 0 (V)C)
VWL Q ((Vx 2 XX
(7)
14) = =V Q8 (WT)
x x (00 (WT)T]
where the activation used in the backward pass is deter- ministically quantized from the full-precision X rather than (X).
(x)(). which is biased we will discuss.
which is right side of Eq. (8). Similarly, the estimation in Eq. (5) for Vw is also unbiased. Given that each linear layer is unbiased, the final gradient calculated with backpropagation is unbiased, which ensures the convergence of SGD, as discussed by (Chen et al., 2020).
3.4. Gradient Bias
We first derive the correct gradient formula with Straight Through Estimator (STE) (Bengio et al., 2013): Given the forward pass in Eq. (3), the correct gradient should be
VxL STE R
( C (VrL) 0 (WT)"
3.5. Impact Analysis of Six Quantizers
VWC
T (w) C (VYC)T (x((x))
(9)
Before making any attempts to improve the training, it is necessary to understand which among the 6 quantizers in

---

Oscillation MXFP4 Training for Vision Transformers
Qunntred Weight
Rate of Change for Weight At Different Stage Training
Rate of Change for Activation At Different Stage of Training
-0.00
Rate Change for Weight
Rate of chanaac for Activation
10"
1100
10-2
-0.745
10-2
Labove Woights
20-5
-0.750
10--
Pull-Precision Training
10-1
Full Precision Training
MXPP4 Training
MXFF4 Training
-0.55
100
200
300
400
500
60
80
20
Reration
Epach
Epoch
Figure 3: Trajectory of some oscillation elements in DeiT- Tiny during the last epoch of training. The top plot shows the change of quantized FP4 value, and the bottom plot shows the oscillating latent weight around the quantization decision threshold thrd -0.75.
Figure 2: Rate of change for weight and activation at differ- ent stages of 90-epoch DeiT-Tiny pre-training. We calculate the average rate for all quantized weights and select a trans- former block to test output activation given fixed input.
Eq. (3,4,5) is the bottleneck. We test the impact of quan- tizers by activating them separately: for the i-th test, we only activate Q(+) while leaving all other matrices in full precision, train the model from scratch, and compute vali- dation accuracy. As shown in Tab. 1, the activation/weight quantizers Q(2)/Q((2) in the forward pass lead to most accu- racy degradation. For example, MXFP4 training on DeïT-T has a 3.98% accuracy loss, while only quantizing the activa- tion/weight in the forward pass accounts for 2.23% / 0.96%, respectively. We reveal in the next section this is due to the instability of low-precision training.
We find that the weight oscillation is the source of this problem. To be clear, we refer to w/S as latent weight, where S is the quantization scale factor of weight element As illustrated in the top plot in Fig. 4, large amount of latent weights lies around the quantization thresholds (the midpoints of two quantized values) at the end of the training process. For these elements, little perturbation on their corresponding master weights will change the quantized values, which results in a giant jump from one quantized value another. This makes the rate of change of the quantized weight matrix much higher than its corresponding master weight, and meanwhile contributes to the instability of activation, which aligns with our finding.
4. Oscillation Phenomenon
4.1. Instability of MXFP4 Training
We tracked several oscillating weight elements during the final epoch of training for a better understanding of this oscillation phenomenon. As shown in Fig. 3. these latent weights are changing with small steps around the quanti- zation threshold thrd -0.75, which is the midpoint of two FP4 values qı = -1, = -0.5. When the latent weight crosses thrd = -0.75 caused by small update, the quantized weight would shift from 91 to q2 (or from 92 to (1). Frequently crossing thrd causes the frequent flipping between qı and q2. Therefore, a direct characterization of oscillating weight is that, the oscillating weight elements will have their latent value stay closely around the quantiza- tion threshold and frequently cross the threshold.
During the final stage of training, the learning-rate (LR) typ- ically approaches zero, so the model can stop exploration and quickly descend to local minimum. However, we find that MXFP4 training cannot converge even with a suf- ficiently small learning rate due to the oscillation between quantization points. To explain this phenomenon, we define rate of change for a tensor X as
To To
X
X
(X)
where í refers to training step, and step O ~ To refers to a short training interval. During pre-training, we can test the rate of change for the master weight W, the quantized weight matrix Q(2((WT) and activation Y at different stages.
4.2. Quantization Confidence of Weight Distribution
To quantitatively assess the severity of the oscillation prob- lem, we define quantization confidence for each weight element w, which measures the normalized distance to the nearest quantization threshold:
As shown in Fig. 2, for full-precision models, the rate of change can gradually decrease to near zero, while for quan- tized models the rate of change would stay high in the final of training, indicating that there are still large changes inside the models.
QuantCCon(ww) = min, ww - thrd,l
Max
5

---

Oscillation- MXFP4 Training for Vision Transformers
Epoch
Epoch
Epoch 90
behind the actual optimization process and is updated slowly, EMA weight is less likely to be affected by oscillations. Consequently, this makes the optimization process much more stable.
6000
Latent Weight
Quant Thresholde
4000
2000
Our EMA quantizer first maintains an EMA weight through- out the training process, When doing quantization to each weight element w with scale factor S and its EMA value WEMA. we first use the latent weight w/S to propose two candidate quantized values Wh: and Uaa as they are the two values that give the smallest MSE. We then use the EMA weight to check which is closer to WEMA. and use this as the quantized value. This algorithm is formalized as Algorithm in Appendix C.
-5
-5
-5
Latert weight wi5
3500
I Ouant Confidence
3000
More Confident
Less Confident
2500
2000
1500
1000
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
10
Quanttor/(e)
Figure 4: The change of latent weight and quantization con- fidence during 90-epoch pre-training of DeiT-Tiny. The top plot shows the distribution of latent weight, and the bottom plot shows the distribution of quantization confidence.
6. Adaptive Ramping Optimizer
Besides smoothing the weight quantization with EMA quan- tizer, another effective approach to reducing oscillations is to manually decrease the update frequency of oscillat- ing weights. Building on this idea, we propose Adaptive Ramping Optimizer (Q-Ramping). which directly locates the frequently oscillating weights according to their updat- ing trajectory, and then adaptively decrease their updating frequency by using a higher gradient accumulation step for these oscillating weights to reduce the oscillation frequency.
where PP4 denotes the quantized FP4 value of w, (thrd,) denotes all the quantization thresholds, and MaxDist(w denotes the maximum possible distance if quantized to FP4 It is ensured that QuantConf(w) e [0, 1].
The rationale behind this metric is that the closer a latent weight is to quantization decision threshold, the more likely it is to oscillate, making it harder for the weight to converge to a stable FP4 value.
6.1. Identifying Oscillating Weights
The first thing is to locate the oscillating weights and quan- tify their degree of oscillation. To achieve this, we would record information about the weight update trajectory for each element. During a training stage with To steps, we sum up updating distance for each master weight element w and its quantized weight uq:
As shown in the bottom plot of Fig. 4, we observe a grad- ual decline in quantization confidence throughout training. This trend indicates an increasing prevalence of oscillation as training progresses. Consequently, effective solutions to mitigate oscillation should be dynamic, adapting to the specific conditions of each stage of the training process.
To
To >
distw Σ
1ww - w' distQ
w
5. EMA Quantizer
And then, we define oscillation ratio Ry for each weight element as
We firstly propose an EMA Quantizer (Q-EMA) to solve the oscillation phenomenon. Since the weight will oscillate between the two possible choices randomly even with small perturbations, we hope to find a better way to choose from these two possible values after quantization.
R == disto /istt
representing the degree of oscillation.
During training, if a weight element we doesn't 't fall into the oscillation process, the master weight w and the quantized weight wo would move with a similar trajectory. In this situation, distQ s distw, so R., would not be too large.
We find that the Exponential Moving Average (EMA) can be used to alleviate the oscillation problem. EMA on weight is determined as:
W WEMA = 3
EMA + (1 - 3) W
(10)
In contrast, for oscillating weight elements, the quantized weight would switch frequently between two discrete quan- tization values and 42. Each switch from to 92 (or from d2 to qi) will increase dist. by lq - 9al. making it relatively large. Meanwhile, the master weight w would be oscillating around the quantization threshold, and the step-
where W' is the BF16 weight. A Typical choice of / is 0.998. Therefore, even when weight makes a very large step, W WEMA only moves slightly. When the weight oscillates between two quantized values, as EMA weight is always left
6

---

Oscillation-Redee MXFP4 Training for Vision Transformers
Table 2: Results on the 90-epoch pretraining of Vision Transformers. We report the Top-1 Accuracy % on validation dataset.
TRAINIGG METHODS
BIT WIDTH QUANTIZATION
DRIT-T
DEIT-S
DEIT-B
SWIN-T
SWIN-S
FULI PRECISION
A16W16G16
63.73
73.33
75.57
78.35
80.44
INT4
A4W4G4
FER-TENSOR
40.14
60.07
68.13
74.22
75.74
MICROSCALING (BASELINE)
A4W4G4
PER-GRE GROUP
58.56
70.10
74.54
76.87
79.45
TETRAJET (OURS)
A4W4G4
PER-GROUP
59.75
71.03
74.91
77.12
79.51
TETRAJET + Q-EMA(OURS)
A4W4G4
PER-GROUP
60.00
72.25
77.32
77.30
79.74
TETRAJET + Q-RAMPING(OURS)
A4W4G4
PER-GROUP
60.31
71.32
75.62
77.33
79.67
size would be < 111 - qal. so in this situation, we would get distw < dist.q and Rw, will be quite large.
2021) based on the official implementation". All the models are trained for 90 epochs on ImageNetIK (Russakovsky et al., 2015) with default training recipes. For Q-EMA & Q- Ramping, we show the insensitivity to their hyperparameter choice in Appendix C.3.
Therefore, the larger Rw. the more frequently and severely the weight element w oscillates, which means that we should put more effort into suppressing the oscillation of w.
We compared our MXFP4 training method, TetraJet, with full-precision training, 4-bit per-tensor quantization method INT4 (Xi et al., 2023), and original Microscaling's MXFP4 training method (Rouhani et al., 2023b). The detailed results are listed in Tab. 2.
6.2. Suppressing Weight Oscillation Adaptively
Based on periodically detecting and quantifying weight oscillation, we propose Adaptive Ramping Optimizer (Q- Ramping) to alleviate the oscillation problem of these weights. We adaptively decrease the updating frequency of these oscillating weights, by setting larger batch-size for them. We also expand their corresponding learning-rate pro- portional to their batch-size. The adapted batch-size would be an integer multiple of the global batch-size, and we would accumulate the gradient for each oscillating weight accord- ing to its own batch-size. This algorithm can he formalized as Algorithm 2 in Appendix C.
As result, our TetraJet can consistently outperform the orig- inal method Microscaling, and we can further improve the performance of MXFP4 training by overcoming oscillation problems in forward pass with Q-EMA Q-Ramping.
7.2. Quantitative Analysis on Oscillation Reduction
To validate our improvements in mitigating oscillation, we analyzed different statistics to show how our methods work in oscillation reduction in real training.
By applying Q-Ramping, the update frequency is reduced for oscillating weights, so that their oscillation frequency is also reduced. Additionally, through larger batch-size and larger learning-rate, the oscillating weights near the quantization thresholds can be updated to place further away from the quantization threshold. Therefore, the weight distribution will have a higher quantization confidence, and the oscillation phenomenon can be alleviated.
Improvement of Training Stability As described in Sec. 4.1, the rate of change for weights and activation can- not converge to zero in MXFP4, which reflects the model cannot converge stably. In Tab. 3, we can see our methods can effectively reduce the instability of both the weight and activation in forward.
7. Experiments
Improvement of Quantization Confidence As described in Sec. 4.2, weight confidence indicates the risk of weight os- cillation. If the confidence is lower at the end of the training, more weights are still oscillating around the quantization threshold, and it is harder for these parameters to converge to decisive values.
7.1. Vision Transformer Pre-Tree Pre-Training
We evaluate our TetraJet training method and oscillation reduction method Q-EMA & Q-Ramping on Vision Trans- formers pre- mre-training. During training, we quantize the for- ward and backward process all the linear layers the At- tention module and the MLP module of transformer blocks.
In Fig. 5, we can see the unique function of Q- D-Ramping in improving quantization confidence. It successfully reduced the weights that are prone to oscillate (those with low confi- dence) by identifying them, reducing their update frequency, and increasing their gradient accumulation steps.
We do pre-training for DeiT-Tiny, DeiT-Small, and DeiT- Base (Touvron et al., 2021) using Facebook's 's training recipe?, and pre-train Swin- Tiny and Swin-Small (Liu et al.,
hhtt:://ghhaa.commmmimmoofffffff Transforrmmr
7

---

MXFP4 Training for Vision Transformers
3.0
Tetrajet (Ours)
f. 2 delt_sman small
m.cck....ll..f2 deit
2.5
Tetrajet Dampen
Tetrajet + Q-Ramping (Ours)
10000
Tetrale Q-Ramping
Tetraj
34000
Tetrajet Q-Rarrping
Osciliating Weights%
Tetrajet + Q-EMA (Ours)
1.5
9000
32000
8000
30000
1.0
7000
28000
0.5
26000
0.0
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
QuantConf(w))
QuantConf(ww)
90
Epoch
Figure 5: Q-Ramping's unique effect on improving the
Figure 6: Q-EMA Q-Ramping's effect on oscillating weights reduction during the whole training process. We present the 90-epoch pre-trate pre-training of DeiT-T.
distribution of quantization confidence of the final model.
Table 3: Effect of Q-EMA & Q-Ramping stabilizing weight and activation at the end of DeiT-T training. r(-) refers to the rate of change for tensors, W is the quantized weights, and Y is the output of 9th transformer block given fixed input.
Table 4: Comparison of our oscillation reduction methods with other methods for DeiT MXFP4-Pretraininng Ima- geNet Classification. We report the top-1 Acc.% of the final model.
DeiT-T
DeiT-S
r(w9)
r(Y)
TetraJet
59.75
71.03
TetraJet TetraJet Dampen
0.0045
0.0401
Tetralet + Dampen
59.75
70.75
0.0044
0.0394
TetraJet + Freeze
16.45
22.04
Tetralet Q-EMA (Ours)
0.0018
0.0251
TetraJet Q-EMA (Ours)
60.00
72.25
Tetralet + Q- R-Ramping (Ours)
0.0028
0.0318
TetraJct Q-Ramping (Ours)
60.31
71.32
Oscillation Reduction throughout the Training We use Oscillation Ratio Rw (defined in Sec. 6.1) to characterize the oscillation problem during the whole training process. We define that those weights with R > are oscillating weights. As shown Fig. 6, both of our methods can effectively reduce the Oscillating Weights. Among them, Q-EMA reduces the most oscillating weights by directly smoothing weight quantization. Q-Ramping also reduces the oscillating level, while method "Dampen" from Nagel et al. (2022) cannot effectively reduce oscillation in MXFP4 pre-training.
in pre-training, and the "Freezing" method would encounter severe degradation when adapted to pre-training tasks.
Stability Improvement We removed weight quantizers in forward to simulate oscillation-free training (set Q identity function), and removed both activation and weight quantizers in forward to simulate MXFP4 training with stable forward process (set Q()) and Q(2) to identity func- tion). Consequently, our stabilization method Q-EMA and Q- Q-Ramping can counteract the influcnce of weight oscilla- tion, and approach a comparable accuracy to training with a full- prelisreoision forward process. Results are listed in Tab, in Appendix B.
7.3. Ablation Study
Training Method We investigate the quantization method in the MXFP4 training. We find that double quantization consistently outperforms Microscaling's incorrect gradient computation. Besides, when we ensure unbiased gradient estimation by double quantization and truncation-free seal- ing, we can get the optimal result with stochastic rounding. The detailed results are listed in Tab. 5 in Appendix B.
8. Conclusion
In this work, we not only proposed a new MXFP4 train- ing method TetruJet for a more accurate 4-bit training in MXFP4 format, but also introduced novel approaches to an- alyzing and resolving the instability of forward pass, which is the bottleneck of MXFP4 training. Extensive experiments revealed that our TetraJet consistently surpasses current 4- bit training methods, and Q-EMA / Q-Ramping can provide additional enhancement with effective oscillation reduction, and even achieve competitive performance compared to full-precision training.
Other Methods on Oscillation Reduction Following the configuration in Nagel et al. (2022), we compared Q-EMA Q- Ramming with their methods. As a result Tab. 4, their "Dampen"" method cannot work well on reducing oscillation
8

---

Reduced MXFP4 Training for Vision Transformers
Acknowledgment
References
The authors gratefully thank Pengle Zhang for the insight- ful discussions. This work was supported by the NSFC Project (No. 62376131), Beijing Natural Science Foun- dation (QY24258), Tsinghua University Initiative Scien- tific Research Program (Student Academic Research Ad- vancement Program), and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
Bengio, Y., Léonard, N., and Courville, A. Estimating propagating gradients through stochastic neurons for conditional computation, 2013.
Chen, J., Gai, Y., Yao, Z., Mahoney, M. W., and Gonzalez, J. E. A statistical framework for low-bitwidth training of deep neural networks. Advances in neural information processing systems,
2020.
Chmiel, B., Banner, R., Hoffer, E., Yaacov, H. B., and
Impact Statement
Soudry, D. Logarithmic unbiased quantization: Practical 4-bit training in deep learning. 2021.
Our MXFP4 low-precision training method enhances AI efficiency, reduces energy consumption, and improves ac- cessibility by lowering hardware costs. This can help bridge technological gaps and promote sustainable Al develop- ment. However, the reduced computational cost could lower barriers malicious uses, such as deepfake generation or automated disinformation. Ensuring that such technologies used ethically and for the benefit of society is essential to maximizing their positive impact.
Courbariaux, M., Bengio, Y., and David, J.-P. Binarycon- nect: Training deep neural networks with binary weights during propagations. Advances in neural information processing systems, 28, 2015.
Kalamkar, D. D., Mudigere, D., Mellempudi, N., Das, D.,
Banerjee, K., Avancha, S., Vooturi, D. T. Jammala- madaka, N., Huang, J., Yuen, H., Yang, J., Park, J., Heinecke, A., Georganas, E., Srinivasan. S., Kundu, A., Smelyanskiy, M., Kaul, B., and Dubey, P. A study of bfloat16 for deep learning training, CoRR, abs/1905. 12322, 2019. URL http:/ http://axvvvorrrge abs abs/19055 abs/1905.1233
Liu, A., Feng. Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., al. Deepseek-v3 2024.
technical report. arXiv preprint ar arXiv:2412.
Liu, S.-Y., Liu, Z., and Cheng, K.-T. Oscillation-free quantization for low-bit vision transformers. In Inter- national Conference on Machine Learning, pp. 21813- 21824. PMLR, 2023.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Micikevicius, P., Stosic, D., Burgess, N., Cornea, M., Dubey, P., Grisenthwaite, R., Ha, S., Heinecke, A., Judd, P., Kamalu, et al. Fp8 formats for deep learning. arXiv preprint arXiv:2209.054: 2022.
Nagel, M., Fournarakis, M., Bondarenko, Y., and Blankevoort, T. Overcoming oscillations in quantization- aware training. In International Conference on Machine Learning, pp. 163 16318-16330. PMLR, 2022.
Narang, S., Diamos, G., Elsen, E.,, Micikevicius, P., Alben, J., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., et al. Mixed precision training. In Int. Conf. on Learning Representation, 2017.

---

SecillationReeduce MXFP4 Training for Vision Transformers
NVIDIA. Nvidia blackwell architecture, 2024a. URL https://resources nviiia com/ en-us-Be en-us-blackin - www.bblcclllllchhtteeeeeeeeeeeee
the 38th International Conference Machine Learn- ing, volume 139 of Proceedings of Machine Learn- ing Research, pp. 10347-10357, PMLR, 18-24 Jul 2021. URL https: https//s//p/ccceiigg.... mlr h v139/touvron211 httm..
Accessed:
2025-01-30.
NVIDIA. Nvidia ıιx blackwell gpu architecture, 2024b. URL https://images nvidia. en/ htt.://mm/oouuttoooon///eeerr//r
Vaswani, A. Attention is all you need. Advances Neural Information Processing Systems, 2017.
seidia-rtx-Baee htts-bllccwel--p-prrrottteuuee pdf. Accessed: 2025-01-30.
nvidia
Wortsman, M., Dettmers, T., Zettlemoyer, L.., Morcos, A.,
Farhadi, A., and Schmidt, L. Stable and low-precision training for large-scale vision-language models. Ad- vances in Neural Information Processing Systems, 36: 10271-10298, 2023.
NVIDIA. Transformer engine, 2024c. URL https:
//github.. www.c///////// /tt./////rranaoooorrrnnnnnnononn Accessed: 2025-01-30.
Peng, H., Wu, K., Wei, Y., Zhao, G., Yang, Y., Liu, Z., Xiong, Y., Yang, Z., Ni, B., Hu, J., et al. Fp8-Im: Training fp8 large language models. arXiv preprint arXiv:2310.18314 2023.
Xi, H., Li, C., Chen, J.. and Zhu, J. Training transform- ers with 4-bit integers. Advances in Neural Information Processing Systems, 36:49146 49168, 2023.
Xi, H., Chen, Y., Zhao, K., Teh, K. J., Chen, J., and Zhu, J. Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization. In Inter- national Conference Machine Learning, pp. 54049 54063. PMLR, 2024.
Rouhani, B. D., Garegrat, N., Savell, T., More, A., Han, K.-N., Zhao, Ritchie amd Hall, M., Klar, J., Chung, E., Yu, Y. Schulte, M., Wittig, R., Bratt, I., Stephens, N., Milanovic, J., Brothers, J., Dubey, P., Cornea, M., Hei- necke, A., Rodriguez, A., Langhammer, M., Deng, S., Naumov, M., Micikevicius, Siu, M.. and Verrilli, C. Ocp microscaling (mx) specification. Technical report, 2023a.
Xi, H., Cai, H., Zhu, L., Lu, Y., Keutzer, K., Chen, J., and Han, S. COAT: Compressing optimizer states and activa- tions for memory-efficient FP8 training. In The Thirteenth International Conference on Learning Representations, 2025. URL https: hwwp//eeeeeeeeeeeeeeeeeeeeeeeeee net /ortmm id=XfKSDgqIR..
Rouhani, B. D., Zhao, R., More, A., Hall, M., Khodamoradi, A., Deng, S., Choudhary, D., Cornea, M., Dellinger, E., Denolf, K., et al. Microscaling data formats for deep learning, arXiv preprint arXiv:2310.1055: 2023b.
Zhu, F., Gong, R., Yu, F., Liu, X., Wang, Y., Li, Z.,, Yang, X., and Yan, J. Towards unified int8 training for convolu- tional neural network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1969-1979, 2020.
Russakovsky, O., Deng. J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3)2211-522 2015. doi: 10.1007/s 12000//11263-----088111
Sun, X., Choi, J., Chen, C.-Y., Wang, N., Venkataramani, S., Srinivasan, V., Cui, X., Zhang, W., and Gopalakr- ishnan, K. Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks. Advances neural information processing systems, 32, 2019.
Sun, X., Wang, N., Chen, C.-Y., Ni. J., Agrawal, A., Cui, X., Venkataramani, S., El Maghraoui, K., Srinivasan, V. V., and Gopalakrishnan, K. Ultra-low precision 4-bit training of deep neural networks. Advances in Neural Information Processing Systems, 33:1796-1880 1807, 2020.
Touvron, H., Cord, M., Douze, M., Massa, F., Sablay- rolles, A., and Jegou, H. Training data-efficient image transformers & distillation through attention. In Meila, M. and Zhang, (eds.), Proceedings of
10

---

Oscillation MXFP4 Training for Vision Transformers
A. Statistics for Measuring Oscillation and Training Instability
In this section, we formally define and explain the statistics we use in this paper to measure weight oscillation and training instability.
A.1. Oscillation Ratio
During a training stage with To steps, we sum up updating distance for each master weight element w and its quantized weight wq = Q(w):
Definition
To Σ lwe
distw =
1=1
distq
luo w/2
We define oscillation ratio Ru: for each weight element, representing the degree of oscillation:
Rw == /
In the Q-Ramping method for pre-training, we set To = 30 to minimize the additional cost of identifying oscillating weights. In the validation experiment (Tab. 6), we set To = 200 to fully validate the oscillation reduction.
Interpretation a weight element w has higher R at a certain stage of training, it means that shows more characteristics of oscillation. The larger Re,, the more frequently and severely the weight element w oscillates, which means that we should put more effort into suppressing the oscillation of
Comparison of Oscillation Ratio and Previous Metric Nagel al. (2022) also define a metric flipping frequency (average frequency of quantization flipping, defined for each weight element) to find out oscillating weights and measure oscillation severity, but it is only suitable for the small learning-rate training (e.g. fine-tuning, near the end of pre-training), because when the learning-rate is relatively large (e.g. the early middle stage of pre-training), the latent weight would be updated with large step size and the quantized weights also change frequently during training, but f would falsely recognize some of them as quantization oscillation. This is also a reason why the "Freeze" method performs badly in pre-training (see the result Tab. 4).
Oscillation Ratio R, overcomes the issue of oscillation detection in the early stage of pre-training. Only the weights that fall into real quantization oscillation would get large Rw:: these weights with small moves around the quantization threshold w is relatively small) but with frequent switch between quantization values (dist q is relatively large).
A.2. Quantization Confidence
Definition To quantitatively assess the severity of the oscillation problem, we define quantization confidence for each weight element which measures the normalized distance to the nearest quantization threshold.
QuantConf(w>) = min, w - thrd, MaxDist(w:FPP)
where FP4 denotes the quantized FP4 value of w. {thrd;} } denotes all the quantization thresholds, and MaxDist(u:PP) denotes the maximum possible distance when quantized to w FP4 It is ensured that QuantConf(ww) E 1].
Interpretation If an element w has less quantization confidence, it is more prone to oscillate, because it is closer to the quantization threshold and little perturbation would make its quantized value switch frequently. If a weight matrix W has more elements with low confidence, we call the weight distribution is less confident, which indicates that the optimization to this weight more unstable. For example, in Fig. 4, weights in Epoch 90 are less confident than weights in Epoch 30.

---

MXFP4 Training for Vision Transformers
A.3. Rate of Change for Weight and Activation
Definition we define rate of change for a tensor X as
1 To
Tb X
X
r(X)
where refers training step, and step To refers to short training interval.
During pre-training, test the rate of change for the master weight W, the quantized weight matrix Q (W ( "T, and activation Y in different stages.
This metric is useful in the end of training. When Learning Rate (LR) is approaching zero to push the model to quickly descend to a local minimum and converge, we expect the rate change for quantized weight and activation can also be near zero to ensure stability of training. However, in Section 4.1, we have found that the rate of change stays high at the end of MXFP4 training.
Interpretation
Therefore, we can decrease the rate of change for quantized weight and output activation of quantized layers, it means effectively improve the training stability. We have shown the results in Tab. 3.
12

---

Oscillation- MXFP4 Training for Vision Transformers
B. More Detailed Results of Ablation Study
Quantization Methods We do an ablation study to compare our training method TetraJet and Microscaling's original training method. Through Tab. we conclude that: (a) Our double quantization corrects the gradient estimation in MXFP4 Linear Layers, and is consistently better than Microscaling's 's original design. (b) As long as we give unbiased gradient estimation, which is guaranteed by double quantization and truncation-free scaling, we reach the optimal strategy with stochastic quantizution in backward. (c) It is necessary to ensure unbiasedness. Only in the unbiased situation, can stochastic quantization exert its advantage.
Table 5: Comparison quantization methods. We report the accuracy the validation set of 90-epoch DeiT-T pre-training.
Backward Quant
XW For Grad Computing
Computation of Shared Scale
Top-1% 1%
Top-5%
Note
Stochastic
Double Quantization
Truncation-Fres Free Scaling
59.75
82.67
TetraJet(unbiased gradient)
Stochastic
Double Quantization
Microscaling's 's Scaling
59.18
82.64
Stochastic
Microscaling's Design
Truncation-Free Free Scaling
56.98
80.60
Stochastic
Microscaling's 's Design
Microscaling's 's Scaling
57.49
81,27
Deterministic
Double Quantization
Truncation- Free Scaling Microscaling's 's Scaling
58.60
82.11
Deterministic
Double Quantization Microscaling's 's Design
59.02
82.18
Deterministic
Truncation-Free Scaling
58.40
81.57
Deterministic
Microscaling' Design
Microscaling's 's Scaling
58.56
81.92
Microscaling
Stability Improvement We simulated oscillation-free training by removing the weight quantizer in forward, and simulated stable forward process by removing both weight & activation quantizers in forward. As result in Tab, 6, our methods Q-EMA&& & Q-Ramping can fully eliminate the negative effects of weight oscillation, and can approach better accuracy with a more stable forward process.
Data Format We study the choice of FP4 format for the forward and backward computation. In Tab. 7, although E3M0 is another possible FP4 format, E2M1 is always better format for weight, activation, and gradient.
Table 6: Ablation study on quantization stability. We re- port the accuracy on validation set of 90-epoch DeïT-B pre-training. WQ: Weight Quantization in forward: AQ: Ac- tivation Quantization forward.
Table 7: MXFP4 Data Format Selection. We report the top-1 Acc.% % of DeïT-T Pre-Training.
Grad
E2M1
E3M0
A&W
Top-1 Acc.%
E2MI
59.75
58.90
E3M0
54.21
53.72
TetraJet TetraJet WQ Tetralet w/o WQ & AQ
74.91
75.16 16
75.86
TetraJet + Q-EMA
77.32
TetraJet + Q-mmpping
75.62
13

---

Oscillation-Rdduee MXFP4 Training for Vision Transformers
C. Detailed Implementation of Q-EMA and Q-Ramping
C.1. Algorithm: EMA Quantizer (Q-EMA)
Algorithm EMA Quantizer for Micro-Block (Q-EMA)
input Weight Block W: EMA weight block W EMA.
output Quantized Weight Block (WFP4, ;8) in MXFP4 format
1: Assume W and W EMA vectors of size 32.
2: + maxisis32 VIl. +- M + E [M = 0)
3: -- log2 Qp
2M
S 2
4: for i f- 1 to 32 do
5: 42 < two nearest MXFP4 values to W.
6: if WEMA: 91
WEMA: Q2 then
7: WPP4 + 91
8: else
W.PL 92
9:
10: end if
11:
end for
Return MXFP4 block (WFP4,.
12:
C.2. Algorithm: Adaptive Ramping Optimizer (Q-Ramping)
Algorithm 2 Adaptive Ramping Algorithm for MXFP4 Training (Q-Ramping)
Hyperparameter: k1, 12.
2: function Osci 1 alat onDetection (Model M, Global Learning-Rate LR, Global Batch-Size BS)
3: Train the model M for To < Tupdste steps calibration dataset without Q-Ramping, to detect oscillating weight.
4: for each weight element in quantized layers do
5: Compute the oscillation ratio Re according the length of trajectory of master weight w & quantized weight wo:
6: LRw + min(k) Kg RRw/ks ] + +,, Nmax) LR:
BS <- min(k2 22 R + 1, Mmax) BS;
8: I/ k1, K2 coefficients for amplifying LR & BS (that is using higher gradient accumulation step).
9: // Niax denotes the maximum amplification factor.
10: end for
11: end function
12: function Model Toddelrraing hhw.:/aaaai...itt.ggggaaammmmm (Initial Model M, Steps T, Learning-Rate LR, Batch-Size BS)
13: for + to do
14:
mod Tupdate = 0 then call Oe000ssillaaaa DacillattooDeteetiiiotk... LR, BS) to adaptively adjust LR.., & BS for each element w;
16: end if
17:
for each weight element in quantized layers do
18:
update w according to LR,,, & BS,, by Customized AdamW;
19:
end for
20: for each parameter W non-quantized layers do
update W by normal AdamW;
end for
23:
end for
24:
end function
14

---

Dacillation-Reduee MXFP4 Training for Vision Transformers
C.3. Selection of Hyperparameter & Insensitity to Hyperparameter
For Q-EMA, the momentum 0.998 for caleulating WEMA a good default choice. For Q-Ramping, 41 16 is good threshold to measure the severity of oscillation, and l2 5 is default ratio for amplifying the Learning Rate & Batch Size (meanwhile, reducing the frequency of oscillation). We can reach better performance through minor tuning. The detailed settings are listed in Tab. 8.
Table 8: Selection of hyperparameter in Q-EMA & Q-Ramping.
DeiT-T
DeiT-S
DeiT-B
Swin-T
Swin-S
Tetralet
59.75
71.03
74.91
77.12
79.51
Tetralet + Q-EMA (default: 0.998)
59.69
71.51
77.18
77.23
79.74
60.00
72.25
77.32
77.30
79.74
Tetralet + Q-EMA (best: 3 tuned)
(3
0.9983)
(3 0.9972)
(3 0.999)
(B = 0.9975)
(B 0.998)
Tetralet Q-Ramping (default: kn = 16, l:2 5)
60.31
71.32
75.62
77.23
79.52
60.31
71.32
75.62
77.33
79.67
Tetralet Q-Ramping (best: k1 16, k2 tuned)
(k2 5)
(k2 = 5)
(k:2 5)
(k2 3)
(k2 4)
We also validate Q-EMA MARiA/QRRauiee O--amping's insensitivity to hyperparameter choice in Tab. 9 & 10.
Table Insensitivity to hyperparameters (Tetralet + Q-EMA) on DeiT-B.
33
0.993
0.995
0.997
0.998
0.999
0.9995
w/o Q-EMA
Accuracy
75.39
76.37
77.23
77.18
77.32
77.30
74.91
Table 10: Insensitivity to hyperparameters (TetraJet Q-Ramping) on DeiT-B.
h1
16
16
16
16
16
16
8
12
16
20
w/o Q-Ramping
l:2
4
5
6
7
8
5
Accuracy
75.35
75.33
75.62
74.96
75.29
75.13
75.19
75.60
75.62
74.85
74.91
C.4. Other Discussion on Q-EMA & Q-Ramping
Q: Why cannot we combine two algorithms?
A: When we use Q-EMA, there are two variables (W & WEMA) that determine the result of weight quantization, training with Q-EMA results in a different MXFP4 training dynamic. Therefore, it is more complicated to identify and track the oscillating weights this situation. Therefore, is not proper to simply combine Q-EMA Ramping.
15


Reviews:
[{"review_id": "W63xSWDlWA", "reviewer": "Reviewer_e9Me", "summary": "The paper introduces TetraJet, a novel training method for Vision Transformers using the MXFP4 low-precision format, which is supported by Nvidia's Blackwell GPUs and offers significant speed improvements. The authors identify weight oscillation as a key issue causing accuracy degradation in MXFP4 training. To address this, they propose two techniques: EMA Quantizer (Q-EMA), which smooths weight quantization using an exponential moving average, and Adaptive Ramping Optimizer (Q-Ramping), which dynamically reduces update frequency for oscillating weights. Their approach achieves over 50% reduction in accuracy degradation compared to the baseline Microscaling method and brings MXFP4 training close to full-precision performance, demonstrating its effectiveness in stabilizing low-precision training.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for valuable comments. Below, we respond to the questions.\n\n> Claims And Evidence & Theoretical Claims:   \n> Lack formal theoretical validation & Leaving some theoretical gaps\n\nThank 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.\n\n> Other Weaknesses:   \n> The study focuses only on Vision Transformers and does not evaluate other architectures (e.g., CNNs, NLP models like LLMs).\n\nThanks 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.\n\nDue 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.\n\n> Experimental Designs Or Analyses:   \n> the study lacks statistical significance testing (e.g., confidence intervals or variance analysis), which would further strengthen the validity of the reported improvements.\n\nThanks 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.\n\n> Question 1:   \n>Have you tested MXFP4 fine-tuning on pre-trained models? Does the oscillation issue persist in fine-tuning settings? \n\nOur 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].\n\nWe 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.\n\n*Table 1. Results of ViT-base Fine-tuning (MS: Microscaling; TJ: TetraJet)*\n| Methods | Acc% (mean±std) |\n| --- | ---|\n| BF16 | 81.49±0.08 |\n| MS (Baseline) | 80.04±0.04|\n| TJ (Ours) | 80.17±0.03 |\n| TJ + Q-EMA (Ours)  | 80.57±0.09 |\n| TJ + Q-Ramping (Ours) | 80.25±0.04 |\n\n> Question 2:   \n> 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?\n\nThanks 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.\nWe 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. \n\nWe 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.\n\n*Table 2: Insensitivity to hyperparameters (TetraJet + Q-EMA) on DeiT-B.*\n\n| $\\beta$ | 0.993 | 0.995 | 0.997 | 0.998 | 0.999 | 0.9995 | w/o Q-EMA |\n| ------- | ----- | ----- | ----- | ----- | ----- | ------ | --------- |\n| Acc%    | 75.39 | 76.37 | 77.23 | 77.18 | 77.32 | 77.30  | 74.91     |\n\n\n*Table 3: Insensitivity to hyperparameters (TetraJet + Q-Ramping) on DeiT-B.*\n| $k_1$ |  16   |  16   |  16   |  16   |  16   |  16   |   8   |  12   |  16   |  20   | w/o Q-Ramping |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| $k_2$ |   3   |   4   |   5   |   6   |   7   |   8   |   5   |   5   |   5   |   5   |               |\n|  Acc  | 75.35 | 75.33 | 75.62 | 74.96 | 75.29 | 75.13 | 75.19 | 75.60 | 75.62 | 74.85 |     74.91     |\n\n[1] https://docs.nvidia.com/cuda/pdf/ptx_isa_8.7.pdf  \n[2] S.-Y. Liu, Z. Liu, and K.-T. Cheng, \"Oscillation-free quantization for low-bit vision transformers,\" ICML, PMLR, 2023.  \n[3] https://github.com/facebookresearch/mae", "claims_and_evidence": "The paper provides strong empirical evidence to support its claims, including extensive experiments on Vision Transformers that demonstrate TetraJet’s superiority over existing MXFP4 training methods. The identification of the weight oscillation problem is backed by quantitative analysis of weight changes, rate of change metrics, and oscillation ratio measurements. The effectiveness of Q-EMA and Q-Ramping is validated through comparative accuracy results, stability improvements, and confidence metrics, showing clear advantages over the baseline. However, while the paper convincingly argues that oscillation is the main source of degradation, it does not theoretically prove why these methods generalize across different architectures or tasks beyond Vision Transformers, leaving room for further validation.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are well-aligned with the problem of low-precision training for Vision Transformers. The authors conduct experiments on ImageNet-1K using established Vision Transformer architectures (DeiT and Swin Transformers), which are widely used benchmarks for image classification. Their evaluation focuses on accuracy degradation, oscillation metrics, and training stability, which are appropriate for assessing the impact of quantization techniques. The use of quantization confidence and oscillation ratio provides insightful analysis beyond standard accuracy metrics. However, the paper primarily evaluates MXFP4 on pre-training tasks, and it would be useful to see results on fine-tuning or downstream applications to confirm generalizability.", "theoretical_claims": "The paper primarily focuses on empirical findings, but it includes some theoretical justifications for its quantization techniques, such as the unbiased gradient estimation from double quantization and truncation-free scaling. The derivations in Section 3.4 align with prior work on Straight-Through Estimators (STE) and unbiased gradient estimation. The quantization confidence metric and the oscillation ratio definition are intuitively reasonable, though they lack formal theoretical validation. While the explanations are convincing, rigorous mathematical proofs on why Q-EMA and Q-Ramping reduce oscillation across different scenarios are not provided, leaving some theoretical gaps.", "experimental_designs_or_analyses": "The experimental design is robust and well-structured, leveraging ImageNet-1K as a benchmark and conducting comprehensive ablation studies to isolate the effects of different quantization components. The impact analysis of six quantizers (Table 1) effectively identifies the activation and weight quantizers in the forward pass as the primary sources of degradation. The oscillation reduction experiments (Figures 2–6) provide strong empirical support for Q-EMA and Q-Ramping. However, the study lacks statistical significance testing (e.g., confidence intervals or variance analysis), which would further strengthen the validity of the reported improvements. Additionally, while Q-Ramping's dynamic adaptation is intuitive, its hyperparameter sensitivity is only briefly discussed, and a more detailed analysis of its tuning could enhance confidence in its robustness.", "relation_to_broader_scientific_literature": "The paper builds on prior work in low-precision training, QAT, and MX quantization (Rouhani et al., 2023) by addressing oscillation issues in MXFP4 training from scratch. It extends findings on weight oscillation in QAT (Nagel et al., 2022; Liu et al., 2023) but uniquely applies EMA-based smoothing and adaptive update strategies to stabilize MXFP4. While related techniques exist in optimization, their use for 4-bit Vision Transformer training is novel. However, its applicability beyond vision models remains unexplored.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "- The study focuses only on Vision Transformers and does not evaluate other architectures (e.g., CNNs, NLP models like LLMs), limiting its broader applicability.", "other_comments_or_suggestions": "None", "questions_for_authors": "- Have you tested MXFP4 fine-tuning on pre-trained models? Does the oscillation issue persist in fine-tuning settings?\n- 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ORdOORZx5O", "reviewer": "Reviewer_KRL2", "summary": "The authors propose a MXF4 training scheme for Vision Transformers. Training at extremely low-bit such as 4-bit formats is challenging and prone to high accuracy loss, mainly due to weight oscillations in the forward pass as identified by the authors. The paper outlines two methods EMA Quantizer (Q-EMA) and Adaptive Ramping Optimizer (Q-Ramping) to resolve this issue.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for valuable comments and the acknowledgment of our contributions. Below we respond to the questions.\n\n> Broader Scientific Literature:   \n> The authors might enhance their context by explicitly citing some recent low-precision training surveys...\n\nThanks for the valuable advice. We have found some surveys on low-precision training [1,2,3]. We will include them in our revised version. \n\n> Other Comments Or Suggestions:  \n> I think the sentence \"Therefore, quantizer …. group shape.\" can be better phrased. It's hard to follow the sentence and what quantizer(*) refer to.\n\nThanks for the valuable advice. We will clarify this by pointing out that they are Quantizers $Q^{(*)}$ in Eq. 3-5.\n\n> Weakness 1:   \n> The analysis primarily focuses on ViT architectures; evaluating more diverse architectures (e.g., CNNs, LLMs) could further validate general applicability.\n\nThanks 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.\n\nDue 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.\n\n> Weakness 2:   \n> Discussion of potential overhead or computational costs associated with the oscillation tracking in Q-Ramping could be expanded.\n\nThanks 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.\n\n> Question 1:  \n> Could you clarify the computational overhead introduced by Q-EMA and Q-Ramping methods in actual training?\n\nThanks for the valuable questions. \nIn 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.\n\n\n> Question 2:  \n> 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?\n\nThanks 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.\n\nWe 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.\n\n*Table 1: Insensitivity to hyperparameters (Q-EMA) on DeiT-B.*\n\n| $\\beta$ | 0.993 | 0.995 | 0.997 | 0.998 | 0.999 | 0.9995 | w/o Q-EMA |\n| ------- | ----- | ----- | ----- | ----- | ----- | ------ | --------- |\n|Acc%|75.39|76.37|77.23|77.18|77.32|77.30|74.91|\n\n*Table 2: Insensitivity to hyperparameters (Q-Ramping) on DeiT-B.*\n|$k_1$|16|16|16|16|16|16|8|12|16|20|w/o Q-Ramping|\n|-|-|-|-|-|-|-|-|-|-|-|-|\n|$k_2$|3|4|5|6|7|8|5|5|5|5||\n|Acc%|75.35|75.33|75.62|74.96|75.29|75.13|75.19|75.60|75.62|74.85|74.91|\n\n*Table 3. Results of 50epoch MAE ViT-base Fine-tuning (MS: Microscaling; TJ: TetraJet)*\n|Methods|Acc%(mean±std)|\n|---|---|\n|BF16|81.49±0.08|\n|MS(Baseline)|80.04±0.04|\n|TJ(Ours)|80.17±0.03|\n|TJ+Q-EMA(Ours)|80.57±0.09|\n|TJ+Q-Ramping(Ours)|80.25±0.04|\n\n\n*Table 4. OLMo-2 [6] 150M Pre-training with 20 Billion tokens*\n| Methods | Perplexity |\n| --- | --- |\n| BF16  | 23.18  |\n| MS (Baseline) | 25.88  |\n| TJ (Ours)| 23.82 |\n\n[1] Wei L et al., \"Advances in the Neural Network Quantization: A Comprehensive Review\". Applied Sciences. 2024.  \n[2] Chitsaz K et al., \"Exploring Quantization for Efficient Pre-Training of Transformer Language Models,\" arXiv:2407.11722, 2024.  \n[3] Kumar T et al., \"Scaling Laws for Precision,\" arXiv:2411.04330, 2024.  \n[4] https://docs.nvidia.com/cuda/pdf/ptx_isa_8.7.pdf  \n[5] OLMo Team et al., \"2 OLMo 2 Furious,\" arXiv:2501.00656, 2024.  \n[6] https://github.com/facebookresearch/mae", "claims_and_evidence": "The authors provide clear evidence through detailed experiments on popular Vision Transformer architectures (DeiT and Swin). They systematically identify the accuracy degradation problem in MXFP4 training due to weight oscillations. The evidence from the ablation studies is thorough, clearly indicating that forward-pass quantizers for activations and weights cause the most degradation. Both proposed solutions, Q-EMA and Q-Ramping, effectively address this oscillation issue, demonstrated through quantitative metrics like rate of change and quantization confidence. \nThe experiments and analyses are solid, with clear ablation studies and comparisons against baseline methods and competitive state-of-the-art methods. I reviewed the supplementary materials", "methods_and_evaluation_criteria": "The proposed method sound reasonable with proper explanation.", "theoretical_claims": "I did not rigorously verify any theoretical proofs, as the primary contributions are experimental and methodological.", "experimental_designs_or_analyses": "Results and analysis sound reasonable.", "relation_to_broader_scientific_literature": "The paper tackles the complex problem of training with MXFP4 format. There are several recent attempts at low-bit training specifically targeting FP8, and INT4 methods. The discussion of prior techniques (e.g., Microscaling, Jetfire, LSQ quantization) is comprehensive and clear. However, the authors might enhance their context by explicitly citing some recent low-precision training surveys, if available, to provide broader context.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n\n- Clearly identifies and addresses a practical issue (weight oscillation) with thoughtful solutions.\n- Extensive empirical validation across multiple ViT architectures demonstrates significant performance improvements.\n- Methods introduced (Q-EMA, Q-Ramping) are innovative yet practical.\n\nWeaknesses:\n\n- The analysis primarily focuses on ViT architectures; evaluating more diverse architectures (e.g., CNNs, LLMs) could further validate general applicability.\n- Discussion of potential overhead or computational costs associated with the oscillation tracking in Q-Ramping could be expanded.", "other_comments_or_suggestions": "I think the sentence \"Therefore, quantizer (1)(3)(5) should use 1 × 32 group shape, and quantizer (2)(4)(6) should use\n32 × 1 group shape.\" can be better phrased. It's hard to follow the sentence and what quantizer(*) refer to.", "questions_for_authors": "- Could you clarify the computational overhead introduced by Q-EMA and Q-Ramping methods in actual training?\n- 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jaL5nAPQbL", "reviewer": "Reviewer_TZFB", "summary": "This paper proposes two methods to train vision transformers with MXFP4-accelerated GEMMs. In the backward pass, the authors add stochastic rounding and scaling to achieve unbiased gradient estimates. In the forward pass, the authors add various EMA-based methods to avoid weight oscillation during quantization. The authors test their method on training smaller vision transformers and find that the proposed methods are sufficient to achieving near-lossless performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for valuable comments and the acknowledgment of our contributions. Below we respond to the questions.\n> Question 1: \n> … 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?\n\nThanks 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. \n\nFor 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.\n\nFor 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).\n\n*Table 1. Oscillation Reduction on DeiT-S 60-epoch pre-training.*\n| Methods | Accuracy |\n| --- | --- |\n| MS (Baseline) | 63.73 |\n| MS+Q-EMA | 64.19 |\n\n*Table 2. Accuracy improvement without additional oscillation reduction techniques (90-epoch pre-train)*\n|Methods|DeiT-T|DeiT-S|DeiT-B|Swin-T|Swin-S|\n|-|-|-|-|-|-|\n|MS(Baseline)|58.56|70.10|74.54|76.87|79.45|\n|TJ(Ours)|59.75|71.03|74.91|77.12|79.51|\n\n> Question 2: \n> Is double quantization actually necessary? \n\nThanks for the valuable question. There are mainly two core issues that necessitate double quantization：\n- **(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.\n- **(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).\n\nIn conclusion, the correctness of optimization goal & the non-square group shape (1x32 / 32x1) make double quantization necessary in our MXFP4 training.\n\n> Question 3:\n> 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.\n\nThanks 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]).\n\nOur 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. \n\n*Table 3. OLMo-2 150M pre-training with 20 Billion tokens*\n|Method|Perplexity|\n|-|-|\n|BF16  | 23.18  |\n|Microscaling(Baseline)|25.88|\n|TetraJet(Ours)|23.82|\n\n> Question 4:\n> What is the cost of storing the EMA components for oscillation reduction?\n\nDuring 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).\n\nIf 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.\n\n[1] Y. Li et al.,\"Q-vit: Accurate and fully quantized low-bit vision transformer,\" NeurIPS, 2022.  \n[2] Y. Liu et al.,\"NoisyQuant: Noisy bias-enhanced post-training activation quantization for vision transformers,\" CVPR, 2023.  \n[3] Z. Wang et al.,\"Quantformer: Learning extremely low-precision vision transformers,\" IEEE TPAMI,2022.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "This is a mostly empirical paper.", "experimental_designs_or_analyses": "The experiments seem reasonable.", "relation_to_broader_scientific_literature": "FP4 training is a relatively unsolved problem. The paper makes a number of references to oscillation reduction in low precision training that seem relevant.", "essential_references_not_discussed": "A number of FP4 LLM training papers have come out since this paper was submitted. The authors may find these works useful and interesting but since they were released around the time of submission or after, it would be unreasonable to have the authors compare against them.", "other_strengths_and_weaknesses": "The paper seems to have two separate methods: unbiased gradient estimation in the backward pass and oscillation reduction in the forward pass. It is not clear to me how orthogonal these two methods are. 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?", "other_comments_or_suggestions": "Is double quantization actually necessary? If you use stochastic rounding to compute unbiased gradients and stochastic rounding is implemented with iid noise, the output should be unbiased regardless of which ways the MX groups go. Perhaps I am misunderstanding what the core issue is that necessitates double quantization.\n\nIs 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.\n\nWhat is the cost of storing the EMA components for oscillation reduction?", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Summary: The paper introduces TetraJet training method for Vision Transformers using MXFP4 low-precision data type. The paper investigates the root cause of degradation in performance and identify weight oscillation as a root cause. The paper further fixes this using EMA method which smoothens weight quantization and Q-Ramping which reduces updates for unstable weights. The paper demonstrates experimental results via training DeiT and Swin transformers.

Review summary:
Reviewers agree this is an important paper for low-precision training and found the paper is well written with good empirical results. The proposed solutions are simple, effective, and supported by ablation studies. Overall, reviewers are satisfied and lean toward acceptance.
Decision: Accept (poster)