Robust Automatic Differentiation of Square-Root Kalman Filters via Gramian Differentials

1. 문제 제기 (Problem)

선형 가우스 상태 공간 모델 (LGSSM) 에서 파라미터 학습을 위해 칼만 필터의 로그 주변 우도 (log-marginal likelihood) 에 대한 기울기를 계산할 때, **제곱근 칼만 필터 (Square-Root Kalman Filter)**의 자동 미분 (Automatic Differentiation, AD) 과정에서 두 가지 주요한 수치적/수학적 문제가 발생합니다.

• 배경: 제곱근 칼만 필터는 수치적 안정성을 위해 공분산 행렬 $P$ 대신 하삼각 행렬 $S$ ($P=SS^\top$) 를 전파합니다. 이 과정의 핵심 연산은 행렬 $M$을 삼각화하는 연산자 $T(M)=L$ (여기서 $LL^\top = MM^\top$) 입니다. 실제로는 $M^\top$의 얇은 QR 분해 (Thin QR) 를 통해 계산됩니다.
• 문제 1 (비유일성): 행렬 $M$의 열 수 ($m$) 가 행 수 ($n$) 보다 큰 경우 ($m > n$), 직교 행렬 $Q$가 무한히 많은 선택지가 존재하여 삼각화 행렬 $L$의 미분이 정의되지 않습니다. 이는 예측 단계나 업데이트 단계에서 자주 발생합니다.
• 문제 2 (특이점): 표준 QR 분해 기반의 야코비안 공식은 $R^{-1}$ (또는 $L^{-1}$) 을 포함합니다. 입력 행렬 $M$이 랭크 결손 (rank-deficient) 인 경우 (예: 공분산 행렬 $U$ 또는 $V$가 특이행렬일 때), 역행렬이 존재하지 않아 기울기 계산이 발산합니다.

2. 방법론 (Methodology)

저자는 필터의 출력 (우도, 평균, 공분산) 이 삼각화 행렬 $L$ 자체보다는 그라미안 (Gramian) $LL^\top = MM^\top$을 통해서만 의존한다는 관찰에 기반한 새로운 미분 기법을 제안합니다.

• 핵심 통찰 (Smooth-Factorization Argument):
  • 필터의 손실 함수 $\ell$은 $L$에 의존하지만, 실제로는 $L$의 그라미안 $\Sigma = LL^\top = MM^\top$에만 의존합니다.
  • $M \mapsto MM^\top$ 매핑은 다항식이므로 모든 곳에서 매끄럽습니다 (smooth). 따라서 $L$의 삼각화 과정이 미분 불가능한 지점이 있더라도, 최종 손실 함수는 입력 $M$에 대해 매끄럽습니다.
• 그라미안 미분 항등식 (Gramian Differential Identity):
  • QR 분해를 직접 미분하는 대신, $LL^\top = MM^\top$ 관계식의 미분인 다음 식을 만족하는 $dL$을 구합니다:
    $$dL L^\top + L dL^\top = dM M^\top + M dM^\top$$
  • 이 식을 만족하는 임의의 $dL$은 그라미안에 의존하는 손실 함수에 대해 정확한 기울기 (JVP) 를 제공합니다.
• 대리 접선 (Surrogate Tangent) 유도:
  • 랭크 결손이 있는 경우를 처리하기 위해 $dL$을 두 부분으로 분해하여 구성합니다.
    1. 열 공간 (Column-space) 성분: $L$의 열 공간에 속하는 섭동에 대해, 모어 - 펜로즈 (Moore-Penrose) 의사역행렬 $L^\dagger$를 사용하여 고유한 하삼각 해를 구합니다.
    2. 영 공간 (Null-space) 성분: $L$의 열 공간 밖의 섭동에 대한 보정 항을 추가하여 전체 식을 만족시킵니다.
  • 이 두 성분을 합친 Definition 2가 제안된 폐쇄형 (closed-form) 접선입니다.

3. 주요 기여 (Key Contributions)

1. 폐쇄형 JVP 공식 도출: 그라미안 의존 손실 함수에 대해 정확하며, 랭크 결손 입력에서도 정의되는 제곱근 필터의 접선 벡터 곱 (JVP) 을 폐쇄형으로 유도했습니다. 이는 기존 QR 기반 공식이 유효한 영역에서는 기존 공식과 일치하고, 그렇지 않은 영역에서는 확장된 해를 제공합니다.
2. 그라미안 충분성 증명: 칼만 필터의 모든 관련 출력 (로그 우도, 필터링된 평균, 필터링된 공분산) 이 $L$이 아닌 $LL^\top$을 통해 입력에 의존함을 증명하여, 제안된 미분 기법의 타당성을 수학적으로 입증했습니다.
3. 자동 전치 (Automatic Transposition) 를 통한 VJP: 제안된 JVP 매핑이 입력 접선 $dM$에 대해 선형임을 증명했습니다. 이를 통해 역방향 모드 (Reverse-mode) 에 필요한 벡터 - 야코비안 곱 (VJP) 을 별도의 유도 없이 선형 연산자의 전치 (transpose) 만으로 자동으로 얻을 수 있음을 보였습니다.
4. 오픈소스 구현: JAX 프레임워크를 기반으로 한 오픈소스 구현을 제공하여, 실제 상태 공간 모델 학습에 즉시 적용 가능하게 했습니다.

4. 실험 결과 (Results)

• 실험 설정: 랭크 1 노이즈에서 풀 랭크 노이즈까지 변화하는 공분산 행렬을 가진 간단한 LGSSM 모델을 사용했습니다.
• 성능:
  • 제안된 방법은 랭크 결손 및 차원 불일치 상황에서도 로그 주변 우도 곡선과 일관된 기울기 (tangents) 를 정확하게 계산했습니다.
  • 반면, 기존 "고전적인" QR 분해 기반 접근법은 랭크 결손 시 기울기 계산이 발산하거나 유효하지 않은 값을 반환하여 실패했습니다.
  • Fig. 1 에서 보듯, 제안된 방법의 기울기 화살표가 우도 곡선의 실제 접선과 완벽하게 일치함을 확인했습니다.

5. 의의 및 결론 (Significance)

• 강건한 학습: 이 연구는 제곱근 칼만 필터를 사용하는 상태 공간 모델의 파라미터 학습을 위해 필수적이었던 수치적 불안정성과 미분 불가능성 문제를 해결했습니다.
• 범용성: 제안된 방법은 랭크 결손 행렬 (rank-deficient matrices) 을 다루는 능력을 갖추어, 노이즈 공분산이 특이행렬인 경우나 필터링 과정에서 공분산이 퇴화하는 상황에서도 안정적인 학습을 가능하게 합니다.
• 실용성: JAX, PyTorch 등 현대적인 자동 미분 프레임워크와 호환되므로, 엔드 - 투 - 엔드 (end-to-end) 상태 공간 모델 학습 파이프라인에 통합하기 용이합니다.

요약하자면, 이 논문은 QR 분해의 미분 한계를 우회하여 그라미안 항등식을 직접 미분함으로써, 제곱근 칼만 필터의 자동 미분을 수학적으로 엄밀하고 수치적으로 강건하게 만드는 획기적인 방법을 제시했습니다.