[논문리뷰] xLSTM : Extended Long Short-Term Memory(작성중)

0. Abstract

• 1990년대에 Long Short-Term Memory(LSTM)의 핵심 아이디어로 constant error carousel과 gating이 도입
• 그러나 병렬화가 가능한 self-attention-Transformer의 등장으로 LSTM을 앞지름
• LSTM의 한계를 극복하면 어떻게 되나?
  • 첫째, 지수 게이팅을 도입 : 적절한 정규화 및 안정화 기술
  • 둘째, LSTM 메모리 구조를 변경
    • (i) sLSTM : 스칼라 메모리, 메모리 mixing
    • (ii) mLSTM(병렬화) : 매트릭스 메모리와 공분산 업데이트 
  • Residual Blcok으로 통합하면 xLSTM block이 생성됨.
• 최신 Transformer 및 State Space Model과 견줄만한 능력

sLSTM

- 기존 Sigmoid Gating -> Exponential Gating으로 변경 ( Sigmoid의 단점 보완 )

- New Memory Mixing : LSTM에도 비슷한 것이 있었음.

 

mLSTM

- 병렬가능

- 행렬 메모리 탑제

 

xLSTM

- sLMTM 및 mLSTM의 혼합

 

 

1. Introduction

 

LSTM figure

LSTM의 성공

• LSTM은 기울기 소실 문제를 위해 도입되었다.
• C의 경우 Long term memory에 해당한다고 생각하면 되고, H의 경우 Short를 담당한다고 생각하면 된다.

LSTM의 한계

• 어떤것을 저장해야 좋을지 결정하는 역량
• 제한된 저장 용량 : Vector에 longterm memory를 저장하여 저장 공간에 한계가 있다.
• 메모리 혼합 : t-1 -> t로 이동해가며 순차적인 시간 흐름에 따라 데이터를 저장하여, 순차적인 처리가 강요된다.

 

2. Extended Long Short-Term Memory

2.1 Review of the Long Short-Term Memory

모두 LSTM에 대해 알고 있다고 간주하고, 아주 간략하게 설명하자면

- 파란색은 gate들이며, t시점의 데이터 및 h_(t-1)에 의해서 결정된다.

- 녹색은 다음 셀에 들어가기 위한 벡터들 이라고 생각하면 된다.

주로 봐야 할 부분은, C_t, H_t이다.

• C_t : C_(t-1)과 Z_t에 의해서 결정된다.
  • 이때 Z_t는 x_t와 h_(t-1)에 의해서 결정되며, tanh가 중간에 있다 -> sigmoid부분이 xLSTM에서 사라짐.
• H_t : 는 C_t에 의해서 결정됨

 

2.2 sLSTM

• 1번 : Sigmoid term을 exponential term으로 바꾸면서, gradient vanishing을 완화 시킴, 그러나 단점으로는 exponential하게 값들이 커질수 있다.
• 2번 : 1번에 exponential하게 커질 수 있는 단점을 normalize term 을 통해서 완화 시킨다
  • 더 넓은 범위의 게이트 값을 사용할 수 있게 하여 정보 흐름을 더 유연하게 제어 가능
  • 역전파 과정에서 더 강한 기울기 신호를 전달해, 장기 의존성 개선 가능
  • 큰 양수 입력값(중요한 정보)에 대하여 큰 값을 출력 가능함 -> 더 중요한 정보를 더 오랫동안 기억 할 수 있게함
• 3번 : 따라서 h_t는 normalize term을 통해 큰 값을 갖지 않는다

 

허나, exponential이 추가되어 overflow발생이 가능하여 stabilize를 아래와 같이 진행한다.

 

New Memory Mixing

• LSTM과 비슷한 역할을 하는 cell이 있다.
• Rz, Ri, Rf, Ro의 recurrent연결을 통해서 멀티헤드 역할을 하게 한다. R들은 대각행렬 구조를 갖고 있다.
• 여러개의 작은 LSTM이 병렬로 작동하는 것과 유사함
• 헤드간에는 정보 교류가 없다.

R들의 모양의 예시는 아래와 같다.

[[R1  0  0  0],
 [0  R2  0  0],
 [0   0 R3  0],
 [0   0  0 R4]]

 

 

2.3 mLSTM

• LSTM의 저장용량 늘리기 : 스칼라 c ∈ R(LSTM) -> 행렬 C∈R^(d×d)(mLSTM)으로 확장
• 그렇다면 어떻게 공간에 저장된 정보들을 불러오나?

왼쪽 맨처음 빨간 박스에는 아무런 정보도 없는 메모리 블록이 있다라고 하자.

첫번째 key value를 outer product를 실행해 d*d 메트릭을 만들고 Memory Block 0에 정보를 더해준다.

 

두번째 key value를 Outer product를 실행해 d*d 메트릭을 만들고 Memory Block 1에 정보를 더해준다.

 

이렇게 할경우 Memory Block에 t시점들의 정보들이 쌓이게 되는데, 그럼 어떻게 해야 특정 시점의 정보를 다시 복원할 수 있을까?

 

파란색 글씨에 정보 복원에 대한 것을 서술하였다.

 

MemoryBlcok에 해당하는 key를 곱해주면, value를 되찾을 수 있음.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

2.4 xLSTM Architecture

트랜스포머 아키텍쳐와 비슷하게 생겼는데, sLSTM부분에서는 NH의 부분이 어떻게 구성되어 있는지 코드를 살펴볼 필요가 있으며, mLSTM부분에서는 mLSTM블록의 정보들이 어떻게 섞이는지 살펴볼 필요가 있다. 논문으로 구체적인 부분을 살펴보기 어려운 부분이 있다.

 

https://github.com/NX-AI/xlstm

GitHub - NX-AI/xlstm: Official repository of the xLSTM.