机器学习笔记(7.3): FlyLoRA

可能更好的阅读体验 Jeefy's Blog: FlyLoRA

基于果蝇的生物学原理，设计出来的这样一个框架。

其核心的启发可以从两个方向展开：

• 生物学逻辑（这里略）
• 基于实验的数学逻辑

一个核心的观察是，对于微调的情况下，多 expert 是优于高 rank 的，也就是说，微调学习的信息是放在高维度就是冗余的，并且对于不同的任务，这个表示之间可能存在难以在相同空间融合的张力。

这也是为什么 Split-LoRA 很可能比单独的 LoRA 表现更好的原因。

那么为什么 Top-k 就好于全量呢？

首先先说为什么 Top-K 是合理的。如图 (a)，对于 LoRA，可以发现激活的秩的分布是很不均衡的，所以可以用 Top-K 拟合全量。
也可以从 SVD 分解的角度来说，对于一个矩阵可以对角化分解为：\(\sum \sigma_i u_i^T v_i\)，那么取 top-K 的 \(\sigma_i\)，其实就可以拟合整个矩阵了。

然后说说为什么这里的 Top-k 会比全量好。通过从梯度的协方差看，可以发现 LoRA 的协方差分布其实除了对角线还是很大的。也就是说，梯度在不同秩之间的关联还是很大的，根据后文的有效秩和信息坍缩理论分析，这里的相关性会使得能表示的信息远低于上限，每次学习到的内容会相互影响。正相关，则会导致表示冗余；负相关，则会导致互相抵消。
而观察 FlyLoRA 的协方差矩阵，可以发现解耦合能力非常好，每个 rank 占据独立的功能方向，参数利用率高。这正符合专家特化的期望，所以采用 Top-K 确实会合理，而且学习效率高于 LoRA。
这个也可以继续关联到 SVD 的核心思想上：将一个矩阵拆分为若干正交方向上表示。这里非对角线协方差接近 0，可以对应这个架构本身接近于将每个秩正交化，更新互不干扰。
这里其实可以做一个新的图，看看 Split-LoRA 的协方差，应该会介于 (b) 和 (c) 之间。因为 Split-LoRA 实际上没有保证路由的这个 W 保持正交，会导致秩之间的梯度互相影响。

补充：协方差衡量的是 “偏离均值的共动” \(Cov(X, Y) = E[(X - E[x]) (Y - E[Y])]\)
协方差为正，表示通向联动；协方差为负，表示反向联动；协方差为零，则表示几乎无关。一般来说，需要通过 \(Cov(X, Y) \div (\sigma_X \sigma_Y)\) 归一化为相关性。

但是 Split-LoRA 最大的问题是，需要训练一个 Gateway，进行 \(top-k\) 选择。这部分的开销是很大的。

\[\begin{aligned} f(x) &= W_0 x + \frac {\alpha} r \sum_{i = 1}^N G(x)_i B_i A_i x \\ G(x) &= {\text top-}k(W_g x) \\ \end{aligned} \]

那么有没有办法天然的进行 gateway 呢？那就是 FlyLoRA 中的这个固定的稀疏矩阵 A，将信息天然映射到多个低秩空间中，这样自然的将信息进行了 gateway 分组。

这样还有一个好处，那就是路由不仅仅是 loss-free 的，而且还不会引入额外的计算。保持了稀疏性带来的高效计算的能力。

Fixed Sparse Random Projection as Implicit Router

换句话来说，现在是：

\[G(x) = {\rm top-}k(y) = {\rm top-}k(Ax) \]

其中 \(A \in \mathbb R^{r\times n}\) 的每一行只有 p 个非零的，独立从 \(N(0, \frac {1}{r^2})\) 采样的元素。

只有这 \(top-k\) 个元素对应的单秩 expert 会被激活：

\[f(x) = B \times G(x) \]

但是类比 MoE 的思路，我们还需要一个负载均衡的措施，使得专家被更均匀的分配。

具体来说，就是维护一个偏置：

\[f(x) = B \times {\rm top-}k(Ax + d) \]

其中 \(d_i \leftarrow d_i + u \cdot {\rm sign}(\bar c_i - c_i)\)，其中 \(u\) 是学习率，\(\bar c_i\) 是期望激活率。

---

为什么要有负载均衡。可以从如下角度解释。

1. 过拟合和路由坍缩

MoE 的路由存在：被选择次数越多 \(\to\) 反向梯度越多 \(\to\) 拟合越好 \(\to\) 选择次数越多 的正反馈循环，会导致路由坍缩的问题。这也是 Switch Transformer 中会出现的一个问题。
在 FlyLoRA 中，这个叙述需要被修改为：分布是静态的，只依赖 A，不依赖 B，范数的差异会导致始终路由到某几行，其他的始终保持初始化。也就是路由坍缩的情况。

2. 有效秩和信号坍缩

如果所有信号都涌向同一个秩，那么由于能表示的信息有限，必定存在更多的压缩，导致信号无用。而其他的部分不会被激活，这个有效的秩自然就会小于期望的秩，导致信号被学习的不完整。

3. 专家特化和初始噪声

当只有一个专家在某个领域时，特化自然不会发生；当有多个专家在某个领域时，那么自然会分化为更细的领域。而且在初始化时，可能存在某个秩投影天然就对数据分布亲和，出来的权值很大，所以需要增加偏置修正其权值分布。

---

一些理论分析。

• Theorem 3.1

这样构造的 \(A\)，使得：

也就是说，这个路由和利用欧式距离路由的方法非常的相似，每一个秩管辖了一片邻域，而一片邻域可以认为语义相似，所需要的能力相似。

• Theorem 3.3

top-k （\(k \lt r\)）的路由，使得梯度更新的协方差以 \(\frac {k^2}{r^2}\) 的关系压缩，也就是说，\(k\) 越小，梯度协方差压缩越强。
核心证明就是因为掩码的存在，每次更新都会强制让无关的被忽略，不会贡献协方差。

• Theorem 3.4

本文的另一个核心贡献：不同任务上的 FlyLoRA 合并后的表现，依然很不错。

\[W' = W_0 + \frac 1 t \sum_{i = 1}^t B_i A_i \]

为什么会很不错？因为 A 的近似正交性，使得多个 LoRA 组件之间的梯度几乎没有重合。
并且这个冻结的 A 之间，基本上也是正交的：

\[{\mathbb E}[A_i A_j] = 0 \]

并且声称 \(n\) 越大，正交的概率越大。Polynomially decaying correlations：

\[{\mathbb P}(\| A_i A_j^T \|_2 \ge \epsilon r) \le \frac {p^2} {n r^2 \epsilon^2} \]

并且呢，基于这个正交性：

\[\langle B_i A_i, B_j A_j \rangle_F \approx 0 \]

也就是说，这个解耦性非常强，两个 LoRA 对于激活的更新近似正交。

为什么正交了，就没啥抵消了呢？
关键在于正交 \(\to\) 范数可分解 \(\to\) 更新不重叠、不抵消

那为什么能力提升也是线性叠加的？
考虑位移足够小的情况下，在参数空间中更新就是近线性的，可以通过线性相加得到。
如果线性近似不满足（更新太大）那么直接相加就是不可行的，所以才是 “微调” 时使用。