기술 요약: FlashAR

문제 정의

대규모 자기회귀 (AR) 모델은 시각 데이터를 이산 토큰 시퀀스로 처리함으로써 이미지 생성 분야에서 최첨단 성능을 달성했습니다. 그러나 표준 AR 모델은 래스터 스캔 디코딩에 의존하여 토큰을 엄격하게 왼쪽에서 오른쪽, 위에서 아래로 순차적으로 생성합니다. 이 순차적 종속성으로 인해 추론 지연 시간이 이미지 해상도에 따라 2 차 함수적으로 ($O(HW)$) 증가하여 고해상도 생성이 지나치게 느려지고 현대 GPU 의 병렬화 능력을 효과적으로 활용하지 못하게 합니다.

기존 가속화 전략은 상당한 트레이드오프에 직면해 있습니다:

1. 패러다임 재설계: 토큰 예측 순서를 변경하는 방법 (예: VAR, NAR) 이나 다중 스케일 생성을 채택하는 접근법은 종종 처음부터 훈련을 필요로 하므로, 방대한 양의 기존 사전 훈련된 래스터-AR 모델 저장소와 호환되지 않습니다.
2. 훈련 후 적응: 이산 확산 적응과 같은 접근법은 원래 예측 목적 함수를 수정하여 훈련 - 추론 간격을 생성하고, 인과적 모델링에 내재된 세밀한 구조적 일관성을 저하시킬 수 있습니다.

이 논문이 다루는 핵심 과제는 다음과 같습니다: 사전 훈련된 래스터 스캔 AR 모델을 학습된 생성적 사전 지식 (generative prior) 을 보존하고 훈련 후 오버헤드를 최소화하면서 고품질 병렬 생성기로 변환할 수 있는가?

방법론: FlashAR

FlashAR 는 사전 훈련된 래스터 스캔 AR 모델을 양방향 다음 토큰 예측 구조를 사용하여 병렬 생성기로 변환하는 경량 훈련 후 적응 프레임워크입니다. 이 프레임워크는 세 가지 핵심 기술 요소를 도입합니다:

1. 이중 헤드 디코딩을 위한 중간 분기

최종 레이어에 새로운 예측 헤드를 연결하는 것 (이는 원래 수평 래스터 스캔 목적 함수에 편향됨) 대신, FlashAR 는 중간 레이어($m$) 에서 네트워크를 분기합니다.

• 아키텍처: 사전 훈련된 백본은 공유 트렁크 ($F_{0:m-1}$) 와 두 개의 병렬 경로로 분할됩니다.
  • 수평 헤드: 행 단위 예측을 위해 원래 상위 레이어 ($F_{m:L-1}$) 를 재사용합니다.
  • 수직 헤드: 열 단위 예측을 위해 상위 레이어의 클론으로 초기화된 학습 가능한 블록 ($\tilde{F}_{m:L-1}$) 을 활용합니다.
• 근거: 중간 특징은 더 풍부한 의미적 추상화를 유지하며 최종 레이어 표현보다 특정 왼쪽에서 오른쪽 방향에 덜 특화되어 있습니다. 여기서 분기함으로써 수직 경로가 필요한 컨텍스트에 접근하면서도 임계 경로 깊이를 확장하지 않아 런타임 효율성을 유지할 수 있습니다.

2. 학습 가능한 퓨전 게이트

수평 및 수직 예측은 상호 보완적인 종속성을 포착하지만, 그 상대적 중요도는 공간적 위치에 따라 다릅니다 (예: 수평 가장자리 대 수직 가장자리). 정적 평균은 불충분하며 컨텍스트 충돌을 유발할 수 있습니다.

• 메커니즘: FlashAR 는 두 개의 이전 상태 ($h^H_{p,q-1}$ 및 $h^V_{p-1,q}$) 의 컨텍스트에 기반하여 예측을 동적으로 결합하는 학습 가능한 퓨전 게이트를 사용합니다.
• 공식: 퓨트된 로그이트 $z_{p,q}$는 $g_{p,q} z^H_{p,q-1} + (1 - g_{p,q}) z^V_{p-1,q}$로 계산되며, 여기서 게이트 $g_{p,q}$는 MLP 와 시그모이드 함수를 통해 유도됩니다. 이를 통해 모델은 방향성 단서를 적응적으로 가중할 수 있습니다.

3. 2 단계 안정화 적응 파이프라인

급속한 수렴을 보장하고 사전 훈련된 매니폴드의 파괴적 망각을 방지하기 위해 FlashAR 는 2 단계 훈련 전략을 사용합니다:

• 1 단계 (분기 초기화): 백본과 수평 헤드는 고정됩니다. 수직 헤드와 퓨전 게이트만 훈련됩니다. 이를 통해 새로운 수직 경로가 확립된 수평 사전 지식을 방해하지 않고 예측 능력을 신속하게 습득할 수 있습니다.
• 2 단계 (공동 미세 조정): 전체 네트워크 (백본, 두 개의 헤드, 게이트) 가 새로운 대각선 병렬 디코딩 패러다임을 위해 직교하는 공간 표현을 조화시키기 위해 공동으로 미세 조정됩니다.

4. 하드웨어 인식 추론

이론적 병렬성을 실제 시간 단축으로 전환하기 위해 저자는 다음을 구현합니다:

• FlexAttention: 밀집된 제로 패딩으로 인한 메모리 조각화를 방지하기 위해 매우 희소한 2 차 근접 마스크를 동적으로 컴파일합니다.
• 배치된 KV-Cache: 대각선 단계에서 모든 토큰을 동시에 업데이트하여 커널 시작 오버헤드를 분산합니다.

주요 기여

1. FlashAR 프레임워크: 처음부터 사전 훈련할 필요 없이 사전 훈련된 래스터-AR 모델을 병렬 생성기로 변환하는 훈련 후 적응 방법.
2. 이중 헤드 중간 분기: 중간 레이어에서 분기함으로써 수평 사전 지식 보존과 직교 수직 예측 도입 사이의 긴장을 해결하는 새로운 아키텍처.
3. 동적 퓨전: 수평 및 수직 신호를 적응적으로 결합하여 공간적 불일치를 해결하는 학습 가능한 게이트.
4. 효율적인 훈련 파이프라인: 최소한의 데이터 요구 사항으로 빠른 수렴을 보장하는 2 단계 적응 프로세스.

실험 결과

저자는 LlamaGen(ImageNet 256×256) 과 Emu3.5(Text-to-Image 512×512) 에서 FlashAR 를 평가했습니다.

• 가속화: FlashAR 는 표준 AR 베이스라인 대비 512×512 이미지 생성에서 최대 22.9 배의 실제 시간 가속을 달성합니다. LlamaGen 의 경우 447.2 이미지/초의 처리량을 기록했습니다 (표준 AR 의 117.9 대비).
• 데이터 효율성: 적응에는 원래 훈련 데이터의 **0.05%**만 필요합니다 (예: Emu3.5 의 경우 80K 토큰 대 150B 사전 훈련 토큰).
• 품질 보존:
  • LlamaGen: FlashAR-L 은 FID 3.16 및 IS 289.0 을 달성하여 BlockDiffusion(FID 4.55) 을 능가했으며, 처음부터 훈련된 NAR-L 보다도 Inception Score 에서 더 높은 성능을 보였습니다.
  • Emu3.5: FlashAR 는 디코딩 단계를 1024 에서 63 으로 줄였습니다. GenEval 점수는 미미하게 감소했습니다 (80.48 → 80.29) 이며, "Colors" 및 "Position" 지표에서 베이스라인을 능가했습니다.
• 수렴: FlashAR 는 BlockDiffusion 보다 훨씬 빠르게 수렴하여 BlockDiffusion 의 75 에포크 대비 25 에포크 만에 우수한 FID 를 달성했습니다.

의의 및 주장

이 논문은 FlashAR 가 병렬 디코딩의 효율성과 사전 훈련된 자기회귀 모델의 충실도 사이의 간극을 성공적으로 메웠다고 주장합니다. 원래 훈련 목적 함수에 대한 수정을 최소화함으로써 FlashAR 는 기존 대규모 모델의 강력한 생성 능력을 보존합니다. 저자는 이 접근법이 처음부터 새로운 병렬 아키텍처를 훈련하는 데 따른 prohibitive 비용이나 목적 함수 변경 적응 방법에서 흔히 발생하는 품질 저하 없이 실시간 고해상도 이미지 생성을 위한 실용적인 해결책을 제공한다고 강조합니다.

이 연구는 대각선 병렬 디코딩이 중간 분기와 동적 퓨전을 통해 사전 훈련된 모델의 고유한 편향을 존중하도록 모델 아키텍처를 신중하게 적응시킨다면 자기회귀 생성을 위한 실현 가능하고 효율적인 패러다임임을 입증합니다.