Revealing and Mitigating the Local Pattern Shortcuts of Mamba

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: You, Wangjie, Tang, Zecheng, Li, Juntao, Yao, Lili, Zhang, Min
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866913557230125056
author You, Wangjie
Tang, Zecheng
Li, Juntao
Yao, Lili
Zhang, Min
author_facet You, Wangjie
Tang, Zecheng
Li, Juntao
Yao, Lili
Zhang, Min
contents Large language models (LLMs) have advanced significantly due to the attention mechanism, but their quadratic complexity and linear memory demands limit their performance on long-context tasks. Recently, researchers introduced Mamba, an advanced model built upon State Space Models(SSMs) that offers linear complexity and constant memory. Although Mamba is reported to match or surpass the performance of attention-based models, our analysis reveals a performance gap: Mamba excels in tasks that involve localized key information but faces challenges with tasks that require handling distributed key information. Our controlled experiments suggest that this inconsistency arises from Mamba's reliance on local pattern shortcuts, which enable the model to remember local key information within its limited memory but hinder its ability to retain more dispersed information. Therefore, we introduce a global selection module into the Mamba model to address this issue. Experiments on both existing and proposed synthetic tasks, as well as real-world tasks, demonstrate the effectiveness of our method. Notably, with the introduction of only 4M extra parameters, our approach enables the Mamba model(130M) to achieve a significant improvement on tasks with distributed information, increasing its performance from 0 to 80.54 points.
format Preprint
id arxiv_https___arxiv_org_abs_2410_15678
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Revealing and Mitigating the Local Pattern Shortcuts of Mamba
You, Wangjie
Tang, Zecheng
Li, Juntao
Yao, Lili
Zhang, Min
Computation and Language
Artificial Intelligence
Large language models (LLMs) have advanced significantly due to the attention mechanism, but their quadratic complexity and linear memory demands limit their performance on long-context tasks. Recently, researchers introduced Mamba, an advanced model built upon State Space Models(SSMs) that offers linear complexity and constant memory. Although Mamba is reported to match or surpass the performance of attention-based models, our analysis reveals a performance gap: Mamba excels in tasks that involve localized key information but faces challenges with tasks that require handling distributed key information. Our controlled experiments suggest that this inconsistency arises from Mamba's reliance on local pattern shortcuts, which enable the model to remember local key information within its limited memory but hinder its ability to retain more dispersed information. Therefore, we introduce a global selection module into the Mamba model to address this issue. Experiments on both existing and proposed synthetic tasks, as well as real-world tasks, demonstrate the effectiveness of our method. Notably, with the introduction of only 4M extra parameters, our approach enables the Mamba model(130M) to achieve a significant improvement on tasks with distributed information, increasing its performance from 0 to 80.54 points.
title Revealing and Mitigating the Local Pattern Shortcuts of Mamba
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2410.15678