1. 시작하며

워! 드디어 intern students 명단에 이름을 올렸습니다. 이 연구실에서 진행하고 있는 움직임이 정말 신선하고 마음에 들어 정착하고 싶은 마음이 굳혀졌습니다. 오늘은 연구실 세미나에서 진행했던 hungry hungry hippos에 관련된 논문 리뷰를 진행하려고 합니다. 앞의 이니셜을 취해서 H3라고 불리는데요. transformer의 attention 기법과 RNN의 recurrency 속성을 교묘하게 활용해서 language modeling 분야에서 performance(perplexity)와 speed를 모두 개선시킨 논문이라고 할 수 있겠습니다.

2. paper review!

슬라이드 링크: https://docs.google.com/presentation/d/11rgLHIqnl60H84L6wjS82d452bmiZ0ofFLNlXJUu_Cw/edit?usp=sharing

바로 리뷰 시작하겠습니다! hungry Hippos 라는 게임이 있는데요. 왼쪽 부분에 그림으로 첨부했습니다. 이 논문에서 위 게임과 H3 (hungry hungry Hippos)의 관련성은 크게 없는 것으로 보입니다. Catchy하게 작명한 것 같네요.

위의 두 논문을 주로 참고했습니다. 처음으로는 하이에나 논문을 배당 받았었는데요. 결국 하이에나 아키텍쳐가 H3의 일반화된 구조임을 깨닫고 H3부터 리뷰하기로 마음 먹었습니다. 참고로 계속 등장하는 단어 "HiPPO"는 사실 하마가 아니라 "high order polynomial projection operators"를 나타냅니다. 이는 이 state space model 계열의 논문에서 근간으로 하고 있는 이론적 배경입니다. recurrent 구조에서 memory loss가 발생하지 않도록 HiPPO를 활용한 수학적 기교를 부렸습니다. 자세한 내용과 증명은 HiPPO 논문을 참고 부탁드려요. (저도 추후에 리뷰해볼까 계획 중 입니다.)

목차는 위와 같습니다. H3의 contribution인 H3 layer과 flashconv로 향하는 여정이라고 생각하면 됩니다.

보시다시피 하이에나는 H3의 메커니즘의 일반화된 형태라고 볼 수 있습니다. matrix form에 대한 유도는 논문의 부록에 자세하게 나와있습니다.

전달드렸다시피, H3 paper에서 주요 contribution은 h3 layer의 제안 그리고 state space model에 대한 모든 모델에 대해서 범용으로 적용할 수 있는 flashconv를 창조했다는 것에 있습니다. H3 layer의 그림에서 볼 수 있다시피 transformer의 attention module을 겨냥한 것으로 추론해볼 수 있습니다. flashconv 역시 그림이 직관적인데요. recurrent 속성을 활용해 input을 slicing하고 next state을 교묘하게 passing 해줘서 8K sequence 길이를 넘어서까지도 아무런 문제 없이 output을 계산할 수 있음을 확인할 수 있습니다.

langauge domain에서 state space model이 transformer 기반의 모델을 넘어설 수 있을까? 에서 질문을 출발해 H3가 탄생했습니다. 트랜스포머로부터 과거로부터 토큰을 기억하는 능력(casuality)과 token들을 비교하는 능력(multiplicative interaction)을 답습하여 state space model 기반의 H3에 적용시켰습니다. 노란색 줄은 H3 layer 그 자체를 나타내는데요. 여기서 SSM이 HiPPO를 이론적 배경으로 두고 있고 두 개의 SSM이 등장하기 때문에 Hungry Hungry hippos라는 이름을 작명하게 된 것 같네요.

H3에 대한 실험결과를 살펴보겠습니다. langauge modeling에 관련된 실험을 해서 perplexity를 성능의 척도로 활용했습니다. perplexity가 낮을수록 고성능입니다.

보시다시피 같은 parameter 숫자에서 H3 모델의 성능이 transformer 기반의 모델을 능가함을 알 수 있었습니다.

Superglue dataset에 대해서도 실험을 진행했습니다.

H3 모델의 성능이 transformer 기반의 모델을 능가함을 쉽게 확인할 수 있습니다.

H3 모델이 recurrent한 속성을 가지고 있다보니까 inference 측면에서도 높은 스피드를 확인할 수 있습니다.

SSM에 대해서 간단하게 짚고 넘어가보겠습니다. control 분야에서 흔하게 등장하는 system dynamics가 등장하고 있습니다. u만 관측가능한 sequence이며 x는 hidden property를 지닙니다. u를 활용해 y (output)을 추론할 수 있습니다. 한편, 이론적 근간에는 HiPPO가 있습니다. 이는 sequence history를 보존해주는 이론적 장치가 있어서 RNN의 단점을 보완하면서도 transformer의 attention 모듈의 특성을 잘 활용할 수 있습니다. matrix A~D는 learnable 합니다.

y_i 를 전개해보면 위와 같은 결과를 얻을 수 있습니다. 물론, convolution 연산은 frequecy domain에서 단순 곱셈 (convolution theory)과 같으므로 시간을 훨씬 save할 수 있습니다.

ssm을 recurrent module로 나타내봤습니다. 좀 더 직관적인 이해를 돋기 위해서 제가 직접 그린 그림이므로 논문에는 등장하지 않는 점 전달드립니다. SSM은 그 자체로 casuality의 속성(미래를 참조하지 아니하고 과거만을 참조하므로)을 지닙니다.

한편, H3는 겉보기엔 state space model이지만 그 안을 드려다보면 linear transformer에서 엄청난 insight를 얻었다고 볼 수 있습니다. query와 key를 활용해 attention map을 직접 연산하지 않고 이 둘을 decouple하여 (K, V) 부터 연산을 진행해주면 계산복잡도 측면에서 훨씬 이득임을 파악할 수 있습니다. 이들의 식 또한 좌(original transformer), 우(linear transformer)로 다름을 알 수 있습니다. 오른쪽에 등장한 식만 보더라도 적절히 변형시켜 SSM과 매칭 시킬 수 있을 것 같지 않나요?

위와 같이 치환해서 식을 정리해보면 SSM과 완벽히 동일한 식임을 알 수 있습니다.

이제 다시한번 변형을 거칩니다. 곧이 곧대로 linear transformer의 식을 활용하는 것이 아닌 저자들의 구미에 알맞게 식을 변형시켜 H3 layer의 식을 얻어냈습니다. 색깔로 칠한 바운딩 박스의 부분끼리 대응되는 것을 쉽게 관찰할 수 있죠.

한편, A 가 shift matrix일 때 마치 x가 sliding window를 지닌 메모리의 역할을 하고 있음을 확인해볼 수 있습니다. 이를 활용해 key의 값을 적절히 변형시켜 중간값으로 활용하고 있습니다.

H3의 전체적인 알고리즘은 다음과 같습니다. 지금까지의 이야기에서 multihead만 추가된 형태입니다.

그 뒤로, flashconv라는 계산법도 저자들이 제안을 했는데요. 이는 8K의 길이 이상을 지닌 sequence에 대한 forward 연산을 계산할 수 있는 방법론을 제시합니다. 어떤 SSM 기반의 모델에 대해서도 적용할 수 있기때문에 유의미하다고 할 수 있습니다.

알고리즘 자체는 정말 간단합니다. 5번과 6번 라인이 제가 표기한 두 식에서 출발했음을 알 수 있죠.

다만, 계산 검증을 진행하던 중, M_xy의 식이 논문에서 제시된 형태가 맞을까하는 의구심이 들긴했는데요. 이에 관련해서 의견이 있으시다면 메일, 댓글 부탁드립니다.

아무튼 간단한 아이디어 스케치를 살펴보기위해 y chunk의 i번째 index만 뽑아서 살펴보겠습니다. 우리는 what to show!의 윗부분만을 확인하면 이 방법론이 맞다는 것을 검증할 수 있죠.

그래서 검증해보니 식의 형태가 비슷하게 나왔습니다. 물론, 제가 앞서 change한 M_xy를 기준으로 대입하면 완벽히 그 결과가 일치하고 논문에서 제시한 M_xy를 활용하면 한칸씩 밀리는 인상을 받았습니다만, 더욱 검증을 해보고 renew 하도록 하겠습니다.

x에 관한 식도 i번째 element만 picking 해서 살펴보면 자명하게 성립함을 검증할 수 있습니다.

이를 활용한 speed up이 기존 transformer의 5.8배까지도 된다고하니 솔깃한 이슈인 것 같네요.

이것으로 리뷰 마치겠습니다. 감사합니다.

3. 마치며

앞으로는 state space model의 근간이 되는 hippo에 관련된 논문을 살펴보고 그의 파생상품들을 정리해가면서 review할 예정입니다. 연구실에서는 transformer 기반으로 조합 최적화 문제를 해결해나가는 과정에 있는데 transformer를 능가할 수 있는! 장래가 촉망되는! state space model들이 힘을 보탤수 있으면 좋을 것 같습니다~