Comprendiendo y Mitigando los Cuellos de Botella de los Modelos de Espacio de Estados a través del Prisma de la Recencia y el Sobreajuste.

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

December 31, 2024

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

cs.AI

Resumen

Los Modelos Estructurados de Espacio de Estados (SSMs, por sus siglas en inglés) han surgido como alternativas a los transformadores. Si bien los SSMs suelen considerarse efectivos para capturar dependencias de largas secuencias, demostramos rigurosamente que están inherentemente limitados por un fuerte sesgo de recencia. Nuestros estudios empíricos también revelan que este sesgo afecta la capacidad de los modelos para recordar información distante e introduce problemas de robustez. Nuestros experimentos de escalado descubrieron que estructuras más profundas en los SSMs pueden facilitar el aprendizaje de contextos largos. Sin embargo, un análisis teórico posterior revela que a medida que los SSMs aumentan en profundidad, muestran otra tendencia inevitable hacia el exceso de suavizado, por ejemplo, las representaciones de tokens se vuelven cada vez más indistinguibles. Este dilema fundamental entre recencia y exceso de suavizado obstaculiza la escalabilidad de los SSMs existentes. Inspirados en nuestros hallazgos teóricos, proponemos polarizar dos canales de las matrices de transición de estados en los SSMs, estableciéndolos en cero y uno, respectivamente, abordando simultáneamente el sesgo de recencia y el exceso de suavizado. Los experimentos demuestran que nuestra técnica de polarización mejora consistentemente la precisión de recordación asociativa de tokens de largo alcance y permite a los SSMs beneficiarse aún más de arquitecturas más profundas. Todos los códigos fuente se encuentran disponibles en 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.