xLSTM: Extended Long Short-Term Memory

요약 1990년대에 Long Short-Term Memory(LSTM)의 핵심 아이디어로 지속적인 오류 회전목마와 게이팅이 도입되었습니다. 그 이후로 LSTM은 시간의 시험을 견디며 많은 딥러닝 성공 사례에 기여했으며, 특히 첫 번째 대형 언어 모델(LLM)을 구성했습니다. 그러나 병렬화 가능한 자기 주의를 핵심으로 하는 트랜스포머 기술의 출현은 새로운 시대의 도래를 알렸고, 대규모에서 LSTM을 능가했습니다. 이제 우리는 간단한 질문을 제기합니다: 최신 LLM의 기술을 활용하되 LSTM의 알려진 한계를 완화하면서 LSTM을 수십억 개의 매개변수로 확장하면 언어 모델링에서 얼마나 멀리 갈 수 있을까요? 첫째, 적절한 정규화 및 안정화 기법을 사용한 지수 게이팅을 도입합니다. 둘째, LSTM 메모리 구조를 수정하여 (i) 스칼라 메모리, 스칼라 업데이트 및 새로운 메모리 믹싱을 갖춘 sLSTM, (ii) 행렬 메모리 및 공분산 업데이트 규칙을 갖춘 완전 병렬화 가능한 mLSTM을 얻습니다. 이러한 LSTM 확장을 잔차 블록 백본에 통합하여 xLSTM 블록을 만들고, 이를 잔차적으로 쌓아 xLSTM 아키텍처를 구성합니다. 지수 게이팅과 수정된 메모리 구조는 xLSTM의 성능과 확장면에서 최첨단 트랜스포머 및 상태 공간 모델과 비교하여 유리하게 수행할 수 있도록 합니다.

 

그림 1: 확장된 LSTM(xLSTM) 계열. 왼쪽에서 오른쪽으로: 1. 지속적인 오류 회전목마와 게이팅을 갖춘 원래의 LSTM 메모리 셀. 2. 지수 게이팅을 도입한 새로운 sLSTM 및 mLSTM 메모리 셀. sLSTM은 새로운 메모리 믹싱 기법을 제공합니다. mLSTM은 새로운 행렬 메모리 셀 상태와 공분산 업데이트 규칙을 갖춘 완전 병렬화가 가능합니다. 3. 잔차 블록 내의 mLSTM 및 sLSTM이 xLSTM 블록을 형성합니다. 4. 쌓인 xLSTM 블록이 xLSTM 아키텍처를 구성합니다.

 

1 Introduction

Long Short-Term Memory (LSTM) 아이디어 (Hochreiter, 1991; Hochreiter & Schmidhuber, 1997a, b), 즉 지속적인 오류 회전목마와 게이팅은 점차 사라지는 그래디언트 문제를 극복하기 위해 도입되었습니다 (Hochreiter, 1991; Hochreiter et al., 2000):

LSTM은 다양한 분야에서 성공적으로 적용되었습니다 (Hochreiter et al., 2001, 2007; Schmidhuber, 2015) 및 2017년 트랜스포머의 등장까지 텍스트 생성에서 우위를 점했습니다 (Vaswani et al., 2017). LSTM의 효과는 텍스트 생성 (Graves, 2013; Karpathy, 2015), 시퀀스-시퀀스 학습 (Sutskever et al., 2014), 시퀀스-시퀀스 번역 (Sutskever et al., 2014), 이미지 캡션 생성 (Karpathy & Fei-Fei, 2015; Hossain et al., 2019), 코드 생성 (Yin et al., 2017), 악보 생성 (Sturm et al., 2016), 비강수 모델링 (Karpathy, 2015), 강수량-유출 모델링 (Kratzert et al., 2019), 홍수 경고 시스템의 수문학적 모델링 (Nearing et al., 2022) 등 다양한 시퀀스 관련 작업에서 입증되었습니다. 강화 학습에서도 LSTM은 최고의 성능을 발휘하는 시퀀스 모델로, 스타크래프트 II의 AlphaStar 모델 (Vinyals et al., 2017), 도타 2의 OpenAI Five 모델 (Berner et al., 2019), 핵 융합 제어를 위한 자기학적 제어 모델 (Degrave et al., 2022) 등이 있습니다. LSTM은 학습 추상화에서도 우수하여, 예를 들어 의미 정보를 추출하여 메모리 셀에 저장하고 (Karpathy, 2015), 더 많은 신경 뉴런의 활성화로 쉽게 확장되는 능력 (Lake et al., 2019), 언어 모델링 (Bau et al., 2019), 감정 분석 (Radford et al., 2017) 등에 사용되었습니다. LSTM은 최근에도 활발히 사용되고 있으며 (Degrave et al., 2022; Nearing et al., 2022), 시간의 시험을 견뎌냈습니다.

그림 2: LSTM의 한계. 왼쪽: 평균 제곱 오차(MSE) 측면에서 최근접 이웃 검색 문제. 참조 벡터가 주어지면, 시퀀스 끝에서 연결된 값을 반환하기 위해 가장 유사한 벡터를 찾기 위해 시퀀스를 순차적으로 스캔합니다. LSTM은 더 유사한 벡터를 찾았을 때 저장된 값을 수정하는 데 어려움을 겪습니다. 우리의 새로운 xLSTM은 지수 게이팅을 통해 이 한계를 극복합니다. 오른쪽: 희귀 토큰 예측. Wikitext-103에서 토큰 빈도별 버킷으로 토큰 예측의 당혹도(PPL)를 나타냅니다. LSTM은 제한된 저장 용량 때문에 희귀 토큰 예측에서 성능이 떨어지지만, 우리의 새로운 xLSTM은 행렬 메모리를 통해 이 문제를 해결합니다.

 

그들의 엄청난 성공에도 불구하고, LSTM에는 세 가지 주요 한계가 있습니다: (i) 저장 결정의 수정 불능. 이 한계를 최근접 이웃 검색 문제를 통해 예시합니다 (부록 B 참조). 참조 벡터가 주어지면, 가장 유사한 벡터를 찾기 위해 시퀀스를 순차적으로 스캔해야 하며, 시퀀스 끝에서 연결된 값을 제공해야 합니다. 그림 2의 왼쪽 패널은 이 작업에서의 평균 제곱 오차를 보여줍니다. LSTM은 더 유사한 벡터를 찾았을 때 저장된 값을 수정하는 데 어려움을 겪지만, 우리의 새로운 xLSTM은 지수 게이팅으로 이 한계를 보완합니다. (ii) 제한된 저장 용량, 즉 정보는 스칼라 셀 상태로 압축되어야 합니다. 이 한계를 희귀 토큰 예측을 통해 예시합니다. 그림 2의 오른쪽 패널은 Wikitext103(Merity et al., 2017)에서 토큰 빈도별 버킷에 대한 토큰 예측의 당혹도를 보여줍니다. LSTM은 제한된 저장 용량 때문에 희귀 토큰에서 성능이 떨어집니다. 우리의 새로운 xLSTM은 행렬 메모리를 통해 이 문제를 해결합니다. (iii) 메모리 혼합으로 인한 병렬 처리 불가능성, 즉 한 시점에서 다음 시점으로 히든 상태 간의 히든-히든 연결은 순차 처리를 강제합니다. LSTM의 이러한 한계는 언어 모델링에서 트랜스포머(Vaswani et al., 2017)의 출현을 이끌었습니다. 이러한 한계를 극복하고 LSTM을 현재의 대형 언어 모델 크기로 확장하면 언어 모델링에서 어떤 성능을 달성할 수 있을까요?

 

2. 확장된 장단기 메모리(LSTM)
LSTM의 한계를 극복하기 위해 확장된 장단기 메모리(xLSTM)는 Equation (1)의 LSTM 아이디어에 두 가지 주요 수정을 도입합니다. 이러한 수정 사항, 즉 지수 게이팅과 새로운 메모리 구조는 LSTM 계열을 두 가지 멤버로 확장합니다: (i) 스칼라 메모리, 스칼라 업데이트 및 메모리 혼합을 갖춘 새로운 sLSTM(섹션 2.2 참조) 및 (ii) 행렬 메모리와 공분산(외적) 업데이트 규칙을 갖춘 완전 병렬화가 가능한 새로운 mLSTM(섹션 2.3 참조). sLSTM과 mLSTM 모두 지수 게이팅을 통해 LSTM을 향상시킵니다. 병렬 처리를 가능하게 하기 위해 mLSTM은 메모리 혼합, 즉 히든-히든 순환 연결을 포기합니다. sLSTM과 mLSTM은 모두 다중 메모리 셀로 확장할 수 있으며, sLSTM은 셀 간의 메모리 혼합을 특징으로 합니다. 또한 sLSTM은 각 헤드 내에서만 셀 간의 메모리 혼합을 하되, 헤드 간에는 메모리 혼합이 없는 다중 헤드를 가질 수 있습니다. sLSTM의 헤드를 도입하고 지수 게이팅을 결합함으로써 새로운 방식의 메모리 혼합을 확립합니다. mLSTM의 경우 다중 헤드와 다중 셀이 동일한 역할을 합니다. 이러한 새로운 LSTM 변형을 잔차 블록 모듈에 통합하면 xLSTM 블록이 생성됩니다(섹션 2.4 참조). 이러한 xLSTM 블록을 아키텍처에 잔차적으로 쌓으면 xLSTM 아키텍처가 됩니다(섹션 2.4 참조). xLSTM 아키텍처와 그 구성 요소는 그림 1을 참조하십시오.

 

2.1 장단기 메모리(LSTM) 리뷰

원래의 LSTM 아이디어는 스칼라 메모리 셀을 중심 처리 및 저장 단위로 도입하여 소멸하는 그래디언트 문제를 피하고자 했습니다 (Hochreiter, 1991; Hochreiter & Schmidhuber, 1997a, b). 지속적인 오류 회전목마와 게이팅을 통해 셀 상태 업데이트를 제어했습니다. 메모리 셀에는 세 개의 게이트가 포함됩니다: 입력 게이트, 출력 게이트, 포겟 게이트. 포겟 게이트는 Gers et al. (2000)에 의해 도입되었습니다. LSTM 메모리 셀 업데이트 규칙은 다음과 같습니다:

2.2 sLSTM

LSTM이 저장 결정을 수정할 수 있는 능력을 부여하기 위해, 우리는 정규화와 안정화와 함께 지수 게이트(빨간색)를 도입합니다. 특히, 입력 및 포겟 게이트는 지수 활성화 함수를 가질 수 있습니다. 정규화를 위해, 우리는 모든 미래 포겟 게이트 곱의 결과를 합산하는 정규화 상태를 도입합니다.

sLSTM의 순방향 패스는 다음과 같습니다:

2.3 mLSTM

2.4 xLSTM 모델

그림 3: xLSTM 블록. 왼쪽: 잔차 sLSTM 블록으로, 업 프로젝션 후 (트랜스포머와 유사): 입력은 선택적 컨볼루션과 함께 sLSTM에 전달되고, 그 다음에 게이팅된 MLP가 이어집니다. 오른쪽: 잔차 mLSTM 블록으로, 업 프로젝션 전 (상태 공간 모델과 유사): mLSTM은 컨볼루션, 학습 가능한 스킵 연결, 그리고 요소별로 작동하는 출력 게이트를 통해 두 개의 MLP 내부에 포함됩니다. 자세한 내용은 부록의 그림 9와 그림 10을 참조하십시오.

 

xLSTM 블록. xLSTM 블록은 과거를 고차원 공간에서 비선형적으로 요약하여 서로 다른 히스토리나 컨텍스트를 더 잘 분리해야 합니다. 히스토리를 분리하는 것은 다음 시퀀스 요소, 예를 들어 다음 토큰을 정확하게 예측하기 위한 전제 조건입니다. 우리는 고차원 공간에서 비선형적으로 임베딩된 패턴이 원래 공간보다 선형적으로 분리될 가능성이 더 크다는 Cover의 정리에 의존합니다 (Cover, 1965). 우리는 두 가지 잔차 블록 아키텍처를 고려합니다: (i) 잔차 블록은 업 프로젝션 후(트랜스포머와 유사) 과거를 원래 공간에서 비선형적으로 요약한 다음, 고차원 공간으로 선형 매핑하고 비선형 활성화 함수를 적용한 후 다시 원래 공간으로 선형 매핑합니다; 그림 3의 왼쪽 패널과 그림 1의 세 번째 열을 참조하십시오. 더 자세한 버전은 부록의 그림 9에 나와 있습니다. (ii) 잔차 블록은 업 프로젝션 전(상태 공간 모델과 유사) 고차원 공간으로 선형 매핑하고, 고차원 공간에서 과거를 비선형적으로 요약한 후 다시 원래 공간으로 선형 매핑합니다. sLSTM을 포함하는 xLSTM 블록의 경우, 주로 업 프로젝션 후 블록을 사용합니다. mLSTM을 포함하는 xLSTM 블록의 경우, 메모리 용량이 고차원 공간에서 더 커지기 때문에 업 프로젝션 전 블록을 사용합니다. 그림 3의 왼쪽 패널과 그림 1의 세 번째 열, 또는 자세한 내용은 부록의 그림 9를 참조하십시오.

 

xLSTM 아키텍처. xLSTM 아키텍처는 빌딩 블록을 잔차적으로 쌓아 구성됩니다 (Srivastava et al., 2015; He et al., 2016). 우리는 현대의 대형 언어 모델에서 사용되는 가장 일반적인 preLayerNorm (Ba et al., 2016b) 잔차 백본을 사용합니다. 그림 1의 마지막 열을 참조하십시오.

 

2.5 메모리 및 속도 고려 사항

3 관련 연구

선형 어텐션
트랜스포머의 문맥 길이에 대한 이차 복잡성을 극복하고 어텐션을 문맥 길이에 대해 선형으로 만드는 여러 방법이 제안되었습니다. Synthesizer는 토큰 간 상호작용 없이 합성 어텐션 가중치를 학습합니다 (Tay et al., 2020). Linformer는 저랭크 행렬을 사용하여 자기 어텐션을 실현하며 심지어 이를 선형적으로 근사화합니다 (Wang et al., 2020). Linear Transformer는 어텐션 메커니즘을 선형화합니다 (Katharopoulos et al., 2020). Performer는 양의 직교 랜덤 특징 접근법으로 어텐션 소프트맥스를 선형적으로 근사화합니다 (Choromanski et al., 2021). Structured Global Convolution (SGConv) (Li et al., 2022)와 Hyena Hierarchy (Poli et al., 2023)에서는 어텐션을 빠른 긴 컨볼루션으로 대체했습니다.

상태 공간 모델(SSM)
최근 상태 공간 모델(SSM)은 문맥 길이에 대해 선형적이며 트랜스포머와 비교하여 유망한 성능을 보여주어 매우 인기를 끌고 있습니다. 처음 제안된 모델 중 하나는 Structured State Space 시퀀스 모델(S4) (Gu et al., 2021)이고, 그 뒤를 이어 Diagonal State Space (DSS) 모델 (Gupta et al., 2022), Gated State Space (GSS) 모델 (Mehta et al., 2022), S5 모델 (Smith et al., 2022), Bidirectional Gated SSM (BiGS) (Wang et al., 2022), H3 모델 (Fu et al., 2023), 그리고 Mamba (Gu & Dao, 2023)가 제안되었습니다.

순환 신경망(RNN)
순환 신경망(RNN)은 문맥 길이에 대해 선형적이기 때문에 트랜스포머와 어텐션을 대체할 수 있는 방안으로 제안되었습니다. 깊은 선형 순환 유닛(LRU)을 가진 RNN은 언어 모델링에서 유망한 결과를 보여주었습니다 (Orvieto et al., 2023; De et al., 2024). Hierarchically Gated Linear RNN (HGRN) (Qin et al., 2023)과 HGRN2 (Qin et al., 2024)도 마찬가지입니다. 대형 언어 모델링에 잘 알려진 RNN 접근 방식은 RWKV (Peng et al., 2023, 2024)로, 트랜스포머에 필적하는 성능을 보여줍니다.

게이팅
LSTM의 주요 아이디어 중 하나인 게이팅은 최근 많은 접근 방식에서 재발견되고 재해석되었습니다. 게이팅은 HGRN (Qin et al., 2023), HGRN2 (Qin et al., 2024), Gated Linear Attention (GLA) (Yang et al., 2023), Gated State Space (GSS) 모델 (Mehta et al., 2022), Bidirectional Gated SSM (BiGS) (Wang et al., 2022), Moving Average Equipped Gated Attention (MEGA) (Ma et al., 2022), RWKV (Peng et al., 2023), 그리고 Mamba (Gu & Dao, 2023)에서 사용되었습니다.

공분산 업데이트 규칙
저장 용량을 향상시키기 위해, 우리는 mLSTM 셀에 공분산 업데이트 규칙을 가진 행렬 메모리를 장착했습니다. 이러한 업데이트 메커니즘을 기반으로 하는 다른 방법에는 Fast Weight Programmers (Schmidhuber, 1992; Schlag et al., 2021), RWKV-5 및 RWKV-6 (Peng et al., 2024), Retention (Sun et al., 2023), Linear Transformer (Katharopoulos et al., 2020), 그리고 HGRN2 (Qin et al., 2024) 등이 있습니다.

가장 관련된 연구
개념적으로 xLSTM에 가장 가까운 모델은 Retention (Sun et al., 2023), RWKV (Peng et al., 2023, 2024), 그리고 HGRN2 (Qin et al., 2024)입니다. 이 모델들은 행렬 메모리 및/또는 게이팅 개념을 공유합니다. 그러나 새로운 sLSTM과 달리, 이러한 접근 방식은 메모리 혼합을 허용하지 않습니다. 메모리 혼합은 상태 추적 문제를 해결할 수 있게 하며, 따라서 LSTM은 상태 공간 모델(SSM) 및 트랜스포머보다 표현력이 더 큽니다 (Merrill et al., 2024; Delétang et al., 2023). 상태 추적은 코드 평가나 긴 서사에서 엔티티 추적을 위해 필요합니다.

잔차적 쌓기 아키텍처
대부분의 현대 대형 딥러닝 모델처럼, xLSTM 아키텍처는 빌딩 블록을 잔차적으로 쌓아서 구성됩니다 (Srivastava et al., 2015; He et al., 2016). 이 구성은 딥 컨볼루션 네트워크 (He et al., 2016)와 트랜스포머 (Vaswani et al., 2017)를 가능하게 했습니다. 트랜스포머는 GPT-3 (Brown et al., 2020), ChatGPT (Schulman et al., 2022), GPT-4 (Achiam et al., 2023), Megatron-LM (Shoeybi et al., 2019), Gopher (Rae et al., 2021), ERNIE 3.0 Titan (Wang et al., 2021), GLaM (Du et al., 2021), Chinese M6 (Lin et al., 2021), 다국어 AlexaTM 20B(Soltan et al., 2022), OPT (Zhang et al., 2022), Chinchilla (Hoffmann et al., 2022), BLOOM (Scao et al., 2022), GLM-130B (Zeng et al., 2022), LaMDA (Thoppilan et al., 2022), PaLM (Chowdhery et al., 2022), Llama (Touvron et al., 2023), Gemini (Google, 2023; Reid et al., 2024)와 같은 대형 언어 모델의 궁극적인 힘입니다.

 

4 실험

이 섹션에서는 xLSTM을 실험적으로 평가하고 언어 모델링에 중점을 두고 기존 방법과 비교합니다. 섹션 4.1에서는 합성 작업에서 xLSTM의 특정 능력을 조사합니다. 섹션 4.2에서는 SlimPajama(Soboleva et al., 2023)에서 150억 개의 토큰으로 학습된 다양한 현재 언어 모델링 방법의 검증 세트 당혹도(perplexity)를 비교합니다. 동일한 데이터 세트에서 xLSTM에 대한 제거 연구(ablation studies)를 수행합니다. 그런 다음, Kaplan et al. (2020) 및 Brown et al. (2020)와 유사하게 다양한 방법의 확장 동작을 평가합니다. 섹션 4.3에서는 더 철저한 언어 모델링 실험을 수행합니다. SlimPajama(Soboleva et al., 2023)에서 3000억 개의 토큰으로 학습된 후, 섹션 4.2에서 가장 성능이 좋은 방법들과 xLSTM을 비교합니다. 먼저, 방법들이 더 긴 문맥으로 외삽하는 성능을 평가하고, 둘째로 검증 당혹도와 다운스트림 작업(Sutawika et al., 2024)에서의 성능을 테스트합니다. 셋째로, PALOMA 언어 벤치마크 데이터셋(Magnusson et al., 2023)의 571개 텍스트 도메인에서 방법들을 평가합니다. 넷째로, 이번에는 20배 더 많은 학습 데이터를 사용하여 다양한 방법의 확장 동작을 다시 평가합니다. 모든 실험에서, mLSTM 기반 xLSTM 블록 대 sLSTM 기반 xLSTM 블록의 비율을 xLSTM[a:b]로 표시합니다. 예를 들어, xLSTM[7:1]은 여덟 개의 블록 중 일곱 개가 mLSTM 기반 블록이고, 하나가 sLSTM 기반 블록임을 의미합니다. 총 블록 수가 48인 경우, 이는 6개의 sLSTM 기반 블록과 42개의 mLSTM 기반 블록으로 변환됩니다. 또한 모든 실험에서 mLSTM과 sLSTM에 각각 사전 및 사후 업 프로젝션 블록을 사용합니다.

 

4.1 합성 작업 및 롱 레인지 아레나

먼저, 형식 언어에서 메모리 혼합과 함께 xLSTM의 새로운 지수 게이팅의 효과를 테스트합니다 (Delétang et al., 2023). 그런 다음, 다중 쿼리 연상 회상 작업에서 xLSTM의 새로운 행렬 메모리의 효과를 평가합니다 (Arora et al., 2023). 마지막으로, 롱 레인지 아레나에서 긴 시퀀스를 처리하는 xLSTM의 성능을 평가합니다 (Tay et al., 2021).

그림 4: 메모리 혼합과 함께 xLSTM의 지수 게이팅 테스트. 상태 추적이 필요한 형식 언어 작업을 해결하는 다양한 모델의 스케일된 정확도로 결과가 제공됩니다. 다양한 작업은 촘스키 계층별로 그룹화되어 있습니다

 

xLSTM의 메모리 혼합과 지수 게이팅 테스트
xLSTM의 새로운 지수 게이팅과 메모리 혼합을 테스트하여 상태 추적 문제를 해결할 수 있는지 확인합니다 (Merrill et al., 2024; Merrill & Sabharwal, 2023). 우리는 Delétang et al. (2023)의 형식 언어 작업을 구현하고 확장하여 길이 외삽을 위한 다중 길이 훈련을 가능하게 합니다. 모든 작업에 대한 자세한 설명과 확장된 결과는 부록 B.1.1을 참조하십시오. 우리는 xLSTM을 트랜스포머, 상태 공간 모델, 순환 신경망(RNN)을 포함한 다른 방법들과 비교합니다. 테스트된 방법들의 정확도는 작업에 관련된 토큰에서 평가됩니다. 정확도는 0(랜덤)에서 1(완벽)까지의 범위로 조정됩니다. 이러한 작업에서 다음 방법들의 2-블록 아키텍처를 비교합니다: xLSTM[0:1] (즉, sLSTM만), xLSTM[1:0] (즉, mLSTM만), xLSTM[1:1], Llama, Mamba, RWKV, Retention, Hyena, LSTM, 그리고 트랜스포머 블록 내의 LSTM (LSTM (Block)). 이 실험의 결과는 그림 4에 나와 있습니다. 메모리 혼합(상태 추적 없음) 없이 트랜스포머나 상태 공간 모델과 같은 모델은 예를 들어 패리티 작업과 같은 정규 문법을 해결할 수 없습니다. 이 결과는 트랜스포머와 상태 공간 모델이 RNN보다 근본적으로 덜 강력하다는 연구 결과와 일치합니다 (Merrill et al., 2024; Merrill & Sabharwal, 2023; Delétang et al., 2023).

 

그림 5: 컨텍스트 길이 2048에서 다중 쿼리 연상 회상 작업에 대한 다양한 모델의 메모리 용량 테스트. 각 패널은 다른 수의 키-값 쌍에 전념합니다. x축은 모델 크기를, y축은 검증 정확도를 나타냅니다.

 

연상 회상 작업에서 xLSTM의 메모리 용량 테스트
이 실험에서는 다중 쿼리 연상 회상 작업에서 xLSTM의 새로운 행렬 메모리를 메모리 용량 측면에서 테스트합니다 (Arora et al., 2023): 각 시퀀스마다 나중에 검색을 위해 큰 어휘에서 키-값 쌍이 무작위로 선택되어야 합니다. 원래 작업의 난이도를 높이기 위해, 키-값 쌍의 수를 최대 256까지 늘리고 문맥 길이를 최대 2048까지 확장합니다. 따라서 다양한 모델의 메모리 용량에 대한 광범위한 테스트를 수행할 수 있습니다. 우리는 Llama, Mamba, RWKV-5, RWKV-6, xLSTM[1:1] 및 xLSTM[1:0]의 2-블록 아키텍처를 비교합니다. 모델은 쌍을 회상하는 정확도로 평가됩니다. 트랜스포머(e.g., Llama)는 코딩 차원에서 지수적인 메모리를 가지고 있기 때문에 이 작업의 골드 스탠다드가 됩니다 (Ramsauer et al., 2021). 결과는 그림 5에 나와 있습니다. xLSTM[1:1]은 모든 비-트랜스포머 모델 중에서 가장 우수한 성능을 보이며, 작은 모델에서도 마찬가지입니다. 흥미롭게도, sLSTM 블록은 메모리 용량을 감소시키지 않고 오히려 이를 활용하여 256개의 키-값 쌍과 같은 가장 어려운 작업에서도 유리합니다. 추가 결과는 부록 B.1.2에 나와 있으며, 외삽 분석은 훈련 중에 본 것보다 더 긴 문맥으로 외삽할 때 xLSTM의 향상된 메모리 용량이 유지됨을 나타냅니다.

 

롱 레인지 아레나에서 xLSTM의 긴 문맥 처리 능력 테스트
긴 시퀀스와 큰 문맥에서 xLSTM의 성능을 평가하기 위해 롱 레인지 아레나(Tay et al., 2021)에서 다양한 방법을 비교합니다. xLSTM은 모든 작업에서 일관되게 강력한 성능을 보여주어, xLSTM 아키텍처가 긴 문맥 문제의 다양한 측면을 처리하는 데 매우 효율적임을 시사합니다. 자세한 내용은 부록 B.1.3을 참조하십시오.

 

4.2 방법 비교 및 제거 연구

이 논문의 주요 질문은 새로운 LSTM 변형을 확장할 때 언어 모델링에서 무엇을 달성할 수 있는가입니다. 따라서 우리는 xLSTM, 트랜스포머, 상태 공간 모델(SSM) 및 다른 방법들을 SlimPajama에서 150억 개의 토큰으로 자동 회귀 언어 모델링 설정에서 학습시킵니다. 우리는 학습된 모델들을 검증 세트에서 비교합니다. 마지막으로, xLSTM에 대한 제거 연구를 수행합니다.

xLSTM을 다른 방법들과 비교하기
비교를 위해, 우리는 SlimPajama(Soboleva et al., 2023)에서 150억 개의 토큰으로 모델들을 학습시킵니다. 학습된 모델들은 검증 세트에서의 퍼플렉시티(perplexity)로 평가됩니다. 우리는 다음과 같은 방법들을 비교합니다: xLSTM (우리의 새로운 방법), GPT-3 (트랜스포머) (Brown et al., 2020), Llama (트랜스포머) (Touvron et al., 2023), H3 (SSM) (Fu et al., 2023), Mamba (SSM) (Gu & Dao, 2023), RWKV-4 (RNN) (Peng et al., 2023), RWKV-5 (RNN) (Peng et al., 2024), RWKV-6 (RNN) (Peng et al., 2024), GLA (선형 트랜스포머) (Yang et al., 2023), HGRN (RNN) (Qin et al., 2023), HGRN2 (RNN) (Qin et al., 2024), RetNet (선형 트랜스포머) (Sun et al., 2023), Hyena (선형 트랜스포머) (Poli et al., 2023), xLSTM[1:0], 그리고 xLSTM[7:1] (섹션 4 참조). 모델들은 혼합 정밀도로 학습되었으며, RWKV-5, RWKV-6, GLA, HGRN, HGRN2는 참조 구현에서 혼합 정밀도 학습을 지원하지 않았습니다. 우리는 방법들을 (a) 트랜스포머, (b) 상태 공간 모델(SSM), (c) 순환 신경망(RNN) 및 선형 트랜스포머로 분류합니다. 선형 트랜스포머는 트랜스포머 어텐션 메커니즘을 대체하는 선형 방법입니다. 모델들은 크기가 350M 파라미터인 GPT-3 모델과 일치하며, 즉, 임베딩 차원 1024 및 24개의 잔차 블록을 갖습니다. GPT-3만 토큰과 출력 임베딩에 공유 가중치를 사용하므로 파라미터가 더 적습니다. 표 1의 결과는 xLSTM이 검증 퍼플렉시티에서 기존의 모든 방법을 능가함을 보여줍니다. 자세한 내용은 부록 B.2를 참조하십시오. 그림 6은 이 실험에 대한 확장 동작을 보여주며, xLSTM이 더 큰 모델에서도 유리하게 수행될 것임을 나타냅니다.

 

표 1: SlimPajama에서 150억 개의 토큰으로 학습한 후 다음 토큰 예측에 대한 방법 비교. 모델 클래스 내에서 최고의 검증 퍼플렉시티(즉, 트랜스포머, LSTM, SSM, RNN, 선형 트랜스포머)는 밑줄로 표시되며, 전체에서 가장 좋은 결과는 굵게 표시됩니다. 각 모델 클래스에서 최고의 성능을 보이는 방법은 나중에 섹션 4.3에서 대형 언어 모델(LLM) 학습에 사용됩니다. 새로운 메모리를 갖춘 xLSTM(xLSTM[1:0] 및 xLSTM[7:1])이 가장 우수한 성능을 보입니다.

 

제거 연구
표 1과 그림 6은 xLSTM이 SlimPajama에서 150억 개의 토큰으로 학습할 때 언어 모델링에서 뛰어난 성과를 달성함을 보여줍니다. 따라서 xLSTM의 어떤 요소가 바닐라 LSTM 성능을 개선하는지에 대한 질문이 자연스럽게 제기되며, 이는 개별 새로운 xLSTM 구성 요소에 대한 제거 연구를 촉발합니다. 이를 위해, 우리는 바닐라 LSTM 아키텍처를 단계별로 xLSTM 아키텍처로 변환합니다. 먼저, LSTM 레이어를 pre-LayerNorm 잔차 백본에 통합하고, 그 다음 업 프로젝션 블록을 추가하며, 그런 다음 지수 게이팅을 추가하고, 마지막으로 행렬 메모리를 추가합니다. 결과는 표 2(상단)에 나와 있습니다. 제거 연구는 지수 게이팅과 행렬 메모리가 강력한 성능 개선에 기여함을 나타냅니다. 또한, 게이팅은 RNN 및 상태 공간 모델에서 항상 중요한 주제이므로, 우리는 다양한 게이팅 메커니즘을 제거하여 연구했습니다. 표 2(하단)에서, 각 게이트를 학습 가능하게 하고 입력에 의해 영향을 받도록 하는 것이 점진적인 긍정적 효과를 가져온다는 결론을 내립니다. 개별 백본 구성 요소에 대한 추가 연구는 부록 B.2에서 논의됩니다.

 

그림 6: SlimPajama에서 150억 개의 토큰으로 학습한 후 다음 토큰 예측에 대한 방법 비교. 각 모델 클래스(표 1 참조)의 최고의 방법에 대한 검증 퍼플렉시티 성능 측정이 보고됩니다. 27억 단계에서 xLSTM[7:1]의 성능 저하는 초기 학습 수렴이 느려서 특히 덜 학습된 모델로 이어진 것입니다. xLSTM은 모든 크기에서 최고의 방법입니다.

 

표 2: 제거 연구. 상단: 새로운 xLSTM 구성 요소에 대한 제거 연구는 지수 게이팅과 행렬 메모리가 바닐라 LSTM 대비 xLSTM의 강력한 성능 개선에 기여함을 보여줍니다. 하단: 다양한 게이팅 기법에 대한 제거 연구. 우리는 시그모이드 포겟 게이트와 지수 입력 게이트가 있는 xLSTM[1:0]을 고려합니다. 바이어스 초기화 ∞는 포겟 게이트가 1로 설정됨을 의미하며, [3,6]은 해당 구간에서 등간격으로 값을 취함을 나타내고, N(0,0.1)은 평균이 0이고 표준편차가 0.1인 가우시안에서 무작위로 선택된 값을 나타냅니다. PPL은 검증 퍼플렉시티를 나타냅니다. 첫 두 줄은 선형화된 어텐션과 유사한 모델에 해당하며, 네 번째 줄은 Retention, 다섯 번째 줄은 RWKV-5, 여섯 번째 줄은 RWKV-6에 해당합니다. 게이트의 입력 종속성은 더 나은 성능으로 이어집니다.

 

4.3 xLSTM을 대형 언어 모델(LLM)로 사용하기

이 연구는 대규모 언어 모델링 실험에서 xLSTM의 잠재력을 테스트하여 xLSTM을 LLM으로 사용합니다. 이를 위해 학습 데이터의 양을 증가시키고 SlimPajama에서 3000억 개의 토큰으로 학습합니다. 동일한 양의 토큰은 예를 들어 Mamba (Gu & Dao, 2023) 및 Griffin (De et al., 2024)에서도 사용됩니다. 우리는 섹션 4.2의 모델 비교에서 각각의 방법 클래스에서 최고의 성능을 보인 xLSTM, RWKV-4, Llama, 및 Mamba를 비교합니다. 우리는 다양한 모델 크기(125M, 350M, 760M, 1.3B)를 학습하고 모든 모델의 길이 외삽 능력을 테스트하며 검증 세트에서 성능을 평가합니다. 우리는 다운스트림 작업에서의 성능을 평가하고, PALOMA 벤치마크의 471개 텍스트 도메인에서 언어 모델링 성능을 테스트하며, 마지막으로 확장 법칙 동작을 조사합니다.

시퀀스 길이 외삽
먼저, xLSTM, RWKV-4, Llama, 및 Mamba의 1.3B 크기의 대형 모델에 대해 시퀀스 길이 외삽을 테스트합니다. 모든 모델은 문맥 길이 2048로 학습된 후 문맥 길이 최대 16384까지 테스트됩니다. 결과는 그림 7을 참조하십시오. 다른 방법과 달리, xLSTM 모델은 더 긴 문맥에서도 낮은 퍼플렉시티를 유지합니다.

 

그림 7: 언어 모델링에서 시퀀스 외삽. SlimPajama에서 3000억 개의 토큰으로 학습한 후 SlimPajama 검증 세트에서 다음 토큰 예측을 수행하는 xLSTM, RWKV-4, Llama, 및 Mamba의 1.3B 크기 대형 모델 비교입니다. 모델은 문맥 길이 2048로 학습되었고, 문맥 길이 최대 16384까지 테스트됩니다. 왼쪽: 다양한 문맥 길이에서 평가된 토큰 퍼플렉시티. 다른 방법과 달리, xLSTM 모델은 더 긴 문맥에서도 낮은 퍼플렉시티를 유지합니다. 오른쪽: 검증 퍼플렉시티(PPL) 측면에서 긴 문맥 크기로 외삽할 때의 예측 품질. xLSTM은 최고의 PPL 값을 제공합니다(가장 좋은 값은 굵게, 두 번째로 좋은 값은 밑줄로 표시).

 

검증 퍼플렉시티와 다운스트림 작업
둘째, 모든 모델 크기에 대해, xLSTM, RWKV-4, Llama, 및 Mamba 모델의 다음 토큰 예측 성능과 상식 추론을 측정하는 다운스트림 작업에서의 성능을 평가합니다. 표 3의 세 번째 열에는 다양한 방법의 검증 세트 퍼플렉시티가 나와 있습니다. 검증 세트 퍼플렉시티 측면에서 xLSTM[1:0]과 xLSTM[7:1]이 모든 모델 크기에서 최고의 모델입니다. 표 3의 다른 열에는 다운스트림 작업에서의 성능이 나와 있습니다. 대부분의 작업과 모든 모델 크기에서 xLSTM이 최고의 방법입니다 — ARC 작업에서는 경우에 따라 Mamba가 최고의 방법일 수 있습니다. 자세한 내용은 부록 B.3을 참조하십시오.

PALOMA 언어 작업에서의 성능
셋째, 모든 모델 크기에 대해, PALOMA 언어 작업(Magnusson et al., 2023)에서 xLSTM, RWKV-4, Llama, 및 Mamba 모델의 다음 토큰 예측 성능을 테스트합니다. 우리는 nytimes.com부터 Reddit의 r/depression까지 571개의 텍스트 도메인에서 다음 토큰 예측 퍼플렉시티로 성능을 측정합니다. 표 4는 언어 모델링(첫 7열)과 세분화된 도메인 벤치마크(마지막 5열)로 그룹화된 토큰 예측 퍼플렉시티를 보여줍니다. 이러한 언어 작업에서 xLSTM[1:0]이 xLSTM[7:1]보다 더 좋은 성능을 보입니다. xLSTM[1:0]은 571개 텍스트 도메인 중 568개(99.5%)에서 Mamba보다 낮은 퍼플렉시티를, 571개 중 486개(85.1%)에서 Llama보다 낮은 퍼플렉시티를, 571개 중 570개(99.8%)에서 RWKV-4보다 낮은 퍼플렉시티를 가집니다. 자세한 내용은 부록 B.3을 참조하십시오.

 

표 3: 검증 세트 퍼플렉시티와 다운스트림 작업

SlimPajama에서 3000억 개의 토큰으로 학습한 후, 다음 토큰 예측 및 다운스트림 작업에서 xLSTM, RWKV-4, Llama, 및 Mamba를 비교합니다. 모델 크기는 125M, 250M, 760M, 및 1.3B입니다. 첫 번째 열은 방법들을, 두 번째 열은 실제 파라미터 수를 보여줍니다. 세 번째 열은 검증 세트 퍼플렉시티를 나열하고, 나머지 열은 다운스트림 작업에서의 성능을 보여줍니다. 모델 크기별 최고의 모델은 굵게 표시되고, 두 번째로 좋은 모델은 밑줄로 표시됩니다. 대부분의 작업과 모든 모델 크기에서 xLSTM이 최고의 방법입니다 — ARC 작업에서는 경우에 따라 Mamba가 최고의 방법입니다. 검증 세트 퍼플렉시티 측면에서 xLSTM[1:0]과 xLSTM[7:1]이 두 가지 최고의 모델입니다.

 

표 4: PALOMA 언어 모델링 작업에서의 성능

SlimPajama에서 3000억 개의 토큰으로 학습한 후, PALOMA 언어 벤치마크에서 다음 토큰 예측 퍼플렉시티로 xLSTM, RWKV-4, Llama, 및 Mamba를 비교합니다. 모델 크기는 125M, 250M, 760M, 및 1.3B입니다. 두 번째 열은 실제 파라미터 수를 보여줍니다. 571개 텍스트 도메인은 언어 모델링(첫 7열)과 세분화된 도메인 벤치마크(추가 5열)로 그룹화됩니다. 마지막 열은 모든 작업에서의 평균 퍼플렉시티를 보여줍니다. 모델 크기별 최고의 모델은 굵게 표시되고, 두 번째로 좋은 모델은 밑줄로 표시됩니다. xLSTM은 최고의 성능을 제공합니다.

 

확장 법칙
넷째, 성능을 더 큰 모델 크기로 외삽할 수 있게 하는 멱법칙 확장 동작을 평가합니다(Kaplan et al., 2020; Brown et al., 2020). 그림 8은 확장 동작을 보여줍니다. 모든 모델은 유사한 확장 동작을 보이지만, 서로 다른 오프셋을 가집니다. RWKV-4가 가장 낮은 성능을 보이고, 그 다음으로 Llama와 Mamba가 뒤따릅니다. xLSTM은 Mamba보다 더 나은 성능을 보이며, Mamba가 Llama에 비해 가지는 마진과 유사한 마진을 가집니다. 확장 동작은 더 큰 모델에 대해서도 xLSTM이 트랜스포머와 상태 공간 모델에 비해 계속해서 유리한 성능을 보일 것임을 나타냅니다.

그림 8: 확장 법칙. SlimPajama에서 3000억 개의 토큰으로 학습한 후 SlimPajama 검증 세트에서 xLSTM, RWKV-4, Llama, 및 Mamba의 다음 토큰 예측 퍼플렉시티. 모델 크기는 125M, 350M, 760M, 및 1.3B입니다. 각 모델 클래스에서 최고의 모델이 선택되었습니다(표 1 참조). 확장 법칙은 더 큰 모델에 대해서도 xLSTM이 잘 수행될 것임을 나타냅니다.

 

5. 제한사항

(i) mLSTM과 달리, sLSTM의 메모리 혼합은 병렬화 가능한 작업을 금지하므로 빠른 병렬 구현을 허용하지 않습니다. 그럼에도 불구하고 우리는 sLSTM에 대해 빠른 CUDA 커널을 개발했으며, 현재 우리의 병렬 mLSTM 구현보다 약 1.5배 느립니다.

(ii) mLSTM에 대한 CUDA 커널은 최적화되지 않았으므로 현재 구현은 FlashAttention 또는 Mamba에서 사용된 것보다 약 4배 느립니다. FlashAttention 관점에서 더 빠른 CUDA 요소를 얻을 수 있습니다. 여전히, 메모리 업데이트 및 검색은 파라미터를 타이지 않고 표준 행렬 연산을 통해 병렬 처리될 수 있으므로 실제 시간 오버헤드는 복잡한 메모리 소요가 적습니다.

(iii) mLSTM의 행렬 메모리는 높은 계산 복잡도를 가지며 d×dd \times dd×d 행렬을 처리해야 합니다. 그러나 행렬 메모리는 파라미터를 타이지 않고 표준 행렬 연산을 사용하여 병렬 처리될 수 있으므로 실제 시간 오버헤드는 복잡한 메모리 소요가 적습니다.

(iv) 포겟 게이트의 초기화는 신중하게 선택해야 합니다. 행렬 메모리가 시퀀스 길이와 독립적이므로 시퀀스 길이가 증가해도 메모리를 유지하지 않아도 됩니다. 그럼에도 불구하고 이는 최대 16k의 컨텍스트 길이에서는 제한사항으로 보이지 않습니다(섹션 4.3 참조).

(v) 대규모 언어 실험을 위한 계산 비용이 비싸기 때문에, 특히 대형 xLSTM 아키텍처에서는 아키텍처나 하이퍼파라미터를 완전히 최적화하지 않았습니다. 우리는 xLSTM이 완전한 잠재력에 도달하기 위해 광범위한 최적화 과정이 필요할 것으로 예상합니다.

 

결론

우리는 간단한 질문에 부분적으로 답변했습니다: LSTM을 수십억 개의 파라미터로 확장하면 언어 모델링에서 얼마나 멀리 갈 수 있을까요? 현재까지 우리는 "최소한 트랜스포머나 상태 공간 모델과 같은 최신 기술만큼 멀리 갈 수 있다"고 답할 수 있습니다. 우리는 메모리 혼합과 새로운 메모리 구조를 갖춘 지수 게이팅을 통해 LSTM을 xLSTM으로 개선했습니다. xLSTM 모델은 트랜스포머와 상태 공간 모델과 같은 최신 방법들과 비교할 때 언어 모델링에서 유리한 성능을 보입니다. 확장 법칙은 더 큰 xLSTM 모델이 트랜스포머 기술로 구축된 현재의 대형 언어 모델에 진지한 경쟁자가 될 것임을 나타냅니다. xLSTM은 강화 학습, 시계열 예측 또는 물리 시스템 모델링과 같은 다른 딥러닝 분야에 상당한 영향을 미칠 잠재력을 가지고 있습니다.