Retrievit: 트랜스포머, 상태 공간 모델, 그리고 하이브리드 아키텍처의 문맥 내 검색 능력

Retrievit: In-context Retrieval Capabilities of Transformers, State Space Models, and Hybrid Architectures

트랜스포머는 문맥 내 검색 능력에서 뛰어난 성능을 보이지만, 입력 시퀀스 길이에 따라 2차 복잡도를 갖는다는 단점이 있습니다. 반면, 상태 공간 모델(SSM)은 효율적인 선형 시간 처리를 제공하지만, 검색 능력은 제한적입니다. 본 연구에서는 트랜스포머와 SSM을 결합한 하이브리드 아키텍처가 두 가지 인공적인 문맥 내 검색 작업에서 두 모델의 장점을 모두 활용할 수 있는지 조사합니다. 첫 번째 작업은 n-gram 검색으로, 모델이 입력 시퀀스 내에서 쿼리에 이어지는 n-gram을 식별하고 재현해야 합니다. 두 번째 작업은 위치 검색으로, 모델에게 단일 쿼리 토큰을 제시하고, 시퀀스 내에서 해당 요소를 찾아내고, 해당 요소의 위치 인덱스를 출력하도록 합니다. 통제된 실험 환경에서, 트랜스포머, SSM, 그리고 하이브리드 아키텍처의 데이터 효율성, 길이 일반화 능력, 외부 데이터에 대한 강건성, 그리고 학습된 표현을 평가합니다. 연구 결과, 하이브리드 모델은 SSM보다 우수한 성능을 보였으며, 데이터 효율성과 정보가 밀집된 문맥 검색에서의 확장성 측면에서 트랜스포머와 동등하거나 더 나은 성능을 보였습니다. 그러나, 트랜스포머는 위치 검색 작업에서 여전히 우수한 성능을 유지했습니다. 표현 분석을 통해, SSM 기반 모델은 인접한 위치를 나타내는 토큰들이 임베딩 공간에서 서로 가깝게 위치하는, 위치 정보를 고려한 임베딩을 생성한다는 것을 발견했습니다. 이러한 특성은 트랜스포머에서는 나타나지 않으며, 이는 SSM 및 하이브리드 모델의 다양한 검색 작업에서의 강점과 한계를 설명합니다. 본 연구는 작업 요구 사항에 따른 아키텍처 선택에 대한 지침을 제공하며, 트랜스포머와 SSM, 그리고 하이브리드 모델이 위치 관계를 학습하는 방식의 근본적인 차이를 밝혀냅니다.

Transformers excel at in-context retrieval but suffer from quadratic complexity with sequence length, while State Space Models (SSMs) offer efficient linear-time processing but have limited retrieval capabilities. We investigate whether hybrid architectures combining Transformers and SSMs can achieve the best of both worlds on two synthetic in-context retrieval tasks. The first task, n-gram retrieval, requires the model to identify and reproduce an n-gram that succeeds the query within the input sequence. The second task, position retrieval, presents the model with a single query token and requires it to perform a two-hop associative lookup: first locating the corresponding element in the sequence, and then outputting its positional index. Under controlled experimental conditions, we assess data efficiency, length generalization, robustness to out of domain training examples, and learned representations across Transformers, SSMs, and hybrid architectures. We find that hybrid models outperform SSMs and match or exceed Transformers in data efficiency and extrapolation for information-dense context retrieval. However, Transformers maintain superiority in position retrieval tasks. Through representation analysis, we discover that SSM-based models develop locality-aware embeddings where tokens representing adjacent positions become neighbors in embedding space, forming interpretable structures. This emergent property, absent in Transformers, explains both the strengths and limitations of SSMs and hybrids for different retrieval tasks. Our findings provide principled guidance for architecture selection based on task requirements and reveal fundamental differences in how Transformers and SSMs, and hybrid models learn positional associations.