기술 요약: jNO – 신경 연산자 및 기초 모델 훈련을 위한 JAX 라이브러리

문제 제기

신경 연산자 및 물리 정보 신경망 (PINNs) 과 관련하여 과학적 머신러닝 (SciML) 의 현재 상황은 심각한 소프트웨어 분열로 고통받고 있습니다. 최첨단 기초 모델 (예: Poseidon, Walrus, PDEformer2) 과 연산자 학습 도구는 주로 PyTorch 생태계 내의 별도의 저장소에서 유지 관리됩니다. 이러한 분리는 서로 다른 모델 계열 간의 체계적인 비교, 전이 학습, 그리고 재현 가능한 미세 조정을 위한 장벽을 조성합니다. 또한, 실제 워크플로는 종종 단절되어 있어 사용자가 기하학, 손실 구성, 모델 정의, 훈련 논리에 대한 별도의 추상화를 관리해야 합니다. 감독된 연산자 학습과 물리 정보 잔차를 결합하거나 사전 훈련된 백본을 새로운 편미분 방정식 (PDE) 설정으로 이전하려는 시도가 있을 때 이러한 분열은 더욱 심화됩니다.

방법론

jNO(jax Neural Operators) 는 전체 훈련 워크플로를 단일 추적 프로그램으로 취급하는 통합된 JAX 네이티브 라이브러리를 제공함으로써 이러한 과제를 해결합니다. 핵심 방법론은 **추적 (tracing)**과 **지연 실행 (deferred execution)**을 중심으로 한 상징적 도메인 특화 언어 (DSL) 에 의존합니다.

1. 추적 우선 설계

eager 실행 모델과 달리 jNO 는 도메인, 모델 호출, 잔차, 감독 손실 및 진단을 하나의 상징적 언어로 작성하는 지연 계산 그래프를 구성합니다.

• 상징적 노드: 시스템은 변수, 텐서, 모델 및 미분 연산자를 나타내기 위해 Placeholder 노드를 사용합니다. 산술 및 미분 연산은 즉시 실행하는 대신 표현 트리를 구축하기 위해 표준 JAX 원시 연산을 오버로드합니다.
• 통합 파이프라인: 사용자는 도메인과 변수를 정의하고 미분 표현으로 모델을 구성한 후 단일 인터페이스를 통해 모든 것을 해결합니다. 이를 통해 코드 구조를 변경하지 않고도 연산자 회귀, 메시 인식 잔차 평가, PDE 제약 훈련 간에 원활한 전환이 가능합니다.
• 최적화: 실행 전에 시스템은 추적된 트리에서 공통 부분 표현 제거 (CSE) 를 수행하여 구조적으로 동일한 서브트리를 단일 XLA 프로그램으로 정규화합니다. 이는 목적 함수, 연산자 정의 및 물리 제약을 통합합니다.

2. 도메인 및 메시 워크플로우

domain 클래스는 기하학 정의와 추적된 표현을 연결하는 메시 인식 데이터 조직 계층으로 작용합니다.

• 통합: PyGmsh/Gmsh를 통한 메시 생성 및 meshio를 통한 로딩을 지원합니다.
• 구조: 데이터는 태그가 지정된 점 집합을 위한 메시 풀과 런타임 컨텍스트 사전으로 조직됩니다. 각 배치 요소가 서로 다른 PDE 매개변수화 또는 강제 장을 나타낼 수 있는 배치 훈련을 지원합니다.
• 적응성: 시스템은 기호 방정식을 변경하지 않고 훈련이 점 집합을 업데이트할 수 있도록 적응형 점 관리 후크를 포함합니다.

3. 유한 요소법 (FEM) 통합

jNO 는 FEAX 백엔드를 통해 약형 (weak form) 조립으로 확장된 추적 프로그래밍 모델을 제공합니다.

• 통합 DSL: 별도의 FEM 프론트엔드 대신 jNO 는 메시 기반 약형, 변분 PINN 및 연산자 학습에 동일한 상징적 계층을 사용합니다.
• 조립: 상징적 약형은 백엔드 중립 표현으로 낮아집니다. 목표에 따라 시스템은 vpinn(변분 PINN), fem_system(선형 정상), fem_residual(비선형 정상), 또는 fem_time(과도) 로 디스패치할 수 있습니다.
• 미분 가능성: 통합은 조립된 유한 요소 연산자를 JAX 생태계 내에 유지하여 조립, 잔차 평가, 야코비안 구성과 같은 엔드투엔드 미분 가능 워크플로우를 JAX 변환과 호환되도록 합니다.

4. 모델 인터페이스 및 기초 모델

jNO 는 신경 연산자 아키텍처를 추적된 프로그램 내의 일반 연산으로 취급합니다.

• 기초 모델: 외부 foundax 저장소에서 유지 관리되는 PDE 기초 모델 (Poseidon, Walrus, Morph 등) 의 JAX 네이티브 변환을 지원합니다. 이러한 모델은 jno.nn.wrap()을 통해 래핑되고 임의의 추적된 표현과 구성될 수 있습니다.
• 세밀한 제어: 사용자는 매개변수 수준에서 모델, 옵티마이저 및 학습률 동작을 구성할 수 있습니다. 기능으로는 Equinox 및 Flax 모델 모두를 위한 선택적 매개변수 마스킹, 동결/동결 해제 및 LoRA(저랭크 적응) 지원이 포함됩니다.

5. 실행 및 훈련

jno.core 모듈은 추적된 제약 조건을 단일 JIT 컴파일된 훈련 프로그램으로 컴파일합니다.

• 성능: XLA 를 활용하여 공격적인 부분 표현 제거, 메모리 최적화 및 2 축 장치 메쉬를 통한 다중 장치 병렬화를 실현합니다.
• 하이퍼파라미터 튜닝: ArchSpace 인터페이스는 아키텍처 및 훈련 매개변수에 대한 그리드 검색 및 비경사 최적화 (Nevergrad 를 통해) 를 지원하며, 훈련 루프에 직접 통합됩니다.
• 지속성: 실험은 jno.save(무결성을 위한 RSA 서명 지원) 를 사용하여 직렬화하여 솔버 상태, 도메인 객체 및 컴파일된 추론 래퍼를 보존할 수 있습니다.

주요 기여

1. 통합된 추적 언어: jNO 는 JAX 내 단일 상징적 인터페이스에서 신경 연산자, 물리 정보 잔차 및 기초 모델 미세 조정을 표현하는 최초의 통합 프레임워크입니다.
2. 생태계 통합: 다양한 PDE 기초 모델 계열 (예: Poseidon, Walrus) 을 JAX 로 변환 및 통합하는 통합 백엔드로 작용하여 별도의 코드베이스를 유지 관리하지 않고도 나란히 평가 및 구성을 가능하게 합니다.
3. FEM 및 연산자 통합: 추적된 DSL 을 약형 조립으로 확장함으로써 jNO 는 프로그래밍 모델을 변경하지 않고도 메시 기반 약형, 변분 PINN 및 연산자 모델을 구성할 수 있게 합니다.
4. 매개변수 수준 제어: 이 라이브러리는 상징적 흐름 내에서 선택적 마스킹 및 LoRA 적응을 포함한 훈련 역학에 대한 세밀한 제어를 제공합니다.
5. 재현성 및 이식성: 이 라이브러리는 PyPI 에 패키징되어 있으며 자동화된 테스트를 포함하고 검증 가능한 워크플로우를 위한 서명된 아티팩트 공유를 지원합니다.

결과 및 주장

이 논문은 특정 PDE 데이터셋에 대한 새로운 벤치마크 결과를 제시하지는 않지만 아키텍처 및 워크플로우 기여에 중점을 둡니다:

• 워크플로우 효율성: 이 논문은 추적된 설계가 연산자 학습, 잔차 평가 및 PDE 제약 훈련 간 전환 시 코드 구조 변경의 필요성을 제거한다고 주장합니다.
• 성능: 단일 XLA 그래프로 컴파일함으로써 jNO 는 PyTorch 중심 대안과 비교하여 컴파일 오버헤드를 줄이고 수치적 안정성을 향상시키면서도 프로토타이핑의 단순성을 유지한다고 주장합니다.
• 상호 운용성: FEAX 의 통합은 JAX 솔버 (예: Diffrax, Optimistix) 또는 외부 라이브러리 (예: SciPy) 와 호환되는 미분 가능 유한 요소 워크플로우를 가능하게 합니다.

중요성

저자들은 jNO 를 SciML 생태계의 분열을 줄이기 위한 중요한 단계로 위치시킵니다. 그 중요성은 다음과 같습니다:

• 장벽 완화: 통합된 JAX 네이티브 스택을 제공함으로써 연구자들에게 고급 신경 연산자 및 기초 모델 워크플로우를 더 접근 가능하게 만듭니다.
• 하이브리드 워크플로우 활성화: 훈련 제어 및 매개변수 수준 최적화를 표준화함으로써 앙상블 방법, 점진적 전이 학습 및 다중 모델 제약과 같이 이전에 어려웠던 워크플로우를 가능하게 합니다.
• 미래 연구의 기반: 기초 모델을 단일 프레임워크로 통합함으로써 이전에 다양한 저장소로 인해 방해받았던 체계적인 비교 및 구성 워크플로우를 가능하게 합니다.

이 논문은 jNO 가 JAX 생태계 내에서 이러한 고급 워크플로우를 더 접근 가능하고 효율적으로 만들어 커뮤니티를 위한 재사용 가능한 연구 소프트웨어 스택으로 기능함으로써 신뢰할 수 있는 연구적 중요성을 제공한다고 결론지었습니다.