최근성과 과도한 평활화의 시각을 통해 상태 공간 모델의 병목 현상을 이해하고 완화하는 방법

Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing

December 31, 2024

저자: Peihao Wang, Ruisi Cai, Yuehao Wang, Jiajun Zhu, Pragya Srivastava, Zhangyang Wang, Pan Li

cs.AI

초록

구조화된 상태 공간 모델(SSM)은 트랜스포머에 대안으로 등장했습니다. SSM은 종종 장기 의존성을 포착하는 데 효과적이라고 여겨지지만, 우리는 엄격히 증명합니다. 그들은 강한 최근성 편향으로 인해 본질적으로 제한되어 있다는 것을. 우리의 경험적 연구는 또한 이 편향이 모델이 먼 정보를 회상하는 능력을 손상시키고 견고성 문제를 도입한다는 것을 밝혀냅니다. 우리의 확장 실험은 그 후 SSM의 깊은 구조가 장거리 맥락을 학습하는 데 도움이 된다는 것을 발견했습니다. 그러나, 이어지는 이론적 분석은 SSM이 깊어질수록 더 매끄러워지는 또 다른 불가피한 경향을 보인다는 것을 밝혀냅니다. 예를 들어, 토큰 표현이 점점 구별하기 어려워집니다. 최근성과 과도한 매끄러움 사이의 이 기본적인 딜레마는 기존 SSM의 확장 가능성을 방해합니다. 우리의 이론적 발견을 영감받아, 우리는 SSM의 상태 전이 행렬의 두 채널을 극성화하는 것을 제안합니다. 각각을 제로와 원으로 설정하여 동시에 최근성 편향과 과도한 매끄러움에 대처합니다. 실험은 우리의 극성화 기술이 장거리 토큰의 연상 회상 정확도를 일관되게 향상시키고 SSM을 깊은 아키텍처로부터 더 많은 혜택을 누릴 수 있도록 만든다는 것을 입증합니다. 모든 소스 코드는 https://github.com/VITA-Group/SSM-Bottleneck에서 공개되었습니다.

English

Structured State Space Models (SSMs) have emerged as alternatives to transformers. While SSMs are often regarded as effective in capturing long-sequence dependencies, we rigorously demonstrate that they are inherently limited by strong recency bias. Our empirical studies also reveal that this bias impairs the models' ability to recall distant information and introduces robustness issues. Our scaling experiments then discovered that deeper structures in SSMs can facilitate the learning of long contexts. However, subsequent theoretical analysis reveals that as SSMs increase in depth, they exhibit another inevitable tendency toward over-smoothing, e.g., token representations becoming increasingly indistinguishable. This fundamental dilemma between recency and over-smoothing hinders the scalability of existing SSMs. Inspired by our theoretical findings, we propose to polarize two channels of the state transition matrices in SSMs, setting them to zero and one, respectively, simultaneously addressing recency bias and over-smoothing. Experiments demonstrate that our polarization technique consistently enhances the associative recall accuracy of long-range tokens and unlocks SSMs to benefit further from deeper architectures. All source codes are released at https://github.com/VITA-Group/SSM-Bottleneck.