ReQAT: Achieving Full-Precision Reasoning Accuracy with 4-bit Floating-Point Quantization-Aware Training

Summary

The paper diagnoses why FP4 quantization degrades chain-of-thought reasoning in distilled Large Reasoning Models and proposes ReQAT, a training-time recipe tailored to that diagnosis. The central conceptual move is entropy-selective: rather than treat all tokens equivalently, the authors argue — through a mixed-precision routing experiment, a logit-noise sensitivity study, and tail-mass measurements — that FP4 failures concentrate at low-entropy "symbolic commitments" (digits, operators) where quantization rarely flips the top-1 token but reliably inflates the probability of non-top-1 alternatives, so sampling errors cascade through long reasoning traces. Three matching interventions follow. Trace-Aligned QAT (TAQ) is a two-stage FT→QAT pipeline that reuses the same reasoning traces so learning is reallocated onto the same low-entropy positions. Selective Entropy Minimization (SEM) is an auxiliary loss that soft-weights entropy sharpening toward the lowest-entropy tokens. Q-FIT is a small calibration that jointly picks a pre-RoPE paired scale and a post-RoPE shift for the KV cache. The reported outcome — matching or slightly exceeding BF16 full-FT accuracy under NVFP4 W4A4KV4 while retaining roughly 3× serving throughput — is a striking positive result if it holds up, and links quantization-aware training to a specific, testable statistical property of reasoning traces.

The main conceptual limitation is that the mechanism story is stated more strongly than the evidence bears. The margin over BF16 full FT is very small (65.94 vs 65.46 on AIME-120, well within seed variance on a 120-problem set, yet the paper reports no per-seed variance despite advertising 8-seed averaging), the "cascade" causal story is partially undermined by the qualitative example the authors themselves display (the model self-corrects but still fails), and interpretive claims — "monotonic" component gains, "consistent" SEM improvements, "up to 3.9× speedup" — read stronger than the tables and heterogeneous cross-platform configurations support. Empirical scope is also narrow: two DeepSeek-distilled models, math + code only, with several supporting analyses (mixed-precision routing, noise injection, tail-mass, gradient-contribution ratio) not yet reproducible from the released code.

Strengths

• Diagnosis before remedy. The paper first establishes, through three complementary experiments (mixed-precision routing, logit-noise perturbation, and tail-mass vs top-1 mismatch), that FP4 damage is concentrated at low-entropy positions and manifests as non-top-1 sampling inflation rather than argmax flips. Each of the three interventions (TAQ, SEM, Q-FIT) then maps cleanly onto one prong of the diagnosis.
• Trace-alignment mechanism. The FT vs QAT vs FT+QAT entropy-change analysis (Fig. 4) and the low-entropy gradient-contribution ratio $C_{\mathrm{low}}$ (Fig. 7) together provide a plausible learning-dynamics story for why reusing the same reasoning traces is important — not just that it helps, but through what channel.
• Q-FIT is a principled joint calibration. Rather than committing to pre-RoPE scaling or post-RoPE shifting alone, Q-FIT chooses per layer between them via a small $(\alpha_s, \alpha_m)$ grid search, and the visualization in Fig. 5 makes the layer-specific rationale legible.
• Honest disclosures. The paper is explicit about several inconvenient details a weaker treatment might hide: DGX Spark uses a W4A4KV8 fallback, MMLU shows only marginal gains, greedy decoding reduces the PTQ gap, and multi-epoch TAQ alone does not close the FP4 gap.
• Breadth of ablations. Component-wise ablations (Tables 4, 5, 12, 13, 15–17) cover trace alignment, Q-FIT design, SEM soft vs binary mask, SEM strength, multi-epoch TAQ, decoding, and rotation baselines.

Weaknesses

• Headline gain is over an un-fine-tuned baseline, not over BF16 full FT. The advertised "+9.1%" is 65.94 − 56.83 (vs the BF16 baseline). The margin over the matched BF16 full-FT comparison — the one implied by "surpasses BF16 fine-tuning accuracy" — is +0.48 pp. Fig. 1(d) and Fig. 3 juxtapose "+9.1%" with the "Full-FT" label, inviting misreading.
• ~0.5 pp margin is likely within seed noise. AIME-120 is 120 problems; per-seed variance and paired significance tests are not reported despite Appendix D.1 stating 8-seed averaging. The categorical "surpasses" claim is not statistically supported by what is shown.
• Monotonicity claim is overstated. The paper asserts accuracy increases monotonically as TAQ, Q-FIT, SEM are added; Table 1 (NVFP4 W4A4KV4) and Table 2 (GSM8K, MATH-500) contain non-monotonic entries and SEM regressions.
• Cascade story vs its own example. In Fig. 9, the low-entropy-noise trace detects and self-corrects the injected symbolic error but still fails. The narrative "later steps are built upon an incorrect intermediate result" therefore does not explain the shown failure; some other downstream corruption is doing the work.
• KV cache format inconsistency. The paper defines NVFP4 as E2M1 throughout but adopts E1M2 for the KV cache. Whether the reported hardware-native throughput uses the same E1M2 KV format (i.e., is realizable on Blackwell tensor cores) is not explained.
• Cross-platform throughput claim. The abstract's "3.9× on DGX Spark, 3.1× on B200" is a W4A4KV8 number on one platform and W4A4KV4 on the other; the two are not measuring the same quantization system as the accuracy tables.
• Compute vs bandwidth attribution. The 1.8–1.9× speedup at 2K output tokens is ascribed "primarily to compute benefit of 4-bit GEMM" even though autoregressive decode is typically bandwidth-bound and 4-bit weights/KV alone cut bytes moved by ~4×.
• Sample-count confound in response-length analysis. Table 6's per-cell sample counts vary by nearly 3× across methods within the same length bin (PTQ 96 vs QAT 295 in 24–32K), so accuracy is not compared on matched problems; the response-length story is confounded by selection.
• Underspecified block rotation for MXFP4. The MXFP4 block-Hadamard rotation is described only as "a special case within Q-FIT," which conflates two distinct mechanisms; the main text does not state the block size, target tensors, or interaction with the MXFP4 group size.
• Missing equations for $s_t$ and $C_{\mathrm{low}}$. Section 6's mechanism analysis introduces "the per-token gradient magnitude $s_t$ and the gradient contribution ratio $C_{\mathrm{low}}$" but the equations themselves are not displayed.
• Narrow model coverage. Both main models are DeepSeek-R1 distillations from a single teacher; the framework's generality to non-distilled LRMs (Qwen3 reasoning, gpt-oss, Llama-3-instruct-style reasoning models) is untested.
• Ablation baseline ambiguity. Table 16 includes a Q-FIT-off row even though the prose says "with Q-FIT enabled," leaving the SEM baseline ambiguous (~+1.6 pp vs ~+0.5 pp).
• Q-FIT vs RHT framing hides that "no transformation" is best at 140M. Table 17 bolds the no-transformation column at 140M (60.32) above Q-FIT (59.79); the prose reports only the Q-FIT-vs-RHT comparison.
• Overreach in qualitative narrative. "In all tested cases, the model fails to recover" is a universal quantifier attached to a small qualitative case study and contradicts the quantitative Fig. 2(e), which shows a reduced (nonzero) accuracy.

Reproducibility & code

• Training pipeline present, but ablation switches not exposed. The full ReQAT recipe — Stage-1 BF16 SFT, Q-FIT calibration via methods/daroc/run_daroc.py, Stage-2 QAT+SEM via training/reqat/sft.py — is shipped and its hyperparameters (LR 1e-5, cosine-with-min-LR, warmup 0.03, per-device batch 1, grad-accum 2, max seq 25K, bfloat16) match Table 11. No result files or expected numeric targets ship with the release, and the individual component ablations (Q-FIT scale-only vs shift-only vs E1M2/E2M1 KV; SEM soft vs binary mask; Table 4 misaligned Stage-2 dataset) require code edits rather than CLI flags.
• Throughput evidence not reproducible from released code. The K-shift in the released repo runs as a Python KShiftCallback, and the README states that the TensorRT-LLM plugin needed for hardware-native NVFP4 inference is not included. No trtllm-bench configs, engine build recipes, output-length sweeps, or throughput-parsing utilities are shipped, so the up-to-3.9× / 3.1× headline throughput numbers rest entirely on the authors' bench.
• Section 3 analyses have no shipped code. The mixed-precision BF16/FP4 dual-forward decoder, the multiplicative-Gaussian logit-noise perturbation harness with entropy-quantile selection, and the tail-mass / top-1 mismatch computation with 20-bin quantization are not implemented anywhere in the release. Since these three analyses ground the paper's core diagnosis, a reproducer must re-implement them from scratch.
• Mechanism analysis (Fig. 4, Fig. 7) unsupported. No intermediate checkpoint saves (training configs set save_strategy: 'no'), no per-baseline-bin entropy-change script, and no per-token embedding-gradient hook exist in the release, so Fig. 4's entropy-change trajectories and Fig. 7's $C_{\mathrm{low}}$ trajectory require significant custom instrumentation.
• PTQ / QAT / QAD baselines not shipped. The vLLM custom registry contains only ReQAT-style fake-quantized Qwen2 and Llama classes; AWQ, QuaRot, FlatQuant, and QAD comparisons rely on external re-implementations not included.
• MMLU evaluation not wired. inference.py supports GSM8K, MATH-500, AIME variants, and LiveCodeBench, but not MMLU; the Table 14 neutrality claim is not directly reproducible.
• Seed-averaging claim not scripted. Appendix D.1 states 8-seed averaging but scripts/inference/inference.sh hard-codes seed=42 and no aggregation script is shipped — a real concern given the small AIME margins.
• Llama-8B multi-epoch recipe not scripted; dataset naming misleading. The shipped dataset generator selects ~5,562 examples despite the folder name 89k-math-sft, and the paper's 70M-token Llama-8B subset with multi-epoch schedule (Tables 2, 9, 13, 15) has no dedicated generator or driver.

Recommended Changes

Essential

• Reframe the "+9.1%" and "surpasses BF16 FT" headline. In the abstract, intro, Fig. 1, and Fig. 3, make explicit that +9.1% is measured against the un-fine-tuned BF16 baseline and that the margin over BF16 full FT is ~0.5 pp on AIME-120. (Addresses Headline gain is over an un-fine-tuned baseline in Weaknesses.)
• Report per-seed variance and a paired significance test for the AIME comparisons. The 8-seed averaging in Appendix D.1 must be surfaced as error bars or standard deviations for at least Tables 1–3 and Fig. 1; without this the "surpasses" claim is not statistically supported. (Addresses ~0.5 pp margin is likely within seed noise.)
• Qualify "monotonic" and "consistent" claims. Replace "Accuracy increases monotonically" with an accurate qualifier (e.g., "on AIME"), and revise the SEM $\lambda$ ablation prose to acknowledge that $\lambda = 0.03$ regresses on GSM-8K and MATH-500. (Addresses Monotonicity claim and non-uniform $\lambda$ behaviour.)
• Replace or explain the low-entropy qualitative example. Either show an example where the low-entropy error is genuinely uncorrected and propagates, or soften the "later steps are built upon an incorrect intermediate result" narrative to match the self-correction shown. (Addresses Cascade story vs its own example.)
• Clarify the KV-cache format story end-to-end. State whether the throughput numbers were obtained with E1M2 or E2M1 KV, and, if E1M2 is not natively supported by Blackwell tensor cores, explain how the accuracy and throughput systems match. (Addresses KV cache format inconsistency.)
• Rewrite the throughput headline to name the KV precision per platform. Make it explicit that the DGX Spark number is W4A4KV8 and the B200 number is W4A4KV4, and that they are not two measurements of the same system. (Addresses Cross-platform throughput claim.)
• Release the TRT-LLM K-shift plugin (or a strong caveat). Ship the plugin or clearly state that the throughput results depend on an unreleased inference kernel, and that the shipped Python KShiftCallback is a training/simulation path only. (Addresses Throughput evidence not reproducible.)
• Ship the Section 3 analysis code. Add the dual BF16/FP4 mixed-precision decoder (Fig. 2(c–d), Fig. 8), the logit-noise perturbation harness with entropy-quantile selection (Fig. 2(e), Fig. 9), and the tail-mass / top-1 mismatch bin script (Fig. 2(f)). These support the paper's core diagnosis and should not require re-implementation. (Addresses Section 3 analyses have no shipped code.)
• Ship the mechanism-analysis code and intermediate checkpoints. Add the per-token embedding-gradient hook and $C_{\mathrm{low}}$ logger for Fig. 7, the per-bin entropy-change script for Fig. 4, and set save_strategy so the 35M / 70M / 140M / 280M snapshots are recoverable. (Addresses Mechanism analysis unsupported.)
• Display the equations for $s_t$ and $C_{\mathrm{low}}$. The current text ends in a colon with no formulas, breaking the mechanism story. (Addresses Missing equations.)

Suggested

• Broaden model coverage. Evaluate on at least one non-R1-distilled LRM (e.g., Qwen3 reasoning or gpt-oss) to test whether the low-entropy failure mechanism generalizes beyond R1 distillation. (Addresses Narrow model coverage.)
• Stratify the response-length analysis by matched problems. Report per-problem outcome across methods rather than per-bin averages so Table 6 controls for the fact that different methods generate different length distributions. (Addresses Sample-count confound.)
• Decompose the compute-vs-bandwidth contribution to speedup at 2K tokens. Add a weight-only quantization row on the same hardware; if 4-bit weights alone reproduce the 1.8–1.9× number, then the "primarily compute" attribution needs revision. (Addresses Compute vs bandwidth attribution.)
• Reconcile the "4–5% Q-FIT overhead" figure with the observed 0.8–2.6% end-to-end drop. State the measurement scope (per-layer vs end-to-end) explicitly. (Addresses Q-FIT overhead figure.)
• Rework the MXFP4 rotation description. Give the block-Hadamard rotation its own name, state the block size, target tensors, and interaction with the MXFP4 group size in the main text — do not fold it under "a special case within Q-FIT." (Addresses Underspecified block rotation and its adjacent conceptual conflation.)
• Fix Table 16 baseline framing. Either drop the Q-FIT-off row or explicitly annotate which row is the SEM baseline the ablation compares against. (Addresses Ablation baseline ambiguity.)
• Show the full ordering in Table 17. Report Q-FIT alongside RHT and no transformation, and note that at 140M no transformation wins. (Addresses Q-FIT vs RHT framing.)
• Soften "in all tested cases". Replace with "in the cases we examined" so the qualitative appendix and the quantitative Fig. 2(e) do not contradict each other. (Addresses Overreach in qualitative narrative.)
• Note the AIME evaluation-set switch in Section 3.2. The logit-noise study uses AIME-2025 alone; flag this explicitly to distinguish it from AIME-120 numbers elsewhere.
• Ship missing pieces of the reproducibility surface. Expose Q-FIT / SEM ablation flags as CLI options, add an NVFP4 W4A4KV4 + block-Hadamard config for Table 17, script the misaligned Stage-2 dataset for Table 4, wire MMLU into inference.py, add a seed sweep and aggregation script, and ship the 70M-token Llama-8B subset generator + multi-epoch driver. Rename 89k-math-sft to reflect that ~5,562 examples are actually kept. (Addresses the corresponding items in Reproducibility & code.)