[论文解读] From block-Toeplitz matrices to differential equations on graphs: towards a general theory for scalable masked Transformers

Krzysztof Choromański, Han Lin+8 more|arXiv (Cornell University)|Jul 16, 2021

Quantum optics and atomic interactions被引用 7

一句话总结

本文通过将掩码建模为图上的拓扑调制，提出了一种用于可扩展掩码注意力的通用理论框架，利用分块托普利茨矩阵结构和快速算法实现次二次计算。该框架统一并扩展了先前的方法（如线性因果注意力和对数线性 RPE-注意力），引入了新型可扩展机制，适用于 d 维 RPE 和图核掩码，并具备可证明的效率。

ABSTRACT

In this paper we provide, to the best of our knowledge, the first comprehensive approach for incorporating various masking mechanisms into Transformers architectures in a scalable way. We show that recent results on linear causal attention (Choromanski et al., 2021) and log-linear RPE-attention (Luo et al., 2021) are special cases of this general mechanism. However by casting the problem as a topological (graph-based) modulation of unmasked attention, we obtain several results unknown before, including efficient d-dimensional RPE-masking and graph-kernel masking. We leverage many mathematical techniques ranging from spectral analysis through dynamic programming and random walks to new algorithms for solving Markov processes on graphs. We provide a corresponding empirical evaluation.

研究动机与目标

• 开发一种通用且可扩展的理论，用于在不引入二次复杂度的前提下，将多种掩码机制整合到 Transformer 架构中。
• 在单一理论框架下统一现有的高效注意力机制（如线性因果注意力和对数线性 RPE-注意力）。
• 将可扩展注意力的适用范围扩展至新场景，包括 d 维位置编码和基于图的核掩码。
• 基于图上的谱分析、动态规划和马尔可夫过程，为高效掩码注意力计算提供数学基础。
• 通过新模型 Graph Kernel Attention Transformer (GKAT) 的实证验证，对比当前最先进的 GNN 模型。

提出的方法

• 将掩码形式化为未掩码注意力的拓扑（基于图的）调制，其中掩码由图结构或距离函数导出。
• 通过掩码矩阵支持次二次矩阵-向量乘法的性质，表征高效掩码，从而实现可扩展的低秩注意力。
• 利用 d 维网格（如图像、视频）中的分块托普利茨结构，将相对位置编码推广至 1D 序列之外的场景。
• 应用快速傅里叶变换（FFT）和汉克尔矩阵技术，加速基于图的掩码的矩阵-向量乘积，尤其适用于最短路径和扩散核。
• 设计基于树的分治递归算法，通过树结构图的自底向上和自顶向下遍历，高效计算掩码注意力。
• 通过建模修改后的图拓扑和在树分解上的高效计算，将填充和序列打包整合到框架中。

实验结果

研究问题

• RQ1 在何种条件下，掩码机制可被整合到注意力机制中并实现次二次时间复杂度？
• RQ2 能否利用结构化矩阵性质，使 d 维相对位置编码在图像和视频建模中实现可扩展性？
• RQ3 基于图核的注意力掩码（如最短路径或扩散核）能否在低秩注意力设置中实现高效计算？
• RQ4 能否将现有的高效注意力机制（如因果注意力和对数线性 RPE）统一在一个单一的理论框架下？
• RQ5 能否从该理论中推导出一种新型可扩展注意力机制——Graph Kernel Attention Transformer (GKAT)，并优于现有 GNN 模型？

主要发现

• 本文建立了可扩展掩码注意力的一般条件：若掩码矩阵支持次二次矩阵-向量乘法，则掩码低秩注意力可实现次二次时间复杂度。
• 该框架将 1D RPE 和因果注意力作为 d 维 RPE 和基于图的掩码的特例，使 2D 和 3D 输入（如图像和视频）的高效注意力成为可能。
• 提出了一种新颖的 d 维 RPE 机制，通过 FFT 实现快速计算，支持视觉和视频任务中的可扩展注意力。
• 基于最短路径距离和扩散核的图核掩码，通过 FFT 和树分解技术被证明可高效计算。
• 针对树结构图的所提算法在计算掩码注意力时达到 $O(|V( au)| imes \text{diam}( au))$ 时间复杂度，结合递归分解与循环移位操作。
• Graph Kernel Attention Transformer (GKAT) 模型在实证评估中优于九种 SOTA GNN 模型，验证了所提框架的实际有效性。