OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers

Core Methodology§

OrbitQuant addresses the challenge of post-training quantization (PTQ) for diffusion transformers (DiTs), where activation statistics vary across timesteps, prompts, and guidance branches. Instead of data-dependent calibration, it applies a randomized permuted block-Hadamard (RPBH) rotation before quantization. This rotation makes each activation coordinate follow a fixed, known marginal distribution (approximately Gaussian) regardless of input, enabling a single Lloyd-Max optimal scalar codebook to be used for all layers, timesteps, and prompts.

Mathematical Formulation§

For a weight matrix $W \in \mathbb{R}^{m \times n}$ and activation vector $x \in \mathbb{R}^{n}$, standard linear layer computes $y = Wx$. OrbitQuant applies an orthogonal rotation $R$ (RPBH) to the activations before quantization and absorbs its inverse into weights:

$$ \hat{y} = Q_w(W R^{-1}) \cdot Q_a(R x) $$

where $Q_w$ and $Q_a$ are weight and activation quantizers. The rotation $R$ is designed so that each element of $R x$ has a distribution close to a fixed standard Gaussian with zero mean and unit variance, independent of the original input. Therefore, a single codebook learned via the Lloyd-Max algorithm on a standard Gaussian suffices for all inputs.

Implementation Details§

• RPBH Rotation: A block-Hadamard matrix is randomly permuted and applied in chunks. It is an orthogonal transform that can be computed efficiently via fast Walsh-Hadamard transform.
• Quantization: Weights are quantized offline to low-bit integers (e.g., 2-4 bits) using standard min-max or percentile clipping, after absorbing the inverse rotation. Activations are rotated online then quantized with a fixed, precomputed Lloyd-Max codebook.
• Codebook Design: The Lloyd-Max codebook is optimized for a zero-mean, unit-variance Gaussian, minimizing MSE. Since the rotated activations match this distribution, no calibration data is needed.

Code Snippet§

import torch
import torch.nn.functional as F

def rpbh_rotation(x, block_size=256):
    # Simplified: apply random permutation and block Hadamard
    n = x.shape[-1]
    H = torch.hadamard(n // block_size).float()  # Hadamard matrix
    perm = torch.randperm(n)
    x_perm = x[..., perm]
    x_rot = torch.cat([F.linear(x_perm[..., i*block_size:(i+1)*block_size], H) 
                       for i in range(n // block_size)], dim=-1)
    return x_rot

class OrbitQuantLinear(torch.nn.Linear):
    def __init__(self, in_features, out_features, bit_width=4):
        super().__init__(in_features, out_features, bias=False)
        self.bit_width = bit_width
        # Precompute inverse rotation and absorb into weights
        R_inv = rpbh_rotation(torch.eye(in_features)).inverse()
        self.weight.data = self.weight.data @ R_inv.t()
        # Quantize weights offline
        self.weight = quantize_weight(self.weight, bit_width)
        # Precompute activation codebook for standard normal
        self.codebook = lloyd_max_gaussian(bit_width)

    def forward(self, x):
        # Rotate activations
        x_rot = rpbh_rotation(x)
        # Quantize activations with fixed codebook
        x_q = quantize_activation(x_rot, self.codebook)
        return F.linear(x_q, self.weight)

Results & Significance§

OrbitQuant sets new state-of-the-art PTQ results on FLUX.1, Z-Image-Turbo, Wan 2.1, and CogVideoX at various bit-widths (W8A8 down to W2A4). It is data-agnostic and transfers directly from image to video models without retuning. The key insight is that a fixed rotation stabilizes activation distributions, eliminating the need for calibration data and making PTQ truly plug-and-play.