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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Wang, Peihao, Cai, Ruisi, Wang, Yuehao, Zhu, Jiajun, Srivastava, Pragya, Wang, Zhangyang, Li, Pan
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866915190477422592
author Wang, Peihao
Cai, Ruisi
Wang, Yuehao
Zhu, Jiajun
Srivastava, Pragya
Wang, Zhangyang
Li, Pan
author_facet Wang, Peihao
Cai, Ruisi
Wang, Yuehao
Zhu, Jiajun
Srivastava, Pragya
Wang, Zhangyang
Li, Pan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2501_00658
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing
Wang, Peihao
Cai, Ruisi
Wang, Yuehao
Zhu, Jiajun
Srivastava, Pragya
Wang, Zhangyang
Li, Pan
Machine Learning
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.
title Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing
topic Machine Learning
url https://arxiv.org/abs/2501.00658