기술 요약: Wasserstein 및 Kalman–Wasserstein KL 발산을 통한 잘 정의된(Well-Posed) KL 정규화 제어

1. 문제 정의

쿨백-라이블러(Kullback-Leibler, KL) 발산은 선형 해결 가능 마르코프 결정 과정(linearly-solvable MDP), 경로 적분 제어(path-integral control), 신뢰 영역 방법(trust-region methods) 등 강화 학습(RL) 및 최적 제어에서 표준적인 정규화 도구로 사용된다. 그러나 고전적인 KL 발산은 저노이즈(low-noise) 영역에서 치명적인 병리 현상을 보인다. 제어된 역학(controlled dynamics)을 가우시안 프로세스 노이즈가 매우 작은 수동적 역학(passive dynamics)과 비교할 때, KL 페널티는 노이즈 공분산에 반비례하는 계수를 가진 이차 제어 비용(quadratic control cost)으로 축소된다. 노이즈 공분산이 소멸함에 따라(결정론적 한계), 이 페널티는 무한대로 발산하여 정규화된 제어 문제를 특이(singular)하게 만든다. 이러한 "저분산 불안정성(low-variance instability)"은 최적 피드백 이득(optimal feedback gains)을 소멸시켜, 결정론적 또는 준결정론적 환경에서 불량한 폐루프 성능과 불안정성을 초래한다.

또한, 고전적인 KL 발산은 통계적 매니폴드에는 내재되어 있으나 기저 상태 공간(state space)의 기하학적 구조는 무시하는 피셔-라오(Fisher–Rao) 기하학을 통해 정의된다. 결과적으로 KL은 서포트 불일치(support mismatch, 예: $x \neq y$인 경우 $KL(\delta_x | \delta_y) = \infty$)에 대해 불연속적으로 반응하며, 상태 간의 거리를 인코딩하지 못한다.

2. 방법론

저자들은 "상태 공간 인식형(state-space-aware)" KL 아날로그를 구축하기 위해 통합된 정보 기하학적 프레임워크를 제안한다. 핵심 방법론은 KL 발산의 역학적 공식화에서 피셔-라오 기하학을 운송 기반 기하학(transport-based geometries)으로 대체하는 것이다.

2.1 기하학적 프레임워크

본 논문은 확률 밀도 매니폴드 상의 리만 기하학 $G$에 기반한 일반적인 발산 클래스 $D_G$를 정의한다.

• 역학적 공식화: 발산은 $G$-쌍대 측지선(G-dual geodesics)을 따르는 시간 가중 작용 적분(time-weighted action integral)으로 정의된다:
  $$D_G(\mu | \nu) = \inf \left\{ \int_0^1 t \|\dot{\gamma}(t)\|_{\gamma(t)}^2 dt : \gamma(0)=\mu, \gamma(1)=\nu \right\}$$
  여기서 $\gamma$는 $G$-쌍대 측지선이다.
• 기하학적 선택: 저자들은 피셔-라오 메트릭($G=FR$)을 다음으로 대체하는 것을 조사한다:
  • Wasserstein-2 ($G=O$): 상태 공간 기하학을 인코딩하는 오토(Otto) 메트릭.
  • 선형화된 Wasserstein ($G=\dot{H}^{-1}$): 평평한(flat) 구조.
  • **Kalman–Wasserstein ($G=KW$):** 정규화 파라미터$\lambda$와 공분산 정보를 포함하는 메트릭.
  • Stein 기하학: 쌍선형 커널(bilinear kernel) 사용.

2.2 이론적 유도

• 쌍대 측지선: 저자들은 $G$-쌍대 측지선이 퍼텐셜 에너지 범함수의 메트릭 경사 흐름(metric gradient flows)임을 규명한다.
• 폐쇄형 표현식(Closed-Form Expressions): 이들은 타원 분포(가우시안 및 스튜던트 $t$-분포를 포함하는 클래스) 사이의 명시적인 발산 공식을 유도한다.
  • Wasserstein의 경우($WKL$), 공분산이 퇴화하더라도 발산이 유한하다.
• **Kalman–Wasserstein의 경우($KWKL$):** 이 발산은 정규화 파라미터$\lambda$를 통해 "노이즈 플로어(noise floor)"를 도입하며, 이는 고전적 KL과 Wasserstein KL 사이를 보간한다.

2.3 제어에의 적용

제안된 발산들은 가우시안 프로세스 노이즈를 가진 선형 이차 제어(LQR) 문제의 정규화 항으로 적용된다.

• 정규화 항: 제어된 전이 분포 $p_t$와 제어되지 않은 분포 $p_0$ 사이의 발산은 가산적인 이차 제어 페널티 $u_t^\top R_\circ u_t$를 생성한다.
• 제어 페널티:
  • 고전적 KL ($FR$):$R_{FR} = B^\top \Sigma_w^{-1} B$($\Sigma_w \to 0$일 때 특이함).
  • Wasserstein ($O$): $R_O = B^\top B$ (노이즈와 무관함).
  • Kalman–Wasserstein ($KW$):$R_{KW} = B^\top (\Sigma_w + \lambda I)^{-1} B$($\Sigma_w \to 0$일 때도 유한함).

3. 주요 기여

1. 일반적 구성: 임의의 기하학에 대한 $G$-쌍대 측지선과 그에 따른 발산 $D_G$를 도입하여, 피셔-라오에 대한 고전적 KL과 운송 기하학에 대한 상태 공간 인식형 변형을 모두 회복한다.
2. 경사 흐름 규명: $G$-쌍대 측지선이 퍼텐셜 에너지의 메트릭 경사 흐름에 대응함을 증명하여, 발산을 기하학적 최적화 및 평균장 신경망(mean-field NN) 역학과 연결한다.
3. 명시적 공식 유도: Wasserstein, 선형화된 Wasserstein, 그리고 Kalman–Wasserstein 기하학 사이의 타원 분포 간 $D_G$에 대한 폐쇄형 표현식을 유도한다. 이 공식들은 운송 기반 기하학이 고전적 KL에 결여된 상태 공간 정보를 어떻게 포함하는지를 명시적으로 보여준다.
4. 잘 정의된 제어(Well-Posed Control): LQR 설정에서 Wasserstein 및 Kalman–Wasserstein 정규화 항이 프로세스 노이즈가 사라질 때도 유한하게 유지됨을 입증하여, 고전적 KL이 실패하는 지점에서 잘 정의된 최적 제어 문제를 제공한다.
5. 경험적 검증: 더블 인테그레이터(double integrator) 및 카트-폴(cart-pole) 시스템 분석을 통해, 제안된 방법이 비자명한(nontrivial) 피드백을 보존하고 저노이즈 영역에서 KL 정규화 제어보다 더 나은 폐루프 성능을 달성함을 보여준다.

4. 결과

• 이론적 한계: 노이즈가 소멸하는 극한($\Sigma_w \to 0$)에서, 고전적 KL 정규화($F_{FR}$)의 최적 피드백 이득은 0으로 수렴하여 사실상 제어를 제거한다. 반면, Kalman–Wasserstein 이득($F_{KW}$)은 0이 아닌 극한(구체적으로 $F_O$)으로 수렴하여 안정적인 피드백을 보존한다.
• 보간(Interpolation): Kalman–Wasserstein 발산($D_{KW}$)은 파라미터 $\lambda$에 의해 제어되는 고전적 KL(고노이즈)과 Wasserstein KL(저노이즈) 사이의 매끄러운 보간 역할을 한다.
• 안정성: 더블 인테그레이터 실험에서 노이즈가 감소함에 따라, KL 정규화 제어기는 안정성 경계(스펙트럼 반지름 $\to 1$)에 접근하여 큰 진동을 일으킨다. 반면 WKL 및 KWKL 제어기는 스펙트럼 반지름을 1보다 훨씬 낮은 수준으로 유지하여 안정성을 보장한다.
• 비등방성 노이즈(Anisotropic Noise): 비등방성 노이즈 시나리오에서 KWKL은 방향별로 적응한다. 즉, 저노이즈 차원은 WKL처럼(노이즈 규모를 무시) 취급하고, 고노이즈 차원은 KL처럼(노이즈 규모에 따라 스케일링) 취급한다. 반면 KL은 저노이즈 차원에서 제어를 완전히 억제한다.

5. 의의 및 주장

본 논문은 제안된 상태 공간 인식형 발산이 RL 알고리즘의 전면적인 개편 없이도 KL 정규화의 저노이즈 퇴화 문제를 해결하는 원칙적인 해결책을 제공한다고 주장한다. 피셔-라오 기하학을 운송 기반 기하학으로 대체함으로써 저자들은 다음을 달성한다:

• 강건성(Robustness): 결정론적 한계에서도 제어 문제가 잘 정의된 상태를 유지한다.
• 기하학적 일관성(Geometric Consistency): 정규화 항이 기저 상태 공간 기하학을 존중하며, 서포트 불일치에 대한 KL의 불연속적 동작을 피한다.
• 실용적 유용성(Practical Utility): 유도된 폐쇄형 표현식은 연속 제어에 사용되는 표준 LQR 및 가우시안 정책 프레임워크에 즉각적으로 통합될 수 있다.

저자들은 이 연구를 완전한 딥 RL 알고리즘이라기보다는 KL 유형의 발산을 구축하기 위한 기하학적 프레임워크로 포지셔닝한다. 이들은 본 연구가 다루기 쉬운 LQR 설정에서 검증되었으나, 다음 단계는 모델 프리(model-free) 및 모델 기반(model-based) 강화 학습에서의 경험적 평가가 될 것이라고 제언한다. 본 연구는 NEITALG 프로젝트 하에 유럽 연구 위원회(ERC)의 지원을 받았다.