초고해상도 이미지 생성을 위한 계층적 로컬 어텐션 UltraGen

초록

UltraGen은 고해상도 텍스트‑투‑이미지 확산 모델의 메모리·연산 병목을 해결하기 위해, 고해상도 라티스를 고정 크기 로컬 윈도우로 나누고 저해상도 라티스로 전역 의미를 전달한다. 윈도우‑우선 토큰 순열과 LoRA 기반 경량 적응을 통해 기존 모델을 재학습 없이 8K 이상 해상도로 확장하면서 10배 이상 속도 향상과 메모리 절감을 달성한다.

상세 분석

UltraGen은 기존 확산 모델이 1K × 1K 수준에서 급격히 증가하는 O(N²) 어텐션 비용 때문에 초고해상도 이미지 생성에 한계가 있다는 점을 정확히 짚는다. 이를 해소하기 위해 제안된 핵심 아이디어는 ‘계층적 로컬 어텐션’이다. 고해상도 라티스 X를 하드웨어에 맞춘 고정 크기(예: 16 × 16 토큰) 윈도우로 분할하고, 각 윈도우 내부에서만 완전 어텐션을 수행한다. 이때 윈도우‑우선 순서로 토큰을 재배열(window‑first permutation)하면, GPU 메모리 레이아웃이 2D 윈도우와 일치해 플래시 어텐션 같은 고성능 커널을 그대로 활용할 수 있다. 복잡도는 O(N·l²)로 거의 선형에 가깝게 감소하고, 윈도우 크기 l을 16으로 고정함으로써 연산 효율과 디테일 보존 사이의 트레이드오프를 최적화한다.

전역 의미 전달을 위해 저해상도 라티스 X_lr를 별도로 생성한다. 이 라티스는 ρ = 4 정도의 스케일링 비율로 고해상도 좌표에 매핑되며, RoPE(회전 위치 임베딩)를 스케일링해 동일한 상대 위치 정보를 제공한다. X_lr 토큰은 전역 어텐션 마스크를 통해 서로 완전 연결되면서, 각 고해상도 윈도우는 자신의 로컬 토큰과 인접 토큰, 그리고 대응되는 X_lr 토큰을 동시에 참조한다. 결과적으로 전역 구조와 레이아웃은 저해상도 라티스가 담당하고, 미세 텍스처는 로컬 윈도우가 담당하는 명확한 역할 분리가 이루어진다.

LoRA(Low‑Rank Adaptation) 모듈은 기존 사전학습된 MMDiT 가중치를 거의 그대로 유지하면서, X_lr 경로에만 가벼운 저차원 변환을 삽입한다. 이렇게 하면 고해상도와 저해상도 경로 사이의 정보 흐름을 효율적으로 연결하면서도 추가 파라미터는 수십만 개에 불과해 메모리 오버헤드가 최소화된다.

실험에서는 FLUX·1·2와 같은 최신 텍스트‑투‑이미지 모델을 UltraGen 프레임워크에 그대로 적용해 4K, 8K, 10K 등 다양한 해상도로 확장하였다. FID, IS, CLIP‑Score 등 정량 지표에서 기존 밀집 어텐션 기반 재학습 방법을 크게 앞섰으며, 메모리 사용량은 10 GB 이상 절감되고 추론 속도는 10배 이상 가속화되었다. 특히, 훈련 데이터는 256 × 256 ~ 1K × 1K 범위만 사용했음에도 불구하고, 초고해상도에서 구조적 일관성과 세밀한 디테일을 동시에 유지한다는 점이 주목할 만하다.

한계점으로는 윈도우 경계에서 발생할 수 있는 미세한 아티팩트와, 저해상도 가이드가 지나치게 강하면 전역 구조가 고정돼 창의적 변형이 제한될 가능성이 있다. 저해상도 라티스의 해상도와 스케일링 비율을 동적으로 조절하거나, 다중 스케일 윈도우를 혼합하는 방식이 향후 연구 방향으로 제시된다.