구형 인코더를 이용한 이미지 생성

기존 디퓨전 모델은 이미지를 생성할 때 수십 번의 반복 계산이 필요해 속도가 느리고 비용이 많이 든다. 이 논문은 이미지를 구(Sphere) 형태의 공간에 고르게 배치하는 새로운 방식을 통해, 단 한 번 혹은 아주 적은 횟수의 계산만으로도 디퓨전 모델 수준의 고화질 이미지를 생성할 수 있음을 입증했다.

기존의 VAE(Variational Autoencoder)는 이미지 데이터를 가우시안 분포(정규 분포)에 맞추려 노력한다. 하지만 이 과정에서 실제 데이터가 없는 빈 공간인 'Posterior Hole'이 생겨, 임의의 노이즈를 넣었을 때 이상한 이미지가 생성되는 한계가 있다. 디퓨전 모델은 이 빈 공간을 메우기 위해 수많은 단계를 거쳐 노이즈를 제거하지만, 연산량이 너무 많다.

이 논문은 데이터를 가우시안 분포가 아닌 '구(Sphere)'의 표면에 고르게 펼쳐놓는 아이디어를 제시한다. 구는 경계가 있고 회전 대칭성을 가지므로, 데이터들을 서로 밀어내게 만들면 구 표면 전체에 균일하게 배치할 수 있다. 이는 마치 자석들이 서로 밀어내며 구 표면을 빽빽하게 채우는 것과 유사한 원리이다.

이렇게 구 표면 전체가 이미지 데이터로 빽빽하게 채워지면, 구 위의 어떤 점을 선택하더라도 유효한 이미지가 생성된다. 결과적으로 디퓨전처럼 여러 번 깎아낼 필요 없이, 구 위의 점 하나를 찍어 바로 선명한 이미지를 얻을 수 있으며, 이는 생성 속도를 획기적으로 높이는 결과로 이어진다.

인코더 $E$는 입력 이미지 $x$를 잠재 표현 $z$로 매핑한 후, RMS Normalization을 통해 반지름이 $\sqrt{L}$인 구 표면 위의 벡터 $v$로 투영한다. [입력 벡터 $z$의 각 원소를 제곱 평균의 제곱근으로 나누어 $z$를 갱신 $\to$ 벡터의 길이는 항상 일정한 값($\sqrt{L}$)이 됨 $\to$ 모든 데이터가 구의 중심에서 같은 거리에 있는 표면에 위치함]

학습 시에는 $v$에 무작위 노이즈 $\sigma \cdot e$를 더한 후 다시 구 표면으로 투영하여 $v_{NOISY}$를 만든다. 디코더 $D$는 이 노이즈 섞인 벡터로부터 원본 이미지를 복원하도록 학습된다. [원본 벡터 $v$와 노이즈 벡터 $e$를 더한 후 다시 RMS Normalization 수행 $\to$ 구 표면 위의 새로운 점 $v_{NOISY}$ 생성 $\to$ 디코더가 이를 입력받아 원본 이미지 $x$를 재구성함]

손실 함수는 세 가지 요소의 가중치 합으로 구성된다. 픽셀 단위 복원 손실($L_{pix-recon}$), 인접한 잠재 벡터가 유사한 이미지를 만들도록 하는 픽셀 일관성 손실($L_{pix-con}$), 그리고 노이즈가 섞인 이미지의 잠재 벡터가 원본과 일치하도록 유도하는 잠재 일관성 손실($L_{lat-con}$)을 사용하여 잠재 공간의 구조를 안정화한다.

추론 시에는 구 위에서 무작위 벡터 $e$를 샘플링하여 디코더에 통과시키면 이미지가 생성된다. 더 높은 품질을 위해 생성된 이미지를 다시 인코딩-디코딩하는 반복 과정을 2~4회 수행할 수 있으며, 이때 Classifier-Free Guidance를 적용하여 조건 일치성을 높인다.

CIFAR-10 데이터셋에서 단 1단계(1-step) 생성만으로 FID 18.68을 기록했으며, 4단계에서는 2.72까지 개선되어 수천 단계의 DDPM(3.17)보다 우수한 성능을 보였다. 이는 적은 연산량으로도 기존 모델을 능가할 수 있음을 시사한다.

ImageNet 256x256 벤치마크에서 Sphere-XL 모델은 4단계 추론으로 gFID 4.02를 달성했다. 이는 수백 단계가 필요한 ADM-G(4.59)보다 우수하며, 수천 단계의 SID(2.44)와 경쟁 가능한 수준의 효율성을 입증한 결과이다.

Ablation Study를 통해 잠재 공간의 각도 $\alpha$가 생성 품질에 결정적인 영향을 미침을 확인했다. ImageNet의 경우 $\alpha=85^\circ$ 부근에서 가장 사실적인 디테일이 생성되었으며, 각도가 너무 작으면 이미지가 흐릿해지고 너무 크면 구조적 일관성이 떨어지는 경향을 보였다.

아키텍처는 표준 ViT(Vision Transformer)를 기반으로 하며, 인코더 끝과 디코더 시작 부분에 4개 층의 MLP-Mixer를 추가하여 토큰 간의 전역적 특징 혼합을 강화했다. 이는 단순한 선형 레이어보다 풍부한 특징 표현을 가능하게 한다.

잠재 공간의 크기 $L$은 이미지 크기에 따라 조절되며, ImageNet 256 해상도의 경우 $32^2 \times 64$ 크기를 사용하여 기존 VAE 대비 훨씬 낮은 압축률(3.0)을 유지한다. 이는 정보 손실을 최소화하여 복원 품질을 높이는 핵심 요인이다.

Classifier-Free Guidance(CFG)를 픽셀 공간과 잠재 공간 모두에서 적용할 수 있는 'Combo' 방식을 제안했다. 픽셀 공간에서의 CFG는 시각적 선명도를 높이고, 잠재 공간에서의 CFG는 의미적 일치성을 강화하여 두 방식의 장점을 결합했다.

학습 시 노이즈 강도 $r$을 조절하는 Jittering 전략을 사용하여 인코더가 구 표면 전체를 조밀하게 덮도록 유도했다. 이는 특정 지점에만 데이터가 몰리는 현상을 방지하여 잠재 공간의 연속성을 확보한다.

인코더와 디코더 모두에 파라미터를 할당해야 하므로 전체 모델 크기가 커질 수 있으며, 학습 시 일관성 손실 계산을 위해 인코더를 두 번 통과해야 하는 연산 부담이 존재한다.