[논문 리뷰] FIT: Far-reaching Interleaved Transformers

ting chen, Lala Li|arXiv (Cornell University)|2023. 05. 22.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

FIT는 그룹화된 데이터 토큰과 적응형 잠재 토큰을 사용하여 국소적이고 전역적인 자기주의 메커니즘을 번갈아 적용함으로써 장거리 시퀀스에서 효율적이고 이차적 복잡도의 어텐션을 달성하는 새로운 트랜스포머 아키텍처를 제안한다. 국소적 윈도우드 어텐션과 학습된 잠재 토큰의 소수 집합에 대한 전역 어텐션을 조합함으로써 FIT는 계산 복잡도를 O(L⁴/³)로 감소시키며, 16GB 메모리 내에서 160K토큰 입력(예: 6400×6400 이미지)에 대해 모델 병렬화나 최적화 없이 엔드 투 엔드 학습을 가능하게 하여 고해상도 이미지 작업에서 최신 기술 수준의 성능을 입증한다.

ABSTRACT

We present FIT: a transformer-based architecture with efficient self-attention and adaptive computation. Unlike original transformers, which operate on a single sequence of data tokens, we divide the data tokens into groups, with each group being a shorter sequence of tokens. We employ two types of transformer layers: local layers operate on data tokens within each group, while global layers operate on a smaller set of introduced latent tokens. These layers, comprising the same set of self-attention and feed-forward layers as standard transformers, are interleaved, and cross-attention is used to facilitate information exchange between data and latent tokens within the same group. The attention complexity is $O(n^2)$ locally within each group of size $n$, but can reach $O(L^{{4}/{3}})$ globally for sequence length of $L$. The efficiency can be further enhanced by relying more on global layers that perform adaptive computation using a smaller set of latent tokens. FIT is a versatile architecture and can function as an encoder, diffusion decoder, or autoregressive decoder. We provide initial evidence demonstrating its effectiveness in high-resolution image understanding and generation tasks. Notably, FIT exhibits potential in performing end-to-end training on gigabit-scale data, such as 6400$ imes$6400 images, or 160K tokens (after patch tokenization), within a memory capacity of 16GB, without requiring specific optimizations or model parallelism.

연구 동기 및 목표

• 장거리 시퀀스를 처리할 때 표준 트랜스포머의 이차적 계산 복잡도 문제를 해결하기 위해.
• 성능을 훼손하지 않으면서 모델 병렬화 없이도 장거리 시퀀스에서 효율적이고 확장 가능한 어텐션을 가능하게 하기 위해.
• 다양한 비전 작업을 위한 인코더, 디퓨전 디코더 또는 자동회귀 디코더로 기능하는 통합 아키텍처를 설계하기 위해.
• 적응형 잠재 토큰을 활용한 국소적 및 전역 트랜스포머 레이어의 번갈아 적용이 장거리 의존성 모델링을 향상시키는지 탐색하기 위해.

제안 방법

• 입력 데이터를 토큰 그룹으로 나누고, 각 그룹은 윈도우드 컨텍스트 내에서 국소 자기주의 어텐션을 처리한다.
• 각 그룹에 대해 전역적 맥락을 나타내는 소수의 학습 가능한 잠재 토큰을 도입하여 어텐션 복잡도를 감소시킨다.
• 국소적 및 전역적 트랜스포머 레이어를 번갈아 적용하고, 교차 어텐션을 통해 데이터 토큰과 잠재 토큰 간의 이중 방향 정보 흐름을 가능하게 한다.
• 단일 순방향 전파에서 데이터 토큰과 잠재 토큰을 반복적으로 업데이트함으로써 국소적 및 전역적 표현 간의 깊은 통합을 보장한다.
• 국소적 및 전역적 레이어 모두에서 표준 피드포워드 네트워크와 어텐션 메커니즘을 유지하며, 입력 범위와 어텐션 패턴을 제외하고는 동일하다.
• 각 그룹 내에서 교차 어텐션을 활용해 데이터 토큰에서 잠재 토큰으로, 그리고 그 반대로 정보를 라우팅함으로써 적응형 계산을 가능하게 한다.

실험 결과

연구 질문

• RQ1 잠재 토큰을 활용한 국소적 및 전역 트랜스포머 레이어의 번갈아 적용이 표준 또는 희소 어텐션 메커니즘보다 더 나은 장거리 모델링을 달성할 수 있는가?
• RQ2 그룹당 잠재 토큰의 수가 고해상도 이미지 작업에서 모델 성능과 효율성에 어떤 영향을 미치는가?
• RQ3 국소적 및 전역 레이어를 번갈아 적용하는 것이 순차적으로 스택하거나 국소 레이어만 사용하는 것보다 성능을 향상시키는가?
• RQ4 FIT는 모델 병렬화나 특수 최적화 없이도 기가비트 규모의 입력(예: 6400×6400 이미지)에 대해 엔드 투 엔드 학습을 수행할 수 있는가?
• RQ5 FIT는 비전 분야에서 다양한 자동회귀, 디퓨전 및 인코더 작업에 걸쳐 얼마나 유연하게 활용될 수 있는가?

주요 결과

• FIT는 153M 파라미터로 ImageNet-64×64 자동회귀 모델링에서 최신 기술 수준에 근접한 성능을 달성하여, 1픽셀당 3.42비트를 기록하며 더 큰 모델과 유사한 성능을 보였다.
• 잠재 토큰의 수를 늘릴수록 성능이 향상되며(낮은 nll 및 bpd), 파라미터 수나 학습 속도에 미치는 영향은 최소한으로 유지되어 효율적인 확장이 가능함을 시사한다.
• 국소적 및 전역 레이어를 번갈아 적용하는 것이 순차적 스택이나 국소 레이어만 사용하는 것보다 뛰어난 성능(예: ImageNet에서 3.49 bpd)을 제공하며, 학습 비용 증가가 거의 없음.
• FIT는 모델 병렬화나 특수 최적화 없이도 16GB 메모리 내에서 160K토큰 시퀀스(예: 6400×6400 이미지)에 대해 엔드 투 엔드 학습을 가능하게 한다.
• 이 아키텍처는 다양한 작업에 대해 강력한 일반화 능력을 보이며, 고해상도 비전 벤치마크에서 인코더, 디퓨전 디코더, 자동회귀 디코더로 모두 효과적으로 기능한다.