GaLore

The problem with Lora

Lora给出了一种非常 Parameter Efficient 的方法，通过更新两个额外的低秩矩阵来sft。

$$W' = W + BAx$$

其中，$B \in \mathcal{R}^{m \times r}$，$A \in \mathcal{R}^{r \times n}$，$r \ll min(m, n)$

但是，参数矩阵并不满足低秩假设。根据线性代数的知识，我们知道 $rank(AB) \leq min(rank(A), rank(B))$，因此有:

$$rank(BA) \leq r$$

倘若原矩阵 $rank(W) >r$，则 lora 无论如何也无法很好地近似 full parameter 的更新。

对 $rank(AB) \leq min(rank(A), rank(B))$ 的一种证明:

首先，$rank(A) = dim(Col(A)) = dim(Row(A)) = number(\text{pivots})$

考虑 $AB$ 的每一列 $AB_{[:, j]}$，我们能发现 $AB_{[:, j]} = \sum_{i=0}^{m-1} B_{[i, j]}A_{[:, i]}$

即 $AB$ 的每一列都是 $A$ 的列线性组合

所以 $rank(AB) = dim(Col(AB)) \leq dim(Col(A)) = rank(A)$， 直接转置再应用一下就可以得到完整的结论了。

Low-Rank Property of Weight Gradient

However，Galore 这篇文章指出，虽然参数矩阵不见得是 low rank 的，不过可以证明，对于一类被称为 Reversible network 的网络结构，它们的梯度矩阵是低秩的。

Reversible network

Definition

一个网络 $\mathcal{N}$ 执行映射 $y = \mathcal{N}(x)$ 被称为 Reversible network 当：

• 存在 $L(x, W)$ 使得 $y = L(x, W)x$，并且 $g_x = L^\top(x,W)g_y$

其中 $g_x = \frac{\partial L}{\partial x}$，$g_y = \frac{\partial L}{\partial y}$

常见的不带偏置项bias的线性层，ReLU，leaky ReLU等激活函数，是Reversible network

定义这类网络是为了能够更加定量地分析各层梯度的数值表示，而不至于始终把梯度当作黑箱的形式

性质

• 线性性质：$\alpha_1 \mathcal{N}_1(x) + \alpha_2 \mathcal{N}_2(x)$ 仍然是 Reversible network

• $\mathcal{N}_2(\mathcal{N}_1(x))$ 仍然是 Reversible network

• $\mathcal{N}(x) = \frac{\partial \mathcal{N}(x)}{\partial x} x$

第三个性质可以展开一点说，我们先说一下它的证明：

根据链式法则，我们知道 $g_x = (\frac{\partial y}{\partial x})^\top g_y$

与 $g_x = L^\top(x,W)g_y$ 对比后，我们知道 $(\frac{\partial y}{\partial x})^\top = L^\top(x,W)$，所以由 $y=L(x, W)x$，可得

$$y = \frac{\partial y}{\partial x}x，\text{即} \mathcal{N}(x) = \frac{\partial \mathcal{N}(x)}{\partial x} x$$

这在数学上也是欧拉齐次函数定理的一种形式。欧拉齐次函数定理指出，一个函数 $f(x)$ 是 $m$ 阶齐次函数（即满足 $f(tx) = t^mf(x)$）的充要条件是 $\nabla f(x) \cdot x = mf(x) $。显然，Reversible Network 即是满足 $m=1$ 的一阶齐次函数。

既然 $\mathcal{N}(x)$ 是一阶齐次函数，那么 $\mathcal{N}(x)$ 满足 $\mathcal{N}(tx) = tN(x), \forall t > 0$，对其两边求导，得到：

$$\begin{aligned} t \frac{\text{d}\mathcal{N}(tx)}{\text{d} x} &= t \frac{\text{d}\mathcal{N}(x)}{\text{d} x} \\ K(tx) &= K(x) \end{aligned}$$

所以一阶齐次函数对应的雅可比矩阵 $K(x) = \frac{\partial \mathcal{N}(x)}{\partial x}$ 必然是一个零阶齐次函数。

零阶齐次性并不意味着 $K(x)$ 与 $x$ 无关，有可能存在导函数阶跃的情况。

比如说Relu是一个符合要求的 Reversible Network，它的雅可比矩阵是:

$$H(x)= \left\{ \begin{aligned} 1, \quad & x>0 \\ 0, \quad & x<0 \\ \end{aligned} \right.$$

显然与 $x$ 有关

定理

设 $K_i$ 是第 $i$ 层的雅可比矩阵，也即 $K_l(x) = \frac{\partial N_l(f_{l-1})}{\partial f_{l-1}}$，所以

$$\partial \mathcal{N}(x) = \partial \mathcal{N_L}(\mathcal{N_{L-1}(\cdots \mathcal{N_1}(x))}) = K_L(x)K_{L-1}(x)\cdots K_1(x) \partial x$$

我们可以展示一下，当损失函数为MSE loss时 $\varphi := \frac{1}{2}||y-f_L||_2^2$

$$\begin{aligned} \text{d}\varphi &= (\mathcal{N}(x)-y)\top\text{d}\mathcal{N}(x) \\ &= (\mathcal{N}(x)-y)^\top K_L(x)K_{L-1}(x)\cdots K_{l+1}(x) \text{d}f_l \\ \end{aligned}$$

因为 $f_l = W_lf_{l-1}$，所以 $\text{d}f_l = \text{d}W_lf_{l-1} + W_l\text{d}f_{l-1}$

我们的目标是分析到 $W_l$ 的梯度，因此可以忽略后一部分的微分

$$\begin{aligned} \text{d}\varphi &= (\mathcal{N}(x)-y)^\top K_L(x)K_{L-1}(x)\cdots K_L(x) \text{d}W_lf_{l-1} + \text{与d}W_l\text{无关的部分} \\ \end{aligned}$$

令 $J_l := K_L(x) \cdots K_{l+1}(x)$

由 $\text{d} \mathcal{N}(x) = K_L(x)K_{L-1}(x)\cdots K_{l+1}(x) \text{d}f_l$ 与 $\mathcal{N}(x) = \frac{\partial \mathcal{N}(x)}{\partial x} x$，可得

$$\begin{aligned} \mathcal{N}(x) &= K_L(x)K_{L-1}(x)\cdots K_{l+1}(x) f_l \\ &= J_l W_l f_{l-1} \end{aligned}$$

所以

$$\begin{aligned} \text{d}\varphi &= (J_l W_l f_{l-1}-y)^\top J_l\text{d}W_lf_{l-1} \end{aligned}$$

因为 $\text{d}\varphi$ 是标量，因此我们可以用迹来表示结果。同时利用迹的循环不变性（$\text{tr}(ABC) = \text{tr}(BCA) = \text{tr}(CAB)$），我们可以更方便的调整结果的形式。

$$\begin{aligned} \text{d}\varphi &= \text{tr}((J_l W_l f_{l-1}-y)^\top J_l\text{d}W_lf_{l-1}) \\ &= \text{tr}(f_{l-1}(J_l W_l f_{l-1}-y)^\top J_l\text{d}W_l) \\ \end{aligned}$$

而由矩阵梯度的定义，$(G_l)_{ij} = \frac{\partial \varphi}{\partial (W_l)_{ij}}$，有

$$\begin{aligned} \text{d}\varphi &= \sum_i \sum_j (G_l)_{ij} (\text{d}W_l)_{ij} \\ \end{aligned}$$

这就是 Frobenius 内积，它有一个重要的性质，可以用迹来代替计算

$$\begin{aligned} \text{d}\varphi &= \langle G_l, \text{d}W_l \rangle_F \\ &= \text{tr}(G_l^\top \text{d}W_l) \end{aligned}$$

进行比较，我们可以得到：

$$\begin{aligned} G_l^\top &= f_{l-1}(J_l W_l f_{l-1}-y)^\top J_l \\ G_l &= J_l^\top (J_l W_l f_{l-1}-y) f_{l-1}^\top \\ &= J_l^\top J_l W_l f_{l-1} f_{l-1}^\top - J_l^\top y f_{l-1}^\top \end{aligned}$$

考虑到梯度 $G_l$ 的方向是向上的，而常说的梯度下降的方向是相反的，因此我们记 $\hat{G_l} = -G_l = J_l^\top y f_{l-1}^\top - J_l^\top J_l W_l f_{l-1} f_{l-1}^\top$ 表示真实的梯度下降的方向。

其它损失函数也一样，变化一下最后一层的梯度就可以了。

突然发现 MSE loss 和 CE loss 的梯度都是 $\mathcal{N}(x)-y$ 啊

作者也推导了softmax层的结果，得到的结果是 $\hat{G_l} = (J_lP_1^{\bot}y-\gamma K^{-1}J_l^\top P_1^{\bot}J_lW_lf_{l-1})f_{l-1}^\top$，其中 $P_{1}^\bot := I - \frac{1}{K} 1 1^\top$，$K$ 是softmax的维度

总而言之，言而总之，除了 Attention 层不满足 Reverse Network 的定义（Attention层有 $\mathcal{N}_1(x) \cdot \mathcal{N}_2(x)$ 这样的积性操作），若我们暂且不考虑 Attention 层，则剩下的所有层的梯度似乎都可以表示为统一形式：

$$\hat{G_l} = A - BW_lC$$

Positive Semi-definite

对于 $\hat{G_l} = J_l^\top y f_{l-1}^\top - J_l^\top J_l W_l f_{l-1} f_{l-1}^\top$ 的形式

$$\left\{ \begin{aligned} B &= J_l^\top J_l \\ C &= f_{l-1} f_{l-1}^\top \end{aligned} \right.$$

而根据线性代数的知识，我们知道这其实就是 $B$ 和 $C$ 都是半正定(Positive Semi-definite, PSD)矩阵的充要条件，因为 $x^\top B^\top B x = (Bx)^\top Bx = ||Bx||^2 \geq 0$

而到了softmax层的结果，我们只需额外证明 $J_l^\top P_1^{\bot}J_l$ 是一个半正定矩阵。根据合同变换，我们知道如果一个矩阵 $A$ 是半正定矩阵，那对任意矩阵 $P$，$P^\top AP$ 也是半正定的。所以我们只要证明 $P_1^{\bot}$ 是一个PSD即可

我们先来求 $Y = \mathbf{1}\mathbf{1}^T$ 的特征值。因为这是一个全1矩阵，也就意味着秩为1，因此它必然有 $K-1$ 个数值为0的特征值。

剩下一个我们简单地根据求特征值定义 $Ax = \lambda x$，能求得最后一个特征值为 $K$

我们要求 $P = I - \frac{1}{K}Y$ 的特征值，根据定义：

$$Pv = (I-\frac{1}{K}Y)v = Iv - \frac{1}{K}(Yv) = v - \frac{1}{K}(\lambda_Y v) = (1-\frac{\lambda_Y}{K})v$$

所以 $P$ 和 $Y$ 的特征值存在一个简单的关系: $\lambda_P = (1-\frac{\lambda_Y}{K})$，代入可求得 $P$ 的特征值有1个0和 $K-1$ 个1

所有特征值非负，所以 $P$ 是一个半正定矩阵。

综上，我们讨论的Large Language Model中的梯度，表达式中的 $B$ 和 $C$ 总是半正定的