LoRA 与 QLoRA

LoRA

LoRA (Low-Rank Adaptation): 保留原来的权重, 只学习一个由两个小矩阵表示的更新量.

令 $y = W_0x$. Fine-tuning (FT) 时, 将权重 $W_0$ 更新为 $W_0 + \Delta W$, 其中:

$\Delta W = sBA.$

各矩阵的维度为:

$W_0 \in \mathbb{R}^{d_{\mathrm{out}} \times d_{\mathrm{in}}}, \qquad B \in \mathbb{R}^{d_{\mathrm{out}} \times r}, \qquad A \in \mathbb{R}^{r \times d_{\mathrm{in}}}.$

$r \ll d_{\mathrm{out}}, d_{\mathrm{in}}$, 因而大大减少需要更新的参数量.

$y_2 = (W_0 + \Delta W)x = W_0x + sB(Ax).$

其中 $W_0$ 固定, FT 时实际更新的参数是 $A, B$.

对于序列输入 $x \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{seq}}}$:

$Ax \in \mathbb{R}^{r \times d_{\mathrm{seq}}}.$

rank 的证明

$\operatorname{rank}(BA) \leq \min\{\operatorname{rank}(A), \operatorname{rank}(B)\} \leq r.$

把 $A$ 的列向量记为 $a_1, a_2, \ldots, a_n$, 则:

$BA = [Ba_1, Ba_2, \ldots, Ba_n].$

$BA$ 的每一列都是 $B$ 的列向量的线性组合, 因而:

$\operatorname{rank}(BA) \leq \operatorname{rank}(B).$

对应地, 从行空间也可得到 $\operatorname{rank}(BA) \leq \operatorname{rank}(A)$.

例子: LoRA 生图模型

冻结 Foundation Model 的 $W_0$, 用特定画风的 dataset 训练一组 $A, B$:

$Y = W_0X + sB(AX).$

这样可以用较低成本适配不同画风. 示意图见下方手写原稿.

类似思路: DeepSeek 的 MLA.

QLoRA

1. NF4

$\text{4-bit weight} \xrightarrow{\mathrm{dequant}} \text{BF16 approximate weight} \xrightarrow{\mathrm{GEMM}} \text{output}.$

4-bit 可以表示 $2^4 = 16$ 个数据点. 如何分布这 16 个点? 如果零附近数据多, 两端数据少, 就可以考虑 Quantile Quantization (QQ).

Gaussian 假设与分位数

假设权重分布为:

$W \sim \mathcal{N}(0, \sigma^2).$

令 $Z \sim \mathcal{N}(0, 1)$, 则分位数满足:

$Q_W(p) = \sigma Q_Z(p).$

$Q_Z(\cdot)$ 表示标准正态分布的分位数函数: 给定分位 $p$, 返回对应的随机变量取值.

在 CDF 的概率轴上等分, 映射到数值轴后, 零附近的区间更窄, 两端的区间更宽. 原稿中的 CDF 图画的就是这个关系.

分位数采样与端点问题

一个直观的构造是:

$\widetilde q_i = Q_Z\!\left(\frac{i}{2^k}\right).$

但要避开两个端点:

$Q_Z(0) = -\infty, \qquad Q_Z(1) = +\infty.$

因此, 对单个分位数样点, 可以考虑用 $2^k + 1$ 作分母, 取 $i = 1, \ldots, 16$.

原文中的形式则是:

$q_i = \frac{1}{2}\left[ Q_Z\!\left(\frac{i}{2^k + 1}\right) + Q_Z\!\left(\frac{i + 1}{2^k + 1}\right) \right].$

这里仍有一个疑问: 如果采样包含 $i = 0$, 或包含使 $(i + 1)/(2^k + 1) = 1$ 的端点, 仍然会遇到无穷值. 原文的采样索引与端点该怎样处理?

这个式子是在用区间端点的均值代替区间内的值, 而且均值取在随机变量的数值空间, 并非分位数的概率空间.

量化还有一个要求: $0 \to 0$, 即必须有 exact-zero. 这与零值和 sparsity 有关, 同时也要考虑对称性问题.

归一化与匹配

令:

$a = \max |W|.$

把权重归一化到 $[-1, 1]$, 再与归一化后的分位数匹配:

$\frac{W}{a} \in [-1, 1] \quad \longleftrightarrow \quad \frac{Q_Z(p_i)}{\max_i Q_Z(p_i)}.$

2. Double Quantization (DQ)

每个权重块还需要额外保存一个 32-bit 的量化常数 $a_j$. 如果一个块包含 64 个 weight, 平均位宽就会达到 4.5 bit. 因此, 可以进一步压缩 $a_j$.

对于包含 64 个 weight 的块 $W_j$, 用量化常数 $a_j$ 和量化后的数值 $q$ 表示:

$W_j \approx a_j q, \qquad a_j = \max |W_j|.$

$q$ 来自前面的 QQ 过程, $a_j$ 是 quantization constant.

对所有 $a_j$ 再做量化. 先减去均值 $\mu$, 再缩放:

$r_j = \frac{a_j - \mu}{b}.$

将 $r_j$ 从 32-bit 量化到 8-bit, 得到 $\widehat r_j$, 再重建:

$\widehat a_j = \mu + b\widehat r_j.$

其中 $\mu$ 和 $b$ 用 32-bit 保存. 最终得到约 4.127 bit/param.

手写原稿