[Paper short review] Self-attention Does Not Need O(n^2) Memory 논문 리뷰

@ 굵은 글씨는 중요한 내용, 빨간 글씨는 내가 추가한 내용

 

https://arxiv.org/pdf/2112.05682v1.pdf

 

• 정리: self-attention의 연산 과정에 트릭을 사용하여 공간 복잡도를 기존 O(n^2)에서 O(log n)로 줄인 알고리즘 제안

---

 

0. Abstract

Sequence 길이와 관련하여 O(1) 메모리가 필요한 attention 알고리즘과 O(log n) 메모리가 필요한 self-attention 알고리즘을 제안한다. 기존 self-attnetion은 O(n^2) 메모리를 필요로 한다. 본 논문에서 제안하는 알고리즘도 시간 복잡도는 여전히 O(n^2)이지만 현재 가장 큰 문제인 메모리 문제를 완화한다. 따라서 attention의 메모리 요구 사항을 줄이면 더욱 긴 시퀀스를 처리할 수 있다. 또한 메모리 효율성을 유지하면서 기능을 차별화하는 방법을 시연한다. 시퀀스 길이가 16384의 경우, self attention의 메모리 overhead는 추론에서 59배, differentiation에서는 32배 감소한다.

Overhead

 

어떤 처리를 하기 위해 들어가는 간접적인 처리 시간 혹은 메모리 등을 의미한다.

EX) A라는 task의 실행시간은 5초인데, 안정성을 위해 부가적인 B라는 task를 추가하여 최종 실행시간이 15초가 걸렸을 경우, overhead는 10초이다.

 

1. Introduction

Attention은 현대 신경 아키텍처에서 널리 사용된다. 특히 nlp에서 많은 변화를 가져왔고, 이후 연구된 transformer 아키텍처의 핵심이기도 하다. 이는 여러 분야에 영향을 미치고 있다.

길이가 n인 경우 query=q, key=k, value=v라 하면, attention은 다음과 같이 정의된다.

따라서 단일 query에 대한 attention의 결과는 value와의 가중합이며, 여기서 가중치는 query와 key의 dot-product의 softmax 이다.

이러한 attention을 구현하려면 먼저 모든 i에 대해 si를 계산하고 기억해야 하므로 각 query에 대해 O(n) 시간 및 메모리 복잡도를 갖는다. Transformer는 self-attention을 사용하므로 시간 및 메모리 복잡도는 O(n^2)이다. 

많은 연구에서 self-attention의 n^2 시간 및 메모리 복잡도는 기존 attention 메커니즘 및 아키텍처의 연구 동기 중 하나였다. GPU와 TPU와 같은 최신 가속기 하드웨어는 딥러닝의 응용 분야에 대해 메모리 제약이 있는 반면 컴퓨팅은 상대적으로 저렴하다. 특히, transformer의 공간 복잡성은 큰? 문제이다.

 

본 연구에서는 각각 상수 메모리가 필요한 attention과 로그 메모리만 필요한 self-attention을 제안한다. 기본 알고리즘은 매우 간단하지만 수치적으로 실현 가능하려면 트릭이 필요하다. 또한 TPU에서 효율적으로 실행되며 self-attention을 위해 O(√ n) 메모리가 필요한 JAX을 구현한다.

 

우리의 방법은 여전히 self-attention에서는 O(n^2) 시간 복잡도를, single attention에서는 O(n) 시간 복잡도를 갖고, 최근 활발히 연구되고 있는 효율적인 long-context attention mechanism은 attention의 대안으로 떠오르고 있다. 하지만 attention을 위한 메모리 효율 알고리즘을 사용하면 아키텍처 선택을 재고하거나 더 길고 밀도 높은 attention이 필요한 새로운 데이터 세트로 확장할 수 있다. 입력에 대해 메모리 효율적인 attention을 통해 이러한 제한을 해결할 수 있다.

 

2. Algorithm

먼저, 우리는 single query로 attention에 대한 알고리즘을 제안하고, 이를 self-attention까지 확장한다. 우리는 Σe^sj에 의한 분할이 분포 법칙을 사용하여 attention의 맨 끝으로 이동할 수 있음을 관찰한다.

이는 상수 메모리로 계산할 수 있다.: 이 알고리즘의 메모리 overhead는 벡터 v와 스칼라 s로 구성되며, 둘 다 0으로 초기화된다. Query q, key k, value v가 주어지면 key와 value를 순서대로 처리한다. key와 value 쌍 ki, vi가 주어지면 si = dot(q, ki)를 계산하고 v* ← v* + + vi e^si와 s* ← s* + + e si를 업데이트한다. 모든 key와 value를 처리한 후 v*/s*를 분할하여 최종 결과를 얻는다. 

공간 복잡성 분석은 입력이 특정 순서로 제공된다고 가정한다. 먼저 query를 읽은 다음 key와 value의 쌍을 나열한다. 만약 입력이 다른 순서로 제공된다면, 인덱스를 시퀀스에 추가로 저장해야 하며, 대신 O(log n) 메모리가 필요하다. 

이 알고리즘을 self-attention으로 확장하기 위해 모든 query에 대한 결과를 순차적으로 계산한다. 이것은 query 목록에 하나의 추가 인덱스만 필요하므로 O(log n) 메모리 복잡성이 발생한다. 이 연산은 쿼리 수, 즉, O(n)의 크기에 선형인 출력을 생성하는데, 이는 공간 복잡도로 계산되지 않는다.

 

3. Numerical Stability

 

4. An Implementation For GPUs/TPUs

5. Empirical Analysis

5.1 Inference

 

5.2 Differentiation

 

5.3 Training

Transformer 아키텍처에 표준 attention 모듈과 메모리 효율적인 attention 모듈을 사용하여 WMT en-de 변환 실험을 실행했다. 학습 내내, 두 방법은 거의 동일하게 행동했다. 10만 번의 학습 단계를 거친 후, 평가 정확도는 메모리 효율적인 구현의 경우 62.69에 달했고 표준 구현의 경우 62.59에 도달했다. 이는 메모리 효율적인 자기 주의 구현이 기존 구현을 대체하는 데 사용될 수 있음을 보여준다. 아래 그림은 두 모델이 매우 유사한 BLEU 점수를 얻었음을 보여준다.

 

6. Conclusion

이 논문은 (self) attnetion의 메모리 요구량을 줄이기 위한 방법을 제안한다. 본 논문을 통해 attention이 본질적으로 memory-hungry가 아니라는 인식을 높여 현재 아키텍처의 설계를 다시 살펴보길 바란다.

 

My Discussion

Why this paper?

transformer를 이미지에 적용하는 프로세스를 생각하면 계속 해상도에 따른 연산 문제가 막히는데 관련 논문이여서 보기로함

 

--> 수학적 트릭을? 사용하여 소프트맥스 연산을 뒤로 옮기는 것만으로 공간 복잡도를 완화 한것이 신기하다

--> 구현부분을 자세히 봐야될 듯