Comprendre et atténuer les goulots d'étranglement des modèles d'espace d'états à travers le prisme de la récence et de la sur-lissage.

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

December 31, 2024

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

cs.AI

Résumé

Les Modèles d'Espace d'État Structuré (MES) ont émergé en tant qu'alternatives aux transformeurs. Alors que les MES sont souvent considérés comme efficaces pour capturer les dépendances de longues séquences, nous démontrons rigoureusement qu'ils sont intrinsèquement limités par un fort biais de récence. Nos études empiriques révèlent également que ce biais nuit à la capacité des modèles de se souvenir d'informations distantes et introduit des problèmes de robustesse. Nos expériences d'échelle ont ensuite découvert que des structures plus profondes dans les MES peuvent faciliter l'apprentissage de contextes longs. Cependant, une analyse théorique ultérieure révèle qu'à mesure que les MES augmentent en profondeur, ils présentent une autre tendance inévitable vers un lissage excessif, par exemple, les représentations de jetons devenant de plus en plus indiscernables. Ce dilemme fondamental entre la récence et le lissage excessif entrave la scalabilité des MES existants. Inspirés par nos découvertes théoriques, nous proposons de polariser deux canaux des matrices de transition d'état dans les MES, en les fixant à zéro et un respectivement, adressant simultanément le biais de récence et le lissage excessif. Les expériences démontrent que notre technique de polarisation améliore de manière constante la précision de rappel associatif des jetons à longue portée et permet aux MES de bénéficier davantage de structures plus profondes. Tous les codes sources sont disponibles sur 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.