A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

본 논문은 감쇠 작용을 최소화함으로써 시간적 스냅샷에서 2 차 인구 역학을 학습하는 새로운 프레임워크이자 알고리즘인 Wasserstein Lagrangian Mechanics(WLM)를 소개하며, 이를 통해 주기성, 와류 역학, 무리 짓기와 같은 복잡한 행동을 정확하게 모델링하는 데 있어 기울기 흐름의 한계를 극복합니다.

핵심

시간 경과에 따른 도시 광장의 군중 이동을 타임랩스 영상으로 보고 있다고 상상해 보세요. 오후 1 시, 1 시 5 분, 1 시 10 분에 각자가 어디에 있는지 snapshot 을 보게 됩니다. 당신의 목표는 그들이 왜 그렇게 움직이는지 파악하고, 1 시 15 분에 그들이 어디에 있을지 예측하는 것입니다.

지난 10 년간 과학자들은 이 문제를 해결하기 위해 군중을 언덕을 굴러 내려가는 공처럼 가정해 왔습니다. 그들은 군중이 항상 가장 낮은 에너지 지점 (예: 계곡) 을 찾아 그곳으로 미끄러지다가 멈춘다고 생각했습니다. 이를 **기울기 흐름 (Gradient Flow)**이라고 부릅니다.

문제점:
실제 생활은 단순히 언덕을 굴러 내려가는 것만이 아닙니다. 때로는 군중이 소용돌이처럼 원을 그리며 회전하고, 때로는 앞뒤로 진동하며, 때로는 목표에 도달한 후에도 계속 움직입니다. 기존의 언덕을 굴러 내려가는 모델은 이러한 움직임을 설명할 수 없습니다. 마치 미끄러지는 바위의 물리학만으로 회전하는 팽이를 설명하려는 것과 같습니다.

새로운 아이디어: 인구 역학 (Population Mechanics)
이 논문의 저자들은 군중을 바라보는 새로운 방식을 제안합니다. 단순히 언덕을 미끄러져 내려가는 것으로만 보지 않고, 전체 군중을 단일하고 거대하며 복잡한 물체로 취급하여 물리 법칙 (특히 뉴턴의 법칙, 하지만 사물 집단에 적용된) 을 따르게 합니다.

그들은 이를 **와서슈타인 라그랑주 역학 (Wasserstein Lagrangian Mechanics, WLM)**이라고 부릅니다.

다음은 비유를 사용하여 작동 방식을 간단히 설명한 것입니다:

1. '작용 (Action)' 원리 (가장 효율적인 경로)

A 지점에서 B 지점으로 이동하려는 등산객이라고 상상해 보세요. 당신은 무작위로 배회하지 않고, 가장 적은 '노력' (또는 '작용') 이 필요한 경로를 선택합니다.

• 기존 방법: 군중은 이용 가능한 가장 가파른 경사를 따라 미끄러집니다.
• 새로운 방법 (WLM): 군중은 자신의 위치와 속도를 모두 고려하여 가능한 가장 효율적인 경로를 선택합니다. 마치 제동만 걸고 멈추는 것이 아니라, 관성을 이용해 부드럽게 코너를 드리프트하는 자동차와 같습니다.

2. '위치 에너지' 지도

물리학에서 물체는 '위치 에너지' (예: 언덕을 굴러 내려가고 싶은 공) 에 기반하여 움직입니다.

• 저자들은 군중을 위한 특별한 '지도'를 만들었습니다. 이 지도는 사람들이 어디에 서 있는지에 관한 것뿐만 아니라, 전체 집단의 형태에 관한 것입니다.
• 만약 한곳에 군중이 너무 빽빽하게 모이면 '에너지'가 상승하여 군중은 자연스럽게 퍼집니다. 너무 멀리 떨어져 있으면 에너지가 변하여 그들이 모일 수도 있습니다.
• WLM 의 마법은 이 지도를 snapshot 들로부터 직접 학습한다는 점입니다. 인간이 규칙을 알려줄 필요가 없으며, 군중의 움직임을 관찰함으로써 '지형'을 파악합니다.

3. '관성' 학습 (왜 즉시 멈추지 않는지)

이것이 가장 큰 업그레이드입니다.

• 기존 방법 (기울기 흐름): 군중이 목표에 도달하면 즉시 멈춥니다. 벽에 부딪히면 그냥 죽어버리는 브레이크가 없는 자동차와 같습니다.
• 새로운 방법 (WLM): 군중은 관성을 가집니다. 그들이 원형으로 빠르게 움직이고 있다면, '언덕'이 평평해지더라도 그 원형 운동을 계속 유지합니다. 그들은 과잉 도달하고, 뒤로 흔들리며, 진동할 수 있습니다. 이를 통해 모델은 다음과 같은 복잡한 행동을 예측할 수 있습니다:
  • 소용돌이 (Vortices): 배수구에서 소용돌이치는 물.
  • 군집 비행 (Flocking): 무리를 지어 나는 새들 (머머레이션).
  • 세포 발달: 배아 성장 중 형태를 바꾸고 이동하는 세포.

컴퓨터가 학습하는 방법 (블랙박스 코치)

저자들은 물리 코치처럼 작동하는 컴퓨터 프로그램 (신경망) 을 구축했습니다.

1. 입력: snapshot 을 봅니다 (예: "여기가 오후 1 시, 1 시 5 분, 1 시 10 분의 군중입니다").
2. 추측: '게임의 규칙' (위치 에너지 지도와 마찰/항력의 정도) 을 추측합니다.
3. 시뮬레이션: 그 규칙에 기반하여 군중이 앞으로 움직이는 가상 시뮬레이션을 실행합니다.
4. 확인: 시뮬레이션 결과를 다음 실제 snapshot(1 시 15 분) 과 비교합니다.
5. 조정: 시뮬레이션이 틀렸다면 코치는 규칙을 수정하고 다시 시도합니다.

결국 코치는 그 특정 군중을 지배하는 정확한 '운동 법칙'을 학습하게 됩니다.

무엇을 테스트했는가

이 논문은 이 '코치'를 세 가지 매우 다른 유형의 군중에서 테스트했습니다:

1. 해양 소용돌이: 멕시코 만의 소용돌이치는 물. 기존 방법들은 소용돌이를 예측하는 데 어려움을 겪었지만, WLM 은 정확하게 맞췄습니다.
2. 배아 세포: 발달 중인 배아에서 분열하고 이동하는 세포. WLM 은 움직임이 복잡하고 혼란스러움에도 불구하고 다음에 세포가 어디에 있을지 예측할 수 있었습니다.
3. 보이드 (Boids, 새들): 군집하는 새들의 컴퓨터 시뮬레이션. 새들은 간단한 규칙 (충돌하지 않기, 가까이 있기, 무리와 함께 비행하기) 을 따릅니다. 기존 방법들은 새들이 단순히 언덕을 미끄러져 내려가는 것으로 생각하여 완전히 실패했습니다. WLM 은 '군집 물리학'을 학습하여 복잡한 루프를 돌고 있을 때조차 새들의 미래 움직임을 예측할 수 있었습니다.

결론

이 논문은 분자, 세포, 또는 동물의 집단을 단순히 언덕을 미끄러져 내려가는 집단이 아니라, 운동량과 관성을 가진 단일 역학적 시스템으로 취급함으로써 그들의 움직임을 훨씬 더 잘 이해하고, 예측하며, 공백을 메울 수 있다고 주장합니다.

이는 모든 사람이 단순히 직선으로 걷는다고 가정하여 춤을 예측하는 것과, 실제로는 운동량, 회전, 리듬을 가진 왈츠를 추고 있음을 깨닫는 것의 차이입니다.

기술 요약

기술적 요약: 라그랑주 행동에의 호소: 시간적 스냅샷으로부터 인구 역학 학습

문제 제기

분자 확산부터 세포 분화와 동물 무리 형성까지, 생물학적 및 물리학적 시스템 전반의 인구 역학은 알려지지 않은 힘에 의해 지배됩니다. 역사적으로 기계 학습과 계산 생물학은 이러한 역학을 주로 와서슈타인 (Wasserstein) 기울기 흐름을 사용하여 모델링해 왔습니다. 수학적으로 엄밀하고 다루기 쉽다는 장점이 있지만, 기울기 흐름은 자유 에너지를 최소화하므로 순수한 소산적이고 비주기적인 역학으로 제한됩니다. 결과적으로 주기성, 진동, 또는 보존적 운동과 같은 핵심 역학적 특성을 포착하지 못합니다.

기존의 대안들인 흐름 매칭 (flow matching) 이나 행동 매칭 (action matching) 은 보간 품질을 향상시키지만, 일반적으로 참조 과정을 명시하지 않는 한 관찰된 시간 범위를 넘어선 역학을 예측할 수 없습니다. 핵심적인 과제는 사전에 지정된 참조 과정이나 라그랑주 함수에 의존하지 않으면서도 관찰되지 않은 주변 분포를 보간하고 미래 상태를 예측할 수 있는 더 표현력 있는 역학 클래스를 추론하는 것입니다.

방법론: Wasserstein Lagrangian Mechanics (WLM)

저자들은 **Wasserstein Lagrangian Mechanics (WLM)**을 제안합니다. 이는 확률 측도 공간 $\mathcal{P}_2(\mathbb{R}^d)$에서 정의된 인구 수준의 작용 (action) 함수를 최소화함으로써 인구 역학 학습을 감쇠된 Wasserstein 라그랑주 하에서의 인구 수준 역학 추론으로 재정의하는 프레임워크입니다.

1. 이론적 프레임워크

이 연구는 확률 측도 공간 $\mathcal{P}_2(\mathbb{R}^d)$에서 정의된 인구 수준의 작용 함수 $S$를 최소화함으로써 인구 역학을 공식화합니다:
$$S[\rho_t, s_t, t] = \int_0^1 e^{\gamma t} \left( \frac{1}{2} \int \|\nabla s_t\|^2 \rho_t dx - U[\rho_t] \right) dt$$
여기서:

• $\rho_t$는 인구 밀도 (위치) 입니다.
• $s_t$는 운송 속도장 $v_t = \nabla s_t$를 정의하는 퍼텐셜입니다.
• $U[\rho_t]$는 인구 수준의 퍼텐셜 에너지 함수 (예: 엔트로피, 상호작용 커널) 입니다.
• $\gamma \geq 0$는 감쇠 계수입니다.

최소 작용의 원리를 적용하여 저자들은 시스템에 대한 해밀턴 운동 방정식을 유도합니다. 이 방정식들은 구조화된 2 차 역학 클래스를 설명합니다:
$$\frac{d}{dt}x_t = v_t, \quad \frac{d}{dt}v_t = -\nabla \frac{\delta U[\rho_t]}{\delta \rho_t}(x_t) - \gamma v_t$$
이 공식화는 여러 역학 체제를 통합합니다:

• 고전 및 양자 역학: $\gamma = 0$이고 $U$가 선형이거나 피셔 정보 (Fisher information) 항을 포함할 때 복원됩니다.
• 기울기 흐름: 과감쇠 극한 ($\gamma \to \infty$) 이나 정확한 초기 조건을 가진 특정 불안정 보존 시스템으로서 복원됩니다.
• 일반 역학: 표준 기울기 흐름에는 없는 진동 및 비소산적 행동을 허용합니다.

2. 알고리즘: 데이터로부터 역학 학습

이 논문은 라그랑주 함수 $L$이나 퍼텐셜 $U$를 명시하지 않고 관찰된 시간적 스냅샷 (주변 분포) $\{p_{t_i}\}$으로부터 직접 이러한 역학을 학습하는 신경 역학 모델인 WLM을 소개합니다.

• 매개변수화: 퍼텐셜 에너지 함수 $U[\rho_t]$는 입자의 경험적 측도에 작용하는 심층 신경망 $\Psi_\theta$로 근사됩니다. 감쇠 $\gamma$는 고정되거나 학습될 수 있습니다.
• 가속도 추정: Proposition 3.1 을 사용하여, 함수 미분 $\nabla \frac{\delta U}{\delta \rho}$는 자동 미분을 통해 개별 입자 좌표에 대한 신경망 $\Psi_\theta$의 기울기를 계산함으로써 우회됩니다.
• 학습: 모델은 Discretize-Then-Optimize 접근법을 사용하여 학습됩니다.
  1. 2 차 Verlet (Leapfrog) 적분기를 사용하여 초기 조건 $(p_0, v_0)$에서 역학을 시뮬레이션합니다.
  2. 시뮬레이션된 주변 분포와 관찰된 주변 분포 사이의 발산 (예: Sinkhorn 발산) 을 손실로 계산합니다.
  3. 시간 이산화된 궤적을 통해 역전파하여 $\theta$와 $\gamma$를 업데이트합니다.
• 추론: 학습이 완료되면 모델은 관찰되지 않은 주변 분포를 예측하거나 관찰된 시간 점 사이를 보간하기 위해 앞으로 시뮬레이션할 수 있습니다.

주요 기여

1. 이론적 통합: 이 논문은 Wasserstein Lagrangian Mechanics 가 고전 역학, 양자 역학, 기울기 흐름을 특수한 경우로 포함하며, 기울기 흐름 단독보다 훨씬 더 표현력 있는 역학 클래스를 제공함을 증명합니다.
2. 참조 없는 학습: WLM 은 참조 과정이나 사전 정의된 라그랑주 함수를 요구하지 않고 관찰된 주변 분포로부터 직접 2 차 인구 역학을 학습하는 최초의 알고리즘입니다.
3. 예측 능력: 관찰 창 내의 보간으로 제한되는 흐름 매칭이나 행동 매칭 방법과 달리, WLM 의 학습된 해밀턴 역학은 관찰 범위를 넘어선 예측을 가능하게 합니다.
4. 실증적 검증: 이 방법은 다양한 데이터셋에서 검증되어 기울기 흐름 학습, 소용돌이 해류 역학, 배아 세포 분화, 그리고 창발적 무리 행동 (Boids) 에서의 견고성을 입증했습니다.

실험 결과

저자들은 네 가지 다른 작업에서 WLM 을 최신 기준선 (JKONET*, NN-APPEX, CURLY-FM, 다양한 슈뢰딩거 브리지 방법 포함) 과 비교 평가했습니다:

1. 기울기 흐름 SDE: 합성 퍼텐셜 구동 SDE 에서 WLM (학습 가능한 마찰 포함) 은 짝지어진 및 짝지어지지 않은 설정 모두에서 최첨단 성능을 달성하며, 높은 마찰 값 ($\gamma \geq 500$) 을 학습함으로써 기울기 흐름 역학을 성공적으로 복원합니다.
2. 해류 소용돌이 데이터: 해류 입자를 보간하고 예측하는 데 있어 WLM 은 참조 과정을 사용하지 않는 방법들보다 우수한 성능을 보입니다. CURLY-FM(도메인 특화 참조 사용) 다음으로 전체 두 번째로 좋은 성능을 달성하며, 참조 과정 없이 정확하게 예측한 유일한 방법입니다.
3. 배아 scRNA 데이터: 인간 배아 줄기세포 분화 데이터에서 WLM 은 명시적인 시간 특징을 사용하지 않았음에도 불구하고, 모든 테스트된 방법 중 leave-one-out 보간에 대해 가장 낮은 Wasserstein 거리 ($0.704$) 를 달성하여 수많은 최적 수송 및 흐름 기반 기준선들을 능가합니다.
4. Boids (창발적 역학): 상호작용 에이전트 무리에 대한 새로운 테스트 (주기적이고 비기울기 역학을 보임) 에서 WLM 은 근본적인 역학을 성공적으로 학습하고 훈련 데이터보다 50 시간 단계 앞선 무리 행동을 예측합니다. 반면, 기울기 흐름 기준선들은 역학을 적합하지 못하게 하여 확산 행동만 생성합니다.

중요성과 주장

이 논문은 WLM 이 인구 역학 모델링에 대한 관점의 근본적인 전환을 나타낸다고 주장합니다. 1 차 기울기 흐름에서 2 차 라그랑주 역학으로 이동함으로써, 저자들은 에너지 최소화 접근법으로는 보이지 않는 주기성, 진동, 그리고 창발적 집단 행동을 포착할 수 있는 프레임워크를 제공합니다.

그 중요성은 다음과 같은 능력에 있습니다:

• 사전 지식 없이 학습: 특정 참조 과정이나 참조 SDE 를 가정하지 않고 복잡한 역학을 추론합니다.
• 일반화: 통계적 상관관계가 아닌 학습된 물리 법칙 (역학) 에 기반하여 미래 상태를 예측하고 누락된 데이터 점을 보간합니다.
• 통합: 기울기 흐름을 극한 사례로 자연스럽게 포함하면서 보존적 및 진동 시스템으로 확장하는 단일 수학적 프레임워크를 제공합니다.

저자들은 식별 가능성에 대한 근본적인 경로 법칙의 식별 가능성은 여전히 열린 질문이며 시뮬레이션 기반 학습 접근 방식이 계산 비용을 수반한다는 점을 인정합니다. 그러나 그들은 WLM 을 세포, 분자, 생태학적 시스템에서의 과학적 추론을 위한 초기 단계이지만 강력한 방향으로 위치시킵니다.