기술 요약: 확산 샘플링을 통한 데이터 변환 역전 (Inverting Data Transformations via Diffusion Sampling)

문제 정의

본 논문은 일반적인 리 군(Lie group)에서의 변환 역전(transformation inversion) 문제를 다룹니다. 많은 머신러닝 및 과학적 모델링 시나리오에서, 관측된 데이터 $\tilde{x}$는 분포 내(in-distribution) 샘플 $x$에 미지의 변환 $g$를 적용하여 생성됩니다 (즉, $\tilde{x} = g \cdot x$). 목표는 이 관측된 데이터를 원래의 데이터 분포로 매핑하는 역변환 $g^{-1}$를 복구하는 것입니다.

이는 미지의 섭동 변환(nuisance transformation)이 존재하며 솔루션이 일반적으로 유일하지 않기 때문에 **블라인드 역문제(blind inverse problem)**로 정식화됩니다. 기존 방법들은 주로 데이터 공간에서 작동하거나 특정 변환 패밀리(예: 아핀 또는 회전 군)에 특화되어 있는 반면, 본 연구는 이미지 호모그래피($PGL(3, \mathbb{R})$)나 PDE의 리 점 대칭(Lie point symmetries)과 같이 비콤팩트(non-compact)하고 비아벨(non-Abelian)이며 쌍불변 메트릭(bi-invariant metric)이 없는 경우를 포함한 일반적인 리 군 작용을 대상으로 합니다.

주요 응용 사례로 강조된 것은 **테스트 타임 등변성(test-time equivariance)**입니다. 사전 학습된 신경망은 종종 기하학적 변환에 대해 강건함이 부족합니다. 저자들은 입력 데이터를 모델의 학습 분포로 되돌리기 위해 테스트 시점에 이러한 변환을 역전시키는 방법을 제안하며, 이를 통해 특수한 등변 아키텍처를 사용하거나 재학습할 필요 없이 강건성을 향상시킵니다.

방법론: 변환 역전 에너지 확산 (TIED)

제안된 방법인 TIED는 데이터 공간 상의 에너지 함수 $E$에 의해 정의되는 볼츠만 분포를 통해 변환에 대한 사후 분포 $p(g | \tilde{x})$를 모델링합니다. 여기서 에너지는 데이터 사전 확률 $p_x(x)$로부터 유도되며, $E(x) = -\log p_x(x)$입니다. 사후 밀도는 다음과 같습니다:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

이 사후 분포로부터 샘플링하기 위해, TIED는 리 군 상에서 직접 수행되는 확산 과정을 도입합니다. 이 방법론은 에너지 지형(energy landscape)의 거칠고 다봉적인(multimodal) 특성과 매니폴드 상의 샘플링에 따른 기하학적 제약이라는 두 가지 주요 과제를 해결합니다.

1. 리 군에서의 트라이비얼라이즈드 확산 (Trivialized Diffusion on Lie Groups)

임의의 군 원소의 접공간(tangent space)에서 업데이트를 수행하는 대신, TIED는 **트라이비얼라이제이션(trivialization)**을 활용합니다. 업데이트는 리 대수(Lie algebra) $\mathfrak{g}$(항등원에서의 접공간)에서 수행되며, 지수 사상(exponential map)을 통해 현재의 군 원소로 전달됩니다.

• 순방향 과정 (Forward Process): 리 대수에서의 미세 섭동을 추출하고 이를 우곱(right multiplication)을 통해 전달함으로써 군 상에 분산 폭발형(variance-exploding) 확률 미분 방정식(SDE)을 정의합니다:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  여기서 $w^{\mathfrak{g}}_t$는 리 대수에서의 브라운 운동입니다.
• 역방향 과정 (Reverse Process): 사후 분포로부터 샘플을 복구하기 위해 역시간 SDE를 유도합니다. 결정적으로, 드리프트 항은 리 대수의 원소인 트라이비얼라이즈드 스코어(trivialized score) $\nabla_g \log p_t(g_t)$를 포함합니다.

2. 트라이비얼라이즈드 타겟 스코어 항등식 (Trivialized Target Score Identity)

핵심적인 이론적 기여는 새로운 트라이비얼라이즈드 타겟 스코어 항등식입니다. 표준 스코어 매칭은 깨끗한 샘플에 대한 접근을 요구하지만, 이러한 역전 설정에서는 샘플을 사용할 수 없습니다. 대신, TIED는 노이즈 분포의 스코어를 깨끗한 에너지 함수의 그래디언트로 표현합니다.
저자들은 (쌍불 invariant 메트릭을 요구하지 않고) 일반적인 리 군에 대해, 트라이비얼라이즈드 스코어가 노이즈 커널 $k_t$에 대한 몬테카를로 기댓값으로 표현될 수 있음을 증명합니다:
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
여기서 $\lambda(w)$는 역전 하에서 하르 측도(Haar measure)의 변화를 설명합니다. 이를 통해 스코어를 에너지 함수 $E$와 자동 미분을 통한 그 그래디언트만을 사용하여 추정할 수 있습니다.

3. 실제 구현

샘플링 알고리즘(Algorithm 1)은 역방로서 SDE를 이산화합니다. 각 단계에서 다음을 수행합니다:

1. 노이즈 커널 $k_t$에서 노이즈 $w$를 샘플링합니다.
2. 변환된 후보들에 대한 에너지 $E$를 계산합니다.
3. 몬테카를로 항등식을 사용하여 트라이비얼라이즈드 스코어를 추정합니다 (수치적 안정성을 위해 log-sum-exp 트릭을 사용).
4. 지수 사상을 사용하여 군 원소를 업데이트합니다.

이 방법은 VAE에서 유도되거나 사전 학습된 분류기의 예측 신뢰도에서 직접 유도되는 다양한 에너지 정의를 지원합니다.

주요 기여

1. TIED 샘플러: 일반적인 리 군에서 데이터 변환을 역전시키기 위한 새로운 확산 샘플러입니다. 이는 업데이트가 매니폴드 내에 머물도록 보장하며, 리 대수만을 사용하여 비콤팩트하고 비선형적인 작용을 처리합니다.
2. 이론적 프레임워크: 변환 역전이 사전 학습된 모델의 테스트 타임 등변성을 위해 사용될 수 있음을 증명했습니다. 저자들은 변환 역전 사후 분포가 $G$-등변(equivariant)함을 보여줌으로써 이 접근법의 타당성을 입증했습니다.
3. 새로운 스코어 항등식: 쌍불 invariant 메트릭 없이도 작동하는 리 군을 위한 일반화된 타겟 스코어 항등식을 제시하여, 깨끗한 데이터 샘플 대신 에너지 그래디언트를 이용한 효율적인 스코어 기반 샘플링을 가능하게 했습니다.
4. 실험적 검증: 이미지 호모그래피($PGL(3, \mathbb{R})$) 및 PDE 대칭($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$)에 대한 실험을 통해 강력한 베이스라인 대비 일관된 성능 향상을 보여주었습니다.

실험 결과

본 논문은 세 가지 주요 작업에 대해 TIED를 평가합니다:

1. $SO(10)$에서의 합성 샘플링: TIED는 고차원 리 군($dim=45$) 상의 다봉 에너지 분포에서 성공적으로 샘플링하였으며, 단순한 트라이비얼라이즈드 키네틱 랑제방(kinetic Langevin) 샘플러보다 현저히 빠르게 수렴했습니다.
2. 이미지 분류 (MNIST):
   • 설정: 사전 학습된 ResNet18을 미지의 아핀($Aff(2, \mathbb{R})$) 및 호모그래피($PGL(3, \mathbb{R})$) 변환이 적용된 MNIST 이미지에 대해 테스트했습니다.
   • 베이스라인: 최적화 기반 방법(LieLAC, FoCal, ITS) 및 샘플링 방법(Kinetic Langevin)과 비교했습니다.
   • 결과: TIED는 모든 베이스라인을 크게 앞질렀습니다. 예를 들어, 분류기 기반 에너지를 사용한 아핀 변환 환경에서 TIED는 **82.64%**의 정확도를 달알성한 반면, 가공되지 않은 ResNet18은 55.48%, Kinetic Langevin은 **74.72%**를 기록했습니다. 또한 일부 설정에서는 특화된 등변 네트워크(affConv)보다도 높은 성능을 보였습니다.
3. PDE 풀이:
   • 설정: DeepONet 뉴럴 오퍼레이터를 미지의 리 점 대칭 변환 하에 있는 1D 열 방정식(Heat equation) 및 버거스 방정식(Burgers equation)에 대해 테스트했습니다.
   • 결과: TIED는 모든 케이스에서 가장 낮은 상대적 $L_2$ 오차를 달성했으며, LieLAC와 FoCal을 능가했습니다. $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$ 대칭이 있는 1D 열 방정식의 경우, TIED는 오차를 0.043까지 줄였는데, 이는 LieLAC의 0.066 및 보정되지 않은 DeepONet의 0.690과 대조적입니다.

의의 및 주장

저자들은 TIED가 사전 학습된 모델을 테스트 시점에 임의의 리 군 변환에 대해 등변하게 만들 수 있는 **훈련이 필요 없는 파이프라인(training-free pipeline)**을 제공한다고 주장합니다. 에너지 기반 모델(사전 학습된 모델 자체로 근사 가능)을 활용함으로써, 이 방법은 특수한 등변 아키텍처나 비용이 많이 드는 재학습을 피할 수 있습니다.

본 연구의 의의는 다음과 같습니다:

• 범용성: 쌍불 invariant 메트릭이 없는 비콤팩트한 경우를 포함하여 광범위한 리 군 클래스에 적용 가능하며, 이는 기존의 확산 또는 최적화 방법들이 어려움을 겪었던 부분입니다.
• 강건성: 변환된 입력을 훈련 분포로 효과적으로 복구함으로써, 사전 학습된 네트워크의 기하학적 변환에 대한 일반화 성능을 크게 향고시킵니다.
• 효율성: 에너지 그래디언트를 사용하는 트라이비얼라이즈드 스코어 항등식과 결합된 확산 기반 접근 방식은 거친 에너지 지형에서 전통적인 MCMC나 반복적 최적화보다 더 효율적이고 안정적인 샘플링 프로세스를 제공합니다.

결론적으로, 이 논문은 이 접근법이 **등변성을 위한 추론 타임 스케일링(inference-time scaling for equivariance)**의 한 형태임을 보여주며, 대규모 파운데이션 모델 시대에 출처 외(out-of-distribution) 변환에 대한 강건성을 확보하기 위한 실용적인 솔루션을 제시합니다.