OpenAI Sparse Transformer

OpenAI는 텍스트, 이미지, 사운드 시퀀스의 다음 요소를 예측하도록 설계된 심층 신경망인 Sparse Transformer를 개발했습니다. 어텐션 메커니즘에 알고리즘적 개선을 구현함으로써, 이 모델은 이전보다 30배 더 긴 시퀀스에서 패턴을 추출할 수 있으며, 수백 개의 레이어를 사용하여 수만 개의 요소를 모델링할 수 있습니다.

Reducing Memory Constraints with Deep Attention

표준 Transformer 아키텍처는 모든 레이어와 어텐션 헤드에 대해 $N \times N$ 어텐션 행렬이 필요합니다. 이는 원시 오디오나 대형 이미지와 같은 고차원 데이터를 처리할 때 상당한 메모리 병목 현상을 생성합니다. 예를 들어, 64x64x3 픽셀의 ImageNet 이미지는 심층 Transformer(64 레이어, 4 헤드)에서 저장된 행렬을 위해 154 GB의 메모리가 필요하며, 이는 표준 GPU의 12-32 GB 용량을 훨씬 초과합니다.

Sparse Attention and Algorithmic Complexity

단일 어텐션 행렬조차 비실용적인 매우 큰 입력을 처리하기 위해, Sparse Transformer는 sparse attention 패턴을 사용합니다. 이 접근 방식에서는 각 출력 위치가 입력 위치의 일부 집합(예: $\sqrt{N}$ 요소)에서만 가중치를 계산합니다. 이는 알고리즘 복잡도를 $O(N^2)$에서 $O(N \sqrt{N})$으로 줄여줍니다.

Factorized Attention Patterns

모델이 동적이고 전역적인 패턴을 학습하는 능력을 유지하도록 보장하기 위해, OpenAI는 어텐션 행렬의 2차원 인수분해를 구현했습니다. 이를 통해 네트워크는 두 단계의 sparse attention을 통해 모든 위치에 어텐션을 수행할 수 있습니다.

  • Strided Attention: 각 위치는 자신의 행과 열에 어텐션을 수행하며, 효과적으로 전체 어텐션 연산을 인수분해합니다. 이는 이미지와 같이 2차원 구조를 가진 데이터에 특히 유용합니다.
  • Fixed Attention: 네트워크는 고정된 열과 최신 열 요소 다음에 오는 요소들에 어텐션을 수행합니다. 이 패턴은 텍스트와 같이 2차원 구조에 맞지 않는 데이터에 최적화되어 있습니다.

Experimental Results and Performance

Sparse Transformer는 다음 벤치마크 데이터셋에 대해 밀도 추정(density estimation)에서 새로운 SOTA(state-of-the-art) 점수를 달성했습니다:

  • CIFAR-10
  • Enwik8
  • Imagenet 64

OpenAI는 sparse attention이 full attention보다 더 빠르게 작동할 뿐만 아니라 더 낮은 손실(loss)을 달성한다는 것을 관찰했습니다. 연구진은 이것이 dense attention의 근본적인 최적화 문제이거나 sparsity 패턴이 제공하는 유익한 귀납적 편향(inductive bias) 때문일 수 있다고 제안합니다.

Generative Capabilities across Modalities

Image Generation

이 모델은 64x64 ImageNet 데이터에 대한 이미지 완성(image completion) 작업을 통해 전역적 구조를 파악하고 있음을 보여줍니다. 모델이 최대 우도(maximum likelihood) 목적 함수를 사용하여 학습되므로, 조정되지 않은 1.0의 softmax temperature로 생성된 무조건부 샘플(unconditional samples)은 모델이 존재한다고 믿는 이미지의 전체 분포를 반영하며, 이는 때때로 이상하게 보이는 샘플을 생성할 수 있습니다.

Raw Audio Generation

위치 임베딩(position embeddings)을 수정함으로써, Sparse Transformer는 원시 오디오 파형을 생성할 수 있습니다. OpenAI는 원시 클래식 음악 클립으로 모델을 학습시켜 약 5초 분량의 원시 오디오에 해당하는 65,000개의 요소 시퀀스를 생성할 수 있게 했습니다.

Implementation and Limitations

sparse attention의 사용을 Facilitate하기 위해, OpenAI는 query와 key 행렬의 필요한 슬라이싱을 효율적으로 수행하는 일련의 block-sparse GPU 커널을 오픈 소스로 공개했습니다.

이러한 발전에도 불구하고, 연구진은 자기회귀(autoregressive) 시퀀스 생성이 매우 높은 해상도의 이미지나 비디오에는 여전히 비실용적이라고 언급했습니다. 그들은 이러한 최적화된 어텐션 연산이 향후 다중 스케일 모델링 접근 방식의 기본 요소(primitives)로 사용될 수 있다고 제안합니다.

Sources