VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

📝 Paper Summary

Model Compression Diffusion Transformers (DiTs) Post-Training Quantization (PTQ)

VQ4DiT enables extremely low bit-width compression of Diffusion Transformers by simultaneously calibrating both the codebook and vector assignments using a zero-data, block-wise strategy, avoiding gradient inconsistencies.

Core Problem

Traditional Vector Quantization (VQ) methods only calibrate the codebook without adjusting assignments, causing weight sub-vectors to be incorrectly mapped and providing inconsistent gradients that degrade performance at low bit-widths.

Why it matters:

Diffusion Transformers (DiTs) like DiT-XL/2 are computationally expensive (17+ seconds/image on A100) and memory-intensive, hindering edge deployment.
Standard uniform quantization fails at extremely low bit-widths (e.g., 2-bit), causing severe quality degradation.
Existing VQ methods for CNNs rely on time-consuming fine-tuning on full datasets, which is impractical for large generative models.

Concrete Example: When quantizing a DiT model to 2-bit, standard methods result in poor image quality (FID > 48) because sub-vectors with the same assignment may have gradients pointing in different directions, preventing the codebook from converging to an optimal state.

Key Novelty

Simultaneous Codebook and Assignment Calibration

Instead of fixing assignments after K-Means clustering, the method treats assignments as learnable probability distributions (soft assignments) over a candidate set.
A zero-data calibration strategy uses synthetic Gaussian noise inputs to generate outputs, minimizing the error between floating-point and quantized block outputs to tune both codebooks and assignment ratios simultaneously.

Architecture

Overview of the VQ4DiT framework, contrasting it with the standard DiT block structure.

Evaluation Highlights

Achieves 3.32 FID on ImageNet 256x256 with 2-bit weights, effectively matching the full-precision model (2.27 FID) while using significantly less memory.
Outperforms standard uniform quantization (UQ) significantly; at 2-bit, UQ fails completely (FID 260.07) while VQ4DiT maintains high generation quality.
Quantizes a DiT XL/2 model on a single NVIDIA A100 GPU within 20 minutes to 5 hours depending on settings, without requiring real training data.

Breakthrough Assessment

8/10

First successful application of Vector Quantization to Diffusion Transformers at extremely low bit-widths (2-bit) with a novel simultaneous calibration mechanism, addressing a key limitation of prior VQ methods.

⚙️ Technical Details

Problem Definition

Setting: Post-training vector quantization of pre-trained Diffusion Transformer weights without access to the original training dataset.

Inputs: Pre-trained floating-point DiT model weights W.

Outputs: Quantized weights represented by a codebook C and assignments A.

Pipeline Flow

Weight Decomposition (K-Means Clustering)
Candidate Assignment Set Construction
Zero-data Input Generation
Block-wise Calibration (Codebook + Assignments)

System Modules

K-Means Initializer

Decompose weights into initial codebooks and assignments using standard K-Means clustering.

Model or implementation: K-Means algorithm

Candidate Set Constructor

Identify top-n closest codewords for each sub-vector to form a candidate assignment set.

Model or implementation: Euclidean distance calculation

Block-wise Calibrator

Optimize codebook values and assignment ratios to minimize output error between quantized and float blocks.

Model or implementation: Gradient descent optimization

Novel Architectural Elements

Learnable assignment ratios: Treating VQ assignments as a weighted average during calibration rather than fixed indices.
Simultaneous gradient updates for both codebook values and assignment selection probabilities.

Modeling

Base Model: DiT-XL/2 (Diffusion Transformer)

Training Method: Zero-data block-wise calibration via gradient descent

Objective Functions:

Purpose: Minimize reconstruction error of block outputs.

Formally: L_MSE = || d_fp(z) - d_q(z) ||^2
Purpose: Encourage assignment ratios to converge to a one-hot vector (entropy regularization).

Formally: L_R = 1 - max(R)

Adaptation: Post-training quantization (PTQ)

Key Hyperparameters:

codebook_size_k: 256 (2-bit), 64 (3-bit)
sub_vector_dimension_d: 4 (2-bit), 2 (3-bit)
candidate_set_size_n: Not explicitly reported in the paper
+ 1 more
calibration_threshold_lambda: 10^-4

Compute: Single NVIDIA A100 GPU; 20 minutes to 5 hours depending on settings.

Comparison to Prior Work

vs. UQ/GPTQ/AWQ: VQ4DiT uses Vector Quantization instead of scalar/uniform quantization, allowing better reconstruction at low bit-widths.
vs. Standard VQ (K-Means): VQ4DiT calibrates assignments simultaneously with the codebook, whereas standard VQ fixes assignments after clustering.
vs. Q-DiT: VQ4DiT focuses on vector quantization and achieves lower bit-widths (2-bit) efficiently without evolutionary search overhead.

Limitations

Calibration time can extend up to 5 hours for high-precision settings.
Trade-off required between codebook size/dimension and quantization error/memory usage.
Specifics of candidate set size 'n' are not detailed in the main text.

Reproducibility

Code: https://github.com/PikachuDeng/VQ4DiT

Code is publicly available at https://github.com/PikachuDeng/VQ4DiT. The method relies on generated Gaussian noise, removing the need for external datasets. Hyperparameters for calibration (lambda_d, lambda_r) are set to 1.

📊 Experiments & Results

Evaluation Setup

Image generation on ImageNet 256x256 using DiT-XL/2.

Benchmarks:

ImageNet 256x256 (Class-conditional Image Generation)

Metrics:

FID (Fréchet Inception Distance)
sFID (Spatial FID)
IS (Inception Score)
Statistical methodology: Not explicitly reported in the paper

Key Results

Benchmark	Metric	Baseline	This Paper	Δ
Comparison with Full Precision (FP32) and Uniform Quantization (UQ) baselines at 2-bit settings shows VQ4DiT maintains performance where UQ fails.
ImageNet 256x256	FID	2.27	3.32	+1.05
ImageNet 256x256	FID	260.07	3.32	-256.75
ImageNet 256x256	IS	278.24	264.45	-13.79
ImageNet 256x256	FID	48.24	3.32	-44.92
ImageNet 256x256	FID	48.24	3.32	-44.92

Main Takeaways

VQ4DiT achieves state-of-the-art trade-offs between model size and performance for Diffusion Transformers.
Standard vector quantization (K-Means) is insufficient for DiTs at low bit-widths due to inconsistent gradients from fixed assignments.
Simultaneous calibration of codebooks and assignments is critical for recovering model performance after quantization.
The method works effectively without real training data, relying only on synthetic noise for calibration.

📚 Prerequisite Knowledge

Prerequisites

Vector Quantization (VQ) concepts (codebooks, codewords, assignments)
Diffusion Transformers (DiTs) architecture
K-Means clustering
Post-Training Quantization (PTQ)

Key Terms

Vector Quantization (VQ): A compression technique where vectors of weights are mapped to the nearest entry in a fixed-size codebook, replacing the vector with an index.

Codebook: A table of representative vectors (codewords) used to reconstruct the original weights.

Assignment: The index pointing to a specific codeword in the codebook for a given weight sub-vector.

DiT: Diffusion Transformer—a class of diffusion models that replaces the U-Net backbone with a Transformer architecture.

Zero-data Calibration: A calibration method that uses synthetic data (e.g., Gaussian noise) instead of real images to tune quantization parameters.

FID: Fréchet Inception Distance—a metric for evaluating the quality of generated images by comparing the distribution of generated images to real images; lower is better.

sFID: Spatial Fréchet Inception Distance—a variant of FID that captures spatial relationships better.

IS: Inception Score—a metric measuring the diversity and clarity of generated images; higher is better.