U‑REPA: U‑넷과 ViT를 잇는 새로운 표현 정렬 기법

초록

U‑REPA는 기존 Diffusion Transformer와 ViT 사이의 REPA 방식을 U‑넷에 적용하기 위해 설계된 프레임워크이다. 스킵 연결로 인해 중간 단계가 가장 의미 있는 표현을 담고 있음을 발견하고, MLP‑후 업샘플링으로 차원 불일치를 해소한다. 토큰‑단위 유사도 대신 샘플 간 관계를 정규화하는 매니폴드 손실을 도입해 공간 격차를 완화한다. 실험 결과 ImageNet‑256에서 200 epoch(1M step)만에 FID < 1.5를 달성했으며, 기존 REPA 대비 절반 정도의 학습 시간으로 더 나은 품질을 얻었다.

상세 분석

U‑REPA 논문은 Diffusion U‑넷 구조와 Vision Transformer(ViT) 기반 시각 인코더 사이의 표현 정렬 문제를 체계적으로 분석하고 해결책을 제시한다. 첫 번째 핵심 인사이트는 U‑넷의 스킵 연결이 네트워크 전반에 걸쳐 정보 흐름을 재구성한다는 점이다. 이로 인해 전통적인 DiT‑ 기반 REPA가 초기 레이어에 정렬을 두는 것과 달리, U‑넷에서는 중간 단계(다운샘플링과 업샘플링 사이)에서 가장 풍부한 의미 정보를 얻을 수 있다. 저자들은 다양한 깊이의 레이어에 정렬을 적용한 실험을 통해 중간 레이어가 FID와 IS 지표 모두에서 최적임을 확인하였다.

두 번째 문제는 공간 차원 불일치이다. U‑넷의 중간 특징 맵은 해상도가 크게 축소된 반면, ViT는 고정된 패치 토큰 수(N)와 동일한 차원을 유지한다. 이를 해결하기 위해 저자는 세 가지 업스케일링 전략을 비교했으며, “MLP 먼저 → 업스케일” 방식이 FLOPs와 성능 측면에서 가장 효율적임을 입증했다. MLP는 저차원 특징을 고차원 임베딩으로 변환하고, 이후 픽셀 언샤플링을 통해 원본 해상도에 맞춘다.

세 번째로, 토큰‑단위 코사인 유사도 기반 손실이 U‑넷과 ViT 사이의 큰 특징 공간 격차를 충분히 반영하지 못한다는 점을 지적한다. 대신 매니폴드 손실을 도입해 샘플 간 상대적 유사도를 정규화한다. 이 손실은 각 배치 내 샘플 쌍의 거리 행렬을 계산하고, ViT 특징이 제공하는 관계 구조와 U‑넷 특징이 학습하는 관계 구조를 최소화하도록 설계되었다. 결과적으로 두 네트워크 간의 정렬이 더 부드럽고 견고해졌다.

실험에서는 ImageNet‑256 데이터셋을 기준으로 U‑REPA가 200 epoch(1M step) 내에 FID < 1.5를 달성했으며, 기존 REPA(σ‑vae‑ft‑ema) 대비 절반 정도의 학습 에포크만에 더 낮은 FID(1.41)를 기록했다. 또한, FLOPs 증가가 미미하고, CFG(Guidance)와 결합했을 때도 안정적인 수렴을 보였다. 코드와 모델 가중치는 공개되어 재현 가능성을 높였다.

전반적으로 U‑REPA는 (1) 중간 레이어 선택, (2) MLP‑후 업스케일링, (3) 매니폴드 기반 관계 정규화라는 세 가지 핵심 설계를 통해 U‑넷과 ViT 사이의 표현 정렬을 성공적으로 구현했으며, 기존 DiT‑ 기반 REPA보다 빠른 수렴과 높은 이미지 품질을 동시에 달성한 점이 가장 큰 기여라 할 수 있다.