Outlier-Aware Post-Training Quantization for Discrete Graph Diffusion Models

TL;DR

Bit-DGDM is a novel post-training quantization framework specifically designed for Discrete Graph Diffusion Models (DGDMs). By combining sparse-dense activation quantization with an ill-conditioned low-rank weight decomposition, it achieves up to 2.5× inference speedup and 2.8× memory reduction while maintaining 4-bit generation quality comparable to full-precision models.

---

What Problem Does This Paper Address?

Discrete Graph Diffusion Models (DGDMs) have become a state-of-the-art approach for generating complex structures like molecules and proteins because they naturally preserve the sparsity and categorical nature of graph data. However, they are computationally expensive.

Unlike Large Language Models (LLMs), which are often "memory-bound" (limited by how fast you can load massive weights), DGDMs are "computation-bound." They have relatively small parameter counts but require many recursive denoising steps involving high-computation operations.

Expand

Source, p. 1

Existing quantization methods fall short here for two main reasons:

1. Incompatible Assumptions: Quantization for Image Diffusion Models (IDMs) assumes Gaussian noise processes, which don't exist in the discrete world of DGDMs.
2. The Outlier Problem: DGDMs exhibit significant "outliers"—values with extreme magnitudes—in both weights and activations. Standard quantization methods either clip these values (losing critical information) or stretch the quantization range (losing precision for the vast majority of normal values).

Methods

The authors introduce Bit-DGDM, which tackles these dual-sided outliers through a "divide and conquer" strategy:

Expand

Source, p. 4

• Sparse-Dense Activation Quantization: The framework identifies activation outliers and keeps them in a sparse, full-precision matrix ($S_X$), while the remaining "dense" values ($D_X$) are quantized to low-bit (e.g., 4-bit). Crucially, they use a data-free adaptive threshold selection, meaning they can calibrate the quantization process using synthetic random graphs rather than requiring the original training data.
• Ill-conditioned Low-Rank Weight Decomposition: Inspired by Robust PCA, the authors decompose the weight matrices into a low-rank component ($LR^T$) and an $\alpha$-sparsity matrix ($S_W$) that captures weight outliers. This allows for faster matrix multiplication while ensuring that the "ill-conditioned" parts of the matrix (the outliers) don't ruin the quantization of the rest.
• Custom GPU Kernels: To make this math efficient in practice, the authors developed two specialized kernels:
• Equidistributed Kernel: For weights, where outliers are relatively balanced across rows.
• Adaptive Kernel: For activations, where outliers are irregular and require dynamic workload balancing across GPU warps.

Datasets

The authors evaluated Bit-DGDM across a diverse set of benchmarks:

• Molecules: QM9 and MOSES (standard for molecular generation).
• General Graphs: Stochastic Block Model (SBM) and Planar graphs.
• Proteins: The CATH dataset for the "inverse protein folding" task (predicting amino acid sequences from 3D structures).

Key Insights and Ideas

• Computation is the Bottleneck: The paper highlights that for DGDMs, reducing FLOPs is more important than just reducing weight-loading latency. This is why LLM-specific methods like GPTQ don't provide significant speedups for graph models.
• Dual-Sided Outliers: In many domains, outliers appear primarily in activations. In DGDMs, they are prominent in weights too. Bit-DGDM is one of the first to explicitly handle this "dual vulnerability."
• Data-Free Calibration: The discovery that thresholds can be set using uniform random distributions of node/edge types (synthetic graphs) is a major win for privacy and usability, as it removes the need for access to the original training set during quantization.
• The Power of Low-Rank + Sparse: Simply using SVD (Singular Value Decomposition) to compress weights often fails because outliers heavily bias the singular vectors. By treating outliers as a separate sparse component, the low-rank approximation becomes much more stable.

Key Results

• Efficiency: Bit-DGDM achieved a 2.5× speedup and a 2.8× reduction in peak memory usage compared to the FP32 baseline.
• Accuracy: At 4-bit precision (W4A4), Bit-DGDM maintained near-original performance. For example, on the QM9 dataset, it achieved 98.2% validity (vs. 99.0% for FP32).
• Superiority over Baselines: It significantly outperformed other SOTA quantization methods. For instance, on the CATH protein task, Bit-DGDM achieved a recovery rate of 51.6%, while standard Round-To-Nearest (RTN) dropped to 32.8%.

Limitations and Future Work

• Hyperparameter Sensitivity: The method relies on a few key hyperparameters, such as the rank ($r$) of the weight decomposition and the sparsity ratio ($\alpha$). While the authors provide a sensitivity analysis, an automated way to select these per-layer would be a valuable extension.
• Overlapping Outliers: The authors note that when outliers in weights and activations overlap perfectly, performance can still degrade slightly, though their $S_X S_W$ term in the expanded multiplication helps mitigate this.
• Hardware Specificity: The custom kernels are currently optimized for NVIDIA GPUs; porting these to other architectures (like TPUs or mobile NPUs) remains a future challenge.