Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Summary

Motivation

  • Transformer는 성능은 매우 좋으나 내부 구조의 복잡함으로 인해서 계산비용이 기하급수적으로 증가하는 문제가 있다. 이 부분에 대한 개선은 Dot Product를 이용해서 지속해서 개선해왔으나, 여전히 개선이 필요한 분야이다.
  • 연구분야는 1)크게 메모리사용량을 줄이기 위한 방안을 찾는 부분과 2)Sequence Length를 늘려서 맥락을 최대한 이해할 수 있게 하려는 부분인데 계산량 자체를 줄인다면 이 두 부분에 기여할 수 있을 것이라고 보았다.

Approach

  • 계산량이 많이 발생하는 Self Attention 부분을 재해석해서 커널 Function으로 간주하고 내부 수식을 내적으로 바꿈으로써 계산량을 Input Sequence를 $N$이라고 할 때, $O(N^2)$에서 $O(N)$로 단축시켰다.
  • 이 때구체적으로 접근한 부분은 Attention Structure로 아래와 같이 수식으로 계된단다고 할 때, $$A_l(x)=V'=softmax( {{QK^T} \over {\sqrt{D}}}V )$$
  • 임의의 시퀀스 $i$만 놓고 본다고 하면 이렇게 정리할 수 있다. 이 때 $sim(Q,K)$는 $exp({{Q,K^T} \over {\sqrt{D}}})$으로 하면 우리가 알고 있는 Transformer 모델이 된다. 이 때 중요한 것은 위에서 보고 있는 $softmax$가 필수가 아닐 수도 있다는 사실이다. $$V' = {{\sum\limits_{j=1}^Nsim(Q_i,K_j)V_j} \over {\sum\limits_{j=1}^Nsim(Q_i,K_j)}}$$
  • 여기가 핵심인게 그러면 $sim(x,y)$이 하나의 유사도 함수라고 보면 유사도함수는 벡터간의 계산을 스칼라로 뱉어주는게 핵심이 이릍 특정 차원에 사상하고 내적하는 계산이라고 보면 그냥 커널함수라고 봐도 되서 아래와 같이 정리할 수 있게 된다. $$sim(x,y) = \phi(x)^T\phi(y)$$
  • 이 걸 이제 위에 $V'$에 대한 식에 대입해서 정리해 보면  $K$,$V$에 대한 부분을 한번에 계산하고 반복해서 쓰는 형태로 다음과 같이 정리할 수 있고 자연스럽게 계산량을 줄일 수 있게 된다. $$V_i' = {{\phi(Q_i)^T\sum\limits_{j=1}^N\phi(K_j)V_j^T} \over {\phi(Q_i)^T\sum\limits_{j=1}^N\phi(K_j)}}$$
  • Encoder, Decoder에 쓰이는 Multi Head Attention은 구조상 약간의 차이가 있기 때문에 Encoder, Decoder에 모두 Linear 구조로 바꾸어서 개선할 수 있었다.

Experiment

  • Synthetic Task, Image Generation, Speech Recognition Task를 기준으로 실험을 진행하였다.
  • 비교를 위한 Baseline Model은 Softmax Attention을 사용하는 Transformer, Reformer(Kitaev et al., 2020)를 선정 → 최근 모델 동향을 제대로 보지 않아 Reformer는 한 번 훑어봐야 할 듯
  • 데이터셋은 MNIST, CIFAR-10,  80 hour WSJ dataset을 사용하였음
  • fast_transformer 라는 별도의 Pytorch based Framework를 구축해서 사용

Transformers are RNNs

  • Linear Transformer의 Decoder 기준 매번 $i$시점의 입력갑을 이전 $i-i$ 시점의 정보를 함께 넣어준다는 점에 RNN의 Hidden State와 비슷하기 때문에 사실상 Linear Transformer는 RNN과 비슷하다는 주장을 짧게 섹션으로 소개하는데 다소 뜬금없기도 했지만, 그것보다 이 구조 때문에 점화식의 형태로 이전시점의 데이터만 참고하면 될 수 있고 성능개선의 중요한 포인트라는 것은 확실하였다. 마치 Markov Property를 보는 느낌적 느낌.

Thoughts

  • 더 훑어봐야겠지만 우선 테스트에서 NLP가 안보여서 안한 걸까, 못한 걸까는 좀 궁금하다.
  • fast transformer framework의 확장성등이 언급되지 않아서 활용가능성에 대해서 의문이 있지만 확인해보 하지만 UberEats에서 성능 개선을 위해서 사용한 기록(링크)가 있어 우선은 가능한 것으로 보인다.

References

Read more

[책]Reshuffle: Who wins when AI restacks the knowledge economy

[책]Reshuffle: Who wins when AI restacks the knowledge economy

원래는 Amazon에 가서 Personal Knowledge Managment에 관한 책을 사려고 했다. Sketch Your Mind라는 책이었는데, 그 때 이 책 “Reshuffle”을 발견하였다. AI가 어떻게 Knowledge Economy를 흔들 것가? 라는 부제를 훑어보면서 저자가 쓴 다른 책을 보게 되었는데 거기에 내가 좋아했던 책을쓴 저자라는 것을 알게 되었다. 그래서 크게 고민하지 않고 구매를 하고

By Bongho, Lee
[책]올라운드투자, 누군가의 투자일기

[책]올라운드투자, 누군가의 투자일기

“올라운드 투자”라는 제목을 보았을 때는, “올라운드 플레이어”가 생각이 났다. “올라운드”라는 표현을 오랜만에 들어본 까닭이었다. 그럼에도 불구하고 이 책을 고른 것은 저자가 그간 보여준 컨텐츠에 대한 신뢰가 있던 까닭이었다. 컨텐츠를 다양하게 보는 편이지만 깊이가 아주 있지는 않았다. 여기서 깊이라 함은 기존 전문적인 정량적 분석의 내용의 수준을 말하는 것이다.

By Bongho, Lee
내가 놓치고 있던 미래, 먼저 온 미래를 읽고

내가 놓치고 있던 미래, 먼저 온 미래를 읽고

장강명 작가의 책은, 유학시절 읽고 처음이었다. 유학시절 "한국이 싫어서"라는 책은 동기부여가 상당히 되는 책이었다. 한국을 떠나 새로운 정채성을 학생으로서 Build up 해나가고 있던 상황에서 이 책은 제목부터 꽤 솔깃하였다. 물론 결말이 기억날 정도로 인상깊은 책은 아니었지만 말이다. 그렇게 시간이 흘러 장강명 작가의 책은 더 이상 읽지 않던

By Bongho, Lee