기술 요약: 트랜스포머를 통한 라그랑지 입자 역학의 통합 시뮬레이션

문제 제기

물리 현상의 전체적 다양성을 모델링할 수 있는 통합 시뮬레이션 프레임워크를 구축하는 것은 시뮬레이션 과학에서 오랜 기간 지속되어 온 과제입니다. 전통적인 접근법은 각 현상에 대해 서로 다른 이산화 방법 (예: 유체의 오일러 격자, 고체의 라그랑지 메쉬) 과 지배 방정식 (예: 나비어 - 스토크스 방정식, 구성 법칙, 원자적 퍼텐셜) 에 의존합니다. 현대의 엔진은 여러 솔버 백엔드를 결합하지만, 여전히 현상별, 솔버별 패러다임으로 작동하여 현상 간 결합을 어렵게 만들고 단일 파이프라인의 적용 범위를 제한합니다.

기존의 머신러닝 기반 시뮬레이터는 명시적인 편미분 방정식 (PDE) 을 우회하여 데이터에서 직접 역학을 학습함으로써 보완적인 접근법을 제공합니다. 그러나 현재 신경망 시뮬레이터는 일반적으로 특정 현상이나 방정식 클래스에 맞춰 설계되고 훈련됩니다. 또한, 차원 축소 방법은 학습된 기저를 고정된 기하학적 형태나 모양 군에 종속시키는 경우가 많아, 물리 법칙이나 기하학이 변경될 때 모델 재설계가 필요합니다. 솔버별 재설계 없이 다양한 라그랑지 입자 시스템 전반에 걸쳐 일반화할 수 있는 단일 아키텍처가 필요합니다.

방법론

저자들은 직물, 탄성 고체, 뉴턴 및 비뉴턴 유체, 입상 물질, 분자 역학을 단일 고정 아키텍처 내에서 통합하여 시뮬레이션하는 입자 기반 예측 - 보정 트랜스포머를 제안합니다. 이 모델은 입자 시뮬레이션 (예: 위치 기반 역학) 에서 일반적으로 사용되는 예측 - 보정 방식을 따르며, 알려진 외부 힘의 적분과 입자 간 상호작용 학습을 분리합니다.

1. 예측 - 보정 프레임워크

시뮬레이션 단계는 두 단계로 나뉩니다:

• 명시적 예측기 (Explicit Predictor): 시간 단계 $\Delta t$ 동안 알려진 외부 힘 ($\mathbf{F}_t$) 을 적분하여 중간 상태 ($\tilde{\mathbf{X}}_t, \tilde{\mathbf{V}}_t$) 를 생성합니다. 이 단계는 외부 가속도를 고려하지만 입자 간 상호작용은 무시합니다.
• 학습된 보정기 (Learned Corrector): 신경망 $F_\theta$가 중간 상태, 입자 속성 ($\mathbf{C}$), 위상 구조 ($\mathbf{T}$), 그리고 경계 조건 ($\mathbf{X}_b, \mathbf{C}_b$) 을 기반으로 잔차 업데이트 ($\Delta \mathbf{X}_t, \Delta \mathbf{V}_t$) 를 예측합니다. 최종 상태는 $\mathbf{X}_{t+\Delta t} = \tilde{\mathbf{X}}_t + \Delta \mathbf{X}_t$ 및 $\mathbf{V}_{t+\Delta t} = \tilde{\mathbf{V}}_t + \Delta \mathbf{V}_t$로 업데이트됩니다.

2. 아키텍처 구성 요소

보정기는 국소 및 전역 상호작용을 모두 효율적으로 처리하도록 설계된 세 단계로 구성됩니다:

• 입자 토크나이저 (Particle Tokenizer): 전역 어텐션 전에 국소 상호작용을 인코딩합니다. 학습 가능한 커널 (삼선형 보간을 사용하는 3D 그리드로 파라미터화됨) 을 사용하여 세 가지 이웃 영역의 기여도를 집계합니다:

  • 공간 이웃 ($N^s_i$): 상대 변위와 속도에 기반한 입자 - 입자 상호작용.
  • 경계 이웃 ($N^b_i$): 입자 - 경계 상호작용.
  • 위상 이웃 ($N^t_i$): 직물/고체의 경우 정지 형태 인접성에 의해 안내되는 상호작용.
    출력은 집계된 특징을 입자의 자체 상태와 연결한 입자별 토큰 $\mathbf{h}_i$입니다.

• 슈퍼 토큰 인코더 (Super-Token Encoder): $N$개의 입자에 대한 전역 자기 어텐션의 계산 비용을 해결합니다. 멀티그리드 방법에 영감을 받아, 이 모듈은 전역 자기 어텐션과 토큰 병합을 번갈아 수행합니다.

  • 자기 어텐션: 3D 회전 위치 임베딩 (RoPE) 을 사용하여 현재 해상도에서 토큰을 문맥화합니다.
  • 토큰 병합: 가장 코사인 유사도가 높은 토큰을 탐욕적으로 매칭하고 병합하여 각 레이어에서 토큰 수를 반으로 줄입니다. 이는 국소 특징, 속도, 재료 정보를 집계하는 컴팩트한 "슈퍼 토큰" 세트를 생성합니다.
  • 이 과정은 $L_e$개의 레이어에 걸쳐 반복되어 깊은 레이어로 갈수록 시퀀스 길이와 계산 비용을 점진적으로 줄입니다.

• 슈퍼 토큰 디코더 (Super-Token Decoder): $L_d$개의 레이어를 통해 거친 문맥을 전체 입자 해상도로 다시 끌어올립니다.

  • 입자 토큰은 크로스 어텐션을 통해 슈퍼 토큰을 쿼리하여, 입자들이 슈퍼 토큰에 인코딩된 전역 정보에 접근할 수 있게 합니다.
  • 그런 다음 입자 토큰은 자기 어텐션과 피드포워드 네트워크를 통해 정제됩니다.
  • 최종 MLP 는 정제된 토큰을 잔차 위치 및 속도 보정으로 매핑합니다.

3. 훈련

모델은 $W$개의 상태 윈도우에 대한 자기회귀 롤아웃 손실 (autoregressive rollout loss) 을 사용하여 훈련됩니다. 손실 함수는 위치와 속도의 평균 제곱 오차 (MSE) 를 물리 정보 정규화 항목 (예: 비압축성 유체를 위한 발산 제로 제약 조건) 과 결합합니다. 아키텍처는 모든 역학 카테고리에서 고정되며, 훈련 데이터 분포만 변경됩니다.

주요 기여

1. 통합 아키텍처: 네트워크 구조를 수정하지 않고도 6 가지 서로 다른 라그랑지 역학 카테고리 (유체, 직물, 입상 물질, 탄성 고체, 비뉴턴 유체, 분자 역학) 를 시뮬레이션할 수 있는 단일 트랜스포머 기반 모델.
2. 새로운 토큰 메커니즘: 국소 이웃 상호작용을 위한 입자 토크나이저와 점진적 토큰 병합을 통한 전역 통신을 가능하게 하는 슈퍼 토큰 인코더 - 디코더의 결합으로, 표현력과 계산 효율성을 균형 있게 조정.
3. 일반화 및 응용: 보이지 않는 재료, 경계 구성, 초기 조건, 외부 힘에 대한 일반화 입증. 이 프레임워크는 상호작용 제어, 미분 가능한 롤아웃을 통한 역설계, 그리고 실제 세계 조작 데이터 학습을 포함한 하위 작업을 지원합니다.

실험 결과

모델은 확립된 솔버 (DFSPH, GIPC, MPM, VBD, OpenMM) 로 생성된 데이터를 사용하여 6 가지 카테고리에 대해 평가되었습니다.

• 정확도: 제안된 모델은 유체, 직물, 모래 작업 전반에 걸쳐 최신 베이스라인 (GNS, DeepLagrangian, Neural Operators) 보다 위치 및 속도에 대한 더 낮은 롤아웃 평균 제곱 오차 (MSE) 를 달성합니다. 예를 들어, 유체 작업에서 위치 오차는 (스케일링된) $5.17 \times 10^{-1}$로, 장기간 롤아웃에서 떠다니는 입자나 침투와 같은 아티팩트를 보인 베이스라인을 크게 능가했습니다.
• 안정성: 베이스라인이 종종 발산하거나 붕괴되는 긴 시간 범위 (예: 유체의 경우 800 프레임, 직물의 경우 250 프레임) 에서 모델은 안정적인 시뮬레이션을 유지합니다.
• 일반화:
  • 보이지 않는 조건: 모델은 보이지 않는 경계 기하학 (예: 호랑이 모양의 모래 더미) 과 초기 조건을 성공적으로 시뮬레이션합니다.
  • 입자 밀도: 훈련 세트와 다른 입자 샘플링 밀도로 일반화됩니다 (예: 약 7k 개 입자로 훈련, 8k~9k 개로 테스트).
  • 긴 롤아웃: 250 프레임에서 테스트될 때 훈련 시간 범위 (200 프레임) 를 넘어 안정성을 유지합니다.
• 어블레이션 연구: 국소 토크나이저 또는 전역 슈퍼 토큰 통신 중 하나를 제거하면 성능이 저하되어 국소 세부 사항과 전역 결합의 필요성이 확인됩니다. 발산 제로 정규화 항목을 추가하면 유체의 물리적 일관성이 더욱 향상됩니다.

중요성 및 주장

이 논문은 단일 미분 가능 아키텍처가 여러 특수 솔버 (FEM, SPH, MPM) 의 필요성을 대체하는 통합 시뮬레이션 프레임워크를 향한 한 걸음이라고 주장합니다. 각 현상을 공유된 입력 - 출력 형식에서의 서로 다른 훈련 분포로 취급함으로써, 모델은 현상별 솔버 엔지니어링의 필요성을 줄입니다.

저자들은 이 아키텍처가 가능하게 하는 세 가지 구체적인 하위 응용 분야를 강조합니다:

1. 상호작용 제어: 예측기가 임의의 외부 힘을 적분할 수 있는 능력을 활용하여 사용자가 지정한 힘 하에서 변형 가능한 물체의 실시간 조작.
2. 역설계: 신경망 롤아웃을 통한 역전파를 통한 물리 매개변수 (예: 마찰 계수) 의 기울기 기반 최적화.
3. 실제 세계 데이터 학습: 3D 가우시안 스플래팅과 같은 기법을 사용하여 실제 세계 관측에서 추출된 입자 궤적에 직접 훈련하고 고전적 시뮬레이터에 의존하지 않는 능력.

저자들은 주로 첫 번째 인코더 레이어의 전체 해상도 자기 어텐션으로 인해 단일 80GB GPU 에서 현재 구현이 약 50,000 개의 입자로 제한되는 메모리 제약 사항을 인정합니다. 희소 어텐션 변형이 이 용량을 확장할 수 있다고 제안합니다. 또한, 아키텍처가 통합되어 있지만 현재 구현은 역학 카테고리별로 별도의 파라미터 세트를 훈련한다는 점을 지적하며, 향후 목표는 모든 카테고리에서 공유 가중치를 가진 단일 모델을 구축하는 것입니다.