Technical Summary: TASQ - Temporal-Adaptive Bit Sparsification Quantization for Diffusion Models

1. Problem Statement

Diffusion models generate high-quality images through an iterative denoising process, repeatedly inferring the same network to transform noisy inputs into clean samples. While quantization (compressing weights and activations to 4 bits or lower) reduces memory and arithmetic costs, existing methods suffer from a fundamental inefficiency: temporal redundancy.

Current static quantization approaches assign a single weight precision to every denoising step. To preserve image quality, this precision must accommodate the most quantization-sensitive step across the entire trajectory, even though many other steps can tolerate significantly fewer bits. Consequently, the model repeatedly pays a "worst-case" arithmetic cost throughout the denoising process, wasting computation on less sensitive stages. Furthermore, recent attempts at timestep-aware quantization have largely focused on activations or required storing separate weight copies for different stages, leading to memory overhead or failing to exploit temporal variations in weight sensitivity.

2. Methodology: TASQ

The authors propose Temporal-Adaptive Bit Sparsification Quantization (TASQ), a framework that decouples storage precision from operating precision. TASQ allows the model to assign different weight precisions to specific layers at different diffusion steps without storing multiple weight copies.

Core Components

• Shared Weight Representation & LSB Masking: TASQ stores a single shared high-precision weight buffer (e.g., 8-bit). Instead of learning separate weights for each stage, it learns a Temporal-Spatial LSB Mask. This mask determines how many least-significant bits (LSBs) to truncate for each layer and denoising stage.

  • Mathematically, for a stored $n$-bit weight $Q_n$, the operating weight at stage $t$ and layer $l$ is derived by truncating bits based on the mask $M_{t,l}$:
    $$Q^{(t,l)}_{M,n} = Q_n - (1 - M_{t,l}) \cdot \text{LSB}$$
  • This allows the model to dynamically switch between effective precisions (e.g., 6-bit, 4-bit, 3-bit) via simple bit-plane truncation or right-shifting, avoiding the need for separate weight copies or runtime search.

• Training Strategy:

  • LoRA-based QAT: TASQ integrates with Low-Rank Adaptation (LoRA) to bridge the gap between Quantization-Aware Training (QAT) efficiency and performance. The quantization operator is applied to the effective weight formed by merging frozen original weights with trainable low-rank adapters.
  • Farthest-Stage-First Training: Jointly updating shared weights across neighboring timesteps can lead to highly correlated gradients and training plateaus. TASQ addresses this by training stages in a "farthest-stage-first" order (e.g., $0 \to 3 \to 1 \to 2$ for four stages). This exposes the shared Most Significant Bits (MSBs) to diverse parts of the trajectory, improving generalization while maintaining separate LSB masks for each stage.
  • LSB Mask Annealing: Since binary selection is non-differentiable, TASQ uses a continuous gate during training that is annealed to a binary decision, following a sparsification schedule.

• Temporal-Precision Engine (Hardware Co-design): To realize the computational benefits, the authors co-design a Temporal-Precision Engine. In bit-serial General Matrix Multiplication (GEMM), execution cycles scale with the product of weight and activation bit widths.

  • The engine keeps activation bit planes on-chip and streams only the weight planes selected for the current layer and stage.
  • This reduces the number of fetched planes from $2LR$ (standard) to $L+R$, reducing execution cycles by 3.2–5.1$\times$ over standard loops.
  • Crucially, switching precision incurs no measured cycle overhead, as the engine simply reads fewer bit planes from the shared buffer.

3. Key Contributions

1. Decoupling Storage and Operating Precision: TASQ enables weight precision to vary across layers and denoising stages to reduce BitOPs, while storage remains fixed at the worst-case precision required by the most sensitive stage.
2. Temporal-Spatial LSB Mask: A novel mechanism that learns stage- and layer-wise precision over a single shared weight representation, eliminating the need for per-stage weight copies.
3. Temporal-Precision Engine: A hardware-aware execution strategy that streams only selected weight planes, achieving 6.1–7.5$\times$ fewer cycles than standard static 8-bit bit-serial execution while preserving image quality metrics (FID, ImageReward).
4. Training Innovations: The introduction of Farthest-Stage-First Training to stabilize optimization and LSB Mask Annealing to handle non-differentiable bit selection.

4. Experimental Results

The authors evaluated TASQ on PixArt-Σ (0.6B), SANA-1.6B, and SDXL-Turbo (2.6B) using datasets like MJHQ-30K and sDCI.

• Quality vs. Computation: TASQ achieves image generation quality comparable to static quantization (and often superior to other low-bit methods) with significantly reduced computation.
  • On a W4A4 budget, TASQ improves FID by 1.0–3.1 points over static quantization.
  • TASQ retains the quality of static 8-bit quantization while reducing execution cycles by approximately 25–50%.
• Efficiency Gains: Compared to a standard static 8-bit bit-serial execution, TASQ reduces execution cycles by 6.1–7.5$\times$.
• Ablation Studies:
  • Temporal vs. Spatial: Joint temporal-spatial adaptation (TASQ) significantly outperforms spatial-only or temporal-only adaptation, confirming that precision requirements vary both across layers and over time.
  • Initialization: TASQ improves performance even without SVDQuant initialization, proving the gains stem from temporal adaptivity rather than just the starting point.
  • Mask Design: The proposed Continuous LSB Mask (CLM) outperforms standard sigmoid gates by avoiding early attenuation of LSB signals.

5. Significance and Claims

The paper claims that TASQ represents a shift in how diffusion models are quantized, moving from static, worst-case allocation to dynamic, trajectory-aware precision.

• Quality-Compute Trade-off: TASQ demonstrates that temporal weight allocation improves the quality-compute trade-off without requiring multiple model checkpoints or complex runtime search mechanisms.
• Hardware Agnosticism: The method is designed to be hardware-agnostic regarding storage (a single buffer) but relies on precision-scalable hardware (like bit-serial accelerators) to realize the cycle savings. The authors note that on fixed-precision hardware (e.g., standard INT4/INT8 kernels), the cycle savings may not be fully realized if the hardware cannot execute below its native precision.
• Practical Deployment: The approach fixes the memory footprint (determined by the worst-case stage) while allowing inference speed to be optimized based on the specific sensitivity of each denoising step.

The authors conclude that by separating the cost of storage from the cost of computation, TASQ effectively leverages the inherent temporal redundancy in diffusion models to achieve high-quality generation with significantly lower arithmetic costs.