Mamba Neural Operator: Who Wins? Transformers vs. State-Space Models for PDEs

이 논문은 PDE(편미분방정식) 해법에서 기존 트랜스포머의 한계를 극복하고 장기 의존성과 연속 동역학을 더 효과적으로 포착하기 위해 구조화된 상태 공간 모델 (SSM) 과 신경 연산자를 통합한 'Mamba Neural Operator(MNO)'를 제안하며, 이를 통해 PDE 작업의 표현력과 정확도를 획기적으로 향상시켰음을 보여줍니다.

핵심

이 논문은 **"PDE(편미분 방정식) 를 푸는 데, 기존에 최고로 불리던 '트랜스포머 (Transformer)'와 새로 등장한 '맘바 (Mamba)' 중 누가 진짜 승자인가?"**를 탐구한 연구입니다.

결론부터 말씀드리면, 새로운 '맘바 (Mamba)'가 기존 '트랜스포머'를 압도하며 새로운 승자로 떠올랐습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 배경: 왜 PDE 를 풀기 어려울까요?

비유: 거대한 날씨 예보
편미분 방정식 (PDE) 은 열이 퍼지는 현상, 물이 흐르는 현상, 바람이 부는 현상 등 자연계의 복잡한 변화를 수학적으로 설명하는 공식입니다. 이를 푸는 것은 마치 **"내일 서울 전역의 날씨를 1 시간 단위로 정확히 예측하는 것"**과 같습니다.

• 기존 방식 (전통적 수치 해석): 그리드 (그물망) 를 촘촘하게 깔고 하나하나 계산합니다. 정확하지만 계산량이 너무 많아 컴퓨터가 지쳐버립니다.
• 딥러닝 방식 (신경 연산자): 컴퓨터에게 날씨 패턴을 학습시켜, 한 번만 보면 바로 미래를 예측하게 합니다.

2. 이전의 챔피언: 트랜스포머 (Transformer)

비유: "모든 사람을 동시에 부르는 거대한 회의"
최근 PDE 해결의 주역은 '트랜스포머'였습니다.

• 장점: 회의실 (데이터) 에 있는 모든 사람 (픽셀) 이 서로 눈을 마주치며 대화할 수 있게 해줍니다. 멀리 떨어진 사람과도 즉시 소통하므로 (전역적 수용 영역), 복잡한 날씨 패턴을 잘 이해합니다.
• 단점: 비효율적이고 비용이 많이 듭니다. 사람 수가 2 배가 되면 대화 횟수는 4 배가 됩니다 (제곱 복잡도). 또한, 긴 시간 동안의 흐름을 기억하는 데는 한계가 있어, "오래된 기억"이 흐릿해지거나 메모리를 너무 많이 잡아먹습니다.

3. 새로운 도전자: 맘바 (Mamba)

비유: "효율적인 정보 전달자"
'맘바'는 '상태 공간 모델 (SSM)'이라는 새로운 기술을 기반으로 합니다.

• 특징: 모든 사람이 동시에 대화하는 대신, 정보를 한 줄로 깔끔하게 전달하며 기억하는 방식입니다.
• 장점:
  1. 메모리 절약: 긴 시간의 흐름도 기억력 저하 없이 잘 유지합니다.
  2. 빠른 속도: 사람 수가 늘어나도 계산 비용은 선형적으로만 증가합니다 (2 배가 되면 비용도 2 배).
  3. 연속성: 자연 현상은 끊어지지 않고 이어지는데, 맘바는 이 '연속적인 흐름'을 더 잘 이해합니다.

4. 이 연구의 핵심 발견: "Mamba Neural Operator (MNO)"

저자들은 "트랜스포머를 버리고 맘바로 갈아타자"가 아니라, **"트랜스포머의 구조에 맘바의 능력을 심어주자"**고 제안했습니다.

• 이론적 연결: 놀랍게도, 수학적 이론을 증명했습니다. **"신경 연산자 (PDE 를 푸는 AI) 의 작동 원리와 맘바 (상태 공간 모델) 의 수학적 구조는 본질적으로 같다"**는 것입니다. 마치 서로 다른 이름의 두 자동차가 사실은 같은 엔진을 쓰고 있는 것과 같습니다.
• 실제 적용: 기존 트랜스포머 모델 (GNOT, OFormer 등) 의 '주의 (Attention)' 부분을 '맘바'로 교체했습니다.

5. 실험 결과: 누가 이겼나?

저자들은 물리학의 난제들 (지하수 흐름, 홍수 예측, 확산 반응, 항공기 주변의 공기 흐름 등) 을 풀며 비교했습니다.

• 정확도: 맘바를 적용한 모델은 기존 트랜스포머보다 오류가 훨씬 적었습니다. 특히 복잡한 난류 (소용돌이) 나 긴 시간 흐름을 예측할 때 압도적이었습니다.
• 효율성: 트랜스포머는 메모리를 10 배나 더 쓰거나 계산 시간이 10 배 더 걸렸지만, 맘바는 이를 획기적으로 줄였습니다.
• 데이터 양: 데이터가 적을 때도 맘바가 더 잘 학습했습니다.

결론: 트랜스포머는 여전히 훌륭하지만, PDE(자연 현상) 를 푸는 특수 목적에서는 **맘바가 더 효율적이고 정확한 '최고의 도구'**임이 입증되었습니다.

6. 한 줄 요약

"기존의 거대한 회의 (트랜스포머) 는 복잡하고 비싸지만, 새로운 정보 전달 시스템 (맘바) 은 적은 비용으로 더 오래, 더 정확하게 자연의 흐름을 예측합니다. 이제 PDE 해결의 새로운 표준은 '맘바'입니다."

이 연구는 인공지능이 과학적 계산 (기후, 유체 역학 등) 에서 더 빠르고 정확하게 작동할 수 있는 길을 열었다는 점에서 매우 중요합니다.

기술 요약

1. 문제 제기 (Problem Statement)

• 배경: 편미분 방정식 (PDE) 은 열 전달, 유체 역학 (Navier-Stokes), 생물학적 시스템 등 복잡한 물리 현상을 모델링하는 데 필수적입니다. 그러나 해의 존재가 보장되지 않거나 해석적 해를 구하기 어려워 수치적 해법 (유한 요소법, 유한 차분법 등) 이 주로 사용됩니다.
• 기존 방법의 한계:
  • 전통적 수치 해법: 계산 비용과 정확도 간의 트레이드오프가 존재하며, 격자 (grid) 가 조밀해질수록 계산 비용이 기하급수적으로 증가합니다.
  • 물리 정보 신경망 (PINNs): 일반화 능력이 부족하고 계수 변경 시 재학습이 필요합니다.
  • 신경 연산자 (Neural Operators, 예: FNO, DeepONet): 함수 공간 간의 매핑을 학습하여 메시 프리 (mesh-free) 접근이 가능하지만, 여전히 특정 아키텍처의 한계가 존재합니다.
  • Transformer 기반 모델: PDE 모델링에서 장기 의존성 (long-range dependencies) 을 포착하는 능력으로 인해 각광받고 있으나, **이차 복잡도 (quadratic complexity)**로 인해 긴 시간 통합 (long-time integration) 과 고해상도 그리드 처리 시 계산 효율성과 메모리 사용량이 큰 병목 현상을 겪습니다. 또한, 연속적인 동역학을 표현하는 데 있어 한계가 있습니다.

2. 방법론 (Methodology)

저자들은 **Mamba Neural Operator (MNO)**라는 새로운 프레임워크를 제안합니다. 이는 구조화된 상태 공간 모델 (Structured State-Space Models, SSMs) 인 Mamba 아키텍처를 신경 연산자 (Neural Operator) 에 통합한 것입니다.

• 핵심 아이디어:
  • SSM 과 신경 연산자의 이론적 연결: 신경 연산자 레이어가 시간 가변 상태 공간 모델 (time-varying SSMs) 과 구조적으로 유사하다는 것을 수학적으로 증명합니다.
  • 이산화 (Discretisation) 이론: Mamba 가 사용하는 Zero-Order Hold (ZOH) 이산화 방법이 PDE 수치 해법에서 널리 쓰이는 **오일러 방법 (Euler method)**과 동치임을 증명합니다. 특히 ZOH 는 1 차 오일러 방법보다 높은 차수의 정확도 (국소 오차 $O(\Delta^2)$) 를 제공하여 PDE 의 연속 동역학을 더 정밀하게 근사할 수 있음을 보여줍니다.
  • 아키텍처 설계:
    • Bi-Directional Scan (양방향 스캔): 2D 그리드 데이터를 두 가지 다른 경로로 스캔하여 시퀀스로 변환합니다. 이는 CNN 의 국소성 (locality) 과 Vision Transformer 의 전역성 (globality) 을 결합하면서도 선형 복잡도 ($O(N)$) 를 유지합니다.
    • S6/Cross S6 Block: 기존 Transformer 의 Self-Attention 과 Cross-Attention 을 Mamba 의 S6 블록과 새로운 Cross S6 블록으로 대체합니다. Cross S6 블록은 두 개의 독립적인 입력 벡터에서 파라미터를 생성하여 결합된 상태 공간 모델을 통해 정보를 융합합니다.
    • 전체 파이프라인: 입력 패치 확장 (Bi-Directional Scan Expand) $\rightarrow$ S6/Cross S6 블록 처리 $\rightarrow$ 출력 병합 (Bi-Directional Scan Merge) 의 구조를 가집니다.

3. 주요 기여 (Key Contributions)

1. Mamba Neural Operator (MNO) 개념 도입: SSM 프레임워크를 다양한 아키텍처 (Transformer 포함) 에 적용 가능한 통합 신경 연산자 접근법으로 확장했습니다.
2. 이론적 증명:
   • 신경 연산자 레이어의 반복적 업데이트 과정이 시간 가변 SSM 의 숨겨진 상태 업데이트와 구조적으로 동형 (isomorphic) 임을 증명했습니다.
   • ZOH 이산화가 오일러 방법의 일반화된 고차 버전임을 수학적으로 규명하여, Mamba 가 PDE 의 시간적 동역학을 더 정확하게 모델링할 수 있는 이론적 근거를 마련했습니다.
3. 성능 및 효율성 입증: Transformer 기반 모델들을 Mamba 로 대체했을 때, 정확도와 계산 효율성이 동시에 향상됨을 다양한 PDE 벤치마크를 통해 입증했습니다.

4. 실험 결과 (Results)

저자들은 PDEBench 벤치마크 (Darcy Flow, Shallow Water 2D, Diffusion Reaction 2D, Compressible Navier-Stokes 2D) 를 사용하여 GNOT, Galerkin Transformer (G.T.), OFormer 등 최신 Transformer 모델과 비교 실험을 수행했습니다.

• 정확도 (Accuracy):
  • 모든 PDE 작업에서 MNO 는 기존 Transformer 모델 (Softmax Attention, Galerkin Attention) 보다 일관되게 낮은 RMSE, nRMSE, RL2 오차를 기록했습니다.
  • 특히 Galerkin Transformer의 경우 MNO 적용 시 RMSE 가 최대 **89%**까지 감소하는 등 가장 큰 향상을 보였습니다.
  • **Diffusion Reaction (DR2D)**과 같은 비선형성이 강한 문제에서 MNO 는 기존 모델 대비 85% 이상의 오차 감소를 달성했습니다.
  • 고해상도 (512x512) CFD2D 작업에서도 MNO 는 가장 낮은 오차를 기록하며 새로운 State-of-the-Art (SOTA) 를 달성했습니다.
• 계산 효율성 (Efficiency):
  • Softmax Attention 기반 모델에 비해 MNO 는 FLOPs, 추론 시간, GPU 메모리 사용량을 획기적으로 감소시켰습니다 (예: OFormer 의 경우 FLOPs 가 977G 에서 60G 로 감소, 메모리는 4.8GB 에서 1.1GB 로 감소).
  • 선형 복잡도 ($O(N)$) 를 유지하여 고해상도 및 긴 시간 시뮬레이션에 적합합니다.
• 안정성 및 일반화:
  • 장기 안정성 (Long-time Stability): MNO 는 시간 단계가 증가함에 따라 오차가 누적되는 현상을 효과적으로 억제하여 장기 예측에서 더 안정적인 성능을 보였습니다.
  • 데이터 효율성: 데이터 양이 적을 때 (1K 샘플) 도 MNO 는 다른 모델들보다 우수한 일반화 성능을 보여주었습니다.
  • 주변 위치 불일치 (Query Misalignment): 입력과 쿼리 위치가 일치하지 않는 상황에서도 MNO 는 더 강건한 성능을 발휘했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• Transformer 대 Mamba: PDE 모델링 분야에서 Transformer 가 우세했던 기존 패러다임을 전환합니다. Mamba 는 단순히 Transformer 의 보완재가 아니라, 연속 동역학 표현과 장기 의존성 포착, 그리고 계산 효율성을 동시에 만족시키는 우월한 프레임워크임을 입증했습니다.
• 이론과 실용의 결합: 수치 해법 (오일러 방법 등) 과 최신 딥러닝 아키텍처 (Mamba/SSM) 간의 깊은 이론적 연결을 제시함으로써, 물리 기반 머신러닝의 해석 가능성과 신뢰성을 높였습니다.
• 미래 전망: MNO 는 고해상도, 복잡한 경계 조건, 그리고 실시간 PDE 시뮬레이션이 필요한 과학기술 및 공학 분야에서 강력한 도구가 될 것으로 기대됩니다.

요약하자면, 이 논문은 Mamba Neural Operator 를 통해 PDE 해법에서 Transformer 의 한계를 극복하고, 이론적으로 검증된 구조화된 상태 공간 모델을 기반으로 한 더 정확하고 효율적인 신경 연산자 프레임워크를 제시했습니다.