SegQuant

Most post-training quantization methods see a linear layer as one numerical matrix. SegQuant also asks where its inputs come from and where its outputs go. It derives semantic segments from the computation graph, quantizes those segments independently, and preserves polarity-asymmetric activations through a hardware-native dual-scale execution path.

SegLinearDiscovers semantic partitions through graph topology, without hand-written layer lists.

DualScaleAssigns separate resolution to negative and nonnegative activations.

Deployment-awareRetains fixed-width integer GEMM, batching, and fused epilogues.

1. Why ordinary tensor quantization is insufficient

Diffusion models repeatedly apply the same denoising backbone, so a small per-layer error can propagate across many timesteps. DiT-style architectures add another difficulty: operations such as chunk, split, concat, and reshape reveal that one physical linear layer can carry several computational meanings.

For example, an AdaNorm linear layer emits shift, scale, and gate parameters for attention and MLP branches. Their numerical ranges and sensitivity differ. Using one quantization configuration for the entire fused matrix creates segment interference: a large-range segment determines the quantizer and wastes resolution needed by a smaller, more sensitive segment.

Conventional uniform view \[ Y=XW, \qquad \widehat{Y}=Q(X)Q(W). \] This formula hides the graph semantics of \(X\) and \(W\). SegQuant recovers that structure before choosing quantization parameters.

2. SegLinear: semantics from graph topology

SegLinear uses torch.fx symbolic tracing and topological pattern matching. It looks around each linear operator for two reusable patterns, rather than relying on model names or manually curated layers.

Output-segmented quantization

If a linear output is later divided by chunk or split, the graph says that different output regions feed different branches. SegLinear partitions the output dimension and quantizes every submatrix independently:

W=[W_1,W_2,\ldots,W_N], \qquad \widehat{Y}=[\widehat{X}\widehat{W}_1,\widehat{X}\widehat{W}_2,\ldots,\widehat{X}\widehat{W}_N].

Input-segmented quantization

If the input is produced by concat or a structure-changing reshape, the input channels originate from distinct paths. SegLinear partitions both the input and the corresponding rows of the weight matrix:

X=[X_1,X_2,\ldots,X_N], \qquad W=[W_1^{\mathsf T},W_2^{\mathsf T},\ldots,W_N^{\mathsf T}]^{\mathsf T}, \] \[ \widehat{Y}=\sum_{i=1}^{N}\widehat{X}_i\widehat{W}_i.

The segment widths are inferred directly from graph metadata. Each segment may then use SmoothQuant, GPTQ, SVDQuant, or another optimizer with its own calibration. SegLinear therefore separates structural discovery from numerical optimization.

SegQuant's graph-to-kernel workflow

Trace graphBuild a static operator graph with tensor shapes and dependencies.

Detect semanticsMatch concat-to-linear and linear-to-split patterns.

Segment tensorsPartition the relevant input channels or weight regions.

Optimize locallyCalibrate and quantize each coherent segment independently.

Lower to kernelsExecute fixed-width low-bit GEMM with fused reconstruction.

The figure is rendered in HTML and CSS. It summarizes the implementation pipeline without relying on a rasterized paper screenshot.

3. DualScale: preserve polarity without changing bit layout

SiLU and GELU retain many small negative values while allowing a much wider positive range. A single uniform step size is dominated by positive outliers and compresses the negative interval into too few bins. Yet those negative activations carry high-frequency texture and detail in diffusion generation.

DualScale assigns one step size to the negative region and another to the nonnegative region:

Polarity-aware quantization \[ Q_{\mathrm{dual}}(x)= \begin{cases} \operatorname{round}(x/s_-), & x<0,\\ \operatorname{round}(x/s_+), & x\geq0, \end{cases} \] \[ s_-=\frac{|\min(X)|}{q_{\min}}, \qquad s_+=\frac{\max(X)}{q_{\max}}. \]

Writing \(X_+=\max(X,0)\) and \(X_-=\min(X,0)\), the linear result is reconstructed as

Y\approx (s_+s_w)\widehat{X}_+\widehat{W}+(s_-s_w)\widehat{X}_-\widehat{W}.

Although the equation has two products, the implementation submits them as one BatchedGEMM and combines the scaled outputs in a fused epilogue. It therefore uses standard fixed-width integer arithmetic and avoids custom number formats or additional kernel launches.

4. Why the two mechanisms complement each other

SegLinear fixes structural interference

It prevents one semantic branch from determining quantization parameters for another branch.

DualScale fixes polarity interference

It prevents the wide positive range from erasing resolution in a narrow but meaningful negative range.

Optimizer-agnostic

Both mechanisms wrap existing PTQ optimizers instead of replacing them with a model-specific method.

Compiler-compatible

Static graph detection and ordinary low-bit GEMM fit current deployment toolchains.

5. Experimental evidence

The study covers SD3.5, FLUX.1-dev, and SDXL; DiT and UNet backbones; INT8 and INT4 configurations; and COCO, MJHQ-30K, and DCI image sets. The ablation below isolates the two proposed components on SD3.5 with W8A8 quantization.

Ablation on MJHQ-30K, SD3.5 W8A8
Method	FID ↓	Image Reward ↑	LPIPS ↓	PSNR ↑	SSIM ↑
Baseline	23.35	0.877	0.419	11.93	0.536
+ SegLinear	23.36	0.899	0.395	12.03	0.554
+ DualScale	22.61	0.909	0.401	12.14	0.551
+ SegLinear + DualScale	22.54	0.952	0.377	12.50	0.567

Layer-level measurements explain the gain. With GPTQ, the Frobenius quantization error of DiT.11.norm1 context falls from \(3.0176\) to \(1.7637\) after semantic segmentation. The improvement appears across SmoothQuant, SVDQuant, and GPTQ, indicating that the graph partition, rather than one optimizer, is the source of the effect.

Deployment cost on SD3.5, RTX 4090
Configuration	Peak memory (GiB)	Backbone size (GiB)	Per-step latency (ms)
FP16	18.63	3.88	317.18
INT8 baseline	16.73	2.00	278.86
SegQuant-8	16.72	2.02	287.85
INT4 baseline	16.06	1.29	209.78
SegQuant-4	16.12	1.42	212.87

22.54 FIDFull method versus 23.35 for the W8A8 ablation baseline.

0.952 IRImage Reward rises from 0.877 when both components are enabled.

1.5% W4 overheadSegQuant-4 adds 3.09 ms per step over the INT4 baseline while preserving its main efficiency gain.

How to read the results. SegQuant does not claim that semantic handling is free. It adds a small runtime and scale-storage cost. The systems contribution is that this cost stays close to ordinary INT8 or INT4 execution while the quality metrics improve consistently.

6. Broader research direction

SegQuant treats model structure as deployment information. Shiva-DiT continues the broader agenda from a different angle by learning which transformer computations should be retained through residual-based differentiable top-k selection. Both works seek structure-aware efficiency instead of applying the same approximation to every parameter, activation, or token.