Distilling Drifting Transformers with Representation Autoencoders

기술 요약: 표현 오토인코더를 이용한 드리프팅 트랜스포머 증류 (Distilling Drifting Transformers with Representation Autoencoders)

문제 정의

표현 오토인코더(Representation Autoencoders, RAE)는 사전 학습된 DINO 특징으로부터 유도된 의미론적으로 풍부하고 레이블별로 클러스터링된 잠재 공간을 활용함으로써 흐름 기반 생성 모델(flow-based generative models)의 성능을 크게 향상시켰으나, 이 모델들을 효율적인 1-스텝 또는 소수의 스텝을 가진 생성기로 증류하는 것은 여전히 도전적인 과제로 남아 있습니다.

저자들은 근본적인 기하학적 불일치를 식별했습니다. RAE 잠재 공간은 기존 VAE 공간과 비교했을 때 그 기저의 상미분 방정식(ODE) 궤적에서 **심각한 비등방성(anisotropy)**과 **큰 곡률(curvature)**을 보입니다. 기존의 궤적 기반 증류 방법들은 매끄럽고 거의 직선에 가까운 교사(teacher) 궤적을 가정합니다. RAE 공간의 높은 곡률은 이러한 방법들을 불안정하거나 비효율적으로 만듭니다. 또한, 최근 제안된 **드리프팅 모델(Drifting Model)**은 궤적을 직접 맞추는 대신 분포 수준의 접근 방식을 제공하여 곡선 궤적 문제를 피하지만, 원래의 공식은 집중된 특징을 제공하기 위해 별도로 학습된 보조 마스크드 오토인코더(MAE)를 필요로 하며, 이는 상당한 계산 오버헤드를 초래합니다.

방법론

본 논문은 보조 특징 추출기의 필요성을 제거하면서 RAE 잠재 공간에 특화되어 적응된 드리프팅 모델 패러다임인 Drift-RAE를 제안합니다.

1. 기하학 및 역학의 이론적 분석

저자들은 먼저 잠재 공간을 정량적으로 분석하여, RAE 특징이 SD-VAE 특징보다 실질적으로 더 비등방적(낮은 참여 비율(Participation Ratio) 및 스펙트럼 엔트로피)임을 입증했습니다. 이들은 고차원 공간에서 양의 샘플(positive samples)이 지나치게 분산되어 있으면, 드리프팅 필드의 인력(attractive) 성분이 소멸되어 퇴화된 필드로 이어진다는 것을 이론적으로 증명했습니다 (정리 1).

• 핵심 통찰: 원래의 드리프팅 모델은 분산된 특징을 집중시키기 위해 MAE를 필요로 하지만, RAE 잠재 공간은 자연스럽게 의미론적으로 집중되어 있습니다. 이를 통해 드리프팅 모델이 보조 MAE 모듈 없이 직접 RAE 잠재값 위에서 작동할 수 있게 합니다.

2. Drift-RAE 프레임워크

이 방법은 사전 학습된 흐름 모델 $v_\theta$를 1-스텝 생성기 $G_\theta(z) = z - v_\theta(z, 1)$로 증류합니다. ODE 궤적을 일치시키는 대신, 생성기는 생성된 샘플들을 실제 데이터 분포로 밀어넣고 다른 생성된 샘플들로부터 멀어지게 하는 드리프팅 목적 함수를 최소화하도록 학습됩니다.

• 목적 함수: 손실(loss)은 양의 샘플(실제 데이터)과 음의 샘플(생성된 데이터)로부터 유도된 정규화된 드리프팅 필드 $\tilde{V}$를 활용하여 RAE 잠재 표현에 대해 토큰 단위로 계산됩니다.

3. 안정성을 위한 실질적 수정 사항

드리프팅 모델을 Diffusion-GAN 프레임워크와 정렬하고 학습 안정성을 개선하기 위해, 저자들은 세 가지 구체적인 수정 사항을 도입했습니다:

• Softmax 차원: 저자들은 원래의 드리프팅 모델에서 사용된 양방향 소프트맥스 트릭을 제거했습니다. 대신 샘플 인덱스에 대한 단방향 소프트맥스를 유지하며, 이는 근간이 되는 log-sum-exp 포텐셜의 그래디언트 방향과 더 잘 일치합니다.
• 입력 섭동(Input Perturbation): 실제 데이터와 생성된 데이터 사이의 비교차 매니폴드(non-intersecting manifolds)로 인해 발생하는 그래디언트 소실을 방지하기 위해, 생성된 샘플에 드리프팅 필드를 계산하기 전 라플라스 분포(Laplace distribution)에서 샘플링된 노이즈를 더해 섭동을 가합니다.
• 음의 샘플의 부분 분리(Partial Detachment of Negative Samples): GPU 메모리 예산을 초과하지 않으면서 분포 근사를 위해 사용할 수 있는 음의 샘플의 수를 늘리기 위해, 저자들은 일부($N_{neg}$)만 그래디언트를 전달하고 나머지는 분리(detached)되는 더 큰 규모의 음의 샘플 풀($N_{total\_neg}$)을 사용하는 방식을 제안합니다. 이는 생성된 분포의 경험적 근사를 개선합니다.

주요 기여

1. 기하학적 분석: RAE 잠재 공간이 높은 곡률과 비등방성을 가지고 있어 표준 궤적 기반 증류를 불안정하게 만든다는 것을 정량적으로 입증했습니다.
2. 이론적 정당성: 드리프팅 모델이 지나치게 분산된 공간에서는 실패하지만(원래 공식에서 MAE가 필요한 이유), RAE의 자연스럽게 집중된 의미론적 공간에서는 효과적임을 보여주는 증명을 제시했습니다.
3. Drift-RAE: 보조 MAE 특징 추출기 없이도 RAE 공간에서 최첨단 1-스텝 생성을 달성하는 새로운 증류 방법을 제안했습니다.
4. 안정성 강화: Diffusion-GAN 원칙에 부합하도록 드리프팅 목적 함수(소프트맥스 차원, 입력 섭동, 샘플 분리)를 수정하는 것에 대한 이론적 및 경험적 통찰을 제공했습니다.

실험 결과

본 방법은 DiTDH-XL 및 DiTDH-L 백본을 사용하여 ImageNet 256 $\times$ 256 데이터셋에 대해 평가되었습니다.

• 성능: Drift-RAE는 단일 생성 단계에서 FID 1.77(DiTDH-XL 사용 시) 및 2.12(DiTDH-L 사용 시)를 달성했습니다.
• 효율성: 이 모델은 단 10,000번의 증류 단계(약 16 에폭)만으로 이러한 결과에 도달합니다.
• 비교:
  • RAE 공간에서의 기존 증류 방법(예: MF-RAE)보다 우수한 성능을 보입니다.
  • 계산 비용이 많이 드는 보조 MAE를 사용하지 않았음에도 불구하고, 원래의 드리프팅 모델과 대등한 FID 및 향상된 FDDINOv2 점수를 달축했습니다.
  • 다른 잠재 공간(예: 픽셀 공간, SD-VAE)에서 학습된 1-스텝 및 소수 스텝 생성기와 경쟁할 만한 수준을 유지합니다.

의의 및 주장

본 논문은 **드리프팅(Drifting)**이 표현 공간에서 생성 모델을 증류하는 데 있어 효과적이고 유망한 프레임워크를 제공한다고 주장합니다. 주요 의의는 RAE의 내재된 의미론적 집중도를 활용하여 RAE 기반 모델 증류의 불안정성을 해결하고, 이로 인해 보조 특징 추출기에 대한 의존성을 제거했다는 점에 있습니다.

저자들은 이론적 분석이 단순화된 고차원 모델에 의존하고 있다는 점과, RAE 공간에서 (증류 없이) 드리프팅 모델을 처음부터 학습시키는 것은 여전히 미해결 과제라는 점을 언급하며 한계를 겸허히 인정했습니다. 또한, 많은 수의 동일 클래스 양의 샘플이 필요하다는 점이 텍스트-투-이미지 생성에서 확장성 문제를 야기할 수 있음을 강조하며, 이를 향후 연구 방향으로 제시했습니다.