跳转至

LoRaQ: Optimized Low Rank Approximation for 4-bit Quantization

会议: NeurIPS 2026
arXiv: 2604.18117
代码: https://github.com/ROCm/tensorcast
领域: 模型压缩
关键词: 扩散模型量化, 4-bit量化, 低秩逼近, 无数据校准, 旋转优化

一句话总结

针对扩散 Transformer 在 4-bit 激进量化下生成质量严重退化的问题,本文提出无需校准数据的 LoRaQ,直接将量化误差塑造成低秩结构,并首次实现低秩分支与主干全流程 sub-16 bit 混合精度量化,在同等甚至更优计算与显存开销下超越 SVDQuant。

研究背景与动机

随着 DiT(如 PixArt-Σ、FLUX、SANA)成为高分辨率文生图的核心基石,其巨大的参数量与密集的迭代去噪步骤使得端侧及资源受限设备上的低延迟部署面临严峻挑战。后训练量化(PTQ)作为压低显存占用与加速推理的关键技术,在推进至 4-bit(W4A4/W4A8)极端低比特时,极易因去噪轨迹中的微小扰动而产生误差累积,导致生成图像结构坍缩或细节失真。先前前沿方案(如 SVDQuant、SmoothQuant)虽然通过平滑激活离群点或引入全精度低秩旁路(W16A16)显著减小了重构误差,但它们始终依赖高精度辅助分支,这不仅在推理时引入了沉重的 16-bit 激活与权重搬运开销,也违背了现代专用 AI 加速硬件全面迈向 sub-16 bit 混合精度计算的设计趋势。

与此同时,现存主流量化算法在构建低秩重构分支或迁移离群特征时,极度依赖计算繁琐且耗时的数据驱动校准流程。校准阶段往往需要反复采样去噪轨迹,并在每个优化迭代中进行多轮计算密集的 SVD 分解,在百余轮迭代下造成极高的显卡计算与显存峰值开销,极大地限制了千亿/百亿参数扩散模型在边缘设备或本地工作站上的敏捷部署。

面对高精度分支带来的显存/带宽惩罚与重度校准依赖,本文从量化误差空间的内在代数结构切入:既然低秩分支的目标是补偿量化误差,那么无需将低秩分支死锁在 16-bit 高精度上,而是可以直接将权重量化误差显式约束并吸收进一个可被量化的低秩矩阵中。核心 idea:利用量化算子的幂等性质,通过轻量级无数据梯度下降将权重微调与误差吸收统一,并结合无推理开销的正交旋转矩阵消除量化截断损失,首次构建全通道 sub-16 bit 混合精度的低秩量化流水线。

方法详解

整体框架

LoRaQ 的核心目标是将线性层权重矩阵 \(\bm{W}\) 的量化误差约束至低秩子空间,从而利用一个低秩分支来完美吸收残差。在推理时,线性变换由两条并行的低精度分支构成:一条是承载主体权重的残差分支,采用 4-bit(如 SINT4 或 MXINT4/MXFP4)量化;另一条是承载误差补偿的低秩分支,由分解矩阵 \(\bm{L}\) 和 \(\bm{R}\) 组成,通过引入正交旋转矩阵 \(\bm{\Omega}\) 并统一量化到 sub-16 bit(如 4-bit、6-bit 或 8-bit)。整个推理流程不产生任何 16-bit 权重或激活传输。

整体 pipeline 包含四个阶段:首先利用激活平滑技术将输入离群点难度迁移至权重;其次仅使用一次 SVD 分解初始化低秩矩阵,随后在无真实校准数据的设定下通过均方误差目标优化低秩扰动,将残差量化误差塑造成低秩流形;接着求解正交旋转矩阵 \(\bm{\Omega}\) 最小化低秩因子自身的量化离散化损失;最后将旋转矩阵离线融合进低秩矩阵中,形成硬件友好的全低比特推理内核。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入特征 X 与权重 W"] --> B["阶段1:通道平滑离群点迁移<br/>降低输入激活尖峰难度"]
    B --> C["基于幂等算子的无数据低秩误差吸收<br/>单次SVD初始化后梯度优化 L 与 R"]
    C --> D["旋转感知的低秩分支全量化<br/>正交旋转 Omega 最小化分解量化误差"]
    D --> E["面向新一代硬件的 OCP Microscaling (MX) 统一精度适配<br/>离线吸收 Omega 并统一量化到 sub-16 bit"]
    E --> F["残差分支 Q1(P) 与低秩分支 Q2(L)Q2(R) 并行推理"]

关键设计

1. 基于幂等算子的无数据低秩误差吸收:克服不可微量化瓶颈直接优化残差空间

以往基于低秩逼近的量化工作(如 SVDQuant)直接对原始权重做高精度 SVD 分解,忽略了量化算子本身在参数空间引入的离散阶跃几何特征。本文将目标转化为:寻找一个低秩扰动矩阵 \(\bm{D} = \bm{L}\bm{R}\)(其中 \(\bm{L}\in\mathbb{R}^{d\times\rho}\),\(\bm{R}\in\mathbb{R}^{\rho\times n}\) 且 \(\rho \ll \min(d,n)\)),使得扰动后的量化结果与原权重的差异严格收敛至该低秩形态,即 \(Q_1(\bm{W} + \bm{D}) - \bm{W} \approx \bm{L}\bm{R}\)。由于量化取整算子的导数几乎处处为零,直接对扰动求导会陷入梯度消失。本文利用量化算子的幂等性(即对量化网格内的点有 \(Q(Q(\bm{A})) = Q(\bm{A})\)),若设量化目标为 \(\bm{W} + \bm{L}\bm{R}\),则可将非凸优化问题重构为显式逼近目标:

\[ \min_{\bm{L},\bm{R}} \left\| Q_1(\bm{W} + \bm{L}\bm{R}) - \bm{W} - \bm{L}\bm{R} \right\|_F^2 \]

由于在计算对 \(\bm{L}\) 和 \(\bm{R}\) 的反向传播时,视外层被量化矩阵为常数目标,从而能够精准推导出关于 \(\bm{L}\) 与 \(\bm{R}\) 的解析梯度更新公式:\(\nabla_{\bm{L}}\mathcal{L} = -2 (Q_1(\bm{W}+\bm{L}\bm{R}) - \bm{W} - \bm{L}\bm{R}) \bm{R}^T\) 以及 \(\nabla_{\bm{R}}\mathcal{L} = -2 \bm{L}^T (Q_1(\bm{W}+\bm{L}\bm{R}) - \bm{W} - \bm{L}\bm{R})\)。该优化全程无需任何前向推理数据和输入激活,仅用 Adam 迭代 1000 步即可快速收敛。相比 SVDQuant 每轮循环都需调用高耗时的 SVD,LoRaQ 仅在初始化阶段执行一次 SVD(\(\bm{L} = -\bm{L}_0, \bm{R} = \bm{R}_0\)),将 SVD 计算开销降低了整整两个数量级。

2. 旋转感知的低秩分支全量化:消除精度瓶颈并释放高秩表征红利

以往方法坚信低秩补偿分支必须维持 FP16/BF16 精度,导致必须专门编写非对称内核且带宽瓶颈突出。LoRaQ 提出直接对低秩分支实施量化(算子 \(Q_2\)),使其同样运行在 sub-16 bit 精度(如 4-bit、6-bit 或 8-bit)。在固定显存预算 \(\beta\)(每通道分配比特数,如 SVDQuant 的 \(16\text{ bits} \times 32\text{ rank} = 512\text{ bits/channel}\))约束下,将低秩分支的数值位宽 \(n\) 降至 8-bit、6-bit 甚至 4-bit,即可成倍换取更高的有效秩 \(\rho' = \lfloor \beta / n \rfloor\)(例如由 rank 32 提升至 64、86 或 128)。对于参数量大、隐层维度宽的扩散模型(如 SANA-1.6B),更高的秩能够覆盖更丰富的特征子空间,其带来的表征增益远超数值精度的微弱损失。

为进一步抵消低秩矩阵量化算子 \(Q_2\) 带来的截断误差,本文在已求得的最优低秩矩阵 \(\bm{L}^*\) 和 \(\bm{R}^*\) 之间引入一个正交旋转矩阵 \(\bm{\Omega} \in \mathbb{R}^{\rho'\times\rho'}\)。由于正交旋转保持内积不变(\(\bm{\Omega}\bm{\Omega}^T = \bm{I}\)),理论上有 \((\bm{L}^*\bm{\Omega})(\bm{\Omega}^T\bm{R}^*) = \bm{L}^*\bm{R}^*\)。通过定义正交损失:

\[ \mathcal{L}_{\bm{\Omega}} = \left\| Q_2(\bm{L}^*\bm{\Omega}) - \bm{L}^*\bm{\Omega} \right\|_F^2 + \left\| Q_2(\bm{\Omega}^T\bm{R}^*) - \bm{\Omega}^T\bm{R}^* \right\|_F^2 \]

采用 Cayley 变换或梯度下降优化 500 步得到最优旋转 \(\bm{\Omega}^*\)。在部署时,\(\bm{\Omega}^*\) 直接通过离线矩阵乘法吸收进 \(\bm{L}' = \bm{L}^*\bm{\Omega}^*\) 与 \(\bm{R}' = (\bm{\Omega}^*)^T\bm{R}^*\) 中,在推理期完全不产生额外的浮点计算、显存访问或算子调用开销。

3. 面向新一代硬件的 OCP Microscaling (MX) 统一精度适配

为了与现代支持微缩放向量量化的加速芯片(如 AMD CDNA4/MI350 系列)原生协同,LoRaQ 全面适配了开放计算项目定义的 OCP Microscaling(MX)数据格式。相比传统按张量或按块粗粒度缩放的量化方案,MX 格式采用 32 个数值一组的微块(Micro-block),每块配备一个 8-bit 的以 2 为底的指数缩放因子(e8m0 规范),数值乘法可被硬件移位器高效执行。LoRaQ 实现了对 MXFP8 (e4m3)、MXFP6 (e2m3)、MXFP4 (e2m1) 以及 MXINT4/MXINT8 的原生混合精度支持,使得主干残差分支与辅助低秩分支能够同时以 sub-16 bit 激活和权重进行纯低比特矩阵乘累加(MAC),达成真正的端到端轻量化推理。

损失函数 / 训练策略

算法分为两步解耦优化: 1. 低秩吸收优化:以 MSE 损失 \(\mathcal{L}_{\text{abs}} = \frac{1}{dn}\| Q_1(\bm{W} + \bm{L}\bm{R}) - \bm{W} - \bm{L}\bm{R} \|_F^2\) 为目标,固定原始权重 \(\bm{W}\),采用 Adam 优化器更新 \(\bm{L}\) 与 \(\bm{R}\),浮点格式学习率设为 \(10^{-3}\),整型格式设为 \(10^{-4}\),优化步数固定为 1000 步。 2. 正交旋转优化:固定 \(\bm{L}^*\) 与 \(\bm{R}^*\),以旋转误差 \(\mathcal{L}_{\bm{\Omega}}\) 为目标更新正交矩阵 \(\bm{\Omega}\),初始学习率设为 \(0.1\) 至 \(0.5\),迭代 500 步。由于每层仅需几百次微型矩阵运算,单卡数分钟即可完成整个扩散 Transformer 的全模型无数据校准。

实验关键数据

主实验

在 PixArt-Σ (0.6B) 与 SANA-1.6B 两个主流扩散 Transformer 上,作者对比了在 MJHQ-30K 和 sDCI 数据集上的生成质量(FID、Image Reward, IR)与基线保真度(LPIPS、PSNR)。主实验对比基于 SINT4 与 MXINT4 格式展开(生成步数均为 20 步)。

模型与数据集 精度格式 (W-A) 方法 FID (↓) IR (↑) LPIPS (↓) PSNR (↑)
PixArt-Σ (MJHQ) FP16 (16-16) Baseline 16.6 0.944 — —
SINT4 (4-4) VIDIT-Q 412.0 -2.270 0.854 6.44
SINT4 (4-4) SVDQuant 19.2 0.878 0.323 17.6
SINT4 (4-4) LoRaQ (本文) 16.9 0.898 0.309 17.6
MXINT4 (4-4) SVDQuant 18.9 0.738 0.424 15.9
MXINT4 (4-4) LoRaQ (本文) 15.4 0.901 0.339 16.8
PixArt-Σ (sDCI) FP16 (16-16) Baseline 24.8 0.966 — —
SINT4 (4-4) SVDQuant 25.9 0.918 0.352 16.5
SINT4 (4-4) LoRaQ (本文) 24.1 0.919 0.346 16.2
MXINT4 (4-4) SVDQuant 26.1 0.902 0.435 14.7
MXINT4 (4-4) LoRaQ (本文) 23.7 0.943 0.374 15.6
SANA-1.6B (MJHQ) BF16 (16-16) Baseline 16.2 1.100 — —
SINT4 (4-4) RTN 20.5 0.894 0.339 15.3
SINT4 (4-4) SVDQuant 19.3 0.935 0.220 17.8
SINT4 (4-4) LoRaQ (本文) 16.1 1.090 0.182 19.2
SANA-1.6B (sDCI) BF16 (16-16) Baseline 22.4 1.070 — —
SINT4 (4-4) SVDQuant 28.1 0.846 0.242 16.2
SINT4 (4-4) LoRaQ (本文) 21.8 1.060 0.208 17.4

消融实验

消融实验重点验证了:① 低秩参数优化(Optimized LR)与 ② 正交旋转矩阵(Rotations)对消除量化误差的必要性。评估基于 PixArt-Σ 在 MJHQ 数据集上进行(残差分支固定为 MXFP4e2,激活统一采用 8-bit)。

低秩格式 秩 (Rank) 优化低秩矩阵 优化旋转矩阵 KID (\(10^{-3}\), ↓) LPIPS (↓) PSNR (↑)
MXFP8e4 64 ✓ ✓ 1.42 0.316 17.1
64 ✓ ✗ 1.55 0.321 17.0
64 ✗ ✓ 1.83 0.356 16.1
64 ✗ ✗ 1.91 0.362 16.2
MXFP6e2 96 ✓ ✓ 1.29 0.305 17.4
96 ✓ ✗ 1.34 0.308 17.2
96 ✗ ✓ 1.85 0.355 16.3
96 ✗ ✗ 1.91 0.358 16.3
MXFP4e2 128 ✓ ✓ 1.34 0.318 17.1
128 ✓ ✗ 1.41 0.326 16.9
128 ✗ ✓ 1.57 0.342 16.6
128 ✗ ✗ 1.61 0.338 16.7

在固定内存预算下(表 3 探索),将低秩分支从传统的 FP16 (rank=32) 替换为 MXFP8 (rank=64)、MXFP6 (rank=86) 或 MXFP4 (rank=128),图像感知质量普遍反超全精度分支(例如在 MJHQ 上 FID 从 FP16 上的 16.8 提升至 MXFP4e2-rank128 上的 16.1,Image Reward 从 0.934 提升至 0.963)。

关键发现

  • 低秩优化是性能质变的关键:消融数据显示,如果仅使用 SVD 截断而不进行低秩参数梯度迭代(✗ Optimized LR),即使加上旋转,KID 也会剧烈恶化(从 1.42 恶化至 1.83),证明离散量化空间中基于幂等假设的直接优化起决定性作用。
  • 高秩低比特权衡显著优于低秩高精度:在相同的显存位宽预算下,选择“低比特数值 + 高秩维度”(如 4-bit 秩 128)所获得的表征补偿能力,系统性优于 SVDQuant 所采用的“16-bit 数值 + 低秩(秩 32)”。
  • 参数规模越大,高秩收益越突出:在参数量达 1.6B 的 SANA 模型上,LoRaQ 在 MJHQ 上的 FID 达到 16.1,甚至微幅超越了 BF16 原模型的 16.2,而 SVDQuant 的 FID 停留在 19.3,充分证明传统低秩方法对大尺度参数矩阵的欠拟合。

亮点与洞察

  • 将不可导量化算子利用幂等性转化为闭式梯度更新:巧妙利用 \(Q(Q(\cdot)) = Q(\cdot)\) 将残差补偿重构为对低秩因子的 MSE 优化问题,绕过了传统 STE(直通估计器)带来的严重梯度偏差。
  • 将正交旋转与因式分解乘法零开销离线融合:在低秩因子间插入正交旋转矩阵平滑奇异值分布,不仅显著压制了低秩矩阵自身的量化误差,而且在部署阶段直接静态内积融合,做到了真正的零推理延时与零参数膨胀。
  • 打破高精度分支必须依赖的工程惯性:首次提出并论证了低秩补偿分支本身完全可以被激进量化至 4/6/8-bit,为全 sub-16 bit 硬件(如 AMD CDNA4/MI350)铺平了纯整型/微浮点混合加速的技术路径。

局限与展望

  • 端到端加速依赖底层定制融合算子:尽管算法在理论吞吐与访存带宽上大幅削减了数据搬运,但要实现端到端的推理时延断崖式下降,仍高度依赖针对 NPU/GPU 定制开发的高性能双路低比特融合算子。
  • 激活量化仍需前置平滑支撑:虽然权重量化实现了完全的无数据自适应,但在 W4A4 极限设置下,依然需要沿用 SmoothQuant 类依赖激活采样统计的前置通道缩放来抑制激活离群峰值。
  • 未来方向:探索动态自适应分配不同网络层与不同去噪时间步的低秩秩数(Adaptive Rank Allocation),并将无数据优化思想扩展至注意力激活张量的压缩中。

相关工作与启发

  • vs SVDQuant: SVDQuant 强制保留 FP16 辅助分支并依赖大量前向数据进行耗时的多轮迭代 SVD;LoRaQ 仅需单次 SVD 初始化,通过无数据梯度下降吸收误差,且允许低秩分支自身完全量化,在同等显存开销下支持高出 2-4 倍的表征秩。
  • vs SpinQuant & QuIP#: SpinQuant 等工作主要在大模型激活或全量权重上利用旋转矩阵抑制离群点;LoRaQ 则是将旋转矩阵精准引入低秩分解的双因子之间,专门用于消除小规模低秩矩阵量化时的方差扩散,且旋转矩阵完全离线吸收。
  • vs SmoothQuant: SmoothQuant 仅做输入通道与权重维度的数学等效缩放;LoRaQ 将其作为前置互补组件,重点解决缩放后 4-bit 权重量化误差的吸收难题。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出低秩分支自身的量化方案并给出基于幂等性的无数据梯度推导。
  • 实验充分度: ⭐⭐⭐⭐⭐ 在主流 DiT 架构(PixArt-Σ、SANA)上跨多种数据集与 sub-16 bit 精度展开详尽测试与消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推导严密清晰,对量化上下界与正交旋转几何特性的数学阐述通透。
  • 价值: ⭐⭐⭐⭐⭐ 彻底击破了低秩补偿必须保留全精度的思维定势,极具软硬件协同设计落地价值。