Mamba 신경 연산자: 트랜스포머를 넘어선 PDE 솔버

초록

본 논문은 구조화된 상태공간 모델(SSM)인 Mamba를 신경 연산자에 결합한 Mamba Neural Operator(MNO)를 제안한다. MNO는 트랜스포머 기반 연산자와 이론적으로 동등한 구조를 갖지만, 연속적인 동역학과 장거리 상호작용을 더 효율적으로 표현한다. 연속 PDE 문제에 대한 실험에서 기존 트랜스포머와 Fourier Neural Operator 대비 정확도와 메모리 효율이 크게 향상됨을 보인다.

상세 분석

MNO는 먼저 PDE를 함수 공간 간 매핑으로 정의하고, 이를 근사하기 위한 신경망 구조를 설계한다. 핵심 아이디어는 시간‑가변 SSM을 연산자 레이어와 동일시함으로써, 상태 전이 행렬 A와 입력 매핑 B, 출력 매핑 C, D를 학습 가능한 파라미터로 두는 것이다. 이때 Mamba가 제공하는 zero‑order hold(ZOH) 기반 이산화 방식은 연속 시간 동역학을 정확히 보존하면서도 O(N)의 선형 복잡도를 유지한다. 기존 트랜스포머는 전역 어텐션을 통해 O(N²) 메모리·연산 비용을 발생시키며, 연속적인 좌표 입력을 토큰화하는 과정에서 해상도 의존적인 손실이 발생한다. 반면 MNO는 상태공간의 선형 연산을 활용해 연속 좌표를 직접 처리하고, 긴 시계열(또는 고해상도 격자)에서도 장거리 의존성을 자연스럽게 캡처한다.

이론적 기여는 두 부분으로 나뉜다. 첫째, 논문은 연산자 레이어의 업데이트 식을 일반적인 선형 시불변 시스템의 상태 전이 식과 동일시함을 증명한다. 이를 통해 트랜스포머의 어텐션이 실제로는 특정 형태의 상태공간 모델에 해당한다는 새로운 관점을 제시한다. 둘째, Mamba의 파라미터화(A, B, C, D)를 신경 연산자에 매핑함으로써, 기존 FNO와 DeepONet이 갖는 주파수 기반 또는 피처 분리 구조를 넘어서는 표현력을 확보한다.

실험에서는 1D·2D 열전도 방정식, Navier‑Stokes, 반응‑확산 시스템 등 다양한 PDE에 대해 MNO‑Transformer, MNO‑FNO, 순수 트랜스포머, 순수 FNO와 비교하였다. 결과는 MNO가 평균 L2 오차에서 15‑30% 개선을 보이며, 메모리 사용량은 동일 조건에서 40‑60% 감소한다. 특히 긴 시간 통합(시간 스텝 수 > 1000)에서 트랜스포머는 급격히 메모리 초과와 정확도 저하를 겪는 반면, MNO는 안정적인 수렴을 유지한다.

한계점으로는 현재 구현이 아직 고차원(3D 이상) PDE와 복합 경계 조건에 대한 확장성을 충분히 검증하지 않았으며, SSM 파라미터 초기화와 학습 스케줄링이 모델 성능에 민감하다는 점을 들 수 있다. 향후 연구에서는 다중 스케일 상태공간 구조와 하이브리드 어텐션‑SSM 레이어를 결합해 더욱 일반화된 연산자를 설계할 여지가 있다.