# Reformer: The Efficient Transformer

## Motiváció

A Transformer architektúra egyik hátránya, hogy hosszú szekvenciákon (NLP esetében szövegeken) rendkívül idő- és memóriaigényes. Erre a Reformer a következő megoldásokat kínálja:
* Attention közelítő kiszámítása LSH-val (locality-sensitive hashing)
* Megfordítható rétegek (reversible layers)
* Feed-forward rétegek chunkolása

## Áttekintés: Transformer


A Transformer (Vaswani et al. 2017) eredetileg NLP-feladatokra kifejlesztett encoder-decoder architektúra. Minden Transformer réteg egy-egy **multi-head attention** és **feed-forward** alrétegből áll. Részletek:
* Normalizált residual alrétegek
* A decoder bemenete jobbra tolódik az attention maszkolásához (a query "nem láthat a jövőbe")
* Scaled dot-product attention:

$$
Attention(Q,K,V)=softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)
$$

## LSH Attention

Az egyszerűség kedvéért tekintsünk 1-es batch méretet. Dot-product attention esetén a Q, K és V mátrixokat *d_model* dimenziós vektorok lineáris leképzésével kapjuk. Q, K és V [*length*, *d_k*], [*length*, *d_k*], illetve [*length*, *d_v*] alakúak, ahol *length* a bementei szekvencia mérete. Ha tehát *length* nagy, *QK^T* nagyon memóriaigényes. Pl. 64K token esetén *QK^T* 32 bites lebegőpontos számokkal több mint 16GB-ot vesz igénybe.

$softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)$ legtöbb eleme azonban 0-hoz közeli, így nem szükséges a teljes mátrixot kiszámolni: elegendő megtalálni az egymáshoz közeli query-key párokat. Ehhez az attention kiszámítását a következő lépések szerint módosítjuk:

1. K újradefiniálása: $k_i = \frac{q_i}{||q_i||}$
2. *q_i* és *k_i* vektorok hashelése
3. A hashelés megismétlése *n_rounds* alkalommal {*h⁽¹⁾*, *h⁽²⁾*, ...} különböző hash-függvényekkel.
4. A vektorok átrendezese $i \mapsto s_i$ permutációval a hash-értékek szerint
5. A rendezett vektorok chunkolása
6. Attention a chunkokon belül, illetve az előző chunk vektoraival

Magyarázat:

1. Ha különböző lineáris leképzésekkel hozzuk létre Q-t és K-t, különböző számú query és key vektorokat kaphatunk az egyes bucketekben
2. Cross polytope LSH algoritmus (Andoni et al. 2015): hogy *b* hasht kapjunk, létrehozunk egy [*d_k*, *b/2*] alakú gaussi random mátrixot, melyet *R*-rel jelölünk. A hash-értékeket a $h(x) = arg max([xR;−xR])$ képlettel kapjuk, ahol *[u; v]* két vektor konkatenálását jelöli.
3. Az azonos hash-értékhez tartozó bucketek egyesítése növeli annak a valószínűségét, hogy két hasonló vektor egy bucketbe kerüljön
4. Az egy bucketbe kerülő vektorok továbbra is az eredeti rendezés szerint követik egymást, hogy megmaradjon az információ a tokenek sorrendjéről
5. A fentiek szerint rendezett vektorokat *m* mértű chunkokra osztjuk. A gyakorlatban $m = \frac{2l}{n_{bucket}}$
6. Az attention memória- és időkomplexitása csökken: $l^2$ helyett $\left( \frac{4l}{n_c} \right)^2$

Az attention kiszámítása tehát ($\frac{1}{\sqrt{d_k}}$ nélkül):

$$
o_i = \sum_{j\in{\widetilde{\mathcal{P}}_i}}exp(q_i·k_j − m(j,\mathcal{P}_i) − z(i,\mathcal{P}_i))v_j
$$

ahol
$\widetilde{\mathcal{P}}_i = \left \{ j : \left \lfloor \frac{s_i}{m} - 1 \right \rfloor \leq \left \lfloor \frac{s_j}{m} \right \rfloor \leq \left \lfloor \frac{s_i}{m} \right \rfloor \right \}$,
$\mathcal{P}_i=\bigcup_{r=1}^{n_{rounds}}\mathcal{P_i}^{(r)}$,
$\mathcal{P_i}^{(r)} = \left \{ j: h_{(r)}(q_i) = h_{(r)}(q_j) \right \}$
$m(j, \mathcal{P}_i)=\begin{cases}\infty & \text{ha }j \notin \mathcal{P_i} \\ 0 & \text{egyébként} \end{cases}$
$\mathit{z}$ a softmax normalizáló tagja

Legyen $N_{i,j} = \left |\left\{r': j \in \mathcal{P}_i^{(r')}\right\} \right|$. Ekkor

$$
o_i = \sum_{r=1}^{n_{rounds}}exp(z(i, \mathcal{P}_i^{(r)})-z(i, \mathcal{P}_i))\sum_{j \in \widetilde{\mathcal{P}}_i^{(r)}} \frac{1}{N_{i,j}} exp(q_i·k_j − m(j,\mathcal{P}_i^{(r)}) − z(i,\mathcal{P}_i^{(r)}))v_j
$$

$$
o_i = \sum_{r=1}^{n_{rounds}}exp(z(i, \mathcal{P}_i^{(r)})-z(i, \mathcal{P}_i))o_i^{(r)}
$$

$$
o_i^{(r)} = \sum_{j\in{\widetilde{\mathcal{P}}_i^{(r)}}}exp(q_i·k_j − m_{i,j}^{(r)} − z(i,\mathcal{P}_i^{(r)}))v_j
$$

$m_{i,j}^{(r)}$-et definiálhatjuk a következőképpen:

$m_{i,j}^{(r)}=\begin{cases}\infty & \text{ha } j \notin \mathcal{P}_i^{(r)} \\ 10^5 & \text{ha } i = j \\ logN_{i,j} & \text{egyébként}\end{cases}$

Ezzel gátoljuk, hogy $q_i$ önmagára figyeljen, kivéve ha nincs más kontextus.

## Reversible layers

A transzformer rétegek aktivációinak memóriában tartása szintén költéges. A RevNet cikk (Gomez et al. 2017) alapján egy ResNet módosítható úgy, hogy backpropagation során minden réteg aktivációja kiszámolható legyen az azt követő réteg aktivációjából. Így elegendő az utolsó réteg aktivációját a memóriában tartani.

$y_1, y_2$ aktivációkat ehhez a következőképpen számoljuk:

$$
y_1 = x_1 + F(x_2)
$$

$$
y_2 = x_2 + G(y_1)
$$

amiből

$$
x_2 = y_2 - G(y_1)
$$

$$
x_1 = y_1 - F(x_2)
$$

A Transformer esetében

$$
Y_1 = X_1 + Attention(X_2)
$$

$$
Y_2 = X_2 + FeedForward(Y_1)
$$

Az első rétegben $x_1$ és $x_2$ a *d_model* méretű bemenet másolásával jön létre.

## Feed-forward réteg chunkolása

A feed-forward réteg $d_{ff} \geq 4K$ méretű vektorokon működik. A memória kímélése céljából a számítások elvégezhatők chunkokban:

$$
Y_2 = \left[ Y_2^{(1)}; ...; Y_2^{(c)} \right ] = \left[ X_2^{(1)} + FeedForward(Y_1^{(1)}); ...; X_2^{(c)} + FeedForward(Y_1^{(c)}) \right ]
$$

A paraméterek tárolása még így is sok memóriát igényel. Ezért az egyes számítások közben nem használt paramétereket CPU memóriában lehet tárolni. Standard transformerben a CPU-GPU/TPU közötti mozgatás túl lassú lenne, de a Reformer által használt nagy batch- és szekvenciaméret indokolja a mozgatást.

## Eredmények

### Összehasonítás más modelltípusokkal: memória és idő


### Kiértékelés enwik8 és imagenet64 adatokon

A görbék *bits per dim* értékeket mutatnak


### Newstest2014 EnDe fordítás

BLEU score


### Teljesítmény a rétegszám függvényében; attention gyorsasága szekvenciahossz függvényében