[논문 리뷰] xLSTM: Extended Long Short-Term Memory

Maximilian Beck, Korbinian Pöppel외 7명|arXiv (Cornell University)|2024. 05. 07.

Topic Modeling인용 수 86

한 줄 요약

xLSTM은 지수 게이팅과 두 가지 기억 보강 변형(sLSTM: 스칼라 기억, mLSTM: 행렬 기억)을 도입하여 LSTM의 한계를 극복하고, Transformer 및 상태 공간 모델과 비교해 언어 모델링에서 우수한 성능 및 확장성을 보여준다.

ABSTRACT

In the 1990s, the constant error carousel and gating were introduced as the central ideas of the Long Short-Term Memory (LSTM). Since then, LSTMs have stood the test of time and contributed to numerous deep learning success stories, in particular they constituted the first Large Language Models (LLMs). However, the advent of the Transformer technology with parallelizable self-attention at its core marked the dawn of a new era, outpacing LSTMs at scale. We now raise a simple question: How far do we get in language modeling when scaling LSTMs to billions of parameters, leveraging the latest techniques from modern LLMs, but mitigating known limitations of LSTMs? Firstly, we introduce exponential gating with appropriate normalization and stabilization techniques. Secondly, we modify the LSTM memory structure, obtaining: (i) sLSTM with a scalar memory, a scalar update, and new memory mixing, (ii) mLSTM that is fully parallelizable with a matrix memory and a covariance update rule. Integrating these LSTM extensions into residual block backbones yields xLSTM blocks that are then residually stacked into xLSTM architectures. Exponential gating and modified memory structures boost xLSTM capabilities to perform favorably when compared to state-of-the-art Transformers and State Space Models, both in performance and scaling.

연구 동기 및 목표

• 대규모 언어 모델링에서 저장소 재검토, 제한된 기억, 병렬성 부재 등 고전 LSTM의 한계 극복 필요성에 대한 동기 부여.
• 지수 게이팅을 갖춘 두 가지 기억 보강 LSTM 변형(sLSTM 및 mLSTM) 제안.
• 이들 변형을 잔차 아키텍처에 통합한 xLSTM 블록 도입으로 확장 가능한 모델링 구현.
• 합성 과제와 대규모 언어 모델링 실험을 통해 xLSTM이 특정 설정에서 Transformer/SSM 벤치마크를 경쟁하거나 능가할 수 있음을 시연.

제안 방법

• 정규화 및 안정화를 포함한 지수 게이팅을 도입하여 재검토 가능한 저장 결정 가능하게 함.
• 스칼라 기억, 스칼라 업데이트 및 셀/헤드 간 기억 혼합을 가진 sLSTM 개발.
• 높은 용량 저장소 및 병렬 조회를 가능하게 하는 행렬 기억과 공분산 업데이트 규칙을 가진 mLSTM 개발.
• sLSTM 또는 mLSTM을 잔차 블록 백본 내부에 삽입하여 xLSTM 블록 생성(변형에 따라 프로젝션 전후 어느 쪽에 위치하는지 다름).
• 재정렬된 LayerNorm 백본을 사용하여 xLSTM 블록을 쌓아 확장 가능한 모델링을 위한 xLSTM 아키텍처 구성.
• CUDA 최적화 구현 제공 및 Transformer와 비교한 메모리/계산 특성 논의.

실험 결과

연구 질문

• RQ1 지수 게이팅과 새로운 기억 구조가 대규모에서 저장 재검토, 기억 용량, 순차적 병목 현상 등 LSTM의 기본 한계를 바로잡을 수 있는가?
• RQ2 수십억 매개변수로 확장되었을 때 sLSTM과 mLSTM가 언어 모델링 작업에서 Transformer 및 상태 공간 모델과 어떤 차이를 보이는가?
• RQ3 잔차 적층을 가진 xLSTM 아키텍처가 모델 크기 및 데이터 규모에 따라 경쟁력 있는 perplexity와 다운스트림 작업 성능을 제공하는가?

주요 결과

• 지수 게이팅과 매트릭스/스칼라 기억 설계로 xLSTM이 유효성 검증 perplexity에서 여러 기본 모델보다 더 우수한 성능을 보임.
• 합성 및 장거리 작업에서 xLSTM은 전통적인 LSTM 및 일부 Transformer/SSM 변종에 비해 상태 추적과 기억 용량이 향상됨.
• xLSTM[1:0] 및 xLSTM[7:1]이 SlimPajama 실험에서 모델 규모에 관계없이 강력한 검증 perplexity를 달성하고 우호적인 스케일링 동향을 보임.
• 300B 토큰 학습에서 xLSTM 변형은 RWKV, Llama, Mamba에 비해 시퀀스 길이 외삽, 더 긴 컨텍스트 perplexity 및 다운스트림 작업 성능이 더 나은 경향을 유지함.
• ablations에서 지수 게이팅과 행렬 기억이 성능 향상의 주요 원인임을 시사하며, 학습 가능한 입력 의존 게이트가 추가 이점을 제공함.
• xLSTM 아키텍처는 선형에 가까운 메모리 고려 및 병렬 구성 요소를 갖추고 있어 Transformer 주도 체제에 대한 확장 가능한 대안을 제공함.