Inverting Data Transformations via Diffusion Sampling

Motivation

• Known: 기존에는 alignment, image registration, rotation synchronization 등 특정 변환군(affine, homography, orthogonal group 등)에 특화된 방법이나, deterministic canonicalization, equivariant neural network 설계를 통해 변환에 대한 강건성을 확보해왔다. 또한 diffusion model 기반 score-based sampling은 유클리드 공간에서 posterior sampling 및 blind inverse problem 해결에 널리 쓰여왔다.
• Gap: 기존 blind inverse problem 해법들은 데이터 공간에서 샘플링하거나 특정 변환군에 국한되어 일반적인 Lie group action, 비선형 작용, 비콤팩트 군에는 적용하기 어려웠고, canonicalization 기반 방법은 추가적인 equivariant network 학습이나 최적화 기반 접근에 의존해 취약하고 일반 Lie group으로 확장하기 어려웠다.
• Why: 알려지지 않은 기하학적 변환(뷰포인트 변화, 센서 모션, 기준계 변경 등)은 컴퓨터 비전, 의료 영상, 과학적 모델링 전반에서 관측을 심각하게 왜곡시키며, 이를 사전 학습된 신경망의 재학습 없이 test-time에 역변환하여 강건성을 확보할 수 있다면 equivariant architecture 설계의 비용과 확장성 문제를 우회할 수 있다는 점에서 중요하다.
• Approach: 변환에 대한 사후분포를 데이터 공간의 에너지 함수로 정의된 Boltzmann distribution으로 모델링하고, 이 사후분포에서 샘플링하기 위해 Lie algebra에서만 계산이 이루어지는 on-manifold diffusion process를 도입한다.

Achievement

Figure 5. For each PDE, we show an out-of-domain case for DeepONet fθ. From left: true solution, fθ prediction, fθ predi

1. TIED(Transformation-Inverting Energy Diffusion) 제안: 데이터 변환을 역전시키는 새로운 diffusion sampler로, on-manifold 성질을 보장하며(즉, 입력 데이터를 정확히 변환) 일반적인 Lie group과 비선형 작용을 다루면서 Lie algebra 지식만을 요구한다.
2. Test-time equivariance에 대한 이론적 정당화: 데이터 변환의 역변환이 근사 에너지 모델만으로도 사전학습된 모델의 test-time equivariance 확보에 사용될 수 있음을 증명하였다.
3. 어려운 변환군에서의 실증적 성과: 이미지 homography와 PDE에 대한 Lie point symmetry 등 도전적인 군에서 변환을 성공적으로 역전시키며, canonicalization 및 sampling 베이스라인 대비 사전학습 네트워크 성능 개선에서 일관되게 우수함을 입증하였다.

How

Figure 6. An illustration of trivialized forward diffusion on the projective linear group acting on an image.

• 변환 사후분포 p(g|x̃)를 데이터 공간 에너지 함수 E로 정의된 Boltzmann distribution으로 정식화
• score-based diffusion(Song et al., 2020)을 일반 Lie group으로 확장하기 위해 trivialization 기법(Lezcano Casado, 2019; Zhu et al., 2025)을 활용, group 상의 SDE를 정의하되 score는 Lie algebra에서 계산
• 새로운 trivialized target-score identity를 도출하여 효율적인 score-based sampling을 가능케 함
• compactness, bi-invariant metric, linear action 등의 가정 없이 general Lie group에 적용 가능하도록 설계
• 에너지 함수는 pretrained model 자체로부터 근사(Grathwohl et al., 2020)하여 training-free 파이프라인 구성
• 이미지 homography(PGL(3,R))와 PDE의 Lie point symmetry(예: DeepONet)에 적용하여 실험 진행, 강력한 canonicalization 및 sampling 베이스라인과 비교

Originality

• 변환 역전 문제를 일반 Lie group 상의 확률적 posterior sampling 문제로 재정식화하고, 이를 diffusion process로 해결한 최초의 시도 중 하나
• Lie algebra에서만 score를 계산하는 trivialized target-score identity를 새롭게 도출하여, 비콤팩트, 비선형 작용, non-bi-invariant metric을 갖는 general Lie group에도 적용 가능하게 함
• test-time equivariance를 위한 training-free, inference-time scaling 관점의 새로운 프레임워크 제시 (canonicalization/equivariant architecture 대안)

Limitation & Further Study

• 에너지 함수의 근사 품질(pretrained model로부터 유도)에 성능이 크게 의존할 가능성이 있으며, 에너지가 부정확할 경우 posterior sampling의 신뢰성이 저하될 수 있음
• diffusion sampling은 반복적 inference 과정을 필요로 하여 추가적인 test-time 계산 비용이 발생하며, 논문에서 이 비용과 성능 향상 간의 trade-off에 대한 심층 분석이 제한적일 수 있음
• 실험이 이미지 homography와 특정 PDE 문제에 국한되어 있어, 더 고차원적이거나 복잡한 Lie group(예: 무한차원 군, 매우 높은 차원의 데이터)에 대한 확장성 검증이 추가로 필요함
• 다중 모드(multimodal) 에너지 landscape에서의 수렴 보장이나 이론적 sample quality 분석이 더 필요함

Evaluation

총평: 일반 Lie group 상에서의 변환 역전 문제를 diffusion 기반 posterior sampling으로 풀어낸 참신하고 이론적으로 탄탄한 접근으로, test-time equivariance라는 실용적 응용까지 제시한 완성도 높은 연구이다. 다만 에너지 근사 품질 의존성과 계산 비용, 더 넓은 범위의 변환군에 대한 확장성 검증은 향후 과제로 남는다.