과학 AI연구

JAX-PI

PINN이 왜 무너지는지에 대한 6년치 처방을 한 저장소에 모았다 — 손실 가중, 시간 인과 학습, PirateNet, 2차 최적화까지 JAX로

직접 해보기 4빠른 시작

저장소 논문

종류: 오픈소스
라이선스: 비상업 연구용 (Penn)
언어: Python
사실 확인: 2026.10.24

git clone https://github.com/PredictiveIntelligenceLab/jaxpi.git && cd jaxpi && pip install .

TRY IT

직접 해보기

OVERVIEW

개요

무엇인가

JAX-PI는 펜실베이니아대 페르디카리스(Perdikaris) 연구실이 공개한 물리 정보 신경망(PINN) 학습 구현체입니다. PINN은 PDE 잔차와 경계·초기 조건을 손실에 넣어 신경망이 물리적으로 말이 되는 해를 배우게 하는 방법인데, 아이디어보다 훈련이 훨씬 어렵습니다. 이 연구실은 2020년 그래디언트 병리, 같은 해 NTK 분석, 2022년 인과 학습, 2023년 「PINN 훈련 전문가 가이드」, 2024년 PirateNets, 2025년 Gradient Alignment로 그 실패를 하나씩 진단해 왔고, JAX-PI는 그 처방들을 한 코드베이스에 묶은 것입니다.

README가 인용하는 대표 논문이 전문가 가이드(arXiv:2308.08468)이고, 저장소는 단일·다중 GPU 학습과 함께 Allen–Cahn, 이류, 스토크스, 쿠라모토–시바신스키, 리드 구동 공동, 원환면 나비에–스토크스, 실린더 와류, 콜모고로프 흐름, 레일리–테일러 불안정성 같은 벤치마크 예제를 제공합니다.

어떻게 동작하나

예제마다 configs/ 아래 설정 파일로 처방을 켜고 끕니다. 핵심 재료는 넷입니다.

손실 가중

손실 항마다 그래디언트 크기가 수십 배 다른 문제(Type I 충돌)를 weighting.scheme = "grad_norm" 같은 자동 가중으로 맞춘다.

시간 인과 학습

시간 구간 i의 잔차에 앞 구간 누적 잔차로 만든 가중치 w_i = exp(−ε·앞 구간 잔차 합)를 곱한다. 앞 시간이 맞기 전에는 뒤 시간을 배우지 못하게 잠근다. 설정 키는 use_causal, causal_tol.

구조

무작위 가중치 분해(RWF), 푸리에 특성, 정확한 주기 경계 조건, 잔차 적응형 PirateNet.

2차 최적화

손실 항들이 반대 방향을 가리키는 Type II 충돌은 가중만으로 안 풀린다. Gradient Alignment 논문은 SOAP 같은 전처리기가 연속한 업데이트의 방향을 맞춰 준다고 보였다(코드는 pirate 브랜치).

Gradient Alignment 논문이 쓰는 정렬 점수는 정규화한 그래디언트들의 평균 길이로 방향 충돌만 잽니다. 손실이 둘이면 코사인 유사도와 같습니다.

$A(v_1,\dots,v_n)=2\left\lVert \frac{1}{n}\sum_{i=1}^n \frac{v_i}{\lVert v_i\rVert}\right\rVert^2-1$

같은 논문의 벤치마크(상대 L2 오차)에서 격차는 문제가 어려울수록 벌어졌습니다.

벤치마크	Adam	SOAP
Wave	5.15e-5	8.05e-6
Allen-Cahn	2.24e-5	3.48e-6
Lid-driven cavity (Re=5e3)	3.24e-1	3.99e-2
Kolmogorov flow (Re=1e4)	2.04e-1	3.20e-2
Rayleigh-Taylor	7.32e-2	5.22e-3

빠른 시작

README 기준으로 GPU 전용이고, JAX 0.4.36·CUDA 12.4·cuDNN 8.9에서 확인됐습니다. 학습 지표는 Weights & Biases로 기록하므로 계정 설정이 먼저 필요합니다.

bash

pip3 install -U pip
pip3 install --upgrade jax jaxlib

git clone https://github.com/PredictiveIntelligenceLab/jaxpi.git
cd jaxpi
pip install .

cd examples/adv
python3 main.py                            # 기본 설정으로 학습
python3 main.py --config=configs/sota.py   # 다른 설정 파일로 실행
python3 main.py --config.mode=eval         # 평가
CUDA_VISIBLE_DEVICES=0,1 python3 main.py   # 다중 GPU

메모리가 부족하면 --config.batch_size_per_device를 줄입니다. 예를 들어 Allen–Cahn 예제의 기본 설정에는 use_causal = True, causal_tol = 1.0, grad_norm 가중이 이미 들어 있습니다. README가 밝힌 상대 L2 오차는 Allen–Cahn 5.37×10^-5, 이류 6.88×10^-4, 스토크스 8.04×10^-5, 원환면 나비에–스토크스 3.53×10^-1입니다.

언제 쓰고, 언제 피하나

상황	JAX-PI(PINN)가 맞는다	다른 도구를 먼저 본다
데이터	센서가 드물고 물리식은 확실하다	시뮬레이션 데이터가 많다 → 신경 연산자(DeepONet·FNO)
목적	역문제, 파라미터 식별, 디지털 트윈 보정	같은 PDE 가족을 수천 번 빠르게 추론
정확도 기준	물리 제약을 지킨 근사면 충분	고충실도 산업 CFD → 검증된 전통 솔버
연구 단계	PINN 처방을 같은 조건에서 비교·재현	상용 제품에 그대로 탑재(라이선스 확인)

한계와 주의

라이선스가 비상업 연구용이다. 저장소 LICENSE는 펜실베이니아대 이사회가 비영리 연구 목적의 사용·복제·수정을 허락하는 별도 문서이며, 표준 오픈소스 라이선스가 아니다. 상업적 이용은 따로 확인해야 한다.
처방은 브랜치에 흩어져 있다. PirateNets와 Gradient Alignment 코드는 pirate, ActNet은 ActNet 브랜치다.
정밀도가 결과를 바꾼다. 저장소는 2024년 11월 재현성 문제로 JAX 행렬 곱 정밀도 기본값을 highest로 올렸다. FP32에서 실패로 보이던 PINN이 FP64로 바꾸자 오차가 급락한 연구(FP64 is All You Need)와 같은 맥락이다.
인과 학습은 만능이 아니다. 시간 순서를 무시한 학습을 고칠 뿐, 장시간 혼돈 동역학의 오차 누적이나 다중 스케일의 표현 병목은 그대로 남는다. 논문 성과도 modified MLP, 정확한 주기 경계, 시간 창 분할을 함께 쓴 조합에서 나왔다.
ε(인과 파라미터)는 민감하다. 너무 작으면 효과가 없고 너무 크면 학습이 굳는다. 논문은 작은 값에서 키워 가는 스케줄을 쓴다.

DEEP DIVE

더 깊이 읽기

특집 · 2026.03.20 · 45분

PINN은 왜 자꾸 딴길로 샐까 — Gradient Alignment와 2차 최적화로 다시 읽는 물리 AI

특집 · 2026.04.02 · 36분

PINN은 왜 시간 문제에서 자꾸 무너질까 — causal training으로 다시 읽는 학습 실패의 구조

기술 · 2026.04.14 · 32분

왜 PINN은 자꾸 실패할까 — Physics-informed Neural Networks 훈련 전문가 가이드

특집 · 2026.03.13 · 40분

FP64 is All You Need — PINN 실패 모드를 다시 생각하다

함께 볼 기술

과학 AI