[논문 리뷰] Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

Beidi Chen, Tri Dao외 4명|arXiv (Cornell University)|2021. 10. 28.

Advanced Neural Network Applications참고 문헌 66인용 수 9

한 줄 요약

Scatterbrain는 주어진 해시 함수를 사용하여 희소성과 저질서 구조를 동시에 고려하는 이론적으로 탄탄한 방법을 제안한다. 이는 주로 LSH를 통해 희소성을, 커널 특징 맵을 통해 저질서를 구현하며, 기존의 희소 또는 저질서 방법보다 최대 2.1배 낮은 근사 오차를 달성하고, 이미지 모델에서 정확도가 1% 감소하는 조건에서 주의 메모리를 98% 감소시킨다. 이는 언어 및 시각 작업에서 단독으로 사용되는 희소 또는 저질서 방법보다 우수한 성능을 보인다.

ABSTRACT

Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits-all approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature in attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain can achieve 2.1x lower error than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained T2T-ViT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce 98% of attention memory at the cost of only 1% drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up to 4 points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers on language modeling and long-range-arena tasks.

연구 동기 및 목표

• 희소성과 저질서 주의 근사를 통합하여, 현재 서로 다른 영역에서 뛰어난 성능을 보이며 일반적으로 별도로 사용되는 방법을 통합하고자 한다.
• 긴 시퀀스 모델링에서 효율성과 정확도의 상충 관계를 해결하기 위해 두 가지 근사 유형을 결합하고자 한다.
• 이중 계산 문제를 피하기 위해 이론적으로 정당화되고 편향이 없으며 효율적인 방법을 개발하고자 한다.
• 희소성+저질서 근사가 이미지 생성 및 언어 모델링과 같은 다양한 작업에서 개별 방법보다 뛰어난 성능을 보임을 경험적으로 검증하고자 한다.
• 기존의 사전 학습된 모델에서 전체 주의를 최소한의 정확도 손실과 함께 쉽게 대체할 수 있도록 하기 위함이다.

제안 방법

• Scatterbrain는 주어진 해시 함수를 사용하여 소프트맥스 주의 행렬의 희소 성분을 식별하고, 전체 행렬을 메모리에 로드하지 않고 근사한다.
• 저질서 성분을 근사하기 위해 커널 특징 맵을 적용하여, 무작위 특징 근사를 통해 효율적인 계산을 가능하게 한다.
• 주의 행렬을 희소 성분과 저질서 성분의 합으로 분해하며, Robust PCA에서 영감을 얻은 구조적 분해를 통해 이중 계산 문제를 방지한다.
• 이론적으로 저질서 기반 기준보다 더 낮은 근사 오차를 보장한다.
• 기존의 효율적인 Transformer 아키텍처(예: Reformer, Performer)와 호환되는 표준 주의 레이어의 플러그인 대체품으로 설계되어 있다.
• 소프트맥스 온도와 주의 분포의 엔트로피에 따라 최적의 희소성 및 저질서 성분 조합을 자동으로 선택한다.

실험 결과

연구 질문

• RQ1 어떤 영역(예: 소프트맥스 엔트로피 기반)에서 희소성, 저질서성 또는 통합된 근사가 주의 행렬에서 가장 잘 작동하는가?
• RQ2 희소성과 저질서 근사를 통합한 방법이 개별 방법보다 더 낮은 근사 오차를 달성할 수 있는가?
• RQ3 이중 계산이나 편향을 유발하지 않으면서도 이론적으로 타당하고 효율적인 방식으로 희소성과 저질서 성분을 통합할 수 있는가?
• RQ4 개선된 근사 오차가 언어 모델링 및 이미지 분류와 같은 최종 작업에서의 종단 간 성능에 어떻게 영향을 미치는가?
• RQ5 Scatterbrain는 전체 주의를 최소한의 정확도 손실과 함께 효과적으로 대체할 수 있는가?

주요 결과

• T2T-ViT 사전 학습 모델에서 Scatterbrain은 정확도가 1% 감소하는 조건에서 주의 메모리를 98% 감소시켰으며, SMYRF 및 로컬 주의보다 뛰어난 성능을 보였다.
• BigGAN 이미지 생성 작업에서 Scatterbrain은 전체 주의를 대체할 때 기준 방법보다 근사 오차가 2.1배 낮았다.
• 언어 모델링 및 Long Range Arena 작업에서 종단 간 학습을 수행한 결과, Scatterbrain은 희소성 또는 저질서 기반 기준보다 퍼플렉서티 최대 4포인트 향상 및 평균 정확도 최대 5포인트 향상했다.
• GLUE 벤치마크에서 BERT의 주의 레이어를 대체할 경우, Scatterbrain은 Performer 및 Smyrf를 포함한 모든 기준보다 9개 작업 중 8개에서 승리했다.
• 시각적 분석 결과, 희소성+저질서 근사가 전반적으로 가장 낮은 오차를 보였으며, 특히 개별 방법이 실패하는 중간 엔트로피 영역에서 뚜렷한 우수성을 보였다.
• 이 방법은 계층적(예: 언어) 및 비계층적(예: 이미지) 작업 모두에서 최고 성능을 기록하여 광범위한 적용 가능성을 입증했다.