DenseMamba: Modelos de Espaço de Estados com Conexão Oculta Densa para Modelos de Linguagem de Grande Escala Eficientes

DenseMamba: State Space Models with Dense Hidden Connection for Efficient Large Language Models

February 26, 2024

Autores: Wei He, Kai Han, Yehui Tang, Chengcheng Wang, Yujie Yang, Tianyu Guo, Yunhe Wang

cs.AI

Resumo

Os grandes modelos de linguagem (LLMs) enfrentam um desafio significativo devido aos requisitos excessivos de computação e memória da arquitetura Transformer, comumente utilizada. Embora os modelos de espaço de estados (SSM) representem um novo tipo de arquitetura de rede fundamental com menor complexidade computacional, seu desempenho ainda não rivaliza plenamente com o dos Transformers. Este artigo introduz o DenseSSM, uma abordagem inovadora para aprimorar o fluxo de informações ocultas entre as camadas nos SSMs. Ao integrar seletivamente os estados ocultos das camadas superficiais em camadas mais profundas, o DenseSSM retém informações detalhadas cruciais para a saída final. As conexões densas do DenseSSM ainda mantêm a paralelizabilidade no treinamento e a eficiência na inferência. O método proposto pode ser amplamente aplicável a diversos tipos de SSM, como RetNet e Mamba. Com tamanho de modelo similar, o DenseSSM alcança melhorias significativas, exemplificado pelo DenseRetNet, que supera o RetNet original com ganhos de até 5% em precisão em benchmarks públicos.

English

Large language models (LLMs) face a daunting challenge due to the excessive computational and memory requirements of the commonly used Transformer architecture. While state space model (SSM) is a new type of foundational network architecture offering lower computational complexity, their performance has yet to fully rival that of Transformers. This paper introduces DenseSSM, a novel approach to enhance the flow of hidden information between layers in SSMs. By selectively integrating shallowlayer hidden states into deeper layers, DenseSSM retains fine-grained information crucial for the final output. Dense connections enhanced DenseSSM still maintains the training parallelizability and inference efficiency. The proposed method can be widely applicable to various SSM types like RetNet and Mamba. With similar model size, DenseSSM achieves significant improvements, exemplified by DenseRetNet outperforming the original RetNet with up to 5% accuracy improvement on public benchmarks.