Mixed Precision Scaling Laws for Quantized Pretraining

Summary (Overview)

• Novel scaling law: The authors develop a scaling law for quantized pretraining with mixed precision across forward and backward passes, incorporating persistent optimization error terms ($\epsilon_f$ and $\epsilon_b$) that existing laws (e.g., Quartet) fail to model.

• Theoretical grounding: The law is motivated by a convergence analysis (Remark 3.1) that decomposes the optimization gap into three components: a vanishing $O(T^{-1/2})$ term, and two persistent $O(1)$ terms from forward quantization error and backward gradient bias.

• Superior prediction accuracy: The proposed law reduces held-out loss prediction MSE by 1.6× to 11× compared to the Quartet law across model size and data extrapolation settings on both Qwen3 and Llama-style experiments.

• Layer-wise extension: The law is extended to mixed precision across Transformer layers via interpolation $L = (L_2 - L_1)p^{\gamma_S} + L_1$, where $\gamma_S$ captures the layer selection strategy.

• Key finding: Unlike prior precision scaling laws, the proposed law predicts that compute-optimal precision increases with compute budget, since persistent optimization error does not diminish with scale.

Introduction and Theoretical Foundation

Background

Scaling laws guide LLM pretraining by characterizing how loss varies with model size, data, and compute (Kaplan et al., 2020; Hoffmann et al., 2022). Quantized training reduces costs by performing forward/backward computations at reduced precision (down to 4 bits with block-scaled formats like NVFP4). However, existing recipes differ in how quantization scales are shared/selected in the forward pass and how gradients are estimated in the backward pass.

Motivation

The authors identify a critical gap: existing scaling laws for quantized pretraining treat quantization as reducing effective parameter count (forward) or effective token count (backward), implying penalties can be offset by scaling. However, quantization also introduces persistent optimization error that does not decrease with more data or parameters—weights can oscillate across rounding boundaries (forward), and gradient estimates can be biased (backward).

Theoretical Decomposition

The loss gap decomposes into:

$$\mathbb{E}_{J,\theta_J}\left[L_N(\hat{\theta}_J)\right] - E = \underbrace{\left(L^*_{N,\psi_f} - E\right)}_{\text{capacity gap}} + \underbrace{\left(\mathbb{E}_{J,\theta_J}\left[L_N(\hat{\theta}_J)\right] - L^*_{N,\psi_f}\right)}_{\text{optimization gap}}$$

where $L^*_{N,\psi_f} := \inf_{\theta \in \mathbb{R}^N} L_N(Q(\theta; \psi_f))$ is the optimal loss of the quantized model and $E$ is the irreducible loss.

Methodology

Proposed Scaling Law

The core scaling law (Eq. 1) is:

$$L = \left[A(\rho_N(\psi_f)N)^{-\alpha} + B(\rho_D(\psi_b)D)^{-\beta}\right]^{\gamma} + \epsilon_f(\psi_f) + \epsilon_b(\psi_b) + E$$

where:

• $\rho_N$, $\rho_D$: rescale parameter count $N$ and token count $D$ based on forward/backward schemes $\psi_f$, $\psi_b$
• $\epsilon_f$, $\epsilon_b \geq 0$: persistent optimization error from forward/backward quantization
• $A$, $B$, $\alpha$, $\beta$, $\gamma$, $E$: fitted constants

Theoretical Justification (Remark 3.1)

For convex loss $L_N$, deterministic forward quantization, and learning rate $\eta = c/\sqrt{T}$:

$$\mathbb{E}_{J,\theta_J}[L_N(\hat{\theta}_J)] - L^*_{N,\psi_f} \leq \underbrace{\frac{1}{2\sqrt{T}}\left[\frac{R_1(N,\psi_f)}{c} + \frac{c}{T}\sum_{t=1}^{T}\mathbb{E}\|\hat{g}_t\|^2\right]}_{(a) \, O(T^{-1/2})} + \underbrace{G\left(\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\|\hat{\theta}_t - \theta_t\|^2\right)^{1/2}}_{(b) \, O(1)} + \underbrace{\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|b_t\|\|\theta_t - q^*\|]}_{(c) \, O(1)}$$

Terms (b) and (c) motivate the persistent error terms $\epsilon_f$ and $\epsilon_b$, respectively.

Example 3.1 (Persistent Forward Error)

For $L_1(\theta) = (\theta - \frac{1}{4})^2$ with quantization $Q(\theta) = \lfloor\theta + \frac{1}{2}\rfloor$ and $T$ divisible by 4, the optimization gap is:

$$\mathbb{E}_J[L_1(\hat{\theta}_J)] - L^*_{1,\psi_f} = \frac{1}{4} \cdot \frac{9}{16} + \frac{3}{4} \cdot \frac{1}{16} - \frac{1}{16} = \frac{1}{8} = \Omega(1)$$

This shows the gap does not vanish as $T \to \infty$.

Experimental Setup

Two data grids were used:

Recipe Decomposition Results

Forward quantization matters more than backward:

Empirical Validation / Results

Main Results (Table 4)

The law reduces MSE by 1.6× to 11× compared to Quartet.

Ablation Results (Table 5)

Removing $\epsilon_f$ increases error by ~3×, confirming forward persistent error is critical.

Layer-Wise Extension Results

Fitted exponents $\gamma_S$ (higher = better strategy):

Keeping first/last blocks in high precision is particularly effective in the backward pass.

Fitted Parameters (Table 10, selected)

Theoretical and Practical Implications

Compute-Optimal Allocation

The cost model (Eq. 4):

$$C(N, D, p_f, p_b) = 2ND\left(\frac{1-p_f}{\lambda_{f,\text{MX8}}} + \frac{p_f}{\lambda_{f,\text{NV4}}}\right) + 4ND\left(\frac{1-p_b}{\lambda_{b,\text{MX8}}} + \frac{p_b}{\lambda_{b,\text{NV4}}}\right)$$

The optimization problem (Eq. 5):

$$\min_{N, D, p_f, p_b} L(N, D, p_f, p_b) \quad \text{s.t.} \quad C(N, D, p_f, p_b) \leq C_0, \quad p_f, p_b \in [0,1]$$

Key Findings

1. Optimal precision increases with budget: As compute budget $C_0$ grows, model/data terms shrink but persistent error remains, so higher precision becomes optimal. This contrasts with prior laws predicting budget-independent precision.

2. Backward pass should use low precision first: The backward pass (2/3 of FLOPs) benefits more from NVFP4 due to: (a) larger compute savings, (b) higher $\gamma_b$ (less loss penalty), (c) ~half the persistent error of forward quantization.

3. Layer placement matters mainly in backward pass: First/last block selection incurs only 33% of the loss gap vs. 70.7% for middle blocks at $p = 0.5$.

Conclusion

The authors develop a scaling law for mixed-precision quantized pretraining that explicitly models persistent optimization error from both forward and backward quantization. Key contributions:

1. Recipe decomposition: Identified forward quantization as more loss-sensitive than backward techniques, with the best FP4 recipe (Quartet II) still leaving a noticeable gap to BF16.

2. Theoretically-grounded law: The law $L = [A(\rho_N N)^{-\alpha} + B(\rho_D D)^{-\beta}]^\gamma + \epsilon_f + \epsilon_b + E$ outperforms existing laws by 1.6×–11× in held-out MSE.

3. Layer-wise extension: The interpolation $L = (L_2 - L_1)p^{\gamma_S} + L_1$ accurately predicts mixed-layer losses, with first/last block selection being particularly effective in the backward pass.

4. Compute-optimal allocation: The law predicts that optimal precision increases with compute budget—a novel finding driven by persistent optimization error that does not diminish with scale.

Future directions include validating on larger models, exploring additional quantization schemes, and extending to other architectures beyond Transformers.