기술 요약: DiffATS: 정렬된 텐서 공간에서의 확산

1. 문제 제기

고해상도 시공간장 (예: 비디오, 기상 데이터, 유체 역학, 편미분방정식 (PDE) 해) 의 직접적인 확산 모델링은 계산적으로 불가능합니다. 표준 확산 모델 (DM) 은 전체 데이터 차원에서 반복적인 함수 평가를 요구하여 막대한 학습 및 샘플링 비용을 초래합니다. 예를 들어, $84 \times 720 \times 1440$ 크기의 텐서로 기상 예보 모델을 학습시키는 데는 수십 개의 TPU 에서 며칠이 소요됩니다.

이를 완화하기 위해 잠재 확산 모델 (LDM) 은 일반적으로 변분 오토인코더 (VAE) 를 사용하여 데이터를 저차원 잠재 공간으로 압축합니다. 그러나 과학 데이터의 경우, 기성 사전 훈련된 오토인코더를 구하기 어렵고, 작업별 오토인코더를 학습시키는 것은 비용이 많이 들며 정교한 최적화가 필요합니다. 이산 코사인 변환 (DCT) 이나 고정된 Tucker 분해에 기반한 오토인코더 없는 원시적 방법들이 존재하지만, 이들은 원시 데이터의 고유한 구조를 최적으로 근사하지 못할 수 있는 사전 정의된 기저에 의존합니다.

본 논문은 다음과 같은 질문을 제기합니다: 사전 훈련된 오토인코더 없이 고유한 저랭크 구조를 활용하여 고해상도 시공간 데이터에 대한 데이터 의존적이며 매개변수 효율적인 원시적 방법을 구축함으로써 생성 모델링을 개선할 수 있는가?

2. 방법론

제안된 프레임워크인 정렬된 텐서 공간에서의 확산 (Diffusion in Aligned Tensor Space, DiffATS) 은 Tucker 분해와 직교 프로크루스테스 (OP) 정렬에 기반한 새로운 텐서 원시적 방법의 구축을 통해 계산 및 표현적 과제를 해결합니다.

2.1 자연 Tucker 원시적 방법의 퇴화

저자들은 확산 학습을 위한 원시적 방법으로 표준 Tucker 분해 인자 (핵 텐서와 모드별 인자 행렬) 를 직접 사용하는 것이 최적이 아님을 규명했습니다. Tucker 분해는 비유일합니다. 단일 텐서는 직교 회전과 관련된 무수히 많은 핵 텐서와 인자 행렬의 조합으로 표현될 수 있습니다. 이러한 "게이지 모호성"은 원시적 방법 공간을 분산시켜 매핑의 연속성을 깨뜨리고 확산 모델의 스코어 함수 학습을 복잡하게 만듭니다.

2.2 직교 프로크루스테스 (OP) 정렬

이러한 퇴화를 해결하기 위해 DiffATS 는 직교 프로크루스테스 (OP) 정렬을 도입합니다.

1. 앵커 선택: 이 방법은 학습 데이터에서 "메도이드 (medoid)" 앵커 행렬을 선택합니다. 이러한 앵커는 인자 행렬의 경험적 분포 내에서 가장 중심적이거나 대표적인 부분 공간을 나타냅니다.
2. 정렬: 각 데이터 샘플에 대해, 인자 행렬을 해당 앵커 행렬에 정렬하기 위해 프로베니우스 노름 거리를 최소화하는 최적의 직교 회전을 찾습니다. 이는 각 Grassmannian 동치 클래스에서 고유한 표준 대표자를 선택합니다.

2.3 Grassmannian 원시적 방법

OP 정렬에 기반하여, 논문은 두 가지 원시적 방법 군을 정의합니다:

• 행렬 Grassmannian 원시적 방법 (MGPs): 저랭크 행렬 (예: 1 차원 PDE 또는 이미지 패치) 의 경우. 원시적 방법은 정렬된 특이 벡터와 변환된 특이 값으로 구성됩니다.
• 텐서 Grassmannian 원시적 방법 (TGPs): 저-다중선형-랭크 텐서 (예: 비디오 또는 2 차원 PDE) 의 경우. 원시적 방법은 핵 텐서와 선택된 모드에 대한 정렬된 인자 행렬로 구성됩니다.

이러한 원시적 방법들은 데이터 적응형 (특정 데이터셋의 구조에서 유도됨) 이며 매개변수 효율적으로, 원래 데이터를 $3.9$배에서$210$배까지 압축합니다.

2.4 이론적 보장

저자들은 특정 "앵커-중첩 (anchor-overlap)" 조건 (데이터의 부분 공간이 앵커와 비자명한 중첩을 가질 때) 하에서 원본 텐서 공간에서 원시적 방법 공간으로의 매핑이 위상동형사상 (homeomorphism) 임을 증명합니다. 이 이론적 결과는 다음을 보장합니다:

• 매핑은 전단사 (일대일이고 전사) 입니다.
• 매핑은 연속적이며 데이터 분포의 국소 위상적 구조를 보존합니다.
• 표현은 비퇴화적이어서 확산 모델이 원시적 방법의 충실한 분포를 학습하도록 보장합니다.

2.5 DiffATS 파이프라인

1. 전처리: 고차원 텐서는 HOOI(고차 직교 반복) 또는 SVD 를 통해 분해됩니다.
2. 정렬: 인자 행렬은 OP 를 통해 메도이드 앵커에 정렬됩니다.
3. 학습: 확산 모델 (DiT 백본 사용) 이 정렬된 원시적 방법 (핵 텐서 및 정렬된 인자) 에 직접 학습됩니다.
4. 추론: 모델이 원시적 방법 공간에서 샘플을 생성한 후, 명시적인 다중선형 재구성 (핵에 정렬된 인자를 곱함) 을 통해 원래 데이터 공간으로 재구성됩니다.

3. 주요 기여

• 새로운 원시적 방법: 회전 모호성을 앵커 기반 OP 정렬을 통해 해결하면서 Tucker 분해의 표현력을 유지하는 데이터 적응형 및 매개변수 효율적인 표현인 MGPs 와 TGPs 의 도입.
• 이론적 기반: 생성 모델링에 사용되는 저랭크 원시적 방법 군에 대해 중첩 조건 하에서 제안된 원시적 방법 매핑이 위상동형사상임을 증명하여, 최초의 공식적인 위상동형사상 보장을 제공.
• 오토인코더 없는 프레임워크: 사전 훈련된 심층 압축 오토인코더에 의존하지 않고 정렬된 원시적 방법 공간에서 직접 학습 및 샘플링하는 종단 간 프레임워크인 DiffATS 제안.
• 경험적 성능: 이미지, 비디오, PDE 해 전반에 걸쳐 기존 오토인코더 없는 기준선 (DCTDiff, SDIFT, AvgPool) 과 경쟁적인 신경 연산자 (FNO) 보다 우수한 성능을 보이며, 동등하거나 더 높은 압축 비율로 우수한 생성 품질 달성.

4. 실험 결과

저자들은 DiffATS 를 여섯 가지 벤치마크 데이터셋에서 평가했습니다:

• 이미지: CelebA-HQ ($1024 \times 1024$).
• 비디오: Moving MNIST.
• PDE: 1 차원 Burgers, 1 차원 반응 - 확산, 2 차원 Burgers, 2 차원 카르만 와류 거리.

주요 발견:

• 생성 품질: DiffATS 는 CelebA-HQ 에서 최상의 Fréchet Inception Distance(FID) 를 달성했습니다 (AvgPool 의 32.89 대비 23.57). 또한 Moving MNIST 에서 최상의 Fréchet Video Distance(FVD) 를 기록했습니다 (AvgPool 의 487.82 대비 415.42).
• PDE 예측: 조건부 생성 작업 (초기 조건을 기반으로 미래 궤적 예측) 에서 DiffATS 는 모든 기준선보다 현저히 우수했습니다. 예를 들어, 1 차원 반응 - 확산 데이터셋에서 최상의 기준선 (SDIFT) 대비 상대 $\ell_2$ 오차를 약 $9$배 감소시켰습니다.
• 압축: 이 방법은 높은 충실도를 유지하면서 $3.9$배에서$210$배까지의 압축 비율을 달성했습니다.
• 절대 실험: OP 정렬을 사용한 데이터 의존적 기저의 사용이 결정적임을 실험이 확인했습니다. 정렬을 제거하거나 공유 기저를 사용할 경우 성능이 크게 저하되었습니다.

5. 중요성과 주장

본 논문은 사전 훈련된 오토인코더를 구할 수 없거나 학습하기 비실용적인 과학 및 고해상도 시공간 데이터에 대해 오토인코더 기반 잠재 확산에 대한 견고한 대안을 DiffATS 가 제공한다고 주장합니다. 이러한 데이터의 고유한 저랭크 구조를 활용하고 기하학적 정렬을 통해 텐서 분해의 비유일성을 해결함으로써, DiffATS 는 효율적이고 고품질의 생성 모델링을 가능하게 합니다.

저자들은 자신의 작업을 생성 작업에서 표준 텐서 분해의 "퇴화" 문제에 대한 해결책으로 위치시키며, 안정적인 저차원 표현을 생성하기 위한 이론적으로 근거 있는 방법을 제공합니다. 그들은 그들의 접근 방식이 Stable Diffusion 과 같은 막대한 사전 훈련된 오토인코더에 의존하는 최첨단 이미지 합성 모델을 능가하는 것을 목표로 하는 것이 아니라, 그러한 자원이 부족하고 저랭크 구조가 지배적인 특징인 시공간 물리장이라는 특정 도메인을 대상으로 한다고 강조합니다.

저자들이 제안한 향후 방향으로는 위상동형사상 보장에 필요한 가정을 완화하기 위해 다중 앵커 OP 정렬과 동적 랭크 선택을 탐구하는 것을 포함합니다.