Comprendere e mitigare i colli di bottiglia dei modelli dello spazio degli stati attraverso la lente della recente e dell'eccessiva levigatura.

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

December 31, 2024

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

cs.AI

Abstract

I modelli strutturati dello spazio degli stati (SSM) sono emersi come alternative ai trasformatori. Mentre gli SSM sono spesso considerati efficaci nel catturare dipendenze di lunghe sequenze, dimostriamo rigorosamente che sono intrinsecamente limitati da un forte bias di recente. I nostri studi empirici rivelano inoltre che questo bias compromette la capacità dei modelli di richiamare informazioni distanti e introduce problemi di robustezza. I nostri esperimenti di scalabilità hanno poi scoperto che strutture più profonde negli SSM possono facilitare l'apprendimento di contesti lunghi. Tuttavia, un'analisi teorica successiva rivela che all'aumentare della profondità degli SSM, essi mostrano un'altra tendenza inevitabile verso un eccessivo levigamento, ad esempio, le rappresentazioni dei token diventano sempre più indistinguibili. Questo dilemma fondamentale tra il bias di recente e l'eccessivo levigamento ostacola la scalabilità degli SSM esistenti. Ispirati dai nostri risultati teorici, proponiamo di polarizzare due canali delle matrici di transizione degli stati negli SSM, impostandoli a zero e uno, rispettivamente, affrontando contemporaneamente il bias di recente e l'eccessivo levigamento. Gli esperimenti dimostrano che la nostra tecnica di polarizzazione migliora costantemente l'accuratezza del richiamo associativo dei token a lungo raggio e sblocca gli SSM per beneficiare ulteriormente da architetture più profonde. Tutti i codici sorgente sono disponibili su 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.