Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method

Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method [1] とは、セルフアテンションの計算コスト削減に Nyström 近似を応用した研究である。NeurIPS 2021 で発表された。なお、この研究に先立ちセルフアテンションの計算コスト削減に Nyström 近似を応用した研究として Nyströmformer [Xiong et al., 2021] があるが、Nyströmformer ではセルフアテンション行列（半正定値行列ではない）に直接 Nyström 近似を適用していたのに対し、本研究では セルフアテンション行列を半正定値行列に拡張してから Nyström 近似を適用しているのが異なる。


参考文献

  1. Yifan Chen, Qi Zeng, Heng Ji, Yun Yang. Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method. In NeurIPS 2021. [Proceedings]
  2. https://github.com/pkuzengqi/Skyformer（2022年3月29日参照）.
    • モデルの実装のリポジトリである。動くコードがあるのと共に、プレ処理済み Long Range Arena へのリンクがある。
  3. 福水 健次. カーネル法入門. 朝倉書店, 2010.
    • 15ページ（初版第7刷）に正定値カーネルの例がある。
  4. https://cookie-box.hatenablog.com/entry/2022/03/31/125645（2022年3月31日参照）.
    • [1] の内容についての筆者の理解を記述している。

アイデア

正定値カーネルでもなければ対称でもないソフトマックスによるセルフアテンションを、ガウシアンカーネルに置き換えて対称行列に拡張してカーネル法の土俵に持ち込んだ上で Nyström 近似を適用している。

近似誤差の議論

元のセルフアテンション行列とそれを近似したセルフアテンション行列の差の行列のスペクトルノルムが上から抑えられることを示している（Theorem 2）。

コード

リポジトリ [2] をクローンして以下を実行することができる。 
from models.model_LRA import ModelForSC, ModelForSCDual
from config import Config

model_config = Config["lra-text"]["model"]
model_config["mixed_precision"] = True
model_config["attn_type"] = "softmax"
model = ModelForSC(model_config)
print(model)

import torch

x = torch.tensor([[0, 1, 2, 3, 4]])
label = torch.tensor([0])
y = model(x, None, label)
print(y)
Transformer