[논문 리뷰] Repeat After Me: Transformers are Better than State Space Models at Copying

Samy Jelassi, David Brandfonbrener외 2명|arXiv (Cornell University)|2024. 02. 01.

Simulation Techniques and ApplicationsDecision Sciences인용 수 3

한 줄 요약

이 논문은 어텐션 메커니즘을 통한 n-그램 저장 및 검색 능력 덕분에, 장기 입력 시퀀스에서 정보를 복사하고 검색하는 데서 트랜스포머가 일반화된 상태공간모델(GSSMs)을 능가함을 보여준다. 이론적으로, 두 층으로 구성된 트랜스포머는 어텐션 헤드 수의 지수함수만큼 긴 시퀀스를 복사할 수 있으며, 반면 GSSMs는 고정된 크기의 잠재 상태 덕분에 본질적으로 제한된다. 이 잠재 상태는 자신의 차원 수를 초월해 더 많은 비트를 인코딩할 수 없다.

ABSTRACT

Transformers are the dominant architecture for sequence modeling, but there is growing interest in models that use a fixed-size latent state that does not depend on the sequence length, which we refer to as "generalized state space models" (GSSMs). In this paper we show that while GSSMs are promising in terms of inference-time efficiency, they are limited compared to transformer models on tasks that require copying from the input context. We start with a theoretical analysis of the simple task of string copying and prove that a two layer transformer can copy strings of exponential length while GSSMs are fundamentally limited by their fixed-size latent state. Empirically, we find that transformers outperform GSSMs in terms of efficiency and generalization on synthetic tasks that require copying the context. Finally, we evaluate pretrained large language models and find that transformer models dramatically outperform state space models at copying and retrieving information from context. Taken together, these results suggest a fundamental gap between transformers and GSSMs on tasks of practical interest.

연구 동기 및 목표

• GSSMs가 효율적임에도 불구하고, 입력 컨텍스트에서 정보를 복사하고 검색하는 데서 트랜스포머를 따라잡을 수 있는지 조사하기 위해.
• 고정된 크기의 잠재 상태로 인해 GSSMs가 장기 시퀀스를 표현하는 데 겪는 이론적 한계를 분석하기 위해.
• 합성 복사 작업에서 트랜스포머와 GSSMs의 일반화 및 효율성을 평가하기 위해.
• 사전학습된 모델들(예: Pythia 대 Mamba)이 제로샷 설정에서 컨텍스트 검색 및 복사 작업에서의 성능을 테스트하기 위해.
• 퍼스펙티브의 차이가 퍼스펙티브 외의 언어 모델 능력에 미치는 실질적 영향을 이해하기 위해.

제안 방법

• GSSMs와 트랜스포머의 공식 정의를 사용해 문자 복사 작업에 대한 이론적 분석을 수행하여, GSSMs가 잠재 상태 크기 이상의 길이의 시퀀스를 복사할 수 없다는 것을 증명한다.
• 어텐션 헤드 수의 지수함수만큼 긴 시퀀스를 복사할 수 있는 두 층의 트랜스포머를 구축하며, 어텐션을 통한 n-그램 저장 및 검색을 활용한다.
• 점차 길어지는 길이의 시퀀스를 복사하도록 훈련시키는 합성 데이터셋에서의 실증적 평가를 수행하며, 훈련 효율성과 길이 일반화 능력을 측정한다.
• 트랜스포머에서 훈련 컨텍스트 창을 초월한 일반화를 향상시키기 위해 하드-알리비(_hard-alibi_) 위치 인코딩을 사용한다.
• 사전학습된 모델들(Pythia 및 Mamba)을 기반으로 컨텍스트 기반 검색 작업을 평가하며, 다양한 문단 길이에서의 F1 점수를 측정한다.
• 훈련된 트랜스포머에서 학습된 어텐션 패턴을 분석하여, 복사 작업에서 n-그램 검색 메커니즘에 의존한다는 것을 확인한다.

(a) Copying: training efficiency.

실험 결과

연구 질문

• RQ1 트랜스포머 모델은 자신의 크기의 지수함수만큼 긴 시퀀스를 복사할 수 있는가? 만약 가능하다면, 어떻게 이를 달성하는가?
• RQ2 고정된 크기의 잠재 상태로 인해 GSSMs가 복사 작업에서 겪는 본질적 표현 한계는 무엇인가?
• RQ3 합성 복사 작업에서 트랜스포머와 GSSMs는 훈련 효율성과 더 긴 시퀀스로의 일반화 능력 측면에서 어떻게 비교되는가?
• RQ4 사전학습된 트랜스포머 모델은 제로샷 컨텍스트 검색 및 복사 작업에서 사전학습된 GSSMs 모델보다 뛰어나게 성능을 내는가?
• RQ5 GSSMs의 고정된 메모리 용량이 장기 입력 컨텍스트를 포함한 자연어 작업에서 성능을 얼마나 떨어뜨리는가?

주요 결과

• 두 층의 트랜스포머는 어텐션 헤드 수의 지수함수만큼 긴 시퀀스를 복사할 수 있으며, 이는 n-그램 저장을 통한 강력한 컨텍스트 검색 메커니즘을 보여준다.
• GSSMs는 복사 작업에서 본질적으로 제한된다: 잠재 상태 크기 이상의 비트를 포함하는 시퀀스를 정확히 복사할 수 없다.
• 합성 복사 작업에서 트랜스포머는 GSSMs보다 훨씬 더 빠르게 훈련되며, 더 긴 시퀀스로의 일반화 능력도 훨씬 뛰어나며, Mamba의 경우 훈련 길이를 초월해 성능이 급격히 떨어진다.
• 사전학습된 Pythia(트랜스포머 기반) 모델은 Mamba(GSSM 기반) 모델보다 제로샷 컨텍스트 룩업 작업에서 더 뛰어난 성능을 보이며, 문단 길이가 길어질수록 Mamba의 F1 점수는 더 급격히 감소한다.
• 트랜스포머에서 학습된 복사 메커니즘은 어텐션 기반의 n-그램 룩업에 의존하며, 이는 이론적 구성과 일치함을 확인한다.
• Mamba는 Pile 데이터셋에서 Pythia보다 낮은 퍼플렉서티를 달성했지만, 컨텍스트 집약적인 작업에서 성능이 열 劣하다. 이는 훈련 손실만으로는 검색 및 복사 작업 성능을 예측할 수 없음을 시사한다.

(b) Copying: length generalization