[논문 리뷰] Efficient Long Sequence Modeling via State Space Augmented Transformer

Simiao Zuo, Xiaodong Liu외 5명|arXiv (Cornell University)|2022. 12. 15.

Topic Modeling인용 수 7

한 줄 요약

이 논문은 장기간 시퀀스 모델링을 향상시키기 위해 첫 번째 레이어에 상태공간모델(SSM)을 통합하여 전반적인 의존성을 포괄하고, 이후 레이어에서는 효율적인 국소적 어텐션을 사용하여 세부적인 국소적 모델링을 수행하는 하이브리드 트랜스포머 아키텍처인 SPADE를 제안한다. 이 방법은 선형 복잡도를 달성하며, Long Range Arena 및 언어 모델링에서 베이스라인을 능가하고, NLU 및 생성 작업을 위한 대규모 사전학습 모델로도 효과적으로 확장된다.

ABSTRACT

Transformer models have achieved superior performance in various natural language processing tasks. However, the quadratic computational cost of the attention mechanism limits its practicality for long sequences. There are existing attention variants that improve the computational efficiency, but they have limited ability to effectively compute global information. In parallel to Transformer models, state space models (SSMs) are tailored for long sequences, but they are not flexible enough to capture complicated local information. We propose SPADE, short for $\underline{ extbf{S}}$tate s$\underline{ extbf{P}}$ace $\underline{ extbf{A}}$ugmente$\underline{ extbf{D}}$ Transform$\underline{ extbf{E}}$r. Specifically, we augment a SSM into the bottom layer of SPADE, and we employ efficient local attention methods for the other layers. The SSM augments global information, which complements the lack of long-range dependency issue in local attention methods. Experimental results on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method. To further demonstrate the scalability of SPADE, we pre-train large encoder-decoder models and present fine-tuning results on natural language understanding and natural language generation tasks.

연구 동기 및 목표

• 표준 트랜스포머의 장기간 시퀀스 모델링에서의 제곱형 계산 비용 문제를 해결하기 위해.
• 국소적 어텐션 메커니즘의 장거리 의존성 포착 능력에 한계가 있음을 극복하기 위해.
• 상태공간모델(SSM)의 전반적 모델링 능력과 효율적 어텐션의 국소적 인덕티브 바이어스를 통합하기 위해.
• 장기간 컨텍스트 NLP 작업을 위한 확장 가능하고 효율적이며 효과적인 아키텍처를 설계하기 위해.
• 하류 NLU 및 생성 작업에서의 사전학습 및 미세조정을 통해 방법의 효과성을 입증하기 위해.

제안 방법

• SPADE는 첫 번째 레이어가 상태공간모델(SSM)인 다층 트랜스포머를 사용하여 시퀀스에 거시적 전역 정보를 통합한다.
• SSM 레이어 다음에는 윈도우 기반 어텐션 등의 효율적 국소 어텐션 레이어가 여러 개 연결되어 국소적 의존성을 보완한다.
• SSM는 선형 시간 및 공간 복잡도로 작동하여 장기간 시퀀스의 효율적 처리를 가능하게 한다.
• 이 방법은 SSM의 전역 인덕티브 바이어스와 어텐션의 국소적 모델링 유연성을 결합하여 상호 보완적인 약점을 보완한다.
• 모델은 자동회귀 및 비자동회귀 시퀀스 모델링 모두를 지원하며, 인코더-디코더 설정도 포함된다.
• 모델은 표준 목적함수로 훈련되고, GLUE 및 요약 벤치마크에서 미세조정된다.

실험 결과

연구 질문

• RQ1 SSM과 국소 어텐션을 융합한 하이브리드 아키텍처가 트랜스포머에서 장거리 의존성 모델링을 향상시킬 수 있는가?
• RQ2 첫 번째 레이어에 SSM을 통합하면 Long Range Arena와 같은 장기간 시퀀스 벤치마크에서 성능 향상이 이루어지는가?
• RQ3 SPADE는 제곱형 어텐션 및 희소 어텐션 베이스라인을 능가하면서도 선형 복잡도를 유지할 수 있는가?
• RQ4 대규모 언어 모델링에서 사전학습하고 하류 작업에서 미세조정했을 때 SPADE의 확장성은 어느 정도인가?
• RQ5 전반적인 성능 향상에서 SSM 레이어와 국소 어텐션 레이어의 기여도는 각각 얼마인가?

주요 결과

• SPADE는 Long Range Arena 벤치마크에서 최고 성능을 기록하며, 전체, 희소, 낮은 랭크 어텐션을 사용하는 모델들을 능가한다.
• 자동회귀 언어 모델링에서 SPADE는 순수 트랜스포머보다 훨씬 빠르며, 더 낮은 퍼플렉서티 점수를 달성한다.
• 모델은 강력한 확장성을 보이며, 사전학습된 인코더-디코더 버전이 GLUE 및 요약 작업에서 뛰어난 성능을 기록한다.
• 제거 실험 결과 SSM 레이어가 장거리 의존성을 포착하는 데 핵심적임을 확인하였고, 국소 어텐션 레이어는 국소 패턴을 보완한다.
• SPADE는 8k 길이의 시퀀스에 대해 GPU 메모리 소비를 80GB 이하로 줄여 장기간 컨텍스트 모델링을 실용적으로 가능하게 한다.
• 구조적 바이어스가 제한된 상황에서도 높은 성능을 유지하며, 위치 인코딩이나 희소 패턴에 의존하는 모델들을 능가한다.