注意力比矩阵分解更好吗?

0. TL;DR

本文提出了一种基于矩阵分解(Matrix Decomposition, MD)的新型全局上下文模块——Hamburger。它通过将全局上下文建模问题转化为低秩恢复问题,并利用矩阵分解的优化算法设计网络架构,从而在性能和计算成本上与自注意力机制相媲美。实验表明,Hamburger在语义分割和图像生成任务中均优于现有的自注意力模块,并且具有更低的计算复杂度和内存占用。

1. 背景介绍

在深度学习中,自注意力机制(尤其是Transformer架构)已经成为捕捉长距离依赖关系的重要工具。然而,自注意力机制是否是建模全局上下文的唯一选择?本文提出了一个有趣的观点:传统的矩阵分解模型在建模全局上下文时可能比自注意力机制更高效。

矩阵分解是一种经典的数学工具,用于将复杂数据分解为若干子矩阵,从而提取数据的潜在结构。本文将矩阵分解与深度学习结合,提出了Hamburger模块。它通过优化算法将输入表示分解为低秩信号子空间,从而去除冗余信息并增强全局上下文。这种方法不仅在理论上具有创新性,而且在实际应用中也表现出色。

2. Hamburger模块

矩阵分解的核心思想是将输入矩阵分解为若干子矩阵的乘积。例如,给定一个矩阵$X∈R^{d×n}$,矩阵分解的目标是找到一个字典矩阵$D∈R^{d×r}$和对应的编码矩阵$C∈R^{r×n}$,使得$X≈DC$。其中,$r≪\min(d,n)$表示低秩约束。

低秩恢复问题可以表示为:

\[\mathop{\min}_{D,C}\|X-DC\|_{F}^{2}+R_{1}(D)+R_{2}(C)\]

其中,$R_{1}(D)$和$R_{2}(C)$是字典和编码的正则化项。

Hamburger模块由两部分组成:线性变换和矩阵分解。具体来说,输入$Z\in R^{d_z×n}$首先通过线性变换$W_l$映射到特征空间,然后通过矩阵分解$M$提取低秩信号子空间,最后通过另一个线性变换$W_u$输出结果。整个过程可以表示为:

\[H(Z)=W_uM(W_lZ)\]

矩阵分解部分可以采用不同的优化算法,例如向量量化(VQ)、概念分解(CD)和非负矩阵分解(NMF)。这些算法通过迭代优化来分解输入矩阵,并恢复低秩信号。

由于矩阵分解的优化算法是迭代的,直接使用反向传播(BPTT)会导致梯度不稳定。为此,本文提出了一种一阶梯度方法,仅使用优化算法的最后一步来近似反向传播。这种方法不仅大大降低了计算复杂度,还避免了梯度爆炸或消失的问题。

3. 实验分析

本文在两个视觉任务上进行了实验:语义分割和图像生成。具体数据集包括:

评估指标包括:

HamburgerPASCAL VOCPASCAL Context数据集上均取得了优异的性能。不仅优于现有的自注意力模块,还在计算复杂度和内存占用上具有显著优势。

实验结果表明,Hamburger在建模全局上下文时具有以下优势:

  1. 计算效率高:矩阵分解的复杂度为 $O(ndr)$,远低于自注意力的 $O(n^2d)$。
  2. 内存友好:避免了生成大矩阵作为中间变量。
  3. 可解释性强:通过低秩恢复和去噪,Hamburger能够学习到更清晰的全局上下文。