기술 요약: 효율적인 확산을 위한 라미네이팅 표현 오토인코더 (FlatDINO)

문제 정의

최근 이미지 생성 분야의 발전은 확산 모델(diffusion models)이 픽셀 공간의 잠재 변수(latents) 대신 자기 지도 학습(SSL) 패치 특징량(예: DINOv2로부터 추출된 특징)에서 직접 작동할 수 있음을 보여주었습니다. RAE(Zheng et al., 2025)로 대표되는 이 접근 방식은 이미 의미론적 구조가 인코딩된 SSL 특징을 활용함으로써, 픽셀 학습 기반 VAE에 내재된 재구성-생성 간의 트레이드오프(reconstruction-generation trade-off)를 우회합니다.

그러나 상당한 계산적 병목 현상이 여전히 남아 있습니다. DINOv2와 같은 표준 SSL 인코더는 고차원 패치 임베딩(예: 256×256 이미지에 대해 768 차원의 256개 패치)의 조밀한 2차원 그리드를 생성합니다. 이 표현은 원래 신호의 크기와 맞먹는 수준이며, 따라서 픽셀 공간 확산과 비교했을 때 계산상의 이점이 없습니다. 이 조밀한 그리드는 인접한 패치들이 상당한 의미론적 내용을 공유하기 때문에 공간적 중복성이 높으며, 이로 인해 이러한 특징 위에서의 확산은 FLOPs 및 메모리 측면에서 불필요하게 비용이 많이 듭니다.

방법론

저자들은 조밀한 2차원 DINOv2 패치 임베딩을 압축된 1차원 연속 토큰 시퀀스로 압축하도록 설계된 변이형 오토인코더(VAE)인 FlatDINO를 제안합니다.

1. 아키텍처

• 인코더(Encoder): 고정된 DINOv2 ViT-B/14(레지스터 포함)로부터 추출된 256개의 패치 임베딩을 입력으로 받는 비전 트랜스포머(ViT-B)입니다. 인코더는 패치 시퀀스 앞에 $T$개의 학습 가능한 레지스터 토큰을 추가합니다. 처리 후에는 원래의 2D 공간 구조를 버리고 오직 이 $T$개의 레지스터 토큰만을 잠재 표현으로 유지합니다.
• 디코더(Decoder): 학습 가능한 레지스터 토큰을 쿼리(query)로 사용하여 압축된 잠재 시퀀스에 어텐션(attention)을 수행하고, 원래의 256개 DINOv2 패치 임베딩을 재구성하는 ViT-L입니다.
• 학습 목표(Training Objective): 모델은 평균 제곱 오차(MSE) 재구성 손실과 KL 발산 페널티를 사용하는 표준 VAE 목표를 사용하여 학습됩니다. KL 가중치($\beta$)는 일관된 정규화 압력을 보장하기 위해 전체 잠재 차원($T \times d$)에 의해 정규화됩니다.
• 잠재 생성(Latent Generation): 생성 모델(구체적으로 flow matching을 사용하는 DiT-XL)이 압축된 FlatDINO 잠재 공간 위에서 학습됩니다. 이 모델은 표준 가우시안 사전 분포(prior)로부터 데이터 분포를 따라 직선 경로를 따라 샘플을 이동시키는 속도장(velocity field)을 학습합니다.

2. 핵심 설계 선택 사항

• 1D 압축: TiTok(Yu et al., 2024a)에서 영감을 받아, 2D 그리드를 1D 시퀀스로 압축합니다. 저자들은 32개의 토큰이 재구성 충실도와 시퀀스 길이 사이의 유리한 절충점을 제공한다는 것을 발견했습니다.
• 특징 차원 축소: 토큰 수는 256개에서 32개로 줄어들지만(8배 감소), 토큰당 특징 차원 또한 압축됩니다(예: 768에서 128로). 이는 전체 차원을 더욱 줄이기 위함입니다.
• Flow Matching: 생성 모델은 표준 확산 대신 flow matching을 활용하며, 학습된 위치 임베딩(positional embeddings)과 함께 1D 잠재 시퀀스에서 직접 작동합니다.

주요 기여

1. FlatDINO: 자기 지도 학습 패치 특징을 1차원 잠재 표현으로 압축하여, 원래의 2D 공간 구조를 효과적으로 제거하면서도 의미론적 내용을 보존하는 최초의 방법입니다.
2. 효율성 이득: FlatDINO 잠재값으로 학습된 확산 모델이 경쟁력 있는 생성 품질을 달면서도 계산 비용을 크게 줄일 수 있음을 입증했습니다. 구체적으로, 8배의 시퀀스 길이 감소는 압축되지 않은 DINOv2 특징 위에서의 확산과 비교했을 때, 순방향 패스(forward pass)당 8배의 FLOPs 감소와 학습 단계당 최대 4.5배의 FLOPs 감소를 가져옵니다.
3. 토큰화 분석: 토큰 절제 연구(token ablation studies)를 통해, 저자들은 모델이 대부분의 토큰에 대해 국소적이고 덩어리 형태(blob-like)의 수용 영역(receptive fields)을 학습하여 이미지를 공간적으로 분할함을 보여줍니다. 또한 특정 임계값(예: 32개에서 16개로) 미만으로 토큰을 줄이면 모델이 덜 효율적인 인코딩 전략(가로 줄무늬 형태)을 채택하게 되어 품질이 급격히 저하되는 "상전이(phase transition)" 현상을 확인했습니다.

결과

실험은 ImageNet 256×256 데이터셋에 대해 DiT-XL 아키텍처를 사용하여 수행되었습니다.

• 생성 품질: FlatDINO 잠재값(32×128)으로 학습된 DiT-XL은 분류기 없는 가이던스(CFG) 적용 시 1.85의 gFID를, 미적용 시 3.34를 달립다. 이는 압축되지 않은 특징을 사용하는 RAE(CFG 적용 시 1.41 gFID 달성) 및 다른 최첨단 잠재 확산 모델들과 비교해도 경쟁력 있는 수준이며, 상당한 압축에도 불구하고 이를 달성했습니다.
• 계산 효율성:
  • 추론(Inference): FlatDINO는 전체 DINOv2 패치에서 작동하는 것보다 순방향 패스당 약 8배 적은 FLOPs를 요구합니다.
  • 학습(Training): 감소된 시퀀스 길이와 인코더 오버헤드를 포함하더라도, FlatDINO는 RAE 베이스라인과 비교하여 DiT-XL 모델의 학습 단계당 4.5배 적은 FLOPs를 요구합니다.
• 재구성(Reconstruction): 32×128 구성은 0.77의 rFID를 달성했습니다(압축되지 않은 RAE 베이스라인의 0.62와 비교). 이는 특징 공간에 대한 높은 충실도의 재구성을 나타냅니다.
• 토큰 절제 연구: 연구 결과 32개의 토큰이 국소적 수용 영역을 유지하기에 충분하다는 것이 밝혀졌습니다. 토큰을 16개로 줄이면 가로 줄무늬 인코딩으로의 질적 변화와 재구성 품질의 급격한 저하가 발생합니다.

의의 및 주장

본 논문은 FlatDINO를 효율적이고 의미론적으로 근거가 있는 이미지 합성을 향한 단계로 규정합니다. 저자들은 2D 패치 그리드에 내재된 공간적 중복성을 활용함으로써, 의미론적으로 조직된 잠재 공간에서 작동하는 것의 수렴 이점을 희생하지 않으면서도 토큰 수를 8배 줄이는 것이 가능하다는 점을 주장합니다.

이 연구는 압축된 의미론적 표현에서의 확산이 자원이 제한된 환경에서 실용적인 접근 방식임을 시사합니다. 저자들은 현재 결과가 예비적이며 모델이 완전히 수렴하지 않았음(RAE는 800 에포크 학습, 본 모델은 600 에포크 학습)을 인정합니다. 생성 품질의 나머지 격차는 불충분한 학습 기간과 압축된 의미론적 잠재값에 특화된 확산 레시피의 필요성 때문이라고 설명합니다. 향차 연구로는 격차를 좁히기 위한 확장된 학습 일정 및 최적화된 SSL 오토인코더에 대한 조사가 제안되었습니다.