[24′ ICML] Parameterized Physics-informed Neural Networks for Parameterized PDEs

1. 들어가며

이번 포스팅에서는 2024년 ICML에 발표된 Parameterized Physics-informed Neural Networks for Parameterized PDEs 논문의 핵심 내용을 정리해보겠습니다. 이번 논문은 P2INN이라고 부르겠습니다.

기존 순수 data driven 방식의 모델들과 달리 PINN은 학습할 모델이 만족해야 하는 함수를 loss function에 추가해주는 방식 등을 사용하여 explicit하게 알려줍니다. 그 결과 기존 방법들보다 훨씬 더 적은 데이터로도 좋은 성능을 내죠. 덕분에 다량의 PDE를 풀어야하는 기존 시뮬레이션을 대체하는 방법중 하나로 PINN이 부각되고 있습니다.

하지만 실제 시뮬레이션을 대체하는 솔루션을 개발하다보면 공통적으로 부딪히는 난관이 있는데요. 바로 파라미터의 변화입니다. 시뮬레이션을 수행한다는건 PDE를 푼다는 의미인데, 이 PDE에는 보통 다양한 파라미터가 있습니다. 이 파라미터에 따라 재현하는 시뮬레이션 상황은 전혀 달라지죠. 보통 실제 문제 상황에서는 재현해야 하는 현상을 맞추기 위한 파라미터를 찾고 (Parameter Estimation), 이후 PINN 모델을 학습합니다. 그러다보니 파라미터가 변경되면 (예를 들어, 배터리의 경우 셀이 업데이트 되는 경우) 데이터셋 제작과 모델 학습을 다시 수행해야 하죠. 실제 고객에게 배포되어 서비스 해야 하는 솔루션 입장에서는 몇일 이상 걸리는 과정을 매번 반복해야 하는 문제가 있습니다.

이에 다양한 파라미터에 대해 동시에 대응할 수 있는 PINN에 대한 다양한 연구가 수행되고 있습니다. P2INN은 이러한 연구의 일환으로 랜덤한 파라미터로 생성한 데이터셋을 학습한 뒤, inference 단계에서는 다양한 파라미터에 대응 가능한 모델을 제안합니다. 이번 포스팅에서는 P2INN의 아이디어와 제안 방법, 실험 결과를 차례로 정리해보겠습니다.

2. Motivation

먼저 저자들은 주요한 두가지 관찰로부터 아이디어를 제기합니다.

첫 번째 관찰은 파라미터 변화에 따른 솔루션의 변화로부터 나오는데요.

위 그림은 Convection과 Reaction 방정식의 파라미터를 변화시켰을때의 해 공간을 표현한 그림입니다. 파라미터가 바뀌면 당연히 그 해 공간이 바뀌지만, 전혀 예측할 수 없게 변하는게 아닌, 연속적으로 변하는 모습을 보여주고 있습니다. 예를 들어 첫 번째 행의 Convection 해 공간을 보면, beta값이 커짐에 따라 해 공간이 천천히 시계 반대 방향으로 회전하는 모습을 보여주고 있습니다. 따라서 저자들은 이러한 관찰을 토대로 파라미터를 저차원 latent space로 매핑하자는 아이디어를 냅니다. 파라미터 전용 encoder를 구성하고, 이 encoder는 파라미터 변화에 따른 representation만을 학습하는거죠. 이 변화가 너무 드라마틱하다면 학습하기 어렵겠지만, 이 정도의 연속적인 변화라면 단순한 모델로도 충분히 mapping 가능하겠죠?

두 번째 관찰은 여러 PDE로부터 나옵니다.

위 그림은 Convection, Diffusion, Reaction 방정식과 그 조합들의 해 공간을 표현한 그림입니다. 마찬가지로 이 세가지 방정식과 그 조합들 사이에서는 해 공간의 변화가 연속적이라는 느낌을 받습니다. 즉 학습해야 하는 PDE가 달라지더라도, 전혀 다른 함수를 학습하는게 아닌, 비슷한 형태의 함수를 학습해야 한다는거죠. 따라서 저자들은 다양한 PDE를 같이 학습하는 일종의 multi task learning 아이디어를 떠올립니다.

3. 제안 방법

이러한 아이디어를 바탕으로 제안한 P2INN 방법의 핵심 내용을 살펴보겠습니다. 크게 두가지로 나누어 살펴볼건데요. 첫 번째는 핵심 아이디어를 구현한 architecture를 살펴보고, 이어서 optional한 방법으로 제안한 fast finetuning 방법을 살펴보겠습니다.

3-1. Architecture

먼저 P2INN에서 학습할 방정식은 다음과 같습니다.

위 식은 convection, diffusion, reaction을 모두 표현해주고 있어요. beta는 convection항을 표현하고 있고, nu는 diffusion을, rho는 reaction 강도를 표현해주고 있죠. 즉 beta, nu, rho에 따라 convection, diffusionm, reaction이 조합된 방정식을 모두 표현할 수 있습니다. P2INN에서는 이 함수를 사용하여 모델이 학습할 데이터셋을 제작합니다.

Architecture는 간단합니다. 결국 P2INN 제안 방법의 핵심은 기존 PINN에 파라미터 전용 encoder를 더해주고, 이렇게 나온 parameter feature를 coordinate feature와 같이 사용해서 추론하자는건데요. 이를 반영한 P2INN architecture는 다음과 같습니다.

왼쪽 상단에 Parameter 전용 Encoder가 추가되었습니다. 그리고 이렇게 나온 parameter feature는 coordinate feature와 함께 concat되어 최종 decoder 입력으로 들어갑니다. 그 외 loss term은 기존 PINN과 동일하게 구성되어 있습니다.

핵심 내용만 간단하게 파이썬 코드로 구현해보면 다음과 같은데요.

import torch
import torch.nn as nn

# PDE 파라미터 인코더: [beta, nu, rho] → latent vector
class PDEEncoder(nn.Module):
    def __init__(self, input_dim=3, latent_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, latent_dim),
            nn.ReLU(),
        )

    def forward(self, mu):  # shape: (B, 3)
        return self.net(mu)

# 좌표 인코더: (x, t) → latent vector
class CoordEncoder(nn.Module):
    def __init__(self, input_dim=2, latent_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, latent_dim),
            nn.ReLU(),
        )

    def forward(self, xt):  # shape: (B, 2)
        return self.net(xt)

# 해석망: [coord_latent, param_latent] → u(x,t)
class SolutionDecoder(nn.Module):
    def __init__(self, input_dim=128, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),  # scalar output u(x,t)
        )

    def forward(self, z):
        return self.net(z)

# 전체 P2INN 모델
class P2INN(nn.Module):
    def __init__(self):
        super().__init__()
        self.coord_encoder = CoordEncoder()
        self.param_encoder = PDEEncoder()
        self.decoder = SolutionDecoder()

    def forward(self, x, t, mu):
        xt = torch.cat([x, t], dim=-1)      # shape: (B, 2)
        coord_z = self.coord_encoder(xt)   # shape: (B, 64)
        param_z = self.param_encoder(mu)   # shape: (B, 64)
        z = torch.cat([coord_z, param_z], dim=-1)  # (B, 128)
        u_pred = self.decoder(z)
        return u_pred

# loss항
def compute_pde_residual(u_pred, x, t, mu, model):
    # 자동미분을 위해 requires_grad 설정
    x.requires_grad_(True)
    t.requires_grad_(True)

    u = model(x, t, mu)

    # ∂u/∂t, ∂u/∂x, ∂²u/∂x² 계산
    grad_u = torch.autograd.grad(u, [x, t], grad_outputs=torch.ones_like(u), create_graph=True)
    u_x, u_t = grad_u[0], grad_u[1]

    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]

    beta, nu, rho = mu[:, 0:1], mu[:, 1:2], mu[:, 2:3]
    residual = u_t + beta * u_x - nu * u_xx - rho * u * (1 - u)
    return torch.mean(residual ** 2)

loss항은 기존 PINN과 마찬가지로 residual의 MSE로 계산해줍니다. P2INN의 forward 부분을 보면 coordinate feature와 parameter feature가 concat되어 최종 decoder 입력으로 들어가는 부분을 확인할 수 있습니다.

3-2. Fast Finetuning

저자들은 여기에서 한발 더 나아가 빠르게 finetuning 할 수 있는 방법도 제안하고 있는데요. 모든 파라미터에 대응 가능하도록 학습했다고 하지만, 사실 내가 당장 풀어야 하는 파라미터셋에 대해 더 잘 푸는게 중요할 수 있잖아요? (대부분 실제 문제 상황에서는 그렇죠) 이런 상황에서 efficient finetuning 할 수 있는 방법을 제안하고 있습니다.

아이디어는 LLM의 LoRA 등에서 나온것으로 보이는데요. 핵심 원리는 decoder 가중치를 SVD (singular value decomposition)한 뒤 diagonal matrix만 finetuning 하자는겁니다.

위 그림은 제안한 fast finetuning 방법을 표현해주고 있습니다. 최근 LLM에서는 다양한 방식의 parameter efficient finetuning 방법들이 연구되고 있고, LoRA는 그중 대표적인 방법인데요. LoRA 계열들에서 아이디어를 착안한 finetuning 방식으로 보입니다.

4. 실험 결과

이렇게 제안한 방법의 실험 결과를 살펴보겠습니다.

4-1. 1D CDR Equations

먼저 본문에서 제안한 CDR 방정식에 대한 실험 결과를 비교해보겠습니다.

6가지 종류의 PDE에 대해 성능을 비교하고 있는데요. PINN과는 압도적인 차이를 보이고 있고, 다른 방법들과 비교해도 모든 실험에서 큰 성능 차이를 보여줍니다.

4-2. Unseen 파라미터 일반화

그럼 과연 학습 데이터에서 보지 못한 파라미터에 대해서도 inference 할 수 있을까요? Reaction PDE로 실험을 진행했는데요, rho는 1~10 범위로만 데이터셋을 만들어 학습하고 interpolation과 extrapolation으로 못본 파라미터를 만들어 성능 비교를 진행했습니다.

실험 결과를 보면 PINN은 rho가 4 이상부터 에러가 급상승하는 반면, P2INN은 모든 파라미터 영역에서 안정적인 모습을 볼 수 있습니다.

4-3. PINN Failure Mode 비교

이번에는 PINN이 학습하기 어려운 고주파 데이터를 P2INN의 extrapolation으로 학습해보겠습니다. convection에서는 beta=30을, reaction에서는 rho=5인 PDE를 만들고 학습해보겠습니다.

실험 결과를 보면 확실히 PINN은 학습 자체가 안되는 모습인데요. beta=20까지 학습한 P2INN은 extrapolation으로도 그럴듯한 해공간을 만들어내는 모습입니다. reaction 방정식도 마찬가지로 PINN은 완전히 실패하는 반면 P2INN은 extrapolation으로 대응 가능한 모습을 보입니다.

4-4. Fast Finetuning

마지막으로 Fast Finetuning 방식의 성능 비교 실험 결과를 살펴보겠습니다.

위 실험 결과를 보면 convection, reaction, diffusion 방정식에 대해 PINN 방식과 다양한 finetuning 방식의 성능을 비교하고 있습니다. 실험을 수행한 세가지 방정식에 대해서는 제안 방법인 SVD 방법의 성능이 가장 좋은 모습을 볼 수 있습니다.

5. 마치며

이번 포스팅에서는 P2INN 논문의 주요 내용을 정리해봤습니다. 실제 시뮬레이션 관련 딥러닝 솔루션을 개발하다보면 결국 기대효과 부분에서 막히는 경우가 많은데요. 현재 찾은 PDE 파라미터에 대해서는 빠르게 동작하는 PINN 기반 딥러닝 모델을 개발할 수 있지만, 파라미터가 바뀌는 경우에 대해서는 바로 대응할 수 없기 때문이죠. 이렇게 바뀐 파라미터에 대해 다시 학습 데이터셋을 제작하고 모델을 재학습 하는 리드타임이 솔루션의 매력을 떨어뜨립니다.

P2INN에서는 이러한 문제를 해결할 수 있는 방법중 하나인 파라미터들을 별도의 encoder로 학습하는 방법을 제안했습니다. 논문에서는 convection, diffusion, reaction 방정식에 대해서만 실험을 수행했지만, 파라미터를 별도로 encoding 한다는 아이디어를 잘 확장하면 다양한 방정식에서도 활용할 수 있을것으로 기대됩니다.

6. 참고 자료

• P2INN paper
• PINN 논문 리뷰

Ffightingseok