딥러닝 기술의 발전은 놀라울 정도로 빠르게 이루어지고 있습니다. 불과 몇 년 전만 해도 RNN(Recurrent Neural Network)이 자연어 처리(NLP)의 대표 모델로 군림했지만, 2017년 구글의 트랜스포머(Transformer) 아키텍처 등장 이후 판도가 크게 바뀌었죠. 그리고 최근에는 트랜스포머를 넘어설 새로운 패러다임을 찾는 움직임이 활발합니다. 그 중심에는 State Space Model, 줄여서 SSM이 있는데요. 과연 SSM은 트랜스포머를 대체할 수 있을까요? 이 글에서는 트랜스포머와 SSM의 특징을 비교 분석하고, NLP의 미래를 조망해 보고자 합니다.

트랜스포머 아키텍처의 장점과 단점

트랜스포머 아키텍처

2017년 구글에서 발표한 트랜스포머(Transformer) 아키텍처는 자연어 처리(NLP) 분야에 일대 혁신을 가져왔습니다. 현재 대부분의 딥러닝 모델에서 트랜스포머가 사용될 정도로 그 영향력은 지대하죠. 트랜스포머의 핵심은 셀프 어텐션(self-attention) 메커니즘에 있습니다. 기존의 RNN(Recurrent Neural Network) 기반 모델들과 달리 어텐션을 통해 병렬 계산이 가능해졌고, 입력 시퀀스의 각 위치 간 관계를 전부 고려함으로써 문맥 이해력을 높일 수 있게 되었습니다. 또한 모델 크기를 빠르게 확장시킬 수 있다는 점도 트랜스포머의 큰 장점으로 꼽힙니다.

그러나 트랜스포머도 완벽하지는 않습니다. 먼저 계산 복잡도 문제를 들 수 있습니다. 트랜스포머는 어텐션 연산 과정에서 막대한 연산량을 필요로 합니다. 이는 곧 메모리 사용량 증가로 이어지죠. 두 번째로 입력 시퀀스의 길이가 길어질수록 장거리 의존성(Long Range Dependency) 문제가 발생할 수 있습니다. 멀리 떨어진 토큰 간의 관계를 파악하기 어려워지는 것입니다.

flashattention과 Memory Efficient Attention

이런 한계를 극복하기 위해 연구자들은 어텐션 메커니즘 자체를 개선하려는 시도를 해왔습니다. 대표적으로 ‘Memory Efficient Attention’이라는 개념이 제시되었는데요, 이는 어텐션 연산을 근사화하거나 분할하는 방식으로 메모리 사용량을 줄이고자 하는 접근법입니다. 그러나 이런 노력에도 불구하고 트랜스포머가 가진 본질적인 한계를 극복하기에는 역부족이라는 지적이 있어왔습니다.

트랜스포머를 대체하려는 노력 : SSM

상태공간방정식 (SSM)

최근에는 트랜스포머를 대체할 수 있는 새로운 패러다임으로 State Space Model(SSM)에 주목하는 연구자들이 늘고 있습니다. SSM은 제어공학이나 신호처리 분야에서 사용되던 모델로, 다변수 입력과 출력을 다룰 수 있고 행렬로 표현 가능하다는 특징이 있죠. 스탠포드대학교의 Christopher Re 교수팀이 SSM을 딥러닝에 접목시키는 연구를 선도하고 있습니다.

SSM을 딥러닝에 적용하기 위해서는 연속형 변수를 이산형으로 변환하는 작업이 필요합니다. 이를 위해 Euler’s Method, Bilinear Method, ZOH Method 등 다양한 이산화 기법들이 사용되고 있습니다. 성공적으로 이산화된 SSM 모델은 트랜스포머 대비 여러 이점을 가질 수 있습니다. 우선 장거리 의존성 문제를 효과적으로 다룰 수 있습니다. 모델 내에서 정보의 흐름이 시간에 따라 연속적으로 이루어지기 때문이죠. 계산 효율성도 개선되고, 메모리 사용량도 크게 줄일 수 있습니다.

Structured State Space Model 성능비교표

실제로 Albert Gu는 HiPPO(High-order Polynomial Projection Operator) 행렬을 활용해 SSM 모델의 장거리 의존성 처리 능력을 대폭 향상시켰습니다. 벤치마크 테스트 결과 트랜스포머의 어텐션 메커니즘보다 훨씬 우수한 성능을 보여주었다고 합니다. 또한 SSM은 두 가지 작동 모드를 지원합니다. Recurrent Mode는 추론 시 빠른 속도를 보여주고, Convolution Mode는 학습 시 병렬처리에 최적화되어 있습니다. 모드를 선택적으로 사용함으로써 속도와 효율성을 동시에 높일 수 있게 된 것이죠. 메모리 사용량 측면에서도 SSM은 트랜스포머보다 큰 이점이 있습니다. 트랜스포머는 입력 시퀀스 길이의 제곱에 비례하는 O(N^2) 메모리 복잡도를 가지는 반면, SSM은 O(N+L) 수준으로 훨씬 적은 메모리를 사용합니다. (N: state 차원, L: 시퀀스 길이)

SSM의 단점과 개선 노력들

하지만 SSM이 만능은 아닙니다. 아직 해결해야 할 숙제들이 남아있습니다. 무엇보다 LTI(Linear Time Invariant) 기반으로 학습한 SSM 모델들의 성능이 트랜스포머에 미치지 못하고 있습니다. LTI는 모든 토큰을 빠짐없이 읽어야 하고, 한 번 읽은 내용을 잊을 수 없다는 특성 때문에 비효율이 발생하는 것으로 보입니다. 예를 들어 웹 문서 내의 불필요한 문구(저작권 표시 등)를 건너뛰지 못하고, 문맥이 바뀌어도 이전 내용을 계속 기억해야만 하는 식입니다.

prefix-sum 알고리즘

이런 문제를 피하기 위해 최근에는 LTI를 과감히 포기하고 Recurrent Mode에 집중하는 시도가 이루어지고 있습니다. 또한 SSM 학습 방법 자체를 개선하려는 노력도 활발합니다. Associate Scan이라는 병렬 계산 기법을 활용해 Recurrent Mode에서도 병렬 처리가 가능하도록 만드는 것이 대표적 사례입니다. Associate Scan은 prefix-sum 연산을 이진 트리 구조로 구현함으로써 병렬화를 달성합니다. 트리의 루트에서 리프까지 값들을 전파하는 과정에서 중간 노드들의 부분합이 동시다발적으로 계산되는 것이죠.

griffin 성능비교표

Griffin 아키텍처에서는 Linear Scan이라는 기법으로 Hidden State를 가상 메모리에 상주시키고 청크 사이즈를 늘리는 방식으로 성능을 끌어올렸다고 합니다. Jax-triton을 계승한 custom Pallas Kernel을 통해 구현되었다고 하네요.

최근에는 RetNet이나 Mamba와 같은 혁신적인 SSM 아키텍처들도 속속 등장하고 있습니다. 이들은 SSM의 장점은 계승하면서도, 트랜스포머의 강점까지 흡수하려 노력하고 있죠. RetNet은 Linear Attention을 SSM에 결합한 형태로, 기존의 SSM과 트랜스포머 모델을 보완하는 위치에 있습니다. Mamba는 ZOH 기법으로 이산화하면서도 연속적 요소를 일부 남겨두는 독특한 아키텍처를 채택했습니다. FlashAttention에서 영감을 받아 파라미터는 HBM에, 히든 스테이트는 GPU 온칩 메모리에 저장하는 메모리 최적화 기법도 도입했다고 합니다.

결론

SSM 한계점

그럼에도 불구하고 현재까지는 대부분의 태스크에서 트랜스포머의 우위가 점쳐지는 것이 사실입니다. 특히 시퀀스의 특정 부분을 정확히 복사하거나 검색하는 능력, 문맥 내 언어 이해력 측면에서는 SSM이 트랜스포머를 따라잡기 어려워 보입니다. 트랜스포머는 어텐션을 통해 이전 토큰들의 조합을 고려할 수 있고, 이를 바탕으로 문맥에 적합한 표현(representation)을 학습할 수 있기 때문이죠.

그러나 SSM의 발전 속도와 잠재력을 고려하면 가까운 미래에 트랜스포머를 대체할 가능성도 충분해 보입니다. 더욱이 지금 이 순간에도 수많은 연구자들이 새로운 아키텍처를 고안하기 위해 열심히 노력하고 있습니다. 앞으로 트랜스포머와 SSM, 그리고 새로운 아키텍처들 사이의 각축전이 더욱 치열해질 것으로 예상됩니다. 어쩌면 10년 뒤의 딥러닝 판도는 지금과는 완전히 다른 모습일지도 모르겠네요. 자연어 처리와 딥러닝 분야의 연구자들에게는 황금 같은 시간이 아닐 수 없습니다. 새로운 혁신을 향한 끝없는 도전이 이어질 예정일겁니다!

레퍼런스

Attention is all you need by Ashish VASWANI, Noam SHAZEER, Niki PARMAR, Jakob USZKOREIT, Llion JONES, Aidan N. GOMEZ, Lukasz KAISER, Illia POLOSUKHIN(2017)

Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context by Dai et. al (2019)

Self-attention does not need O(n^2) Memory by Markus Rabe and Charles Staats (2021)

Memformer: A Memory-Augmented Transformer for Sequence Modeling by Wu et. al (2020)

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Dao et. al (2022)

HiPPO: Recurrent Memory with Optimal Polynomial Projections by Albert GU, Tri DAO, Stefano ERMON, Atri RUDRA, Christopher RÉ (2020)

Efficiently Modeling Long Sequences with Structured State Spaces by Albert GU, Karan GOEL, et Christopher RÉ (2021)

The doctoral thesis of Albert GU (2023)

Mamba-Primer by Alexander Rush (2024)
Prefix Sums and Their Applications by Guy E. Blelloch (2020)

Mamba: Linear-Time Sequence Modeling with Selective State Spaces by Albert Gu and Tri Dao (2024)

Retentive Network: A Successor to Transformer for Large Language Models by Sun et. al (2023)

Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models by De et.al (2024)

IN-CONTEXT LANGUAGE LEARNING :ARCHITECTURES AND A LGORITHMS by Ekin Akyurek, Bailin Wang, Yoon Kim, Jacob Adnreas (2024)
Repeat After Me:Transformers are Better than State Space Models at Copying by Samy Jelassi, David Brandfonbrener, Sham M.Kakade and Eran Malach (2024)