Repeat After Me: Transformers are Better than State Space Models at Copying

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.

한 줄 요약

트랜스포머는 GSSM 대비 컨텍스트 복사 능력에서 근본적 우위를 보인다.

핵심 기여도

이론적으로 2층 트랜스포머는 헤드 수에 비례하는 지수 길이의 문자열을 복사할 수 있음.
GSSM은 고정 크기의 잠재 상태로 인해 문자열 복사에 한계가 있음.
합성 실험에서 트랜스포머가 GSSM보다 복사 효율성과 일반화 능력이 우수함.
사전 학습된 Pythia(2.8B)는 Mamba(2.8B)보다 컨텍스트 복사 및 정보 검색 성능이 뛰어남.

핵심 아이디어

이 연구는 트랜스포머와 GSSM 간 컨텍스트 복사 능력의 근본적 차이를 탐구한다. 트랜스포머는 다중 헤드 어텐션을 통해 n-그램 기반의 저장 및 검색 메커니즘을 구현할 수 있다. 이는 특정 위치의 입력을 추적하고 복사하는 데 활용된다. 반면, GSSM은 고정 크기의 잠재 상태로 인해 입력 길이가 잠재 상태 크기를 초과하면 정확한 복사가 불가능하다는 이론적 한계를 보인다. 이는 트랜스포머가 GSSM보다 컨텍스트 내 임의 위치의 정보를 접근하는 데 더 유리하다는 통찰을 제공한다.

기술적 접근법

**이론적 분석**: 2층 트랜스포머가 헤드 수에 비례하는 지수 길이의 문자열을 복사할 수 있음을 증명.
**실험 모델**: 약 160M 파라미터의 합성 실험에서 트랜스포머가 GSSM보다 복사 학습이 빠르고 긴 입력에 대한 일반화가 우수함 (Figure 1(a), 1(b)).
**실제 모델 비교**: Pythia(410M, 1.4B, 2.8B)와 Mamba(360M, 1.4B, 2.8B)를 사용한 실험에서 트랜스포머가 컨텍스트 복사 및 정보 검색에서 GSSM을 상회함 (Figure 1(c)).
**모델 평가**: 모든 작업에서 64개 배치, QA 작업에서는 50개 질문으로 평가.

주요 결과

합성 복사 작업에서 트랜스포머는 GSSM 대비 더 빠르게 학습하고, 긴 입력에 대한 일반화가 우수함.
Pythia(2.8B)는 Mamba(2.8B) 대비 컨텍스트 복사 및 정보 검색에서 10% 이상의 성능 개선을 보임.
Mamba는 언어 모델링에서 낮은 퍼플렉시티를 보이지만, 컨텍스트 접근 능력은 트랜스포머에 비해 떨어짐.

의의 및 한계

이 연구는 트랜스포머와 GSSM 간 컨텍스트 접근 능력의 근본적 차이를 이론과 실험을 통해 입증한다. 이는 트랜스포머가 GSSM보다 컨텍스트 내 임의 위치의 정보를 접근하는 데 더 유리하다는 점을 강조한다. 그러나 GSSM은 입력 길이에 무관한 메모리 복잡도를 가지므로, 긴 입력 처리나 추론 효율성 측면에서는 여전히 유리하다. 또한, 인간의 언어 처리가 GSSM과 유사하다는 점에서 GSSM의 잠재적 가치도 언급된다.

실용적 활용

이 연구는 컨텍스트 내 정보 접근이 중요한 작업, 예를 들어 문서 검색, 질문 응답, 정보 추출 등에서 트랜스포머가 GSSM보다 우수함을 시사한다. 또한, GSSM에 어텐션 기반 메커니즘을 결합한 하이브리드 아키텍처 개발이 필요하다는 제안을 제공한다.