Longformer - The Long-Document Transformer

在前一篇我們介紹了Reformer的架構與概念，主要就是要解決Transformer記憶體、計算量，還有固定Sequence的問題。然而後續在同年4月，AllenAI也提出了一個新架構-LongFormer，主要也是改善了Transformer固定Sequence的問題，進而提出了該架構以利於訓練長句子與文本的nlp任務，今天我們要探討的paper如下:

Longformer: The Long-Document Transformer

Background

先前傳統的Transformer架構，其中的Attention機制是採用Fully-connect，也就是說今天有n個Query，每一個Query都要與其他n個key(包含自己)去計算Attention-score，這造成整理時間複雜度為 $O(n^2)$。

Longformer為了改善這個問題，除了延續原先的RoBERTa的pretrain，也提出了windowed local-context self-attention，以及結合依據下游任務的end-task motivated global-attention，來改善傳統的Full Attention，因此也做到了訓練長句子的任務。這裡先簡單述說一下該兩種Attention的概念：

Windowed local-context self-attention
擁有特定Window size的local attention主要是用來學習上下文的representation
End task motivated global-attention
而global attention則是依據任務預測來建立完整的sequence

補充說明: Sparse Attention

這裡的windowed local-context self-attention其實是一種Sparse Attention，作者這邊是引用了Generating Long Sequences with Sparse Transformers這篇的重點，該篇提到傳統的Attention是屬於Full Connected，如下圖所示:

我們會看到每一個token都與其他的token去做attention的計算，而Sparse Attention認為每個token只要跟其中幾個token做計算即可，不需要全部token去關注。

而Sparse Attention依據Generating Long Sequences with Sparse Transformers這篇paper，我們可以快速整理一個重要的概念：就是 Local attention + Atrous self attention

local attetion
簡單來說就是只關注特定局部的上下文，如下圖所示:
Atrous self attention
使起源於Atrous Convolution，就是只跟相對距離有關，如下圖所示:

從左邊可以看到可能是依照特定倍數來選擇要關注的token，其餘白色區域則視為不關注的token。

所以快速將這兩者的概念整合起來，就會呈現如下的狀況：

除了關注特定的local 上下文特徵之外，同時也可以關注到較遠的相關距離之token，如此一來一方面可以降低計算量，另一方面也兼具可計算較長的sequence。

Longformer - Attention Pattern

作者在該篇論文特別描述了Attention Pattern，如上圖所示，其中圖(a)就是傳統Transformer Attention的時間複雜度，其中深綠色為Query數量，淺綠色為Query關注Key的數量，因此可以發現Full Attention時間複雜度為 $O(n^2)$，接下來來帶各位讀者看一下其他 Attention的運作原理。

Sliding Window

各位讀者若知道Local Attention，想必能很快理解知道該方法的運作原理。簡單來說，每次在輪到該Query的時候，會根據設定的Window size $w$，前後位移 $w$個位置與Keys來視為這是Query所要關注的範圍 $(p+w, -w)$。

這邊我們假設原先的Sequence 長度為 $n$，window size為 $w$，則該pattern的時間複雜度為 $O(wn)$，然而通常w會比n還要來得小(local attention)，作者這邊也有提到，為了平衡計算效率與model能力，每一層的$w$可能會有所不同。

Dilated Sliding Window

該方法可以想像成Sliding Window再延伸，多了一個參數 $d$，該參數的目的主要是決定window size的縫隙，舉例來說，若Window size $w$為2，$d$也為2，則從token開始算起就會變成是 $(p+w+d, p+w-d)$。也正因為多了縫隙 $d$，所以就能關注更長範圍的上下文資訊。

Global + Slide Attention

這邊先描述Global Attention，這與過往傳統一樣，就是與全部的Key做關注。那為什麼Longformer仍需要Global Attention呢？原因在於下游任務，假設是text-classification，如果採用的 windowed local attetion，model本身並不會知道[CLS]，[SEP]的位置，更無法取得完整的資訊。若是Question-Answer則在一篇文章當中，每一個位置都很有可能是答案所在，所以還是有Global Attention比較適合。所以不同的下游任務會有不同Global Attention定義的方法。

這邊我們要注意一點，作者在Attetion的用法有特別提到，在Longformer這個架構當中，作者們對於lower layer是採用較小的window size去做 windowed slide attetion，越往上到top layer的時候Window size會隨之增加，最後到Global Attention。

In particular, we use small window sizes for the lower layers and in- crease window sizes as we move to higher layers.This allows the top layers to learn higher-level rep- resentation of the entire sequence while having the lower layers capture local information

如此一來，就會efficiency(小window size可抓到局部上下文)與performance(Model下游任務表現)做到一個平衡。

Tensor Virtual Machine

為了讓這樣的Sparse Attention可以應用到目前的DL Framework，作者也對CUDA做了優化，根據paper這邊可以發現經過CUDA優化的Sparse self-attention(Longformer)、和原先的Full Attention，以及尚未CUDA優化的Sparse self-attention(Native Pytorch)在時間、空間複雜度的比較，發現Longformer帶來的效果最好。

Experiments

在這篇paper的實驗當中，作者採用的text8和enwik8來作為訓練資料，然而其中比較我認為要注意的是作者採用了Staged training，作者發現在學習更長的句子之前，必須要使Model先大量的梯度更新來學習local的上下文。所以這邊採用了階段式的學習，第一階段先從window size較小與短句子開始做訓練，每往後增加一個階段，window size和句子就增加一倍，然而learning rate做減半。

在這次的實驗當中，作者共採用了5階段做訓練，第一個stage sequence length是2048，最後一個stage sequence length是23040。
作者在實驗的時候也將Model分為small, large:

small model: 12 layers, 512 hidden size.
large model: 30 layers, 512 hidden size.

可以看到在 Longformer small所帶來的效果是好的，但在Longformer large比TransformerXL還好。

接著作者也展示了在不同的Attention Pattern所帶來的效果，如下:

依據上面的實驗結果，increased window size的效果表現最好，decreased效果最差。表格下半部表現了有沒有使用dilated的比較，發現增加dilated所帶來的效果最好。

Pretraining & Finetuning

Longformer的Pretraine是採用Masked LM，然而接續在RoBERTa後來做訓練。每一個layer是window size是採用512來作Attetion。

我們可以看見當Longformer越訓練到後面，其帶來的效果也越來越好，代表Model本身以及懂得如何應用window來做attention。