Reformulating Neural Operators in $d+1$ Dimensions for Embedding Evolution

이 논문은 보조 함수 차원을 통해 임베딩 진화를 모델링하는 새로운 $d+1$ 차원 신경 연산자(Neural Operator) 프레임워크를 소개하며, 이는 기존의 임베딩 확장 방식이 갖는 계산 비용을 피하면서도 다양한 물리적 벤치마크에서 최첨단 수준의 정확도와 강건성을 달성한다.

핵심

당신이 컴퓨터에게 금속판을 통해 열이 어떻게 퍼지는지, 혹은 폭풍 속에서 물이 어떻게 소용돌이치는지와 같은 복잡한 물리 시스템이 시간이 지남에 따라 어떻게 변하는지를 예측하도록 가르치려 한다고 상상해 보십시오. 인공지능의 세계에서 이러한 문제들을 해결하기 위해 설계된 모델들(이를 **뉴럴 오퍼레이터(Neural Operators)**라고 부릅니다)은 흔히 편미분 방정식(PDEs)이라고 불리는 수학적 규칙에 의해 설명됩니다.

오랫동안, 이러한 문제를 풀기 위해 설계된 AI 모델들은 일종의 "무차별 대입(brute force)" 전략에 의존해 왔습니다. 만약 모델이 충분히 정확하지 않다면, 엔지니어들은 단순히 모델을 더 "뚱뚱하게" 만들기 위해 내부 채널이나 레이어를 더 많이 추가하곤 했습니다. 이는 마치 양동이가 이미 크고 투박함에도 불구하고, 더 많은 물을 운반하기 위해 더 넓은 양동이를 사용하는 것과 같습니다.

이 논문은 물을 운반하는 더 똑똑한 방법을 소개합니다. 단순히 양동이를 더 넓게 만드는 대신, 저자들은 양동 자체에 새로운 차원을 추가하는 방법을 제안합니다.

핵심 아이디어: "그림자" 차원

물리적인 세상(예를 들어 2D 도시 지도)을 평평한 종이 한 장이라고 생각해 보십시오. 기존의 AI 모델들은 이 종이를 위에서 아래로 층층이 내려다보며 그 패턴을 학습하려고 시도합니다.

Haoze Song과 그의 팀은 우리가 단순히 종이를 바라보는 것에 그치지 않고, 그 종이에 그림자나 유령 차원이 붙어 있다고 상상해야 한다고 제안합니다. 그들은 이를 "보조 차원(auxiliary dimension)"(이하 "p-차원"이라 부름)이라고 부릅니다.

• 기존 방식: 3D 물체를 이해하기 위해 2D 사진을 보면서, 세부 사항을 보기 위해 눈을 더 가늘게 뜨고 애쓰는 것(픽셀을 추가하는 것)과 같습니다.
• 새로운 방식 (SKNO): 2D 사진을 가지고 있지만, 동시에 그 사진의 그림자를 옆 벽면에 투영하는 특별한 "그림자 프로젝터"를 가지고 있는 것과 같습니다. 사진과 그 그림자를 함께 연구함으로써, 더 큰 사진을 필요로 하지 않고도 3D 형상을 훨씬 더 잘 이해할 수 있습니다.

이 논문에서 그들은 SKNO(Schrödingerised Kernel Neural Operator)라는 모델을 만듭니다. 이 모델은 데이터가 하나의 추가된 차원이 존재하는 공간에 존재하는 것처럼 취급합니다. 모델은 물리적 지도 위의 데이터만 업데이트하는 것이 아니라, 지도와 그 그림자의 데이터를 동시에 업데이트합니다.

작동 원리: "두 가지 관점" 전략

SKNO의 마법은 이 추가된 차원을 업데이트하는 방식에 있습니다. 저자들은 양자 역학(구체적으로는 슈뢰딩거 방정식, 다만 여기서는 물리 시뮬레이션이 아닌 설계 청사진으로서 사용됨)에서 영감을 얻은 영리한 트릭을 사용합니다.

그들은 두 가지 다른 방식으로 "그림자" 데이터를 동시에 업데이트합니다:

1. 원시 뷰 (The Raw View): 데이터를 있는 그대로 보는 것 (마치 일반 텍스트로 책을 읽는 것과 같습니다).
2. 푸리에 뷰 (The Fourier View): 데이터를 파동과 주파수의 혼합으로 보는 것 (마치 책을 소리의 파동인 음악 악보로 읽는 것과 같습니다).

이 두 가지 "그림자 차원"의 관점을 결합함으로써, 모델은 훨씬 더 효율적으로 복잡한 패턴을 포착할 수 있습니다. 이는 마치 "일반 영어"와 "시적인 영어"를 모두 구사하는 번역가가 있는 것과 같습니다. 그 번역가는 한 가지 언어만 구사하는 사람보다 문장의 뉘앙스를 훨씬 더 잘 이해할 수 있습니다.

결과: 더 빠르고, 더 작고, 더 정확하게

팀은 단순한 열 방정식부터 매우 혼란스러운 3D 유체 폭발(Rayleigh–Taylor instability)에 이르기까지 10가지 이상의 도전적인 물리 문제들을 대상으로 이 새로운 모델을 테스트했습니다.

그들이 발견한 결과는 다음과 같습니다:

• 더 낮은 오차: SKNO는 기존의 최고 모델들(예: FNO, Transolver, DeepONet)보다 일관되게 적은 실수를 기록했습니다.

• 효율성: 이 모델은 더 "뚱뚱해지거나" 비용이 많이 들지 않고도 이러한 결과를 달성했습니다. 실제로, SKNO는 종종 훈련 속도가 더 빨랐고 더 적은 컴퓨팅 자원을 요구했습니다.

• 강건성(Robustness): 모델이 한 번도 본 적 없는 데이터(예: 훈련되지 않은 날의 날씨 패턴을 예측하거나, 훨씬 더 높은 해상도의 데이터를 다룰 때)로 테스트되었을 때도, 경쟁 모델들보다 더 잘 버텨냈습니다. 데이터의 "그리드(grid)" 크기가 변해도 혼란을 겪지 않았습니다.

요약

이 논문은 어려운 물리 문제를 해결하기 위해 단순히 AI 모델을 더 크고 무겁게 만드는 대신, 데이터를 바라보는 방식을 바꿔야 한다고 주장합니다. "그림자 차원"을 추가하고 두 가지 수학적 렌즈(원시 및 주파수 기반)를 통해 데이터를 업데이트함으로써, 모델은 물리 법칙의 근저에 있는 규칙을 더 자연스럽게 학습합니다.

이는 "문제에 더 많은 자원을 쏟아붓는 것"에서 "문제를 바라보는 더 나은 각도를 찾는 것"으로의 전환입니다. 그 결과, 더 정확할 뿐만 아니라 더 우아하고 효율적인 모델이 탄생했습니다.

기술 요약

기술 요약: 임베딩 진화를 위한 $d + 1$ 차원에서의 뉴럴 오퍼레이터 재구성

문제 정의

뉴럴 오퍼레이터(Neural Operators, NOs)는 편미분 방정식(PDE)을 풀기 위해 함수 공간 사이의 매핑을 학습하도록 설계되었습니다. 최근의 발전은 $d$차원 물리 도메인에 대한 커널 파라미터화의 정교화에 집중해 왔으나, 리프팅된 임베딩(lifted embeddings)의 진화 방식은 충분히 탐구되지 않았습니다. 기존 아키텍처들은 임베딩 표현력의 부족을 보완하기 위해 브루트 포스 방식의 스케일링(임베딩 너비 확장 또는 헤드 추가)을 사용하는 경향이 있습니다. 그러나 이러한 전략은 높은 계산 비용을 초래합니다. 임베딩 너비가 넓어질수록 밀집 채널 믹싱(dense channel mixing)은 이차적으로 증가하며, 헤드별 팩토라이제이션(head-wise factorization)은 블록 대각 구조를 유도하여 헤드 간 결합을 약화시킴으로써 이를 부분적으로만 완화할 뿐입니다. 본 논문은 단순히 용량을 늘리는 것이 아니라, 임베딩이 어떻게 진화해야 하는지를 직접 설계하는 데 존재하는 격차를 식별합니다.

방법론

저자들은 보조 함수 차원 $p$를 도입함으로써 뉴럴 오퍼레이터 파이프라인을 $d + 1$ 차원으로 재구성할 것을 제안합니다. 제안된 프레임워크는 임베딩을 물리 도메인 $D_x$ 위에서만 진화시키는 대신, 곱 도메인 $D_x \times D_p$ 상에서 래턴트 스칼라 함수(latent scalar functions)를 진화시킵니다.

일반적 프레임워크

1. 리프팅(Lifting): 입력 필드 $a(x)$는 곱 도메인 상의 스칼라 래턴트 함수 $v_0(x, p)$로 리프팅됩니다. 이는 리프팅 연산자 $P$를 통해 달성되며, 흔히 분리된 선형 맵 $v_0(x, p) = w^\top(p)a(x)$로 구현됩니다.
2. $(d+1)$-차원 진화: 래턴트 함수는 학습 가능한 선형 연산자 $\mathcal{L}$과 비선형 맵 $\sigma$의 시퀀스를 통해 진화합니다. 핵심 구성 요소는 물리적 좌표 $x$와 보조 좌표 $p$ 모두에 작용하는 커널 적분 연산자 $\mathcal{K}$입니다:
   $$\mathcal{K}_l[v_l](x, p) = \int_{D_x} \int_{D_p} \kappa_l(x, y, p, p') v_l(y, p') \, dp' \, dy$$
3. 복구(Recovery): 진화된 함수 $v_L(x, p)$는 $p$에 대한 적분을 통해 출력 도메인으로 매핑됩니다: $u_{pred}(x) = \int_{D_p} \chi(p) v_L(x, p) \, dp$.

슈뢰딩거화된 커널 뉴럴 오퍼레이터 (SKNO)

본 논문은 이 프레임워크를 푸리에 기반 모델인 SKNO로 구체화합니다. 주요 설계 선택 사항은 다음과 같습니다:

• 기저 다양화된 보조 진화(Basis-Diversified Auxiliary Evolution): 각 공간 위치에 대해, $p$ 차원을 따른 신호는 두 가지 구별된 좌표 뷰를 사용하여 업데이트됩니다:
  1. Raw $p$-좌표 믹싱: $p$의 공간 도메인에서의 선형 믹싱.
  2. Fourier-$p$-좌표 믹싱: $p$의 푸리에 도메인에서의 스펙트럴 믹싱.
     이 이중 브랜치 구조($F_p^{-1} \tilde{A}_l F_p + B_l$)는 단순히 동일한 채널 믹싱 경로를 복제하는 것이 아니라, 두 가지 뷰로부터 특징을 포착할 수 있게 합니다.
• 물리 도메인 전파: SKNO는 $(L-1)$개의 글로벌 전파기(Spectral Convolution Operators 사용, $x$의 푸리에 도메인에서 대각화됨)와, 글로벌 스펙트럴 방식에 의해 손실될 수 있는 로컬 정보를 포착하기 위한 하나의 최종 로컬 전파기를 사용합니다.
• 잔차 연결(Residual Connections): 선형 블록에는 학습 및 안정성을 촉진하기 위한 잔차 연결이 포함됩니다.

주요 기여

1. 연산자 수준의 재구성: 저자들은 커널 적분을 통해 물리 및 보조 좌표 모두에서 래턴트 함수를 진화시키도록 NO 파이프라인을 재구성하여, 임베딩 진화를 위한 명시적인 연산자 기반 메커니즘을 확립했습니다.
2. SKNO 아키텍처: 브루트 포스 스케일링 없이 표현력을 향상시키기 위해 기저 다양화된 보조 진화(raw 및 Fourier-$p$ 좌표 믹싱)를 사용하는 슈뢰딩거화된 커널 뉴럴 오퍼레이터를 제안합니다.
3. 포괄적 평가: 1D 선형 방정식부터 고도의 비선형 3D 불안정성까지 10개 이상의 벤치마크를 통해 모델을 평가했습니다.
4. 통제된 분석: 성능 향상이 단순한 파라미터 수 증가가 아니라 아키텍처 설계(기저 다양성)에서 기인함을 입증하기 위해 스케일링 및 절제된(ablated) 베이스라인들과 엄격한 비교를 제공합니다.

실험 결과

1D 열/이류 방정식, 1D Burgers, 2D Darcy Flow, 2D Gray-Scott, 2D/3D Navier-Stokes, 3D Rayleigh-Taylor instability를 포함한 벤치마크 전반에 걸쳐, SKNO는 평가된 베이스라인(DeepONet, FNO, Transolver, CNO) 중 지속적으로 가장 낮은 상대 $L_2$ 오차를 달로 성했습니다.

• 성능 이득: 2D 비압축성 Navier-Stokes ($\nu=10^{-5}$)에서 SKNO는 FNO 대비 상대 $L_2$ 오차를 약 37.1% 감소시켰습니다. 2D Gray-Scott에서는 42.1% 감소했습니다. 3D Rayleigh-Taylor에서 SKNO는 14.3%의 오차 감소를 달성했습니다.
• 용량 효율성: 통제된 실험 결과, SKNO (A+B)는 더 적은 파라미터와 FLOPs를 사용하여 체계적으로 스케일링된 FNO 변형 및 병렬로 쌓인 FNO들을 능가합니다. "B+B" 변형(raw-$p$ 브랜치를 복제한 모델)은 기초 다양화된 "A+B" 변형의 성능을 따라잡지 못했으며, 이는 이중 좌표 뷰의 가치를 확인시켜 줍니다.
• 강건성: SKNO는 혼합 해상도 훈련 및 제로샷 초해상도 추론(예: 128 그리드에서 학습하고 8192에서 테스트) 하에서도 낮은 오차를 유지하며 우수한 해상도 불변성을 보여줍니다. 또한 미경험 시간 영역(unseen temporal regimes)에 대한 강력한 제로샷 일반화 능력을 보입니다.
• 효율성: 추가된 차원에도 불구하고, SKNO는 임베딩 크기에 대해 이차적 복잡도를 갖는 Transolver와 같은 트랜스포머 기반 모델보다 종종 더 우수한 경쟁력 있는 훈련 시간을 유지합니다.

의의 및 주장

본 논문은 보조 도메인 연산자 진화가 브루트 포스 방식의 임베딩 스케일링에 대한 유망한 대안임을 주장합니다. 보조 좌표를 따라 연산자 설계 원칙을 적용함으로써, 모델은 임베딩을 넓히는 데 따르는 과도한 계산 비용 없이 표현력과 근사 능력을 향상시킵니다.

저자들은 "슈뢰딩거화된(Schrödingerised)"이라는 명칭이 PDE 자체에 대한 직접적인 고전적 수치 가속 메커니즘을 주장하기보다는, 보조 좌표를 따른 구조적 연산자 진화에 대한 설계 영감을 의미함을 강조합니다. 결과는 제안된 $d+1$ 차원 설계가 더 낮은 오차, 더 나은 해상도 강건성, 그리고 우수한 용량 효율성을 뒷받는 실증적 증거와 함께, 뉴럴 오퍼레이터 성능을 개선하는 더 직접적이고 효율적인 경로를 제공함을 시사합니다.

결론적으로, 저자들은 향후 연구가 최종 테스트 에러를 넘어 뉴럴 오퍼레이터를 비교할 수 있는 정량적 기준을 개발해야 하며, 특히 서로 다른 집계(aggregation) 설계가 최적화 궤적과 고차원 에러 지형에서의 로컬 미니마 선택에 어떻게 영향을 미치는지 조사해야 한다고 언급합니다.