MoBiQuant: Mixture-of-Bits Quantization for Token-Adaptive Any-Precision LLM¶
会议: NeurIPS 2026
arXiv: 2602.20191
领域: 模型压缩
关键词: 任意精度量化, 混合位宽量化, 大语言模型后训练量化, 异常值迁移, 动态标记路由
一句话总结¶
针对传统后训练量化跨位宽泛化失效与异常值迁移难题,MoBiQuant 提出基于递归残差切片的 MoBiSlice 结构与动态标记级路由器 MoBiRoute,在统一校准参数下实现免重新打包、零内核重启的高吞吐动态任意精度推理。
研究背景与动机¶
大语言模型在边缘设备与动态计算负载场景中部署时,面临着高度波动的算力与内存限制。在系统资源充足时,模型需要以较高精度(例如 4-bit 或 6-bit)推理以保证严苛任务的生成质量;而在并发加剧或内存受限时,则必须即时降低计算精度以满足端到端延迟与吞吐要求。然而,传统的后训练量化(PTQ)方法均假设静态位宽,通常引入一组校准参数专门最小化固定精度下的重构误差。当强行将某一位宽下校准的参数迁移至未见位宽时,模型困惑度剧烈恶化,迫使系统针对不同位宽分别存储多套完整权重,造成高达数倍的存储冗余。
深入分析这一泛化退化现象的本质,可以发现其根源在于精度相关的「异常值迁移」(Outlier Migration)机制。尽管激活值异常值主导量化误差已被广泛证实,但造成高量化误差的特定敏感标记集合在不同位宽间并非静态不变。在较高位宽(如 4-bit)下拟合良好的平滑标记,降低至 3-bit 时往往会跃升为主要误差源;而 4-bit 时的主要异常标记在 3-bit 下反而不再占据主导地位。这一敏感标记分布在不同精度下的剧烈偏移,导致单精度校准参数不可避免地过拟合到局部的异常分布上,从而丧失跨位宽泛化能力。
现有任意精度量化方案在此场景下均面临显著阻碍:基于向量量化(VQ)的方案计算访存极不友好,依赖码本查表导致吞吐受限;而基于标量位切片或位编码的方法往往需要针对每个位宽训练额外的缩放因子,或在切换精度时依赖离线重新打包和重部署特定内核。核心 idea:通过递归残差量化将全精度权重解构为共享低比特基底与多级残差切片,并在校准阶段引入轻量可学习路由器依据标记敏感度自适应分配残差激活状态,以标记自适应机制消除异常值迁移并在推理期通过统一缩放因子与位平面矩阵乘法实现任意精度无缝弹性切换。
方法详解¶
整体框架¶
MoBiQuant 的核心目标是将大语言模型各线性层转换为支持按标记弹性调节有效位宽的动态计算单元,其整体工作流分为权重解构表示(MoBiSlice)与动态路由机制(MoBiRoute)两个核心组件。输入序列的各个标记首先输入 MoBiRoute 轻量路由层,由可学习 MLP 计算标记对各残差切片的敏感度置信度得分;随后通过阈值化生成离散激活掩码,决定该标记在计算中需要累加多少级残差切片;最后由定制的高效位平面矩阵乘法内核,以统一的缩放因子直接在打包位平面上完成加权输出累加。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入标记序列 X"] --> B["MoBiRoute 路由网络<br/>轻量两层 MLP 评估敏感度得分 S"]
B --> C["门控与阈值判定<br/>可调全局阈值 δ 确定激活掩码 G"]
C --> D["MoBiSlice 递归切片选择<br/>共享基底 W1 与各阶残差切片 We"]
D --> E["定制 BMMA 计算内核<br/>位平面高效矩阵乘与移位累加"]
E --> F["输出隐藏状态 Y"]
关键设计¶
1. MoBiSlice 递归残差切片解构:支持单一低位宽内核的任意精度层次重构
传统任意精度量化在位宽切换时往往需要重新打包权重或维护多组缩放因子,不仅带来显著的存储与训练开销,还需频繁重启硬件内核。MoBiSlice 将全精度权重矩阵 \(\mathbf{W}\) 显式分解为 \(E\) 个固定低比特(默认 4 个 2-bit 切片)的级联残差切片 \(\{\mathbf{W}_1, \mathbf{W}_2, \dots, \mathbf{W}_E\}\)。第 1 阶最高有效位(MSB)切片 \(\mathbf{W}_1\) 负责量化全精度矩阵,其余切片递推地对上阶残差 \(\mathbf{R}_e = \mathbf{R}_{e-1} - \mathbf{W}_{e-1}\) 执行向下取整对齐量化:
$\(x_{\mathrm{int}}=\mathrm{clamp}\left(\left\lfloor\frac{x}{s}+z\right\rfloor,\,0,\,2^{b}-1\right),\quad x_{\mathrm{deq}}=s(x_{\mathrm{int}}-z+0.5)\)$
为了消除维护多组缩放因子的开销并保持无偏重构,各级残差切片的缩放因子直接由共享缩放因子严格按 \(s_{e+1} = s_e / 2^{b_e}\) 递减,偏置固定为 \(z_e = 2^{b_e-1}\)。数学证明表明,激活或截断残差切片引入的截断误差期望严格为零且上界小于粗糙量化器半个步长,从而保证任意子集累加均属于保序细化,无需重新打包权重,所有有效精度均复用固定 2-bit GEMM 内核执行计算。
2. MoBiRoute 标记自适应路由门控:动态抑制异常值迁移与平滑退火训练
针对单精度校准容易陷入固定异常值分布的缺陷,MoBiRoute 引入轻量级双层感知机 \(\mathcal{R}(\mathbf{X}, \mathbf{\Theta}_r)\) 生成标记维度的切片需求得分 \(\mathbf{S} \in \mathbb{R}^{T \times E}\)。为了在训练阶段平滑跨越离散选择屏障,设计了带有自适应温度调度的可微 Sigmoid 门控机制:
$\(\mathcal{G}(\mathbf{S}) = \sigma(\tau(t) \cdot \mathbf{S}),\quad \tau(t) = \frac{\ln(L)}{\ln(L) - \ln(t)}\)$
随着训练步数 \(t\) 推进到最终步 \(L\),温度 \(\tau \to \infty\),门控输出无缝退化为离散二值掩码 \(\mathbb{I}(\mathbf{S} > 0)\)。更关键的是,为引导路由器探索全局位宽组合而非过早固化在单一位宽,引入了预算正则化动态衰减目标:
$\(\mathcal{L}_{\mathrm{reg}}(t) = (\mathrm{AvgBits} - b(t)) \cdot \|\mathcal{G}(\mathbf{S})\|_1,\quad b(t) = b_{\mathrm{init}} - (b_{\mathrm{init}} - b)\frac{\ln(t)}{\ln(L)}\)$
该目标自适应惩罚超出当前阶段预算的切片激活,强制路由器在多位宽混合空间充分探索标记对位宽的真实敏感度,从根本上消除了静态校准参数对单一精度异常分布的过拟合。
3. 运行时阈值动态调节与定制 BMMA 高吞吐内核:低开销零重载的弹性推理
在离线校准完毕后,所有切片的门控逻辑与权重均已固定。推理阶段若遇系统负载或延迟预算变化,系统无需重载模型权重或重新执行校准,只需在全局层级调整标量阈值 \(\delta\),即依据 \(\mathcal{G}_\delta(\mathbf{S}) = \mathbb{I}(\mathbf{S} - \delta > 0)\) 调节各标记的切片分配。在底层算子实现层面,不同于以往方法需要复杂的码本查找(VQ)或跨位宽非对齐解包,MoBiQuant 采用位平面主导(bit-major)打包格式,直接在打包位平面上调用硬件级二进制矩阵乘法(BMMA)。针对标记动态路由带来的非连续内存访问与负载不均衡问题,内核实现了路由后标记重排(Token Permutation),将指派至相同切片的标记连续存放以跑满内存带宽,并通过独立 CUDA 流将基础切片与其他切片重叠计算,使路由与打包开销在 4-bit 典型解码延迟中占比压缩至 18% 以下。
损失函数 / 训练策略¶
MoBiQuant 采用逐层重构校准策略,优化参数仅包括量化步长/剪裁范围参数 \(\mathbf{\Theta}_q\) 与轻量路由器参数 \(\mathbf{\Theta}_r\),基模型权重全量冻结。每层分为两阶段: - 阶段 1:首切片稳定化。强制所有标记仅经过第 1 阶基底切片 \(\mathbf{W}_1\) 计算输出 \(\mathbf{Y}_{\mathrm{FS}}\),最小化其与全精度参考输出 \(\mathbf{Y}_{\mathrm{fp}}\) 的平方重构损失 \(\mathcal{L}_{\mathrm{FS}} = \|\mathbf{Y}_{\mathrm{FS}} - \mathbf{Y}_{\mathrm{fp}}\|_2^2\),快速奠定粗粒度主干表征。 - 阶段 2:联合多切片优化。基于阶段 1 参数派生残差切片,联合优化全流重构误差与路由正则损失: $\(\mathcal{L}_{\mathrm{MoBi}} = \|\mathbf{Y}_{\mathrm{q}} - \mathbf{Y}_{\mathrm{fp}}\|_2^2 + \lambda (\mathrm{AvgBits} - b(t)) \|\mathcal{G}(\mathbf{S})\|_1\)$ 校准默认在 128 条 WikiText2 序列上进行 20 个 epoch,批大小为 1,目标位宽 \(b\) 默认设为平衡探索空间的 3.0-bit。在涉及激活值量化与等价变换(LET)场景时,路由器输入通过逆缩放与逆位移重构在原始激活空间,规避输入分布漂移导致的门控饱和。
实验关键数据¶
主实验¶
在主流开源大模型上评估 WikiText2 语言建模困惑度(PPL,越低越好)以及端到端解码吞吐(tokens/sec,越高越好)。在相同平均有效位宽下,MoBiQuant 显著超越了现有任意精度 SOTA 方案,并在高吞吐方面大幅压制静态 VQ 方案。
| 模型 | 位宽 | 方法类型 | 方法 | WikiText2 PPL (↓) | 解码吞吐 tokens/s (↑) |
|---|---|---|---|---|---|
| LLaMA2-7B | 2-bit | Any-precision | AnyPrecisionLLM | 2000.00 | 279 |
| LLaMA2-7B | 2-bit | Any-precision | AnyBCQ | 15.38 | 312 |
| LLaMA2-7B | 2-bit | Any-precision | MoBiQuant | 10.91 | 404 |
| LLaMA2-7B | 3-bit | Any-precision | AnyBCQ | 6.25 | 268 |
| LLaMA2-7B | 3-bit | Any-precision | MoBiQuant | 6.07 | 311 |
| LLaMA2-7B | 4-bit | Any-precision | AnyBCQ | 5.91 | 225 |
| LLaMA2-7B | 4-bit | Any-precision | MoBiQuant | 5.82 | 256 |
| LLaMA3-8B | 2-bit | Any-precision | AnyPrecisionLLM | 1000.00 | 228 |
| LLaMA3-8B | 2-bit | Any-precision | AnyBCQ | 19.01 | 247 |
| LLaMA3-8B | 2-bit | Any-precision | MoBiQuant | 58.12 | 352 |
| LLaMA3-8B | 3-bit | Any-precision | AnyBCQ | 8.08 | 216 |
| LLaMA3-8B | 3-bit | Any-precision | MoBiQuant | 7.97 | 257 |
| LLaMA3-8B | 4-bit | Any-precision | AnyBCQ | 6.84 | 182 |
| LLaMA3-8B | 4-bit | Any-precision | MoBiQuant | 7.31 | 209 |
与固定单精度静态标量量化(Static SQ)基准的对比表明,具有完全弹性推理能力的 MoBiQuant 完全可匹敌甚至超越静态独立调优模型(LLaMA3.2-3B 4-bit PPL: OmniQuant 12.71 vs. MoBiQuant 9.41;LLaMA3-8B 3-bit PPL: OmniQuant 9.11 vs. MoBiQuant 7.97)。
消融实验¶
针对校准数据分布敏感性(以 LLaMA3.2-1B 在 3-bit 下为例)及关键模块设计的消融验证:
| 校准数据集 | 评估集/任务 | OmniQuant 3-bit | MoBiQuant 3-bit | 表现说明 |
|---|---|---|---|---|
| WikiText2 | WikiText2 PPL (↓) | 21.60 | 18.95 | 域内语言建模困惑度降低 2.65 点 |
| WikiText2 | C4 PPL (↓) | 32.86 | 27.92 | 跨域泛化至通用网页文本显著提升 |
| WikiText2 | PTB PPL (↓) | 37.47 | 34.92 | 跨域新闻文本鲁棒性更优 |
| WikiText2 | BoolQ (Acc ↑) | 0.593 | 0.613 | 常识推理准确率提升 2.0% |
| PTB | WikiText2 PPL (↓) | 22.69 | 16.91 | 跨数据集鲁棒性优势大幅扩大(降 5.78 点) |
| C4 | WikiText2 PPL (↓) | 22.61 | 26.36 | 仅在 C4 校准特定组合略逊,总体稳健 |
针对路由正则化衰减调度的消融(比较线性、余弦、指数与对数衰减):在 2.5~3.0 有效比特区间,对数调度(Logarithmic)与指数调度显著优于线性和余弦调度,且对数调度与温度退火天然协同,促成了最平滑的弹性跨位宽收敛。
关键发现¶
- 消除异常值迁移的直接证据:当静态 OmniQuant 在 4-bit 校准而在 3-bit 推理时,标记误差剧烈分化;而 MoBiQuant 路由器的置信度打分与标记的跨位宽误差增量呈现出极强的正相关,敏感标记被精确赋予 4~6 阶切片,非敏感标记回落至 2-bit,使误差分布在各精度下保持稳定一致。
- 块级与层级异构分配行为:虽然全局损失函数针对单一平均位宽优化,但路由器自发涌现出跨层差异化偏好——靠近输入的前几层以及 MLP 模块获得了更高的有效位宽预算,这与经验性混合精度设计(如逐层重要性分配)不谋而合。
- 吞吐与显存双重突破:相比部署多套独立位宽权重的方案,MoBiQuant 压缩了 3.5× 部署显存;相比 AnyPrecisionLLM 与 AnyBCQ,定制 BMMA 内核消除了码本查表与多套反量化参数开销,实现最高 1.34× 的端到端解码吞吐提升。
亮点与洞察¶
- 将混合专家(MoE)思想下放至位宽分配:将模型层面的混合精度或任意精度重构问题,精巧地转化为单 token 视角的位切片动态路由(Mixture-of-Bits),以极小的可微门控参数量解决了量化领域长期存在的跨位宽泛化难题。
- 数学严谨的免校准残差下采样:基于向下取整与中点对齐机制,使各阶残差切片的缩放因子直接以 \(2^{-b_e}\) 级联派生,在数学上证明了残差累加的无偏性与误差界,为硬件高效的位平面移位累加扫清了理论障碍。
- 训练连续、推理纯离散的完美解耦:训练期借助温度退火与对数预算调度引导离散掩码收敛,推理期退化为单参数 \(\delta\) 阈值切分,使云端/端侧系统仅需下发一个标量浮点参数即可完成全网瞬时弹性调度。
局限与展望¶
- 超大规模模型评估仍需验证:论文当前主实验集中在 1B 至 13B 参数量级的模型上,尚未在 70B 及以上超大规模 LLM 上开展系统性评测,其在长上下文超大模型上的训练耗时与门控行为仍需进一步论证。
- 批处理动态特性的工程挑战:当前定制的高吞吐内核主要针对单批次或小批次延迟敏感型边缘推理进行了细致优化;在具有高并发、动态批处理与连续批处理(Continuous Batching)的生产级服务框架(如 vLLM、SGLang)中,不同序列的异构切片激活可能引发 GPU 线程束分化(Warp Divergence),仍需探索更高效的批重排算子。
相关工作与启发¶
- vs AnyPrecisionLLM: 该工作利用低比特种子模型结合向量量化逐步重建高精度,但在硬件部署中重度依赖硬件不友好的码本查找和位转置操作,推理吞吐受阻;MoBiQuant 采用纯标量递归切片和统一位平面 BMMA,在 2-bit 下将困惑度由 2000 降至 10.91 的同时吞吐提升 33.8%。
- vs AnyBCQ: AnyBCQ 为每个支持的精度训练独立的一套缩放因子,造成显存占用和推理内核分支复杂度上升;MoBiQuant 采用共享缩放因子结合标记自适应路由,推理期仅需调节标量阈值 \(\delta\) 即可切换任意精度,无需额外缩放因子存储与重新校准。
- vs OmniQuant / SpinQuant 等静态 PTQ: 传统方法在选定位宽下表现卓越但跨精度迁移时困惑度崩溃;MoBiQuant 可无缝作为此类静态框架的插件,赋予静态基座在 2~6 位的平滑弹性推理能力,且在同等平均位宽下性能持平甚至更优。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次揭示后训练量化跨精度失效的「异常值迁移」本质,并将标记自适应路由引入标量任意精度量化体系。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 LLaMA2/3/3.2 多种尺寸,对比了静态 PTQ、任意精度与 VQ 基线,且包含深入的困惑度、算子延迟、吞吐与消融实验。
- 写作质量: ⭐⭐⭐⭐⭐ 动机与现象剖析清晰明了,数学推导与误差界分析严谨,架构与实验设计环环相扣。
- 价值: ⭐⭐⭐⭐⭐ 为边缘端大模型应对动态算力约束提供了兼具高精度、低显存与高吞吐的工业级落地范式。