기술 요약: 공간적 추측 디코딩 (Spatially Speculative Decoding, SSD)

1. 문제 정의

자기회귀(Autoregressive) 모델은 이미지를 언어 모델링과 유사하게 이산적 토큰의 1D 시퀀스로 취급함으로써 시각적 생성 분야에서 상당한 성공을 거두었습니다. 그러나 이 방식은 이미지의 본질적인 2D 공간적 국소성(spatial locality)을 선형 래스터 스캔(raster-scan) 시퀀스로 평탄화하는 기하학적 타협을 강요합니다.

이러한 평탄화는 추론 과정에서 흔히 "메모리 벽(memory wall)"이라 불리는 심각한 계산 병목 현상을 야기합니다. $n \times n$ 이미지를 생성하려면 $n^2$번의 순차적인 순방향 패스(forward pass)가 필요합니다. 각 패스는 단 하나의 토큰을 예측하기 위해 전체 트랜스포머 파라미터를 다시 로드해야 하므로, 이 과정은 연산량(compute)보다는 메모리 대역폭(memory bandwidth)에 의해 크게 제약됩니다.

표준 추측 디코딩(Speculative Decoding) 및 자코비 기반 반복(Jacobi-based iteration)과 같은 기존 가속 기술들은 자연어 처리(NLP) 기법을 비전 분야로 도입하려고 시도했습니다. 그러나 이러한 방법들은 여전히 1D 순차적 가정에 묶여 있습니다. 시각적 패치(patch)는 엄격한 문법 없이도 높은 국소적 엔트로피를 보이기 때문에, 평탄화된 1D 지평선 상에서 여러 토큰을 미리 예측하는 것은 낮은 초안 수락률(draft acceptance rate)을 초래합니다. 결과적으로, 이러한 방법들은 완만한 속도 향상(1.8배~3.7배)만을 제공하거나, 부자연스러운 독립성 가정을 강제함으로써 이미지 품질을 저하시킵니다.

2. 방법론: 공간적 추측 디코딩 (SSD)

저자들은 이미지의 자연스러운 2D 기하학 구조와 예측 목표를 일치시키는 프레임워크인 **공간적 추측 디코딩 (SSD)**을 제안합니다. SSD는 1D 시퀀스의 즉각적인 다음 토큰만을 예측하는 대신, 인접한 수평 토큰과 바로 아래의 토큰을 동시에 예측합니다.

핵심 메커니즘

1. 2D 공간적 예견 (2D Spatial Anticipation):

   • SSD는 2D 예견을 두 개의 직교하는 1D 예측 스트림, 즉 수평 방향(래스터 스캔을 따름)과 수직 방향(열을 따라 내려감)으로 분해합니다.
   • 모델은 경량 헤드를 사용하여 수평 방향으로 한 행(row) 전체를 초안 작성하고, 그 후 수직 방향으로 후속 행들을 병렬로 초안 작성합니다.
   • 이 접근 방식은 평탄화된 시퀀스 상에서는 오프셋이 크더라도, 특정 토큰이 바로 왼쪽 토큰에 의존하는 정도만큼 바로 위쪽 토큰에도 강하게 의존한다는 관찰 결과를 활용합니다.
   • 복잡도 감소: 공간 블록 전체를 병렬로 초안 작성함으로써, SSD는 $n \times n$ 이미지에 대한 이론적 추론 복잡도를 $O(n^2)$에서 $O(n)$으로 줄입니다.

2. 연속 잠재 공간 예측 (Continuous Latent Space Prediction):

   • 거대한 코드북에 대한 평탄한 확률 분포로 인해 이산적 토큰 ID를 직접 예측하는 것은 어렵고, 이는 낮은 수락률(<5%)으로 이어집니다.
   • 대신 SSD는 마지막 트랜스포머 레이어의 연속적인 잠재 특징(continuous latent features)(구체적으로 최종 RMSNorm 이전의 은닉 상태)을 예측합니다.
   • 경량 예측기 $f_\phi$는 은닉 상태와 토큰 임베딩을 입력받아 특정 공간 오프셋에서의 은닉 상태를 예측합니다.
   • 예측된 은닉 상태는 이후 베이스 모델의 기존 출력 레이어를 통해 디코딩되어 토큰 후보를 얻습니다.

3. 자기 수정으로서의 검증 (Verification as Auto-Correction):

   • 첫 번째 불일치 발생 시 전체 블록을 거부하는 표준 추측 디코딩과 달리, SSD는 검증을 자기 수정 메커니즘으로 취급합니다.
   • 초안이 응집력 있는 공간 블록을 형성하기 때문에, 미세한 표현상의 불일치는 국소적으로 해결될 수 있습니다.
   • 거부된 토큰은 단순히 버려지는 것이 아니라, 시스템이 **잔차 분포(residual distribution)**로부터 샘플링하여 동일한 순방향 패스 내에서 이를 복구합니다.
   • 이 과정은 $r$ 라운드 동안 반복되며, 거부된 위치가 업데이트되고 KV 캐시가 재평가를 위해 롤백됩니다. 이는 블록당 순차적 재생성이 아닌 $r+1$번의 순방향 패스만을 요구합니다.

4. 모듈성 (Modularity):

   • SSD는 플러그 앤 플레이 모듈입니다. 사전 학습된 백본 모델의 수정이 필요하지 않습니다.
   • 자기 증류(self-distillation)를 통해 경량 드래프팅 헤드만을 학습시키므로, FLOPs 오버 overhead를 최소한으로 유지합니다.

3. 주요 기여

• 기하학적 재정렬 (Geometric Realignment): 본 논문은 시각적 데이터의 기저에 깔린 2D 기하학을 존중하는 것이 1D 순차적 디코딩이라는 "계산적으로 비용이 많이 드는 인위적 산물"을 넘어 막대한 계산 효율성을 실현함을 입증합니다.
• 잠재 공간 드래프팅 (Latent Space Drafting): 연속적인 잠재 특징 예측의 도입은 시각 도메인에서 이산적 토큰 예측에 비해 초안 정확도를 크게 향상시킵니다.
• 자기 수정 검증 (Auto-Correcting Verification): 거부된 토큰을 폐기하는 대신 병렬로 복구하는 새로운 검증 전략을 통해 공간적 드래프팅의 이점을 보존합니다.
• 플러그 앤 플레이 가속 (Plug-and-Play Acceleration): 거대한 백본 모델을 재학습하지 않고도 생성 속도를 가속화하는 프레임워크를 제공하며, 이는 이산적 시각 토큰을 생성하는 모든 통합 모델에 적용 가능합니다.

4. 실험 결과

저자들은 DPG-Bench 및 GenEval 벤치마크를 사용하여 세 가지 최첨단 자기회귀 모델인 Janus-Pro-7B, Lumina-mGPT-7B, Emu3-8B에 대해 SSD를 평가했습니다.

• 속도 향상: SSD는 표준 자기회귀 추론과 비교하여 최대 13.3배의 실제 시간(wall-clock) 속도 향상을 달축했습니다.
  • Emu3-8B: 339초 $\to$ 25.5초 (13.27배 속도 향상).
  • Lumina-mGPT-7B: 91.6초 $\to$ 7.5초 (12.19배 속도 향상).
  • Janus-Pro-7B: 7.8초 $\to$ 1.4초 (5.74배 속도 향상).
• 품질: SSD는 높은 충실도를 유지합니다. DPG-Bench 및 GenEval에서의 생성 품질 지표는 베이스라인 자기회귀 모델과 대등한 수준을 유지하며, 심각한 품질 저하를 겪은 1D-MTP 베이스라인 모델들을 크게 상회합니다.
• 비교:
  • 1D-MTP: 약 2.0~2.4배의 속도 향상을 달성했으나 심각한 품질 저하를 초래했습니다 (예: Emu3의 전체 점수가 78.69에서 53.11로 하락).
  • SJD (Jacobi-based): 품질은 보존했으나 1.5~2.9배의 가속에 제한되었습니다.
  • SSD: 품질을 유지하면서 가장 높은 속도 향상을 달성했습니다.

5. 의의 및 주장

본 논문은 현재의 자기회귀 시각 생성의 비효율성이 모델 자체의 내재적 한계가 아니라, 언어 모델링에서 물려받은 1D 순차적 가정의 결과라고 주장합니다. 이 가정을 버리고 이미지의 본질적인 2D 기하학에 디코딩 과정을 정렬함으로써, SSD는 메모리 벽을 극복합니다.

저자들은 이러한 기하학적 전환이 실시간 고해상도 자기회귀 생성 모델을 위한 길을 열어준다고 단언합니다. 결과는 시각적 국소성을 존중하는 것이 단일 백본 내에서 시각과 언어의 통합을 해치지 않으면서도 막대한 계산 절감을 가능하게 함을 시사합니다. 이 방법은 기존의 통합 모델에 적용하여 추론 속도를 경쟁력 있게 만들 수 있는 일반적인 솔루션으로 제시됩니다.