기술 요약: MOSAIK

문제 정의

픽셀 공간 확산 모델(Pixel-space Diffusion Models, PDMs)은 변이형 오토인코더(VAE)에 의해 부과되는 재구성 한계(reconstruction ceiling)를 피함으로써 미세한 고주파 공간 세부 사항을 보존할 수 있다는 점에서 잠재 공간 확산 모델(LDMs)보다 상당한 이점을 제공한다. 그러나 픽셀 공간에서 직접 생성하는 것은 토큰 수를 대폭 증가시킨다. 트랜스포머의 셀프 어텐션 메커니즘은 토큰 수에 대해 이차 복잡도(quadratic complexity)를 가지기 때문에, PDM은 과도한 계산 비용을 발생시킨다.

기존의 효율성 방법들은 디노이징 타임스텝이나 네트워크 깊이에 따라 패치 크기(단일 토큰이 나타내는 공간 영역)를 변화시켜 이를 완화하려고 시도한다. 이러한 접근 방식은 전체 토큰 수를 줄여주기는 하지만, 임의의 단계에서 이미지 전체에 걸쳐 공간적으로 균일한(spatially uniform) 패치 레이아웃을 적용한다. 이러한 균일성은 근본적인 비효율성을 초래한다. 즉, 단순하고 정보량이 적은 영역에 계산량을 과하게 소비하거나, 반대로 미세한 디테일이 중요한 민감한 영역에서의 표현 능력을 제한하게 된다. 본 논문은 시각적 정보 밀도가 본질적으로 불균질(heterogeneous)하며, 현재의 방법들이 결여하고 있는 동적인 영역별 할당 전략이 필요하다고 주장한다.

방법론: MOSAIK

MOSAIK(Multi-patch cOntent-aware Spatial Allocation of Image toKens)은 거칠게 만드는 과정(coarsening)으로 인해 발생하는 예측된 "손상(damage)"(충실도 손실)에 따라 서로 다른 이미지 영역과 디노이징 단계에서 패치 크기를 가변적으로 조절하도록 설계된 프레임워크이다. 이 접근 방식은 PixelDiT 백본을 기반으로 하며, 3단계 훈련 파이프라인을 활용한다:

1단계: 멀티 패치 유창성 (Multi-Patch Fluency)

목표는 모델이 원래의 생성적 사전 지식(generative prior)에 대한 파괴적 망각 없이 다양한 균일 패치 크기($P \in \{16, 32, 64\}$)에서 작동할 수 있도록 하는 것이다.

• 아키텍처: 백본 가중치는 고정된다. DiT 블록에 저계수 적응(Low-Rank Adaptation, LoRA) 모듈이 삽입된다. $p32$와 $p64$를 위한 별도의 토크나이저가 도입되며, 이는 사전 훈련된 $p16$ 투영(projection)을 쌍선형 다운샘플링(bilinear downsampling)으로 합성하여 초기화된다.
• 훈련: 모델은 모든 균일 패치 크기에 대해 고정된 $p16$ 텔러(teacher) 모델의 속도 예측(velocity predictions)을 일치시키도록 최적화되는 자기 증류(self-distillation) 과정을 거친다.

2단계: 혼합 패치 적응 (Mixed-Patch Adaptation)

이 단계는 모델이 단일 순방향 패스(forward pass) 내에서 임의의 불균질한 패치 레이아웃(서로 다른 영역이 서로 다른 패치 크기를 사용하는 경우)을 처리할 수 있게 한다.

• 위치 임베딩: 표준 회전 위치 임베딩(Rotary Positional Embeddings, RoPE)은 균일한 그리드를 가정한다. MOSAIK는 혼합 패치의 공간적 불일치를 해결하기 위해 각 토큰을 정규화된 $p16$ 좌표계 내의 패치 기하학적 중심에 고정한다.
• 훈련: 생성기는 이미지 콘텐츠와 패치 크기의 공간적 배치를 의도적으로 분리하여, 무작위로 샘플링된 불균질한 레이아웃 상에서 훈련된다. 이를 통해 모델은 레이아웃이 시각적 특징과 어떻게 정렬되든 상관없이 어떤 레이아웃이라도 실행할 수 있도록 학습된다.

3단계: 손상 가이드 레이아웃 예측기 (Damage-Guided Layout Predictor)

이 단계는 특정 영역을 거칠게 만들 때 발생하는 추정된 저하(degradation)에 따라 토큰을 동적으로 할당하는 경량 예측기를 도입한다.

• 손상 정의: "손상"은 균일한 $p64$ 레이아웃의 예측값과 고충실도 $p16$ 텔러 사이의 제곱 차이를 중첩되지 않는 $64 \times 64$ 영역에 대해 평균낸 값으로 정의된다.
• 예측기 훈련: FiLM 컨디셔닝을 갖춘 경량 컨볼루션 네트워크가 중간 백본 특징(구체적으로 DiT 블록-6 활성화 함수)과 텍스트 임베딩으로부터 이 손상 맵을 직접 회귀하도록 훈련된다. 이는 추론 시 별도의 순방향 패스를 수행할 필요가 없게 한다.
• 추론 전략: 각 디노이징 단계에서 예측기는 손상 맵을 추정한다. 손상이 높게 예측된 영역에는 더 미세한 패치($p16$)가 할당되는 반면, 손상이 낮은 영역에는 더 거친 패치($p32$ 또는 $p64$)가 할당된다. 할당은 사용자가 지정한 토큰 예산에 의해 제약되며, 추가된 토큰당 가장 높은 손상 감소를 제공하는 리파인먼트(refinement)를 우선시한다.

주요 기여

1. 공간 적응형 패치화 (Spatially Adaptive Patchification): MOSAIK는 픽셀 공간 확산 모델에 불균질한 패치 레이아웃을 적용하는 최초의 프레임워크로, 기존의 효율성 방법들이 가진 공간적 균일성 제약을 넘어선다.
2. 손상 가이드 할당 (Damage-Guided Allocation): 저자들은 패치 할당을 손상 예측 문제로 공식화하였으며, 경량 예측기를 사용하여 거칠게 만드는 과정이 가장 큰 충실도 손실을 일으킬 수 있는 영역을 식별한다.
3. 3단계 훈련 파이프라인: 멀티 패치 유창성을 먼저 확립한 후, 혼합 레이아웃에 적응시키고, 마지막으로 생성기를 재학습하지 않고도 토큰 라우팅을 가이드할 별도의 예측기를 훈련하는 새로운 훈련 전략을 제시한다.
4. 통합 가능성 (Integration Capability): 이 프레임워크는 시간적 효율성 방법(예: 특징 캐싱)과 직교하도록 설계되어 결합된 최적화가 가능하다.

결과

1024x1024 이미지 생성을 위해 GenEval 및 DPG-Bench 벤치마크에서 균일 패치 감소, 시간적 패치 스케줄링(DDiT), 특징 캐싱(TeaCache, TaylorSeer)을 포함한 베이스라인들과 MOSAIK를 비교하는 실험이 수행되었다.

• 효율성 대비 품질: MOSKI는 GenEval(0.74 점수)에서 전체 연산량의 PixelDiT와 대등한 성능을 보이면서도 70%의 FLOPs 감소와 83%의 토큰 수 감소를 달est했다. 또한 DPG-Bench에서도 단 1.0 포인트의 차이(82.5 대 83.5)만을 보였다.
• 시간적 스케줄링과의 비교: 공격적인 계산 환경(예: 60-70% FLOPs 감소)에서 MOSAIK는 DDiT를 지속적으로 능가한다. DDiT는 전역적으로 단일 패치 크기를 적용하는 반면, MOSAIK는 계산 자원을 민감한 영역(예: 피사체의 경계)으로 동적으로 재지정하여 구조적 무결성과 질감, 에지 같은 미세한 디테일을 보존한다.
• 캐싱과의 비교: 매칭된 예산 하에서 MOSAIK는 캐싱 기반 베이스라인(TeaCache, TaylorSeer)보다 우수한 성능을 보이며, 특히 매우 제한적인 환경에서 더욱 그러하다.
• 상보성 (Complementarity): 캐싱 방법과 결합했을 때, MOSAIK는 캐싱 단독으로는 심각한 품질 저하가 발생하는 높은 압축 수준(최대 80% FL한 FLOPs 감소)에서도 안정적인 성능을 유지한다.
• 어블레이션 (Ablation): 손상 가이드 예측기를 무작위 또는 콘텐츠 무관 할당기로 교체하면 성능이 크게 저하되며, 이는 예측된 손상 순위가 공간적 민감도에 대한 진정한 정보를 담고 있음을 입증한다.

의의 및 주장

본 논문은 MOSAIK가 **공간적 불균질성(spatial heterogeneity)**이 시간적 또는 깊이 방향의 적응에만 집중했던 방법론들이 간과해 온 효율적인 생성을 위한 핵심 요소임을 보여준다고 주장한다.

• 최적의 자원 할당: 저자들은 균일한 미세 패치화는 단순한 영역에 계산을 낭비하고, 균일한 거친 패치화는 복잡한 영역의 디테일을 파괴한다고 주장한다. MOSKI는 가장 필요한 곳에 자원을 할당함으로써 이 간극을 메운다.
• 픽셀 공간의 이점: 이 프레임워크는 (VAE 압축이 이미 중복성을 제거한 잠재 모델과 달리) 픽셀 공간 모델에서 가용한 고유한 공간적 중복성을 활용하며, 라우터와 생성기의 엔드투엔드 공동 최적화 없이도 동적으로 거칠게 만들기(coarsening)를 활용할 수 있다.
• 확장성: 이 접근 방식은 중간 정도의 감소부터 공격적인 제약에 이르기까지 광범위한 계산 예산 범위에서 안정적인 성능을 허용하며, 서로 다른 예산 수준을 위해 재학습할 필요가 없다.

논문은 결론적으로 MOSAIK가 픽셀 공간 확산의 새로운 표준을 제시하며, 시간적 스케줄링 및 캐싱에 대해 매우 경쟁력 있는 대안을 제공한다고 밝힌다. 또한 향후 연구가 이러한 원칙을 비디오 확산으로 확장하여 공간과 시간 모두의 중복성을 활용할 수 있음을 시사한다.