【ICML 2023】FlexRound: Learnable Rounding based on Element-wise Division for Post-Training Quantization

FlexRound 解决的是后训练量化（PTQ）里一个很具体的问题：权重已经被训练好，能不能不用微调，只靠一个轻量的重构优化，把权重量化到 4 bit 还不掉点？

它的答案是把舍入方案从”逐元素加法”换成”逐元素除法”。常见做法（AdaRound、AdaQuant、BRECQ）都是在权重上加一个可学习的偏移量，再往最近的网格点上取整；FlexRound 反过来，给每个权重配一个可学习的除法因子，让权重先被缩放再取整。这样不仅每层共享的量化网格步长 s1 可以学，每个权重自己的缩放也能学，而且两者的梯度天然和权重本身的数值挂钩。

本文基于 arXiv:2306.00317（ICML 2023 正式版，PMLR v202）。作者是 NAVER Cloud 的 Jung Hyun Lee、Jeonghoon Kim、Se Jung Kwon、Dongsoo Lee，代码在 onliwad101/FlexRound_LRQ。文中数值均标注为论文报告结果。

先给结论

FlexRound 的核心增量是三件事，按重要性排序：

用除法而不是加法做舍入，把网格步长 s1 和逐权重缩放一起放进同一个可微目标里联合学习。
由除法的导数规则得到一个性质：权重的模越大，它对应的缩放被更新的力度越大（Proposition 3.1），等价于让重要权重有更大自由度去落到更远的量化格点上。
消融显示：如果冻结 s1，FlexRound 就退化成接近 AdaRound；收益主要来自”联合学习步长”这一条，而不是额外加的通道因子。

在 MobileNetV2 这种权重模值大、对量化敏感的网络上，4-bit 权重量化的差距最明显：论文报告 FlexRound 70.82% top-1，AdaRound 69.46%，AdaQuant 直接崩到 15.17%（BRECQ 实验设置，仅量化权重）。

ImageNet 仅权重量化到 4-bit 时，ResNet-18/50 与 MobileNetV2 的 top-1 准确率对比

上图是论文 Table 1 / Table 2 的直观版本：三个模型上 FlexRound 都略高于 AdaRound，而 MobileNetV2 上 AdaQuant 直接崩到 15.17%。ResNet-18 的差距最小，正好对应后面要讲的”权重模值越大，除法舍入越有空间”。

背景：重构式 PTQ 与”加法派”

BRECQ 等工作证明，量化带来的损失退化可以近似成网络输出的二次型，于是逐层/逐块地重构输出，就成了 PTQ 提升精度的通用路径。目标是最小化重构误差：

\min_{\text{可学参数}} \; \lVert \bm{W}\bm{X} - \widehat{\bm{W}}\widetilde{\bm{X}} \rVert_F^2

其中 $\bm{X}$ 是浮点输入， $\widetilde{\bm{X}}$ 是所有前序层已量化后的输入， $\widehat{\bm{W}}$ 是量化权重。

已有方法对这个目标的设计几乎都基于”加法”：

AdaRound： $\widehat{\bm{W}} = s_1\big(\lfloor \bm{W}/s_1\rfloor + h(\bm{V})\big)$ ，用软松弛的 $h(\bm{V})$ 决定每个权重往上还是往下取整，但网格步长 $s_1$ 固定。
AdaQuant： $\widehat{\bm{W}} = s_1\lfloor (\bm{W}+\bm{V})/s_1\rceil$ ，可以同时学步长和偏移，但在 MobileNetV2 这种网络上低比特会严重退化。
BRECQ / QDrop：换的是重构的块单元或激活处理，舍入机制仍是 AdaRound 那一套。

问题在于：AdaRound 永远只能在最近的两个格点里选一个（round up 或 round down），当位宽稍高、格点更密时，这个约束反而不够灵活。

方法：把舍入写成除法

FlexRound 的基础形式是：

\widehat{\bm{W}} = s_1\left\lfloor \frac{\bm{W}}{\bm{S}} \right\rceil

$s_1$ 是整层共享的网格步长（标量）， $\bm{S}$ 是形状与 $\bm{W}$ 相同、逐元素为正且可学习的除法因子。作者把 $\bm{S}$ 再拆成两部分：

\bm{S} = s_1 \odot \bm{S}_2 \odot \bm{s}_3

$\bm{S}_2$ ：与权重同形状的逐权重缩放；
$\bm{s}_3 \in \mathbb{R}_{>0}^{C_{out}\times 1}$ ：逐输出通道的附加因子（因为不同输出通道的统计差异很大）；
对 2D 卷积再补一个 $\bm{s}_4 \in \mathbb{R}_{>0}^{1\times C_{in}\times 1\times 1}$ ，得到 $\bm{S}=s_1\odot\bm{S}_2\odot\bm{s}_3\odot\bm{s}_4$ 。

所以最终（线性层）：

\widehat{\bm{W}} = s_1\left\lfloor \frac{\bm{W}}{s_1 \odot \bm{S}_2 \odot \bm{s}_3} \right\rceil

flowchart LR
  W["预训练权重 W"] --> D["逐元素除法 W / S"]
  S1["s1：全层公共网格步长"] --> S
  S2["S2：逐权重缩放"] --> S
  S3["s3：逐输出通道因子"] --> S
  S["S = s1 ⊙ S2 ⊙ s3"] --> D
  D --> R["round 取整"] --> M["× s1 还原"] --> WH["W_hat"]

round-to-nearest 与 FlexRound 的舍入结果对比

上图用两个权重说明核心差异：A=0.47、B=0.55 时，round-to-nearest 把它们取到 0.4 和 0.6，顺序不变；FlexRound 因为每个权重有独立的缩放因子，可以让 A 取到 0.6、B 取到 0.4，从而改变取整结果的顺序。论文 Figure 1(b) 展示的是同一个现象。

初始化很关键： $\bm{S}_2,\bm{s}_3,\bm{s}_4$ 全部初始化为 1，于是公式退化成标准的 round-to-nearest $s_1\lfloor \bm{W}/s_1\rceil$ 。也就是说，FlexRound 是从最近舍入出发再去学偏移，而不是从随机初始化开始。

论文主推的是 per-tensor uniform PTQ（一层共用一个 $s_1$ ），因为它对硬件最友好；需要时把标量 $s_1$ 换成向量即可退化成 per-channel。作者强调，在 per-tensor 设定下 FlexRound 已经足够接近全精度，这一点对部署端很重要。

为什么”除法”能起作用：reciprocal 性

这是这篇论文最漂亮的部分。用直通估计（STE，straight-through estimator）近似取整的梯度后，重构损失 $\mathcal{L}$ 对缩放因子 $\bm{S}'=\bm{S}_2\odot\bm{s}_3$ （或再乘 $\bm{s}_4$ ）某一项的梯度满足（论文 Proposition 3.1，附录推导）：

\left| \frac{\partial \mathcal{L}}{\partial S'_{(i,j)}} \right| \;\propto\; \left| W_{(i,j)} \cdot \frac{\partial \mathcal{L}}{\partial \widehat{W}_{(i,j)}} \right|

这里 $W_{(i,j)}$ 是原始浮点权重。含义很直接：

权重的模越大，它对应的缩放获得的更新力度越大。 而权重模值常被当作重要性度量（Han 等关于剪枝的工作），所以大权重被允许落到”不止最近两个格点、而是更远格点”的自由度更高。
如果某个权重的重构梯度 $\partial\mathcal{L}/\partial\widehat{W}_{(i,j)}$ 本来就接近 0，那它的缩放梯度也是 0，会被保留在 round-to-nearest 的格点上。

逐权重缩放的梯度幅度随权重模值增大

上图把 Proposition 3.1 画成一条上升趋势线：横轴是权重模值 $|W|$ ，纵轴是缩放被更新的力度；模值越大，缩放可被改动得越多，权重越可能落到更远的格点。论文 Figure 3 用直方图和散点展示了同一趋势。

作者在 MobileNetV2 上统计：第一层卷积约 12.8% 的权重发生了”跨多于一步”的舍入；ResNet-18 因为没有模值大于 1 的权重，这个比例只有约 1.5%。这也解释了为什么 FlexRound 在 MobileNetV2 上的增益明显大于 ResNet。

一个值得注意的推论：位宽越高，FlexRound 越灵活。 AdaRound 无论多少 bit 都只能在相邻两格之间选；FlexRound 可以跨多格，所以论文认为它比 AdaRound 更适合高比特场景（例如 LLM）。这与直觉相反——通常低比特才更难，但 FlexRound 的差异化优势恰恰在网格更密的地方体现。

消融：收益主要来自”联合学习步长”

论文做了两组消融（ImageNet，仅权重量化到 4-bit，BRECQ 设置），结果很有信息量：

设置	ResNet-18	ResNet-50	MobileNetV2
全精度	71.00	76.63	72.62
B + AdaQuant	67.50	72.79	15.17
B + AdaRound	70.18	75.86	69.46
B + FlexRound	70.28	75.95	70.82
消融 1：冻结 `s1`	70.09	75.88	69.47
消融 2：去掉 `s3/s4`	70.22	75.92	70.51

（数值为 top-1 准确率 %，取自论文 Table 1 与 Table 2。）

两个结论：

冻结 $s_1$ 后，FlexRound ≈ AdaRound。 说明”能联合学习网格步长”是核心，只要网格固定，除法带来的额外自由度就被浪费掉了。
去掉通道因子 $s_3/s_4$ 后只略降一点（MobileNetV2 70.51 → 70.82）。所以 $s_3/s_4$ 是锦上添花，主要贡献仍然是”可学步长 + 逐权重除法”。

更低位宽下差距拉大。3-bit 时 MobileNetV2 上 FlexRound 66.87% vs AdaRound 62.51%；2-bit 时 ResNet-50 上 FlexRound 63.67% vs AdaRound 48.47%。

论文报告的实验覆盖

FlexRound 的实验面比较宽，这也是它的卖点之一：

图像分类（ImageNet，ResNet-18/50、MobileNetV2）：仅权重量化、以及权值+激活量化（W4A4、W3A3）。
自然语言理解（GLUE，BERT-Base/Large、GPT-Neo 125M–2.7B，8-bit W/A）：MRPC 上 BERT-Base FlexRound 84.07 对 AdaRound 81.62（全精度 85.05）。
自然语言生成（WikiText2、PTB、WebNLG，GPT-Neo、OPT、GPT-2）。
LLM（LLaMA）：per-channel 权重量化 + per-tensor 激活量化，在 6 个常识推理任务和 WikiText2 上报告接近半精度基线。

论文强调它是第一个在 per-tensor uniform 设定下同时覆盖图像分类、NLU、NLG 的工作，也是最早用逐块输出重构把 LLaMA 这类 LLM 量化到接近半精度的一组实验之一。

局限与需要注意的点

按论文自述和数据来看，有几条边界值得记住：

需要小样本校准数据。 逐块重构要用输入激活；校准集质量和数量会影响结果（论文附录有 sample size 消融）。它比纯 RTN 贵，但比 QAT 便宜得多。
局部最优。 逐层/逐块的连续松弛优化，不保证全局最优；它优化的是每块输出重构，不是端到端任务指标。
对 ResNet 这类权重都在 (-1,1) 的网络，增益有限。 因为 reciprocal 性质需要”大权重”来发力，权重集中在小范围时，FlexRound 的表现接近 AdaRound。
取整仍不可微，靠 STE 近似，训练稳定性和步长 $\beta$ （论文的软松弛参数）等会影响结果。
部署粒度约束。 主推的 per-tensor 最省事；但要复现论文里 per-channel 的 LLM 结果，后端必须支持对应粒度的 weight encoding。

一句话总结：FlexRound 把”舍入”从一个离散的取整动作，改写成一个关于缩放因子的连续优化，并用除法的导数性质让权重大小自然成为重要性先验，从而在 per-tensor 设定下也能逼近全精度。

参考

Lee, J. H., Kim, J., Kwon, S. J., Lee, D. FlexRound: Learnable Rounding based on Element-wise Division for Post-Training Quantization, ICML 2023（PMLR v202, pp. 18913–18939）。
官方代码：onliwad101/FlexRound_LRQ
站内相关：SeqMSE 原理：用输出重构误差选择量化 Encoding
站内相关：HQQ 算法解读：不用校准数据，怎样把大模型权重压到 4 bit