FraQAT: Quantization Aware Training with Fractional bits

TL;DR

FraQAT addresses the challenge of deploying high-quality diffusion-model generative systems on-device by introducing Fractional Quantization-Aware-Training, a curriculum-based QAT that progressively lowers parameter precision from $32$ to $4$ bits using intermediate fractional bit widths. By keeping activations at $INT8$, leveraging a FP32 teacher for distillation, and applying a fractional-bit schedule (e.g., $8\rightarrow7\rightarrow6\rightarrow5.5\rightarrow\cdots\rightarrow4$), FraQAT achieves higher fidelity than traditional QAT and PTQ methods while reducing training time. Across multiple diffusion-model families (SD3.5-Medium, Sana, PixArt, Flux), FraQAT yields $4$–$7\%$ lower FiD scores than the strongest QAT baselines and enables practical on-device deployment, demonstrated by Sana running on a Samsung S25U with $66$ ms forward-time per step and $30.5\%$ latency improvement over a comparable baseline. The work also shows targeted outlier analysis and selective layer training can reduce computational cost without compromising quality, highlighting FraQAT’s potential for mobile privacy-preserving generation and personalized content on edge devices.

Abstract

State-of-the-art (SOTA) generative models have demonstrated impressive capabilities in image synthesis or text generation, often with a large capacity model. However, these large models cannot be deployed on smartphones due to the limited availability of on-board memory and computations. Quantization methods lower the precision of the model parameters, allowing for efficient computations, \eg, in \INT{8}. Although aggressive quantization addresses efficiency and memory constraints, preserving the quality of the model remains a challenge. To retain quality in previous aggressive quantization, we propose a new fractional bits quantization (\short) approach. The novelty is a simple yet effective idea: we progressively reduce the model's precision from 32 to 4 bits per parameter, and exploit the fractional bits during optimization to maintain high generation quality. We show that the \short{} yields improved quality on a variety of diffusion models, including SD3.5-Medium, Sana, \pixart, and FLUX.1-schnell, while achieving $4-7\%$ lower FiD than standard QAT. Finally, we deploy and run Sana on a Samsung S25U, which runs on the Qualcomm SM8750-AB Snapdragon 8 Elite Hexagon Tensor Processor (HTP).

Figures (7)

• Figure 1: FraQAT is a Quantization aware Training (QAT) technique that grants generative models high fidelity at a fraction of training time required. Large text-to-image (T2I) models quantized with FraQAT (W4A8) achieve $16\%$ lower FiD score than the state-of-the-art.
• Figure 2: Bit vs Loss: as we reduce the precision (number of bits), the average knowledge distillation loss increases -- the gap between student and teacher widens. From a quantization perspective, this implies outliers incrementally affect the student model.
• Figure 3: Classic QAT first computes the loss computed at the lower precision ($\circ$), then propagates it back to the original precision and optimize the weights ($\rightarrow$). This results in coarse and noisy gradients. Fractional Quantization Aware Training rely on intermediate precisions (from INT8 to INT4 as training progresses) to incrementally adjust to weights distributions. Consequently, during training Parameters smoothly shift between bins ($\bullet$) thanks to the finer gradients from in intermediate precision.
• Figure 4: Qualitative comparison: FraQAT (d) generates images similar to the original model (a). Prompts are from MJHQ dataset li2024playground.
• Figure 5: Outliers: outliers distribution for activations varies across models. SD3.5-M (left) experience most of its outliers right after Feed Forward layers, while for PixArt-$\Sigma$, most outliers are in Attention layers.