Hybrid Random Features：利用Bagging的思路降低Kernel法的误差

论文地址：

Hybrid Random Features

说明

由于工作的原因平时要看不少论文，加上经常看苏神对论文的解析，所以决定自己也开始写论文解析，每篇解析主要分为以下几个部分：

• 论文的整体思想概述；
• 背景知识（可选）；
• 论文的主体思想，忽略不必要的细节；
• 对该方法的分析和讨论（可选）；

整体概述

Hybrid Random Features是Rethinking Attention with Performers的续作，其主要思想是：

基于蒙特卡洛的Kernel法对Softmax的近似是有误差的，利用Bagging可以减少这个误差。

背景知识

这部分对所需要的背景知识进行介绍。

Attention回顾

首先回顾attention的核心计算公式：

输入：

$$Q\in \mathbb R^{n\times d},K \in \mathbb R^{n\times d} ,V \in \mathbb R^{n\times d}$$

输出：

$$\begin{aligned} O&= \mathrm{SoftMax}\left(Q K^\top \right) V \end{aligned}$$

备注：

• 为了方便讨论，这里忽略了多头的部分；
• SoftMax部分没有除以$\sqrt{d_1}$；
• 只考虑SelfAttention的情形；

时间复杂度分析：

• $Q K^\top :(n, d) ,(d , n) \to (n, n) $
  • 时间复杂度为$O(n^2d)$
• $\mathrm{SoftMax}\left(Q K^\top \right) : (n, n) \to (n, n)$
  • 时间复杂度为$O(n^2)$
• $\mathrm{SoftMax}\left(Q K^\top \right) V : (n, n), (n, d)\to (n, d)$
  • 时间复杂度为$O(n^2 d)$
• 总时间复杂度为：
  • $O(n^2 d)$

可以看到，关于序列长度的二次复杂度成为速度的主要瓶颈，所以一个主要问题就是如何解决这点。

左乘变右乘

现在换个计算方式，假设利用映射$f$将$Q,K$变换为：

$$Q \overset \phi\to Q_1 \in \mathbb R^{n\times d_1}, K\overset \phi \to K_1 \in \mathbb R^{n\times d_1}$$

然后在该变换下，有下式成立：

$$\begin{aligned} O&= \mathrm{SoftMax}\left(Q K^\top \right) V \\ & \approx M Q_1 K_1 ^{\top} V\\ &=M (Q_1 (K_1^{\top} V)) \end{aligned}$$

其中

$$M= \mathrm{diag}\left( 1_n^{\top} Q_1 K_1^T\right)^{-1} \in \mathbb R^{n\times n}$$

为归一化矩阵，类似于SoftMax中的分母部分，注意该矩阵为对角矩阵。

那么此时的时间复杂度为：

• $ K_1^\top V:(d_1 , n), (n, d_1) \to (d_1, d_1) $
  • 时间复杂度为$O(d_1^2n)$
• $(Q_1 (K_1^{\top} V)) : (n, d_1), (d_1, d_1) \to (n, d_1)$
  • 时间复杂度为$O(d_1^2 n)$
• $ 1_n^{\top} Q_1 : (1,n),(n, d_1)\to (1, d_1)$
  • 时间复杂度为$O(d_1 n)$
• $ 1_n^{\top} Q_1 K_1^{\top} : (1, d_1),(d_1, n)\to (1,n)$
  • 时间复杂度为$O(d_1 n)$
• $M (Q_1 (K_1^{\top} V)): (n,n), (n, d_1)\to (n,d_1)$
  • 注意尽管$M$为对角矩阵，所以该操作只要对$(Q_1 (K_1^{\top} V))$的第$i$行乘以$M_{ii}$即可，从而时间复杂度为$O(d_1n)$
• 总时间复杂度为：
  • $O(d_1^2 n)$

所以，利用该方法可以将关于序列长度的二次复杂度降低到一次，所以问题的关键是选择怎样的映射$\phi$。

注意到如果

$$\begin{equation} \phi(Q)^{\top} \phi(K) =Q_1^{\top} K_1\approx \exp\left(Q^{\top } K \right) \tag 1 \end{equation}$$

那么显然就有

$$M (Q_1 (K_1^{\top} V)) \approx \mathrm{SoftMax}\left(Q K^\top \right) V$$

所以一个常见的思路就是找到$f$，使得公式(1)成立。

备注：这种方法也被称为Kernel法。

Performer

Performer的思路就是基于公式(1)，其选择的映射为：

$$\phi(u) = \frac{1}{\sqrt{2 m}} \exp \left(-\frac{\|{u}\|^{2}}{2}\right)\left(\exp \left(\omega_{1}^{\top} {u}\right), \ldots, \exp \left(\omega_{m}^{\top} {u}\right), \exp \left(-\omega_{1}^{\top} {u}\right), \ldots, \exp \left(-\omega_{m}^{\top} {u}\right)\right)^{\top}$$

其中$w_i\sim \mathcal N(0, I_d)$，为iid的多维标准正态分布。

在该论文中，作者证明了无偏估计性

$$\mathbb E_{w\sim\mathcal N(0, I_d)} \left[\phi(Q)^{\top}\phi(K) \right] =\exp\left( Q^{\top } K\right) \tag 2$$

所以利用蒙特卡洛的思想，可以选择适当的$m$，那么就有

$$\phi(Q)^{\top} \phi(K) \approx \exp\left( Q^{\top } K\right)$$

该方法的一个问题是，由于是蒙特卡洛模拟，所以必然有误差，因此一个核心的问题是如何减少误差，原始论文中给出一个方法，在Hybrid Random Features中作者又给出另一个方法——即Bagging。

主体思想

三个核心点

这篇论文有三个核心点：

• Bagging后依然是无偏估计；
• Bagging后的结果依然可以用内积表示；
• Bagging后的结果某种程度上的误差要降低；

细节

无偏估计

由于是使用Bagging的思路，所以必然是许多映射的集成，假设基映射$\mathrm{SM}_k$为

$$\mathrm{SM}_k: \mathbb R^{d} \times \mathbb R^{d} \to \mathbb R, i=1,\ldots ,p+1$$

Bagging的映射为

$${\mathrm{ SM}}(x,y) = \sum_{k=1}^p \lambda_k (x,y) \mathrm{SM}_k (x,y) + \left(1- \sum_{k=1}^p \lambda_k (x,y)\right) \mathrm{SM}_{p+1} (x,y)$$

其中$\lambda_k(x,y)$为系数，形式如下

$$\lambda_k: \mathbb R^{d} \times \mathbb R^{d} \to \mathbb R, k=1,\ldots ,p+1$$

现在假设$\mathrm{SM}_k (x,y)$可以写成内积的期望（无偏估计），即

$$\mathrm{SM}_k (x,y) =\mathbb E_{w \sim \mathcal D} \left[\phi_k (x)^{\top} \phi_k (y) \right] \tag 3$$

记

$$\begin{aligned} \widehat {\mathrm{SM}}_i (x,y) &= \phi_i (x)^{\top} \phi_i (y) \\ {\mathrm{SM}_1}(x,y) &= \sum_{k=1}^p \lambda_k (x,y)\widehat {\mathrm{SM}}_k (x,y) + \left(1- \sum_{k=1}^p \lambda_k (x,y)\right)\widehat {\mathrm{SM}}_{p+1} (x,y) \end{aligned} \tag 4$$

那么

$$\begin{aligned} \mathbb E_{w\sim \mathcal D}[{\mathrm{SM}_1}(x,y)] &= \sum_{k=1}^p \lambda_k (x,y)\mathbb E_{w\sim \mathcal D} \left[\widehat {\mathrm{SM}}_k (x,y) \right] + \left(1- \sum_{k=1}^p \lambda_k (x,y)\right)\mathbb E_{w\sim \mathcal D}\left[\widehat {\mathrm{SM}}_{p+1} (x,y)\right]\\ &\overset{(3),(4)} = \sum_{k=1}^p \lambda_k (x,y) \mathrm{SM}_k (x,y) + \left(1- \sum_{k=1}^p \lambda_k (x,y)\right) \mathrm{SM}_{p+1} (x,y)\\ &={\mathrm{ SM}}(x,y) \end{aligned}$$

因此${\mathrm{SM}_1}(x,y)$是$\mathrm {SM}(x,y)$的无偏估计。

内积表示

现在无偏性解决了，那么问题就来到是否可以用内积表示？对于一些特殊的$\lambda_k$，这点是可以满足的，比如

$$\lambda_k (x,y) =\frac{1}{p+1} ,k=1,\ldots, p$$

那么此时

$$\begin{aligned} \Phi(x)&\triangleq\frac{1}{\sqrt{p+1}}[\phi_1(x)^{\top}, \phi_2(x)^{\top},\ldots, \phi_{p+1}(x)^{\top}]^{\top} \\ \Phi(x)^{\top} \Phi(y) &= \frac 1 {p+1}\sum_{k=1}^{p+1} \phi_k(x)^{\top } \phi_k(y)\\ &= {\mathrm{SM}_1}(x,y) \end{aligned}$$

这个例子显然太特殊了，在论文中作者给出了一个充分条件（这里忽略了常数）：

$$\lambda_k(x,y) =\mathbb E_{w'\sim \mathcal D'} \left[\xi_k (x)^{\top} \eta_k(y) \right] \tag 5$$

记

$$\begin{aligned} \widehat \lambda_k(x,y) &= \xi_k (x)^{\top} \eta_k(y) \\ \widehat{\mathrm{SM}}(x,y) &= \sum_{k=1}^p \widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y) + \left(1- \sum_{k=1}^p \widehat \lambda_k(x,y)\right)\widehat {\mathrm{SM}}_{p+1} (x,y) \end{aligned} \tag 6$$

那么

$$\begin{aligned} \mathbb E_{w\sim \mathcal D,w'\sim \mathcal D'}[\widehat{ \mathrm{SM}}(x,y)] &= \sum_{k=1}^p \mathbb E_{w'\sim \mathcal D'} \left[\widehat \lambda_k(x,y) \right] \mathbb E_{w\sim \mathcal D} \left[\widehat {\mathrm{SM}}_k (x,y) \right] + \left(1- \sum_{k=1}^p \mathbb E_{w'\sim \mathcal D'} \left[\widehat \lambda_k(x,y) \right]\right)\mathbb E_{w\sim \mathcal D}\left[\widehat {\mathrm{SM}}_{p+1} (x,y)\right]\\ &\overset{(3)\sim(6)} = \sum_{k=1}^p \lambda_k (x,y) \mathrm{SM}_k (x,y) + \left(1- \sum_{k=1}^p \lambda_k (x,y)\right) \mathrm{SM}_{p+1} (x,y)\\ &={\mathrm{ SM}}(x,y) \end{aligned}$$

所以这种情形依然保持了无偏估计的特性，下一步只要说明可以用内积表示即可，证明分为两步：

• 第一步，证明$\widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y)$可以用内积表示；
• 第二步，证明$\widehat{\mathrm{SM}}(x,y)$可以用内积表示；

第一步

这里通过两个引理的方式给出证明。

引理1

假设$A,B\in \mathbb R^{m\times n}$，那么存在映射

$$f: \mathbb R^{m\times n} \to \mathbb R^{l}$$

满足

$$f(A)^{\top} f(B)=\mathrm{trace}\left(B^{\top} A\right)$$

证明：

$$\begin{aligned} \mathrm{trace}\left(B^{T} A\right) &= \sum_{i=1}^m \sum_{j=1}^n B_{ij} A_{ij} \end{aligned}$$

那么

$$\begin{aligned} l &=mn\\ f(A)_i &= A_{[i/m] [i\%m]} \end{aligned}$$

即$f$是将$A$按行拼接为一个长向量。

引理2

对于映射

$$\begin{aligned} \eta_i :\mathbb R^{d} &\to \mathbb R^{d_1}, i =1,2\\ \phi_i :\mathbb R^{d} &\to \mathbb R^{d_2}, i =1,2 \end{aligned}$$

存在映射

$$\begin{aligned} \psi_i :\mathbb R^{d} &\to \mathbb R^{d_3}, i =1,2 \end{aligned}$$

满足

$$\begin{aligned} \left(\eta_1(x)^\top \eta_2(y) \right)\left( \phi_1 (x)^\top \phi_2(y) \right) &=\psi_1(x)^\top \psi_2(y) \end{aligned}$$

证明：

$$\begin{aligned} \left(\eta_1(x)^\top \eta_2(y) \right)\left( \phi_1 (x)^\top \phi_2(y) \right) &=\left(\eta_1(y)^\top \eta_2(x) \right)\left( \phi_1 (x)^\top \phi_2(y) \right)\\ &=\mathrm{Trace} \left(\eta_1(y)^\top \eta_2(x) \phi_1 (x)^\top \phi_2(y) \right)\\ &= \mathrm{Trace} \left( \left(\phi_2(y)\eta_1(y)^\top\right) \left(\eta_2(x) \phi_1 (x)^\top \right) \right) \\ &\overset{引理1}= f(x)^\top f(y) \end{aligned}$$

完整证明

注意到

$$\widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y) = \xi_k (x)^{\top} \eta_k(y) \phi_k(x)^{\top} \phi_k (y)$$

所以根据引理2，存在$f_k$，使得

$$\widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y) = f_k(x)^{\top} f_k(y) \tag 7$$

第二步

根据(7)，我们有

$$\begin{aligned} \widehat{\mathrm{SM}}(x,y) &= \sum_{k=1}^p \widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y) + \left(1- \sum_{k=1}^p \widehat \lambda_k(x,y)\right)\widehat {\mathrm{SM}}_{p+1} (x,y)\\ &= \sum_{k=1}^p \widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_k(x,y)-\sum_{k=1}^p \widehat \lambda_k(x,y)\widehat {\mathrm{SM}}_{p+1} (x,y) + \widehat {\mathrm{SM}}_{p+1} (x,y)\\ &= \sum_{k=1}^p f_k(x)^{\top} f_k(y) - \sum_{k=1}^p f_{k,p+1}(x)^{\top}f_{k,p+1}(y) + \phi_{p+1} (x)^{\top} \phi_{p+1} (y) \end{aligned}$$

记

$$\begin{aligned} \Psi(x)&=[f_1(x)^{\top},\ldots,f_p(x)^{\top},if_{1,p+1}(x)^{\top},\ldots,if_{p,p+1}(x)^{\top}, \phi_{p+1}(x)^{\top}]^{\top} \end{aligned}$$

那么

$$\widehat{\mathrm{SM}}(x,y) = \Psi(x)^{\top} \Psi(y)$$

其中

$$i^2=-1$$

降低误差

对于一般的情形，是很难分析出误差的，所以这里作者给了一个特殊的例子，此处$p=1$，两个基本映射为：

$$\begin{equation} \begin{aligned} \phi_{m}^{\mathrm{trig}}({u})&= \frac{1}{\sqrt{m}} \exp \left(\frac{\|{u}\|^{2}}{2}\right)\left(\sin \left(\omega_{1}^{\top} {u}\right), \ldots, \sin \left(\omega_{m}^{\top} {u}\right), \cos \left(\omega_{1}^{\top} {u}\right), \ldots, \cos \left(\omega_{m}^{\top} {u}\right)\right)^{\top} \\ \phi_{m}^{++}({u}) &= \frac{1}{\sqrt{2 m}} \exp \left(-\frac{\|{u}\|^{2}}{2}\right)\left(\exp \left(\omega_{1}^{\top} {u}\right), \ldots, \exp \left(\omega_{m}^{\top} {u}\right), \exp \left(-\omega_{1}^{\top} {u}\right), \ldots, \exp \left(-\omega_{m}^{\top} {u}\right)\right)^{\top} \end{aligned} \tag 8 \end{equation}$$

相应的

$$\xi (z)= \eta(z) =\frac {i}{\sqrt{2 n}} \left(\mathrm{sgn}\left(\tau_{1}^{\top} {z}\right), \ldots, \mathrm{sgn}\left(\tau_{n}^{\top} {z}\right)\right)^{\top}$$

其中

$$\tau_k \sim \mathcal N(0, I_d),k=1,\ldots, d\\ i^2 =-1$$

作者证明了，在这种情形下，Bagging后方法的最大相对误差比单用$\phi_{m}^{\mathrm{trig}}$和$\phi_{m}^{++}$的最大相对误差要低，这也提供了理论保证，这部分见论文第三节。

分析和讨论

这个方法整体来说还是比较清晰，但是作者书写的时候公式符号不太好理解，所以对阅读带来了一定的困难。

还有就是$\xi,\eta$的选择似乎是凑的，不知道是否有更好的方式来选择。