기술 요약: 스패닝 트리 자기회귀 시각적 생성 (STAR)

문제 정의

디코더 전용 트랜스포머(Decoder-only Transformer)에 기반한 자기회귀(Autoregressive, AR) 모델은 언어 생성에서 눈부신 성공을 거두었으며, 차세대 패치 예측(next-patch prediction)을 통해 시각적 생성 분야에도 점점 더 많이 적용되고 있습니다. 그러나 기존의 AR 시각 모델은 일반적으로 고정된 래스터 스캔 순서(raster-scan order)(단방향 의존성)에 의존합니다. 이러한 순서는 높은 샘플링 성능을 제공하지만, 구조적 한계를 수반합니다. 즉, 모델이 미리 정의된 시퀀스의 접두사(prefix)로 나타나지 않는 토큰에는 조건화(condition)될 수 없다는 점입니다. 이러한 경직성은 임의의 이미지 영역을 마스킹하고 남은 컨텍스트를 바탕으로 완성해야 하는 네이티브 이미지 편집(예: 인페인팅)을 저해합니다.

이를 해결하기 위해 최근의 접근 방식들은 모델이 다양한 순서에 노출되도록 훈련 과정에서 토큰 시퀀스의 **무작위 치환(random permutation)**을 사용합니다. 하지만 본 논문은 임의의 시퀀스 순서로 훈련하는 것이 고정된 래스터 스캔 모델에 비해 샘플링 성능을 저하시킨다는 점을 지적합니다. 이는 하나의 트레이드오프를 생성합니다: 고정된 순서는 성능은 높지만 유연성이 부족하고, 무작위 치환은 유연성은 제공하지만 생성 품질과 수렴 속도를 희생합니다.

방법론: 스패닝 트리 자기회귀 (STAR)

저자들은 이미지 패치 위치로 정의된 격자(lattice) 상의 **균등 스패닝 트리(Uniform Spanning Trees, UST)**의 순회 순서를 활용하여 이 간극을 메우는 구조화된 시퀀스 순서 무작로화 전략인 STAR를 제안합니다.

1. 균등 스패닝 트리를 통한 구조적 무작위화

모든 $N!$ 개의 치환 공간에서 샘생하는 대신, STAR는 정규 택시 격자(regular taxicab lattice) $G=(V, E)$ 상의 스패ニング 트리 공간에서 샘플링합니다. 여기서 정점(vertices)은 패치 위치를 나타냅니다.

• 루트 선택: 이미지의 무작위 모서리가 루트 노드 $r$로 선택됩니다. 이러한 선택은 래스터 스캔 순서(모서리에서 시작하는 순서)가 다른 순서보다 우수한 성능을 보이는 경우가 많다는 관찰에 근거하며, 이는 아마도 중심 편향(center bias)(주요 객체는 모서리에 나타날 가능성이 낮음)과 국소성(locality)(인접한 패치를 예측하기가 더 쉬움) 때문일 것입니다.
• 순회 순서: 시퀀스 순서 $\tau$는 샘플링된 스패닝 트리의 너비 우선 탐색(Breadth-First Search, BFS) 순서에 의해 결정됩니다.
• 이론적 근거: 저자들은 사전 정의된 시퀀스의 장점이 사전 지식(중심 편향 및 국소성)을 보존하는 데서 온다고 가설을 세웠습니다. 저자들은 조건부 모델 엔트로피를 측정하여 이를 검증하였으며, 트리 거리상(tree distance) 이전에 예측된 토큰에 가까운 토큰들이 래스터 스캔 동작과 유사하게 낮은 엔트로피를 보임을 확인했습니다. 스패닝 트리의 개수는 지수적으로 증가($\sim \exp(N \cdot z_G)$)하지만, 전체 치환 공간보다는 훨씬 작고 구조화된 부분 집합을 유지합니다.

2. 추론 및 포스트픽스 완성 (이미지 인페인팅)

인페인팅 작업에서 AR 모델의 핵심 과제는 주어진 부분적 관측값(마스킹되지 않은 영역)이 시퀀스 순서에서 **접두사(prefix)**로 나타나도록 보장하여, 마스킹된 영역을 **포스트픽스(postfix)**로서 생성할 수 있게 하는 것입니다.

• 기각 샘플링(Rejection Sampling): 접두사/포스트픽스 분리가 명확한 무작위 치환과 달리, STAR는 비마스킹 영역이 먼저 방문되도록 자연스럽게 배치되는 BFS 순회를 갖는 스패닝 트리를 샘플링해야 합니다.
• 알고리즘:
  1. 비마스킹 영역의 경계(boundaries)를 식별합니다.
  2. 비마스킹 영역의 가장 가까운 경계로부터 맨해튼 거리(Manhattan distance)가 가장 먼 루트 노드 $r$을 선택하여, 비마스킹 영역이 먼저 방문될 확률을 극대화합니다.
  3. 비마스킹 영역에 대한 스패닝 트리를 샘플링합니다. 만약 트리의 최대 깊이가 모든 경계 정점이 마스킹된 정점보다 먼저 방문된다는 조건을 만족하지 못하면, 해당 샘플은 기각됩니다(rejection sampling).
  4. 비마스킹 영역에 대해 유효한 트리를 찾으면, 마스킹된 영역에 대한 스패닝 트리를 샘플링하고 이를 경계 정점들에 연결하여 전체 트리를 완성합니다.
• 효율성: 저자들은 "가장 먼 루트(farthest root)" 선택을 동반한 BFS를 사용하는 것이 DFS(깊이 우선 탐색)에 비해 기각 시행 횟수를 크게 줄여 추론을 효율적으로 만든다는 것을 입증했습니다.

3. 모델 아키텍처

STAR는 표준 디코더 전용 트랜스포머 아키텍처에 최소한의 변경만을 요구합니다. 이는 다음을 포함합니다:

• 인과적 어텐션 마스크(Causal attention masks).
• 이전의 치환 기반 방법들과 유사하게 다음 토큰 위치에 조건화하기 위한 추가 기능(학습 가능한 위치 임베딩 사용).
• 복잡한 아키텍처 변경(예: 양방향 어텐션 또는 특화된 토크나이저)은 필요하지 않습니다.

주요 기여

1. 구조화된 무작위화 전략: AR 시각적 생성을 위한 시퀀스 순서로서 균등 스패닝 트리의 BFS 순회 사용을 제 제안합니다. 이 전략은 시각적 데이터의 고유한 속성(중심 편향, 국소성)을 구조적으로 보존하면서도 유연성을 위한 충분한 무작위성을 도입합니다.
2. 성능-유연성 트레이드오프: BFS 순회 방식의 UST가 임의의 치환과 관련된 성능 저하 없이 (다양한 시퀀스 순서라는) 무작위 치환의 이점을 달성할 수 있음을 입증했습니다. 이는 훈련 복잡도를 증가시키지 않고도 달성되었습니다.
3. 네이티브 이미지 편집 지원: 최소한의 아키텍처 변경만으로 영역 수준의 수정(인페인팅)을 가능하게 하는 메커니즘(BFS를 이용한 기각 샘플링)을 도입하여 기존 AR 모델의 주요 한계를 해결했습니다.

실험 결과

저자들은 ImageNet-1k(256x256 해상도)에서 클래스 조건부 생성 및 이미지 인페인팅 작업을 통해 STAR를 평가했습니다.

• 이미지 생성: STAR 모델(STAR-B, L, XL, XXL)은 FID 및 IS 측면에서 최첨단 모델(예: RAR, DiT, LlamaGen)과 대등한 성능을 달고 있습니다. 특히, STAR-XXL은 FID 1.55를 달성하여 RAR-XXL(1.48) 및 RandAR-XXL(2.15)보다 우수한 성능을 보이면서도 경쟁력 있는 점수를 유지했습니다.
• 이미지 인페인팅: STAR는 인페인팅 작업에서 고정 순서 모델(RAR)과 무작위 치환 모델(RandAR) 모두를 크게 능가합니다.
  • 다양한 마스킹 비율(0.1 ~ 0.9)에서 STAR는 높은 FID 및 IS 점수를 유지합니다.
  • 반면, RAR(고정 래스터 스캔)은 임의의 마스크를 접두사로 자연스럽게 처리할 수 없기 때문에 인페인팅 성능이 크게 하락합니다.
  • RandAR는 인페인팅 능력은 유지하지만, STAR에 비해 생성 품질이 떨어집니다.
• 절제 연구(Ablation Studies):
  • 국소성(Locality): 무작위 치환은 STAR보다 나은 성능을 내지 못했으며, 이는 국소성을 위반하는 시퀀스 순서가 수렴에 해롭다는 것을 시사합니다.
  • 중심 편향(Center Bias): 코너를 루트로 무작위 선택하는 것은 중심 편향이 강한 이미지에서 잘 작동하며, 중심 편향이 약한 이미지에 대해서도 무작위 치환보다 더 잘 일반화됩니다.
  • 엔트로피 분석: STAR로 훈련된 모델은 무작위 치환에 비해 토큰 위치 전반에 걸쳐 더 균일한 규모의 조건부 모델 엔트로피를 보이며, 이는 더 나은 수렴과 상관관계가 있습니다.
  • 오버헤드: 스패닝 트리를 샘플링하는 데 드는 계산 오버헤드는 미미합니다(< 0.0004%의 총 추론 시간). 인페인팅을 위한 기각 샘플링은 매우 효율적이며, 일반적으로 BFS를 사용할 경우 5회 미만의 시행을 필요로 합니다.

의의 및 주장

본 논문은 STAR가 AR 시각적 생성에서의 샘플링 성능과 시퀀스 순서의 유연성 사이의 오랜 트레이드오프에 대한 단순하면서도 효과적인 솔루션을 제공한다고 주장합니다. 균등 스패닝 트리의 구조적 특성을 활용함으로써 STAR는 다음과 같은 이점을 가집니다:

• 디코더 전용 트랜스포머 아키텍처를 보존하여 초기 융합 멀티모달 모델로의 통합을 용이하게 합니다.
• 순서 어닐링(훈련 중 무작위에서 고정 순서로 전환)이나 특화된 토크나이저와 같은 복잡한 훈련 전략이 필요하지 않습니다.
• 생성 품질을 희생하지 않고 **네이티브 이미지 편집(인페인팅)**을 지원합니다.

저자들은 STAR를 미래의 AR 시각 모델링을 위한 "기본 빌딩 블록"으로 자리매김하며, 그 단순함 덕분에 멀티모달 생성 모델에 쉽게 채택되고 확장될 수 있음을 강조합니다. 이 연구는 AR 모델의 성공이 단순히 모델 아키텍처뿐만 아니라, 시각적 데이터의 본질적인 특성에 맞춘 시퀀스 순서의 구조적 무작위화에 달려 있음을 시사합니다.