확산 샘플링을 이용한 데이터 변환 역전 문제 연구

Inverting Data Transformations via Diffusion Sampling

본 연구에서는 일반적인 리 그룹에서의 변환 역전 문제를 다룹니다. 특정 데이터가 알려지지 않은 그룹 요소에 의해 변환되고, 목표는 해당 데이터를 원래 데이터 분포로 되돌리는 역변환을 찾는 것입니다. 이러한 알려지지 않은 변환은 기계 학습 및 과학 모델링 분야에서 광범위하게 나타나며, 관측값에 상당한 왜곡을 초래할 수 있습니다. 우리는 확률적인 관점에서 변환에 대한 사후 분포를 볼츠만 분포로 모델링하며, 이는 데이터 공간 상의 에너지 함수에 의해 정의됩니다. 이 사후 분포에서 샘플링하기 위해, 리 그룹 상에서 모든 업데이트가 다양체 내에서 이루어지도록 하는 확산 과정을 도입합니다. 또한, 관련된 리 대수에서만 계산이 필요합니다. 본 연구에서 제안하는 Transformation-Inverting Energy Diffusion (TIED) 방법은 새로운 형태의 trivialized target-score identity를 활용하여 변환 사후 분포의 효율적인 score-based 샘플링을 가능하게 합니다. 주요 응용 분야로서, 사전 훈련된 신경망의 입력 변환에 대한 강건성을 향상시키는 test-time equivariance 문제를 중점적으로 다룹니다. 이미지 호모그래피 및 PDE 대칭에 대한 실험 결과, TIED는 test time에서 변환된 입력을 훈련 분포로 복원할 수 있으며, 기존의 강력한 정규화 및 샘플링 방법보다 우수한 성능을 보여줍니다. 코드: https://github.com/jw9730/tied

We study the problem of transformation inversion on general Lie groups: a datum is transformed by an unknown group element, and the goal is to recover an inverse transformation that maps it back to the original data distribution. Such unknown transformations arise widely in machine learning and scientific modeling, where they can significantly distort observations. We take a probabilistic view and model the posterior over transformations as a Boltzmann distribution defined by an energy function on data space. To sample from this posterior, we introduce a diffusion process on Lie groups that keeps all updates on-manifold and only requires computations in the associated Lie algebra. Our method, Transformation-Inverting Energy Diffusion (TIED), relies on a new trivialized target-score identity that enables efficient score-based sampling of the transformation posterior. As a key application, we focus on test-time equivariance, where the objective is to improve the robustness of pretrained neural networks to input transformations. Experiments on image homographies and PDE symmetries demonstrate that TIED can restore transformed inputs to the training distribution at test time, showing improved performance over strong canonicalization and sampling baselines. Code is available at https://github.com/jw9730/tied.