推理优化 on Lorn

浅谈低精度浮点数

Wed, 10 Jun 2026 20:25:01 +0800

Motivation

最近在搭建 VLM real quantization baseline。最初的实验环境选在 RTX PRO 6000 Blackwell 上，但在实际部署过程中发现，当前主流推理框架对 Blackwell 设备上传统整数低比特量化的支持还不够成熟，尤其是常见的 W4A8、W4A16 等 INT 量化配置，在 vLLM、SGLang、TensorRT-LLM 等推理栈中的可用性、kernel 覆盖和性能表现，都与 Hopper/Ada 等架构存在明显差异。

进一步调研后发现，Blackwell 架构的低精度推理路线正在发生变化：相比传统 INT4/INT8，硬件与软件生态开始更多地围绕低精度浮点格式展开，例如 MXFP4、NVFP4 等 micro-scaling FP4 格式。这类格式在保持 4-bit 存储与计算效率的同时，通过局部 scale、浮点指数以及更细粒度的动态范围适配，试图缓解 INT 量化中常见的离群值敏感和动态范围不足问题。

与此同时，近两年的低精度量化研究也逐渐从“如何把 INT4 做稳”转向“如何利用硬件原生支持的 FP4 格式”。例如，围绕 MXFP4/NVFP4 的推理量化、FP4 全量训练、FP4 误差分析以及格式感知的旋转/缩放方法，已经在 ICLR、ICML、NeurIPS 以及 arXiv 上陆续出现。因此，我准备写下这篇 blog，作为对低精度浮点量化格式的学习笔记，也作为后续在 Blackwell 设备上构建 VLM real quant baseline 的技术背景梳理。

Note

IEEE 754 标准浮点数

我们先来看看在内存中，浮点数是如何按照 IEEE 754 标准存储的。

本质上，浮点数可以看作是用科学计数法来表示一个数：

$$ V = (-1)^S \times M \times R^E $$

其中 S 表示符号位，M 表示尾数，R 表示基数，E 表示指数。对于计算机中的二进制浮点数而言，基数 R = 2。

以 IEEE 754 单精度浮点数，也就是 float32 为例，它一共占 32 bit，具体划分如下：

符号位 S：1 bit
指数字段 Exponent：8 bit
尾数字段 Fraction：23 bit

为了在有限的位宽内尽可能扩大表示范围并提高精度，IEEE 754 对指数和尾数做了特殊规定：

对于规格化浮点数，尾数的整数部分总是 1，因此这个最高位的 1 不需要显式存储，称为隐藏位。于是，尾数字段虽然只有 23 bit，但实际可以表示 24 bit 的有效数字。
指数字段采用偏置（bias）表示。对于 float32，指数字段占 8 bit，存储值范围为 0 ~ 255，偏置值为 127。对于规格化数，真实指数为：

$$ E = \text{Exponent} - 127 $$

需要注意的是，Exponent = 0 和 Exponent = 255 有特殊含义，分别用于表示 0、非规格化数、无穷大和 NaN。因此，规格化单精度浮点数的真实指数范围是 $-126 \sim 127$。

单精度浮点数在内存中的布局如下：

 bit index
 31 30 23 22 0
 │ │ │ │ │
 ▼ ▼ ▼ ▼ ▼
┌───────┬────────────────────────────┬──────────────────────────────────┐
│ S │ Exponent │ Fraction │
│ 1 bit │ 8 bits │ 23 bits │
└───────┴────────────────────────────┴──────────────────────────────────┘
 │ │ │
 │ │ └─ fraction bits of mantissa
 │ └─ biased exponent, bias = 127
 └─ sign bit: 0 = positive, 1 = negative

例如，我们想用单精度浮点数表示 25.125。

首先将它转换为二进制：

整数部分：25(D) = 11001(B)
小数部分：0.125(D) = 0.001(B)

因此：

25.125(D) = 11001.001(B)

将其写成二进制科学计数法：

11001.001(B) = 1.1001001(B) × 2^4

所以：

S = 0
M = 1.1001001
Fraction = 10010010000000000000000
Exponent = 4 + 127 = 131(D) = 10000011(B)

因此，25.125 的 IEEE 754 单精度表示为：

S | Exponent | Fraction
0 | 10000011 | 10010010000000000000000

完整的 32 bit 表示为：

01000001110010010000000000000000

对应的十六进制表示为：

0x41C90000

浮点数表示会带来一定的精度损失，主要原因是并非所有十进制小数都能被有限位宽的二进制小数精确表示。例如，0.2 转换为二进制后是一个无限循环小数：

0.2(D) = 0.0011001100110011...(B)

由于计算机的位宽有限，存储时只能保留有限位，并按照 IEEE 754 的规则进行舍入，因此会产生精度误差。

在 IEEE 754 单精度浮点数下，常用的有限数表示范围大约为：

-3.4 × 10^38 ~ 3.4 × 10^38

更准确地说，单精度浮点数能够表示的最大有限正数约为：

3.4028235 × 10^38

最小正规格化数约为：

1.17549435 × 10^-38

如果考虑非规格化数，还可以表示更接近 0 的数，但精度会进一步降低。

在如今的 LLM 训练推理过程中，更通用的浮点数格式是 FP16 与 BF16。它们与 FP32 的区别在于：

FP16 的指数位为 5，尾数位为 10
BF16 的指数位为 8，尾数位为 7

我们知道，在总位宽固定的情况下，尾数越多则指数越少，能够表示的范围越小，但相应精度越高。因此 FP16 比 BF16 精度更高，但表示范围更小。

正因如此，BF16 通常用在训练场景（更大的动态范围更利于反向传播中的梯度表示），而 FP16 通常用在推理场景（主要是因为目前很多主流推理框架对 FP16 的 kernel 优化做得更好）。

FP8

FP8 是 Hopper 架构设备开始支持的低精度浮点数，通常有两种表示形式：E4M3 和 E5M2（其中 E 表示指数位数，M 表示尾数位数）。

同样地，E4M3 精度更高、范围更小，而 E5M2 范围更大、精度更低。

在支持 FP8 的硬件（Hopper 架构显卡）上，已经支持 FP8 格式的矩阵乘法（E4M3×E4M3、E5M2×E5M2、E4M3×E5M2、E5M2×E4M3）。

我们接下来从精度和性能两个方面，对比分析 FP8 与 INT8 在量化上的表现。

首先分析二者在精度上的表现。

由浮点数的定义可知，浮点数可以看作是在相邻 2 的幂（$2^E$）之间均匀采样的 $2^M$ 个样本。例如在 E5M2 中，2（一次幂）和 4（二次幂）之间有 4 个样本，4 和 8（三次幂）之间也有 4 个样本；而在 E4M3 中，2 和 4 之间有 8 个样本。由此可知：

E5M2 的量化误差比 E4M3 更大；
FP8 表示的数值是非均匀的，浮点量化的误差会随着数值的增大而增大。

这正是 FP8 与 INT8 量化的最大差别：对于 FP8 而言，越靠近 0 分布越稠密，越远离 0 分布越稀疏。

接下来从数学角度分析 FP8 与 INT8 在量化上的表现。

对于 INT8，量化可以写作：

$$ \hat x = s \cdot \text{round}\left(\frac{x}{s}\right) $$

其中 $s$ 是 scale。如果使用对称 INT8，那么通常有：

$$ q \in [-127, 127] $$

假设一组待量化值的最大绝对值是 $\alpha$，那么 scale 大约是：

$$ s = \frac{\alpha}{127} $$

量化绝对误差满足：

$$ |x - \hat x| \leq \frac{s}{2} = \frac{\alpha}{254} $$

而相对误差为：

$$ \frac{|x - \hat x|}{|x|} \leq \frac{\alpha}{254\,|x|} $$

可以看到，无论从绝对误差还是相对误差来看，INT8 量化都更容易受到 outlier 的影响。

而对于 FP8，假设尾数有 $M$ bit，那么在区间 $[2^e, 2^{e+1})$ 内约有 $2^M$ 个采样点，间隔为：

$$ \Delta_e = \frac{2^e}{2^M} = 2^{e-M} $$

因此量化绝对误差大约满足：

$$ |x - \hat x| \leq \frac{1}{2}\,2^{e-M} $$

而相对误差近似满足：

$$ \frac{|x - \hat x|}{|x|} \leq 2^{-(M+1)} $$

也就是说，虽然 FP8 的绝对误差也会随着数值变大而变大，但相对误差大致稳定。

如果是 E4M3，相对误差不会超过 $6.25\%$；而 E5M2 则不会超过 $12.5\%$。

若令 INT8 的相对量化误差等于 E4M3：

$$ \frac{\alpha}{254\,|x|} = \frac{1}{16} $$

可以得到：

$$ |x| \approx 0.063\,\alpha $$

也就是说，当待量化值与最大值的差距不超过约 16 倍时，INT8 的量化效果往往更好；一旦超出这个范围（即数值相对 $\alpha$ 过小），FP8 则是更优的选择。

而我们知道，权重和激活值虽然都存在 outlier，但二者的数值范围并不相同：权重的 outlier 一般是普通 token 的数十倍，而激活值的 outlier 则不一定，甚至可以达到数百倍。

因此在实际部署时，一个理想的选择是权重用 INT 量化、激活值用 FP8 量化，也就是常说的 W4AFP8（在 TensorRT-LLM 中记作 W4A8）。目前的推理框架普遍推荐这种量化方式，例如 TensorRT-LLM 中：

QUANTIZATION METHODS	PERFORMANCE IMPROVEMENT (BATCH SIZE <= 4)	PERFORMANCE IMPROVEMENT (BATCH SIZE >= 16)	ACCURACY IMPACT	CALIBRATION TIME**
FP8 (W8A8)	Medium	Medium	Very Low	Minutes
Int8 SQ (W8A8)	Medium	Medium	Medium	Minutes
Int8 weight-only (W8A16)	Medium	Low	Low	Not Required
Int4 weight-only (W4A16)	High	Low	High	Not Required
Int4 AWQ (W4A16)	High	Low	Low	Tens of Minutes
Int4 GPTQ	High	Low	Low	Tens of Minutes
Int4-FP8 AWQ (W4A8)	High	Medium	Low	Tens of Minutes

可以看到，W4AFP8 在性能提升最大的同时，还能做到较低的精度影响。

接下来从性能角度比较 FP8 与 INT8。

在正式比较之前，需要先明确一点：低精度 Tensor Core 的核心，并不是把整个矩阵乘法都维持在 8 bit 精度下完成，而是采用 mixed-precision MMA 的方式。也就是说，输入矩阵以 FP8 或 INT8 这样的低精度格式参与乘法，但乘积通常会在更高精度的累加器中累加，最后再根据需要转换为目标输出精度。

这样设计的原因在于，GEMM 中每个输出元素本质上都是一个长度为 K 的点积，需要经历大量乘加操作。如果累加阶段仍然使用过低的精度，那么许多较小的增量会因为舍入而被直接吞掉，累加误差也会随着 K 的增大迅速放大。因此，低精度矩阵乘法的“低精度”主要体现在输入和乘法阶段，而不是简单地意味着整个计算链路都只有 8 bit。

从公开硬件规格来看，在支持 FP8 的 NVIDIA GPU 上，FP8 Tensor Core 和 INT8 Tensor Core 往往具有相同的标称峰值吞吐。这说明在产品级实现中，厂商可以通过专门的数据通路设计、流水线化、tile 级并行以及 mixed-precision 累加机制，使 FP8 和 INT8 在峰值计算能力上达到相同水平。

但这并不意味着 FP8 和 INT8 的底层硬件代价完全相同。INT8 的计算本质上是整数/定点乘加，数据格式规则，乘法和累加逻辑相对简单；而 FP8 虽然同样只有 8 bit，但它包含符号位、指数位和尾数位，计算时需要处理指数对齐、尾数运算、规格化、舍入以及高精度浮点累加等逻辑。因此，从单个 MAC 单元或数据通路的实现复杂度来看，FP8 通常比 INT8 更复杂。

所以更准确的理解是：理论上，FP8 运算单元的硬件实现成本高于 INT8；但在 Hopper、Blackwell 这类专门面向 AI 负载设计的 Tensor Core 中，NVIDIA 通过专门的硬件设计把 FP8 和 INT8 的峰值吞吐做到了相同水平。也就是说，FP8 的优势并不是“硬件更简单”，而是在相近吞吐下提供了更大的动态范围，从而更适合大模型中激活值波动大、outlier 明显的场景。（简单来说：放心大胆地用 FP8 吧，NVIDIA 帮你兜底了。）

FP4 / MXFP4 / NVFP4

在 4-bit 下，想要像 FP8 那样直接使用 FP4 是不现实的，因为 FP4 能表示的数值范围实在太小了。在 OCP 标准中，FP4 一般取 E2M1，它能表示的值大概是：

$$ 0,\ \pm 0.5,\ \pm 1,\ \pm 1.5,\ \pm 2,\ \pm 3,\ \pm 4,\ \pm 6 $$

也就是说，单单使用 FP4，能够表示的数值范围相当有限。因此实际应用时，通常不会只存储一个 FP4 值，还会额外存储一个 scale：

$$ \text{real num} = \text{Encoding of FP4} \times \text{scale} $$

scale 的不同对应着目前两种主流的 FP4 格式（Blackwell 确实支持 FP4 这种数值格式，但没有针对 scale 的专用硬件加速）。

我们先来看 MXFP4，它的 scale 格式是每 32 个值共享一个 E8M0 scale，也就是 2 的幂形式的 scale。

接下来按照与 FP8 相同的思路分析它的量化误差（取 block 内最大绝对值为 $A$）：

$$ \hat x_i = s_{mx} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{mx}}\right) $$

一般而言理想的 scale 是 $\frac{A}{6}$；然而在 MXFP4 中 scale 必须是 E8M0 格式，因此需要将其量化为最近的 2 的幂：

$$ s_{mx} = 2^{\text{round}\left(\log_2 \frac{A}{6}\right)} $$

在 TensorRT 中一般会向上取整。为了计算方便，我们定义：

$$ \rho_{mx} = \frac{s_{mx}}{A/6} $$

由于 E8M0 是 2 的幂，向上取整时有：

$$ 1 \leq \rho_{mx} < 2 $$

由于任意量化器的量化误差不会超过一个量化间隔，因此我们有：

$$ |e_i| \leq s_{mx} = \rho_{mx}\,\frac{A}{6} $$

而 $\rho_{mx}$ 与 block 内最大值和 2 的幂的对齐程度相关：当 $A$ 恰好为 2 的幂时 $\rho_{mx}$ 接近 1，误差较小；当 $A$ 略大于某个 2 的幂时 $\rho_{mx}$ 接近 2，误差较大。

例如 $A/6 = 1.99$ 与 $A/6 = 2.01$：前者会被向上取整为 2，此时 $\rho_{mx}$ 接近 1；后者会被向上取整为 4，此时 $\rho_{mx}$ 约为 2。

若对齐得好，则有：

$$ |e_i| \leq \frac{A}{6} $$

否则为：

$$ |e_i| < \frac{A}{3} $$

除了绝对误差外，我们再看看它的归零阈值（即低于该值会被量化为 0）：

$$ |x| < 0.25\,s_{mx} = \rho_{mx}\,\frac{A}{24} $$

于是有：

$$ \frac{A}{24} \leq \text{zero threshold} \leq \frac{A}{12} $$

大约是 $4.17\% A$ 到 $8.33\% A$。

接下来看看 NVFP4，它与 MXFP4 的区别在于：

block 更小：每 16 个元素共享一个 scale；
scale 更细：NVFP4 的 block scale 采用 FP8 E4M3，而不是 E8M0。

此外，NVFP4 还会额外使用一个全局的（per-tensor）FP32 scale 来避免 overflow。一般而言：

$$ s_{\text{global}} = \frac{\text{global}_{amax}}{6 \times 448}\\[4pt] s_{\text{block}} = \frac{\text{block}_{amax}/6}{s_{\text{global}}} $$

其中 6 是 E2M1 的最大可表示值，448 是 E4M3 的最大可表示值，二者相乘保证 block scale 落在 E4M3 的表示范围内。

同样地，我们可以把 NVFP4 写作：

$$ \hat x_i = s_{\text{global}} \cdot s_{\text{block}} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{\text{global}} \cdot s_{\text{block}}}\right) $$

令

$$ s_{nv} = s_{\text{global}} \cdot s_{\text{block}} $$

那么同样地有：

$$ \rho_{nv} = \frac{s_{nv}}{A/6} $$

又因为此时 scale 为 E4M3，其相邻正规格化数的比例最多约为：

$$ 1 + \frac{1}{8} = \frac{9}{8} $$

因此取向上取整时：

$$ 1 \leq \rho_{nv} \leq \frac{9}{8} $$

于是 NVFP4 的绝对误差上界约为：

$$ |e_i| \leq \rho_{nv}\,\frac{A}{6} \leq \frac{3}{16}A $$

约 $18.75\% A$，这明显好于 MXFP4 在未对齐时的最坏情况。

而 NVFP4 的归零阈值为：

$$ \text{zero threshold} \leq 4.69\% A $$

至于 INT4，可以通过与 INT8 类似的方式分析：对称量化下 $q \in [-7, 7]$，scale 为 $s = A/7$，因此绝对误差上界为 $|e_i| \leq s/2 = A/14$，归零阈值同样为 $A/14 \approx 7.14\% A$。

需要注意的是，INT4 的绝对误差上界（约 $7.14\% A$）看似比 NVFP4（约 $18.75\% A$）更小，但这是均匀分布在整个量程上的误差；而对于神经网络中占绝大多数的小数值，FP4 的误差与数值大小成比例，这些小值上的实际误差远小于 INT4 的均匀误差。再结合 FP4 对长尾异常值更友好的动态范围，可以得出结论：在 4-bit 量化中，面对神经网络这种“大量小值 + 长尾异常值”的分布，NVFP4 的综合表现最好。

而从性能上分析，虽然 MXFP4 由于采用 32 元素 block 和 E8M0 power-of-two scale，理论上具有更低的 scale 存储与缩放复杂度，但在真实的 LLM/VLM 推理中，性能并不只由单个 GEMM 的格式开销决定。NVFP4 通过 16 元素 micro-block 和 E4M3 scale 显著降低了量化误差，使更多层能够稳定进入 FP4 路径，减少混合精度 fallback 和额外补偿。因此，在当前 Blackwell + TensorRT-LLM 生态下，NVFP4 往往是更实际的高性能 FP4 选择；而 MXFP4 更适合在对精度不敏感、或有专门误差补偿方法的场景中追求更高的硬件效率。

此外，目前主流推理框架都对 NVFP4 做了专门优化，是性能—精度权衡下的最佳选择。

Reference

Hessian矩阵系列串讲

Mon, 24 Nov 2025 13:28:52 +0800

定义

假设有一实值函数$f(x_1,x_2,\dots,x_n)$，若$f$的所有二阶偏导数都存在且在定义域里连续，那么我们定义函数$f$的Hessian矩阵为如下一个$n\times n$的方阵:

$$ \mathbf{H} = \left[ \begin{matrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1\partial x_2} & \dots&\frac{\partial^2 f}{\partial x_1 \partial x_n}\\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \dots& \frac{\partial^2 f}{\partial x_2 \partial x_n}\\ \vdots & \vdots & \ddots & \vdots\\ \frac{\partial^2 f}{\partial x_n \partial x_1} & \frac{\partial^2 f}{\partial x_n \partial x_2} & \dots & \frac{\partial^2 f}{\partial x_n^2} \end{matrix} \right] $$

或使用下标标记表示为:

$$ \mathbf{H}_{ij} = \frac{\partial^2 f}{\partial x_i \partial x_j} $$

一般性质

对称性:

对于机器学习中遇到的大多数函数(特别是那些具有连续二阶偏导数的函数),混合偏导数的求导顺序无关紧要。这被称为克莱罗定理或施瓦茨定理,它说明了混合偏导数的相等性:

$$ \frac{\partial^2 f}{\partial x_i\partial x_j} = \frac{\partial^2 f}{\partial x_j \partial x_i} $$

这意味着Hessian矩阵是对称的，即$\mathbf{H} = \mathbf{H}^\top$

正定性与局部最优值的相关性

通过Hessian矩阵的正定性，我们可以判断该点是局部最小值，局部最大值还是鞍点

正定:若Hessian矩阵$H(x^*)$是正定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v>0$)，则函数在$x^*$处有局部最小值。函数在该点周围向所有方向向上弯曲，类似于碗的底部
负定:若Hessian矩阵$H(x^{*})$是负定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v<0$),则函数在$x^*$处有局部最大值。函数在该点周围向所有方向向下弯曲，类似于圆顶的顶部
不定:若Hessian矩阵$H(x^{*})$是不定的,则函数在$x^*$处有鞍点。函数在某些方向向上弯曲，在另一些方向向下弯曲，就像马鞍

Hessian在神经网络中的计算

Hessian矩阵在神经网络计算的许多方面有着重要作用，包括:

⼀些⽤来训练神经⽹络的⾮线性最优化算法是基于误差曲⾯的⼆阶性质的，这些性质

由Hessian矩阵控制(比如牛顿法和拟牛顿法)
对于训练数据的微⼩改变，Hessian矩阵构成了快速重新训练前馈⽹络的算法的基础
Hessian矩阵的逆矩阵⽤来鉴别神经⽹络中最不重要的权值，这是⽹络“剪枝”算法的⼀部分(LeCun的OBD)

对于Hessian矩阵的众多应用而言，一个重要的需要考虑的问题是计算效率。在神经网络中有$W$个参数(包括权值和偏置),那么Hessian矩阵的大小就是$W\times W$,那么计算Hessian矩阵的计算量为$O(W^2)$。这在具有大量参数的神经网络中是难以接受的，因此我们需要进行一些高效的近似。

对角近似

这个方法最早由Yan LeCun在OBD剪枝方法中提出

我们只保留Hessian矩阵的对角线元素$H_{i,i}$，把非对角线元素置为0.(这个方法在很大程度可以可以方便求逆)

对于神经网络中的第j个神经元的输入$a_j$对应权重为$w_{ji}$:

$$ \frac{\partial^2 E_n}{\partial w_{ji}^2} = \frac{\partial^2 E_n}{\partial a_j^2}z_i^2 $$

其中$z_i$是上一层神经元的输出。

而$\frac{\partial^2 E_n}{\partial a_j^2}$可以通过链式法则递归计算(类似于反向传播):

$$ \frac{\partial^2 E_n}{\partial a_j^2} = \underbrace{\frac{\partial}{\partial a_j}[h'(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}]}_{链式法则}=h'(a_j)^2 \sum_{k,k'}w_{kj}w_{k'j}\frac{\partial^2 E_n}{\partial w_k\partial w_{k'}} + h''(a_j)\sum_k w_{kj}\frac{\partial E_n}{\partial a_n} $$

忽略二阶导中的非对角线项$k\neq k'$:

$$ \frac{\partial^2 E_n}{\partial a_j^2}\approx \underbrace{h'(a_j)^2\sum_k w_{kj}^2 \frac{\partial^2 E_n}{\partial a_k^2}}_{链式法则} + \underbrace{h''(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}}_{链式法则} $$

从而一次反向传播便可以计算出来，时间复杂度为$O(W)$

但是如你所见，在这个方法下，Hessian矩阵完全退化为了一个diag，这在神经网络中是不合理的，因为非线性层的引入会让Hessian矩阵的交叉项不为0，更重要的是这种误差会随着网络层数的堆砌不断放大。

因此针对这个问题在对角近似上近年来不断有相关文章发表，下面对我阅读过的一些进行简单介绍:

HesScale:在LeCun的方法的骨架上，针对常见网络的最后一层的softmax+CE结构可以通过一个简单的公式求出diag(H)的精确值:
$$ \frac{\partial^2 E_n}{\partial^2 a_j^2} = p - p\circ p $$
其中p为$\text{softmax}(a_j)$,$\circ$表示向量的逐元素乘积,为方便书写我们将$\{a_1,\dots,a_n\}$写作z

证明如下:

预测$\mathbf{p} = \text{softmax}(z)$,即$p_i = \frac{\exp(z_i)}{\sum_{j = 1}^k \exp(z_j)}$,设目标分布为$t\in \Delta^{K-1},\sum_{i}t_i = 1$，对于单样本而言，其损失可以写作:
$$ \mathcal{L}(z) = -\sum_{j=1}^Kt_i\log p_i = -t^\top z + \log \sum_{j=1}^K \exp(z_j) $$
求一阶导我们有:
$$ \nabla_z \mathcal{L} = -t + \nabla_z \log \sum_{j=1}^K \exp{z_j} = -t + \mathbf{p} $$
求二阶导相当于$\mathbf{p}$对z求导。我们知道$\text{softmax}$的Jocobian为
$$ \frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j),J_{\text{softmax}}(z) = \text{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top $$
其中$\delta_{i,j}$为Kronecker函数，写成矩阵形式，则只有对角线元素为1，其余为0.

因此我们有:
$$ \nabla^2_z\mathcal{L} = \text{diag}(\mathbf{p}) - \mathbf{pp}^\top $$
于是我们可以得到:
$$ \frac{\partial^2 L}{\partial^2 a_j} = p_j(1-p_j) $$
AdaHessian:

具体而言，AdaHessian在进行Hessian矩阵的对角近似的时候利用了Hutchinson估计，它通常也被用来对矩阵的迹进行估计:

对任意矩阵$\mathbf{A}\in\mathbb{R}^{d\times d}$,若随机向量$z = (z_1,\dots,z_d)^\top$满足:
$$ \mathbb{E}[z_i] = 0\\ \mathbb{E}[z_iz_j] = \delta_{ij} $$
那么有:
$$ \mathbb{E}[z\odot (Az)] = \text{diag}(A) $$
下面进行证明:

记估计量$d = z\odot (Az)$的第i个分量:
$$ d_i = z_i(Az)_i = z_i\sum_{j=1}^d A_{ij}z_j = \sum_{j = 1}^d A_{ij}z_iz_j\\ \mathbb{E}[d_i] = \sum_{j=1}^dA_{ij}\mathbb{E}[z_iz_j] = \sum_{j=1}^d A_{ij}\delta_{ij} = A_{ii} $$
由此我们证明了$\mathbb{E}[z\odot(Az)] = \text{diag}(A)$

在神经网络中，相较于直接计算完整的Hessian矩阵以及其OBD方式的对角近似，计算其矩阵向量积(HVP,Hessian Vector product)是并不困难的，它只需要一次反向传播：
$$ Hz = \frac{\partial (g^\top z)}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z + g^\top \frac{\partial z}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z $$
由此，我们可以通过多次在满足Rademacher分布的向量取样计算$z\odot (Az)$的期望就能得到$\text{diag}(A)$的无偏估计，在实际的应用中，只进行一次取样就能得到较为不错的结果。

外积近似

在神经网络应用于回归问题时，通常采用下面形式的平方和误差

$$ E = \frac{1}{2}\sum_{n=1}^N(y_n-t_n)^2 $$

那么Hessian矩阵可以写成如下形式:

$$ H = \nabla\nabla E = \nabla(\sum_{n=1}^N (y_n-t_n)\nabla y_n) = \sum_{n = 1}^N \nabla y_n(\nabla y_n)^\top + \sum_{n=1}^N(y_n-t)\nabla\nabla y_n $$

在网络已经训练好的情况下，输出$y_n$与$t_n$接近，因此第二项可以忽略，由此我们得到了Hessian矩阵的外积近似

$$ H\approx \sum_{n=1}^N b_n b_n^\top $$

其中，$b_n = \nabla y_n = \nabla a_n$(输出单元的激活函数就是恒等函数)。这种方法中的Hessian矩阵可以跟随反向传播算法在$O(W)$个步骤内高效地求出误差函数地一阶导数。再通过简单地乘法就可以在$O(W^2)$步骤内求出矩阵元素。

Hessian矩阵逆的计算

使用外积近似，我们可以提出一个计算Hessian矩阵的逆的高效办法,首先我们有:

$$ \mathbf{H}_N = \sum_{i=1}^n b_n b_n^\top $$

其中，$b_n = \nabla w a_n$时数据点n产生的输出单元激活对梯度的贡献。我们现在推导一个建立Hessian矩阵的顺序步骤，每次处理一个数据点。假设我们已经使用了前L个数据点得到了Hessian矩阵的逆。通过将第L+1个数据点的贡献单独写出来，我们有:

$$ \mathbf{H}_{L+1} = \mathbf{H}_L + b_{L+1}b_{L+1}^\top $$

为了计算Hessian矩阵的逆，我们考虑下面这个矩阵的恒等式:

$$ (M + vv^\top)^{-1} = M^{-1} - \frac{(M^{-1}v)(v^\top M^{-1})}{1+v^\top M^{-1}v} $$

若我们令$H_L = M$,且$b_{L+1} = v$,我们有:

$$ H_{L+1}^{-1} = H_{L}^{-1} - \frac{H_L^{-1}b_{L+1}b_{L+1}^\top H_{L}^{-1}}{1 + b_{L+1}^\top H_L^{-1}b_{L+1}} $$

通过这种方式，数据点可以依次使用，直到L+1=N,整个数据集处理完毕。于是，这个结果表示一个计算Hessian矩阵的逆的算法。这个算法只需对数据集扫描一次。最开始的矩阵$\mathbf{H_0}$被选为$\alpha I$,其中$\alpha$是一个较小的量，从而算法实际找的是$\mathbf{H} + \alpha I$的逆。结果对于$\alpha$的精确值不敏感。

Hessian矩阵的逆计算的K-FAC分解方法：

在特定场景下，Hessian矩阵的逆有一个利用Kronecker分解的近似方法，它利用了在loss为log-like形式下$\mathbb{E}[H] = F$的特点，其中$F$为Fisher信息矩阵:

$$ F = \mathbb{E}[\nabla \log p(x)\nabla \log p(x)^\top] $$

其中上式中做了如下形式的简写:

$$ \nabla \log{p(x)} = \nabla_\theta \log{(p|\theta)}\\ \mathbb{E}[p(x)] = E_{x\sim p(x|\theta)} $$

首先我们有:

$$ \mathbb{E}[\nabla \log p(x)] = \int (\nabla \log p(x))p(x) dx\\ =\int \frac{\nabla \log p(x)}{p(x)}p(x) dx\\ =\int \nabla \log p(x)dx\\ =\nabla 1 = 0 $$

那么当loss为log-like形式下时:

$$ L = -\log p(x) $$

我们求它的Hessian:

$$ \nabla^2 L = -\nabla^2 \log p(x)\\ =-\nabla \frac{\nabla p(x)}{p(x)}\\ =\frac{\nabla p(x)}{p(x)}^\top \frac{\nabla P(x)}{p(x)} - \frac{P(x)^2}{p(x)}\\ =\nabla \log p(x)^\top \nabla \log p(x) - \frac{\nabla^2 P(x)}{p(x)} $$

对该式求期望，我们有:

$$ \mathbb{E}[H]=\mathbb{E}[\nabla \log p(x)^\top \nabla \log p(x)] - \mathbb{E}[\frac{\nabla^2 p(x)}{p(x)}] = F - \int \frac{\nabla^2 p(x)}{p(x)} p(x) dx\\ =F - \int \nabla^2 p(x)dx\\ =F - \nabla^2 \int p(x)\\ =F $$

由于$\theta$本质为所有层$\mathbf{W}$的拼接，我们有:

$$ d\theta = \nabla_{\theta}L(x)\\ F = \mathbb{E}[\nabla L(x)\nabla L(x)^\top] = \mathbb{E}[d\theta d\theta^\top]\\ \theta = [\text{vec}(W_0)^\top,\text{vec}(W_1)^\top,\dots,\text{vec}(W_n)^\top]^\top $$

带入展开得到:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)^\top] $$

令$a_i,g_i$分别为第i层的前向输入和反向传播梯度，由反向传播算法有:

$$ d W_i = g_ia_i^\top $$

带入有:

$$ \text{vec}(dW_i) = \text{vec}(g_ia_i^\top) = a_i\otimes g_i $$

这里可能不太直观，这是因为平常对kronecker积($\otimes$)的接触较少，具体而言，我们这样定义Kronecker积:

$$ A\otimes B = \left[ \begin{matrix} a_{11}B &\dots &a_{1n}B\\ \vdots & \ddots& \vdots\\ a_{m1}B & \dots & a_{mn}B \end{matrix} \right] $$

在这里，由于我们将$dW_i$即$\nabla W_i$第i层参数的梯度进行了向量化，原本的$\nabla W_i$ 是如下形式:

$$ \nabla W_i = g_i a_{i}^\top = \left[\begin{matrix}g_1a_1 & \dots &g_1a_n\\ \vdots & \ddots & \vdots\\ g_ma_1 &\dots&g_ma_n\end{matrix} \right] $$

按列堆叠(vec)就有:

$$ \text{vec}(dW_i) = \left[ \begin{matrix} g_1a_1\\ \vdots\\ g_1a_n\\ \vdots\\ g_ma_n \end{matrix} \right] = a_i\otimes g_i $$

那么我们可以将Fisher矩阵写作如下形式:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)]\\ =\mathbb{E}[(a_i\otimes g_i)(a_j\otimes g_j)^\top]\\ =\mathbb{E}[(a_ia_j^\top)\otimes (g_ig_j^\top)]\\ \approx \mathbb{E}[a_ia_j^\top]\otimes\mathbb{E}[g_ig_j^\top] $$

这个近似相当于我们忽略了$\text{Cov}_\otimes(a_ia_j^\top,g_ig_j^\top)$,这实际上是合理的，尤其是在网络较深的情况下。这里给出一个链接提供一个比较详细的说明。

虽然直接求$F_{ij}$的复杂度仍然不变，但是利用Kronecker积的性质，我们在求逆时可以得到较大的性能提升:

$$ (A\otimes B)^{-1} = A^{-1}\otimes G^{-1} $$

并且在实际运算中，并不是整个网络的Fisher进行计算，而是按层做块对角运算:

$$ \mathbf{F} = \text{blockdiag}(F_1,F_2,\dots),F_l\approx A_l\otimes G_l $$

这样每层求逆的时候只需要求两个小矩阵的逆。

一些统计量的计算：

除了直接对Hessian矩阵的近似计算，我们有时候仅仅需要对Hessian矩阵的统计量进行计算，例如矩阵的迹，最大特征值等。

在计算这些统计量时，我们需要一个重要的算子:Hv,即Hessian矩阵与任意向量的乘积，这个乘积我们在上面已经证明了通过一次简单的反向传播算法可以得到。(PyHessian是一个Python库，它实现了这个算子的高效计算)

顶部若干特征值

利用幂迭代法 + HVP可以高效地求解特征值。

简单来说幂迭代法的流程很简单:

给定矩阵A以及初始非零向量$x_0$,

迭代:
$$ y_{k+1} = Ax_k\\ x_{k+1} = \frac{y_{k+1}}{||y_{k+1}||}\\ \mu_{k+1} = \frac{x_{k+1}Ax_{k+1}^\top}{x_{k+1}x_{k+1}^\top} $$
具体的收敛性证明见教材
矩阵的迹

在上面近似计算中，我们曾用Hutchinson估计对Hessian矩阵进行了对角近似，事实上，我们还可以进行迹的估计。

简单来说，方法一样，从Rademacher分布中取一随机向量$v$,然后有恒等式:
$$ \text{Tr}(H) = \text{Tr}(HI) = \text{Tr}(H\mathbb{E}[vv^\top]) = \mathbb{E}[\text{Tr}(Hvv^\top)] = \mathbb{E}[v^\top H v] $$
Hessian矩阵在神经网络下的特殊结构
本小节内容主要基于Towards Quantifying the Hessian Structure of Neural Networks,B站上有作者的讲解视频[FAI] 港中深张雨舜 | 浅谈神经网络Hessian矩阵的特殊结构

这篇文章主要说明了在神经网络中，Hessian矩阵往往具有近块对角结构，这表明曲率信息(二阶信息)主要在层内耦合，层与层的二阶交互很弱。这间接的说明了Layer-wise的量化/剪枝的合理性。

此外，文章还进一步从理论和实验上阐释了造成这个现象的原因：

静态力量:即使在随机初始化阶段，即在训练开始之前，神经网络的Hessian矩阵就已经呈现出近块对角结构。这种结构的形成与网络的架构设计有关，因此被称为“静态力量”。具体来说，对于线性模型和单隐藏层网络，无论是使用均方误差（MSE）损失还是交叉熵（CE）损失，Hessian矩阵的对角块和非对角块在随机初始化时就已经表现出明显的差异。
动态力量:在训练过程中，Hessian矩阵的结构会进一步发生变化。特别是在使用CE损失时，训练过程会逐渐消除初始时存在于跨层Hessian分量（Hwv）中的“块循环”（block-circulant）模式，而对角块和非对角块的近块对角结构则保持稳定。这种由训练过程引起的结构变化被称为“动态力量”。

另外一点，文章证明了类别数C是影响Hessian矩阵结构的主要因素。(在实验中，隐藏层Hessian的非对角块与对角块的比值以$\frac{1}{\sqrt{C}}$的速度衰减，输出层Hessian的衰减速率为$\frac{1}{C}$).这个结果对于大模型而言是友好的，因为在神经网络中C的大小往往是$1e3\sim 1e4$级别的，这说明大模型的Hessian是具有强对角块的结构的！这一点实际上在传统的方法上人们意识or无意识的用到了(对角近似)，但是更丰富层面以及基于这个发现的在计算的可行性和性能的Trade-off做的工作是比较少的，在优化器那边做的比较多。

此外还有不少文献揭露了Hessian矩阵具有低秩特征谱，即只有少数的特征值显著大，其余大多接近0.