기술 요약: AHPA: 확산 트랜스포머를 위한 적응적 계층적 사전 정렬

1. 문제 제기

확산 트랜스포머 (DiTs) 의 최근 발전은 고충실도 이미지 생성을 개선했으나, 이러한 모델을 처음부터 훈련시키는 것은 여전히 계산적으로 비용이 많이 들며, 의미론적으로 의미 있고 구조적으로 일관된 표현을 학습하기 위해 종종 수백만 번의 반복이 필요합니다. **표현 정렬 (Representation alignment)**은 외부 비전 인코더 (DINOv2 등) 나 내부 VAE 특징에서 얻은 정보 풍부한 시각적 사전 지식과 중간 트랜스포머 특징을 정렬함으로써 이 과정을 가속화하는 패러다임으로 등장했습니다.

그러나 기존 정렬 방법들은 전체 노이즈 제거 궤적에 걸쳐 고정된 감독 목표 또는 고정된 정렬 세분성을 부과한다는 치명적인 한계를 겪습니다. 저자들은 이 "시간 단계 무관 (timestep-agnostic)"한 접근 방식이 신호 대 잡음비 (SNR) 에 따라 표현 감독의 유용한 세분성이 체계적으로 변화하기 때문에 최적이지 않다고 주장합니다:

• 고잡음 영역: 모델은 coarse semantic(거시적 의미) 과 레이아웃 수준의 앵커링을 통해 혜택을 봅니다.
• 저잡음 영역: 모델은 공간적으로 세밀하고 구조적으로 충실한 정교화가 필요합니다.

정적인 단일 수준 감독자는 모델의 진화하는 요구 사항과 표현 불일치를 초래하여, 노이즈 제거 과정이 진행됨에 따라 기울기의 방향 일관성이 감소하고 기울기 간섭이 발생합니다. 또한, 외부 교사 (teacher) 에 의존하는 방법은 상당한 메모리 오버헤드와 훈련 복잡성을 도입하는 반면, 자기지도 학습 방법은 확산 훈련의 비정상적 특성을 포착하지 못하는 단일화된 특징 목표에 의존하는 경향이 있습니다.

2. 방법론: 적응적 계층적 사전 정렬 (AHPA)

표현 불일치를 해결하기 위해 저자들은 AHPA를 제안합니다. 이는 정적 목표 매칭에서 시간 단계 적응적 계층적 사전 스케줄링으로 표현 정렬을 재구성하는 경량 프레임워크입니다.

핵심 구성 요소

1. 계층적 VAE 사전 마이닝:
   고정된 VAE 인코더를 단순한 잠재 압축 모듈로 취급하는 대신, AHPA 는 이를 사전 지식의 다중 스케일 저장소로 활용합니다. 인코더의 중간 레이어는 기능적 그룹으로 분할됩니다:

   • 중간 수준 그룹 ($G_{mid}$): 공간 위상과 부분 수준의 관계 (구조적 발판) 를 포착합니다.
   • 심층 수준 그룹 ($G_{deep}$): 더 넓은 구성과 전역적 의미 추상화를 포괄합니다.
   • 얕은 레이어 ($G_1, G_2$) 는 초기 노이즈 제거 단계에서 과도한 픽셀 수준의 잡음을 유발하는 경우가 많아 제외됩니다.

2. 시간 단계 조건부 동적 라우터:
   경량 다층 퍼셉트론 (MLP) 라우터 $R_\phi$는 현재 시간 단계 $t$를 입력으로 받아 적응적 가중치를 예측합니다. 이는 두 가지 수준의 스케줄링을 수행합니다:

   • 그룹 내 집계: 중간 수준 그룹과 심층 수준 그룹 내의 특징에 가중치를 부여합니다.
   • 그룹 간 조정: $G_{mid}$와 $G_{deep}$의 기여도를 균형 있게 조절합니다.
     라우터는 모델의 현재 노이즈 제거 단계와 정렬 세분성을 동기화하기 위해 이러한 계층적 사전 지식을 동적으로 선택하고 가중합니다.

3. 공동 최적화:
   트랜스포머 백본 ($\theta$) 과 라우터 ($\phi$) 는 공동으로 최적화됩니다. 정렬 목적 함수는 트랜스포머의 투영된 은닉 상태와 집계된 계층적 목표 $z_s(t)$ 사이의 코사인 유사도 손실을 최소화합니다. 중요하게도, 정렬 분기는 추론 (inference) 중에 폐기되므로 추가 추론 비용은 0입니다.

3. 주요 기여

• 비정상 정렬 분석: 저자들은 효과적인 정렬 세분성이 노이즈 제거 시간 단계에 따라 변한다는 점을 규명했습니다. 그들은 기울기 신호 대 잡음비 (G-SNR) 지표를 도입하여 정렬 상태를 정량화했으며, 정적 기준선들이 $t \to 0$으로 갈수록 치명적인 충실도 저하 (기울기 간섭) 를 겪는다는 것을 입증했습니다.
• 계층적 VAE 사전 마이닝: 고정된 VAE 인코더의 중간 레이어가 자연스러운 다중 스케일 사전 은행을 제공한다는 것을 밝혔습니다. 실험 결과, $G_{mid}$와 $G_{deep}$는 상당한 위상 보완성을 보여주며, 심층 특징은 고잡음 영역을 앵커링하고 중간 수준 특징은 저잡음 구조 정교화를 안내하는 것으로 나타났습니다.
• 적응적이고 효율적인 정렬: AHPA 는 외부 교사 없이 적응적 계층적 감독을 달성합니다. 이는 계산 오버헤드를 무시할 정도로만 증가시킵니다 (GFLOPs 약 5% 증가 및 2.5M 파라미터). 또한 기준선과 비교할 수 있는 높은 훈련 처리량을 유지합니다.

4. 실험 결과

실험은 다양한 규모 (B/2, L/2, XL/2) 의 SiT(Scalable Interpolant Transformers) 백본을 사용하여 ImageNet(256×256 및 512×512) 과 MS-COCO 에서 수행되었습니다.

• 수렴 속도: AHPA 는 훨씬 적은 반복 횟수로 경쟁력 있는 수렴을 달성합니다. SiT-XL/2 의 경우, AHPA 는 40 만 번의 반복에서 FID 8.1 에 도달하여 700 만 번 반복으로 훈련된 일반 SiT 의 성능과 일치합니다.
• 성능 대 최첨단 (SOTA):
  • AHPA 는 모든 규모에서 정적 자기 정렬 방법 (SRA 2) 보다 우수합니다.
  • 외부 교사를 사용하는 무거운 방법 (예: DINOv2 를 사용하는 REPA) 과 비교하여 FID 와 구조적 충실도 (sFID) 측면에서 경쟁하거나 능가하며, 외부 교사가 필요 없다는 점이 특징입니다. 예를 들어, ImageNet 256×256 에서 AHPA 는 40 만 단계에서 FID 8.1 을 달성한 반면, REPA 는 8.5 를 기록했습니다.
  • MS-COCO 텍스트-이미지 작업에서 AHPA 는 FID 4.52 를 달성하여 SRA 2 와 REPA 모두를 능가했습니다.
• 효율성: AHPA 는 외부 인코더의 무거운 메모리 및 계산 비용을 피합니다. 총 GFLOPs 를 약 5% 만 증가시키고 7.4 Iter/s 의 훈련 속도를 유지하는 반면 (기준선은 8.1 Iter/s), REPA 와 같은 교사 기반 방법은 약 22% 의 속도 저하를 겪습니다.
• 메커니즘 검증: 학습된 라우팅 정책의 분석은 시간 단계가 감소함에 따라 라우터가 심층 의미 특징에서 중간 수준 구조 특징으로 자연스럽게 초점을 이동시키는 "자율적 인수인계"를 확인시켜 주며, 이는 이론적 비정상 요구 사항을 반영합니다.

5. 중요성 및 주장

이 논문은 AHPA 가 대규모 생성 모델을 가속화하기 위한 원칙적이고, 절약적이며, 자체完結적인 솔루션을 제공한다고 주장합니다. 그 주요 중요성은 확산 과정의 비정상 물리학과 정렬의 세분성을 동적으로 정렬함으로써 정적 안내 패러다임에 내재된 표현 불일치를 해결하는 데 있습니다.

외부 교사 대신 내재적 VAE 계층을 활용함으로써 AHPA 는 효율성과 합성 품질 간의 트레이드오프를 제거합니다. 저자들은 이를 AI 연구의 민주화를 향한 한 걸음으로 위치지으며, 제한된 계산 예산을 가진 연구자들이 외부 비전 인코더의 금지적인 오버헤드나 불필요한 훈련 반복 없이 고충실도 확산 모델을 훈련할 수 있게 한다고 설명합니다. 이 연구는 "비정상 정렬 요구 사항"이 확산 궤적의 근본적인 속성이며, 이를 해결할 때 내재적 계층적 안내의 잠재력을 완전히 unlocked 한다는 것을 시사합니다.