dMX: Differentiable Mixed-Precision Assignment for Low-Precision Floating-Point Formats
dMX employs differentiable gradient-based optimization to assign mixed floating-point precisions, balancing accuracy and efficiency in large language models.
Key Findings
Methodology
This paper introduces a gradient-based continuous parameterization of per-layer floating-point bit-widths, avoiding non-differentiable discrete searches. Using a temperature-controlled annealing schedule, the method gradually discretizes these parameters to hardware-supported MXFP formats. A target-aware regularization guides the average bit-width toward user-defined budgets, balancing model accuracy and deployment cost. Experiments on models like Llama, Qwen3, and SmolLM2 demonstrate that dMX outperforms KL divergence heuristics, achieving Pareto-efficient trade-offs between model quality and bit-width, with stable training dynamics and hardware compatibility.
Key Results
- On Llama 3.2B, Qwen3 1.7B, and SmolLM2 1.7B, dMX reduces perplexity by 5-10 points and improves zero-shot accuracy by 2-3% at target bit-widths around 4.5-6.0, outperforming fixed and greedy methods.
- The approach achieves more stable convergence and smoother Pareto frontiers compared to STE-based rounding, with better alignment to hardware formats.
- Ablation studies confirm the importance of temperature annealing and target regularization, showing their roles in ensuring effective discretization and performance stability.
Significance
This work advances the field of model quantization by enabling end-to-end gradient optimization of floating-point precisions, addressing the longstanding challenge of balancing model accuracy with deployment efficiency. It opens new avenues for hardware-aware, adaptive precision allocation, crucial for deploying large models on resource-constrained devices, and bridges the gap between theoretical optimization and practical hardware constraints.
Technical Contribution
The paper proposes a novel differentiable framework for floating-point bit-width learning, integrating a temperature-controlled annealing mechanism to map continuous parameters to hardware-compatible MXFP formats. It supports multi-format mixed precision, allowing joint optimization of exponent and mantissa bits, and introduces a target-aware regularization to control the average bit-width. This approach surpasses prior integer-only methods by directly modeling floating-point structures and ensuring hardware support, enabling scalable, end-to-end training of mixed-precision models.
Novelty
This is the first work to formulate gradient-based floating-point bit-width optimization within the MXFP standard, leveraging continuous parameterization and temperature annealing to facilitate hardware-compatible discretization. Unlike previous methods limited to integer bit-widths, this approach jointly optimizes exponent and mantissa bits, providing a flexible, end-to-end solution for mixed-precision floating-point quantization, filling a critical gap in the literature.
Limitations
- The method relies on predefined MXFP format sets, limiting flexibility in bit-width choices. Future work should explore broader format support.
- Hyperparameter sensitivity, especially temperature scheduling, requires careful tuning for optimal results.
- Extreme compression near MXFP4 can degrade performance; more sophisticated regularization or adaptive schemes are needed.
Future Work
Future directions include expanding support for more hardware formats, developing adaptive temperature schedules, and integrating hardware-aware feedback for dynamic bit-width adjustment. Extending the framework to multi-task and multi-modal models, as well as exploring automated hyperparameter tuning, will further enhance its applicability and robustness.
AI Executive Summary
Large language models (LLMs) have revolutionized natural language processing, but their deployment faces significant challenges due to enormous memory and computational demands. Quantization, which reduces parameter precision, offers a promising solution, yet uniformly low-bit quantization often leads to unacceptable accuracy loss. To address this, the paper introduces dMX, a differentiable framework for mixed-precision floating-point quantization tailored for MXFP formats supported by modern hardware.
The core innovation lies in representing per-layer floating-point format choices as continuous, learnable parameters. During training, these parameters are optimized via gradient descent, with a temperature-controlled annealing schedule gradually pushing them toward discrete, hardware-compatible formats. This process ensures smooth training dynamics and avoids abrupt performance drops. A target-aware regularization term further guides the average bit-width to meet user-specified constraints, balancing model accuracy against deployment efficiency.
Extensive experiments on models like Llama 3.2B, Qwen3, and SmolLM2 demonstrate that dMX consistently outperforms traditional layer sensitivity heuristics and fixed-precision baselines. Results show significant improvements in perplexity and zero-shot accuracy, with the ability to finely tune the trade-off between model quality and bit-width. The approach's stability and hardware compatibility make it a practical tool for real-world deployment.
This work marks a significant step forward in model compression, enabling end-to-end gradient-based optimization of floating-point precisions. It opens new avenues for hardware-aware, adaptive quantization strategies, crucial for deploying large models efficiently across diverse platforms. Future research will focus on expanding format support, automating hyperparameter tuning, and integrating hardware feedback for dynamic precision adjustment, further bridging the gap between theoretical advances and industrial applications.
Deep Dive
Plain Language Accessible to non-experts
想象你在厨房做饭，每道菜（模型的每一层）需要不同的火候（比特宽度）才能做得最好。以前的方法就像用同样的火力烹饪所有菜，结果有的菜太熟，有的还不熟。现在，有个聪明的厨师助手（dMX），它能根据每道菜的特点，自动调节火力大小（比特宽度），确保每道菜都刚刚好。这个助手在学习过程中会试不同的火力，然后逐渐学会用最合适的火力，最后用硬件支持的火力范围（MXFP格式）完成烹饪。这样，不仅菜做得好，吃得也健康（模型性能高），还节省了很多燃料（存储和计算资源）。这个方法让厨房变得更智能、更高效，菜肴也更美味。
ELI14 Explained like you're 14
想象你在玩一个超级复杂的游戏，里面的角色（模型）非常厉害，但需要很多存储空间和电脑力气。为了让游戏跑得更快，你可以把角色的细节（参数）变得简单一些，但如果太简单，角色就会变得不真实。以前的方法就像用一把尺子把所有角色都缩成一样的大小，效果不好。现在，这个新方法像是有个聪明的机器人助手，它能根据每个角色的不同特点，自动调节细节的程度（比特宽度），让角色既保持真实，又能在普通电脑上顺畅运行。这个机器人助手会在训练阶段学会怎么调节，最后用最适合硬件的细节等级，让游戏既漂亮又不卡顿。是不是很酷？
Abstract
Quantizing large language models (LLMs) to low-precision floating-point representations is central to efficient deployment, yet applying a single bit-width uniformly across all layers is sub-optimal in terms of both performance and accuracy. This work introduces dMX, a differentiable mixed-precision quantization framework for learnable floating-point bit-width assignment. We study its application for the microscaling floating-point (MXFP) family of data types defined by the Open Compute Project (OCP) standard. The per-layer bit-width assignment is formulated as a continuous optimization problem in which each layer's floating-point format format is parameterized by a scalar parameter, folding the multi-variate design space into a single learnable offset. During training this offset takes continuous values, avoiding sudden oscillations between discrete quantization formats. A temperature-based annealing schedule progressively discretizes the learned offsets, ensuring that the final configuration maps to hardware-compatible MXFP formats without abrupt transitions between training and inference behavior. A target-aware regularization term steers the average bit-width toward a user-specified budget, serving as a coarse-grained proxy for inference cost and balancing model quality against deployment efficiency. We performed experiments on different families of LLM, such as Llama, Qwen3, and SmolLM2, evaluating perplexity on WikiText-2 and accuracy on four zero-shot reasoning benchmarks. Across these settings, dMX consistently yields Pareto-dominating models and improves over Kullback-Leibler (KL) divergence-based layer-selection heuristics, efficiently navigating trade-offs between model quality and average bit-width.