본 논문은 장문 시퀀스에 대한 효율적인 대안으로서 상태 공간 모델(SSM)을 제시하지만, 과거 문맥을 검색하는 등의 알고리즘 작업에서는 고정 크기의 순환 상태로 인해 제한적인 능력을 갖는다는 점을 다룹니다. 트랜스포머와 SSM 기반 언어 모델에서 문맥 내 검색이 작동하는 방식을 조사하여 두 모델 모두 Gather-and-Aggregate (G&A) 메커니즘(Gather Head가 문맥에서 관련 정보를 추출하고 Aggregate Head가 단일 표현으로 통합)에 의존함을 발견했습니다. 두 아키텍처 모두에서 G&A는 소수의 헤드에 집중되어 간단한 검색에서조차 중요한 병목 현상을 형성합니다. 예를 들어, 축소된 Llama-3.1-8B에서 단일 Gather 또는 Aggregate Head를 비활성화하면 MMLU에서 정답 문자를 검색하는 능력이 저하되어 정확도가 66%에서 25%(무작위 추측)로 감소합니다. 또한, 이러한 검색 병목 현상은 작업의 제한된 지식 요구 사항을 가릴 수 있습니다. 축소된 모델은 기능하는 G&A 헤드를 사용하여 MMLU에서는 성공하지만 다른 지식 벤치마크에서는 실패합니다. 이 병목 현상은 SSM이 일반적으로 성능이 저조한 GSM8K, BBH 및 대화 이해와 같은 작업에도 마찬가지로 적용됩니다. SSM의 검색 과제는 이러한 헤드에서 나타나며, 효과적인 G&A에 필요한 날카로운 토큰 전환 대신 더 부드러운 어텐션 패턴을 생성합니다. 따라서 트랜스포머와 SSM의 검색 성능 차이는 전체 언어 모델이 아닌 소수의 헤드에 존재합니다. 이는 트랜스포머와 SSM의 성능 차이에 대한 통합적인 설명을 제시하며, 각각의 강점을 결합하는 방법을 보여줍니다. SSM이 소수의 어텐션 레이어와 결합된 사전 훈련된 하이브리드 모델에서 어텐션이 Aggregate Head의 역할을 대신하는 것을 발견했습니다. 마찬가지로, 사전 훈련된 SSM에서 단일 G&A 헤드를 어텐션 변형으로 대체하면 검색 및 벤치마크 점수가 향상됩니다.
