기술 요약: 픽셀 공간 자기회귀 이미지 생성을 위한 병렬 롤아웃 근사 (Parallel Rollout Approximation)

1. 문제 정의

픽셀 공간 자기회귀(Autoregressive, AR) 이미지 생성은 별도의 이산 토큰화나 사전 학습된 토크나이저에 의존하지 않고, 원시 픽셀 패치를 시퀀스로 모델링하는 것을 목표로 합니다. 개념적으로는 매우 직관적이지만, 이 접근 방식은 역사적으로 잠재 공간(latent-space) AR 또는 확산 모델(diffusion models)에 비해 성능을 제한해 온 두 가지 결합된 과제에 직면해 있습니다:

1. 출력 측 과제 (고차원 예측): 단 한 번의 AR 단계에서 고차원의 원시 픽셀 패치(예: 16×16 패치의 경우 768차원 벡터)를 직접 예측하는 것은 매우 어려우며, 이는 큰 단일 단계 생성 오차로 이어집니다.
2. 입력 측 과제 (학습-추론 불일치): 표준 교사 강요(teacher-forced) 학습은 깨끗한 정답(ground-truth) 접두사(prefix)에 모델을 조건화합니다. 그러나 추론 과정에서 모델은 자신이 이전에 생성한 (그리고 잠재적으로 오류가 있는) 픽셀 패치에 조건화되어야 합니다. 이러한 불일치는 오류가 여러 AR 단계를 거쳐 축적되고 전파되는 현상인 노출 편향(exposure bias)을 유발합니다.

$x$-예측(확산 모델에서 효과적임)이나 입력 노이즈 주입과 같은 기존의 완화 전략은 이러한 문제를 부분적으로만 해결할 수 있습니다. 완전한 롤아웃 학습(exact rollout training)은 연속적인 토큰에 대해 순차적인 샘플링을 요구하기 때문에 연속 토큰 기반의 AR에서는 계산 비용이 너무 많이 들어 실질적으로 불가능합니다.

2. 방법론: 병렬 롤아웃 근사 (Parallel Rollout Approximation, PRA)

저자들은 "픽셀 입력, 픽셀 출력(pixel-in, pixel-out)" 인터페이스를 유지하면서 출력 및 입력 문제를 동시에 해결하는 확장 가능한 프레임워크인 **병렬 롤아웃 근사(PRA)**를 제안합니다.

2.1 출력 측: 엔드 투 엔드 중간 타겟 (End-to-End Intermediate Targets)

단일 단계 생성의 난이도를 낮추기 위해, PRA는 고차원 픽셀 패치를 직접 예측하지 않습니다. 대신, 저차원의 중간 상태($z_i \in \mathbb{R}^{d_z}$, 여기서 $d_z \ll d$)를 생성합니다.

• 타겟 구축: 학습 중에 정답 픽셀 토큰 $x_i$와 인과적 접두사 표현 $h_{i-1}$은 경량 투영 네트워크 $g_\psi$를 통해 타겟 중간 상태 $z_i$로 매핑됩니다. 이 타겟은 접두사의 정보를 포함하는 문맥 인식형(context-aware)입니다.
• 디코딩: 인과적 픽셀 디코더 $D_\phi$는 생성된 중간 상태 $\hat{z}_i$(및 이전 상태들 $\hat{z}_{<i}$)를 다시 픽셀 공간으로 매핑하여 출력 토큰 $\hat{x}_i$를 생성합니다.
• 학습: AR 모델은 토큰 수준의 확산 헤드(rectified-flow 목적 함수)를 사용하여 $z_i$를 예측하도록 학습되며, 디코더는 $z_i$로부터 $x_i$를 재구성하도록 학습됩니다. 이를 통해 AR 백본은 저차원 공간에서 작동하면서도 외부 인터페이스는 픽셀 공간을 유지할 수 있습니다.

2.2 입력 측: 병렬 추론 유사 입력 (Parallel Inference-Like Inputs)

순차적인 롤아웃 없이 학습-추론 간의 격차를 메우기 위해, PRA는 추론 유사 픽셀 입력을 병렬로 구축합니다.

• 메커니즘: 깨끗한 정답 픽셀로 학습하는 대신, PRA는 타겟 중간 상태 $z_i$를 (rectified-flow 보간을 사용하여) 섭동(perturb)시켜 노이즈가 섞인 상태 $\tilde{z}_t$를 만듭니다. 이 상태들은 추론 시 사용되는 것과 동일한 픽셀 디코더를 통과하여 재구성된 픽셀 토큰 $\bar{x}_i$를 생성합니다.
• 병렬성: 각 $\bar{x}_i$는 (순차적인 자기회귀 궤적이 아닌) 각각의 대응하는 타겟 상태로부터 독립적으로 구축되므로, 이 과정은 완전히 병렬화가 가능합니다.
• 학습 루프: 이후 AR 모델은 재구성된 추론 유사 픽셀 $\bar{x}_{<i}$의 접두사에 조건화되어 타겟 $z_i$를 예측하도록 학습됩니다. 이는 교사 강요 학습의 효율성을 유지하면서도 추론 시의 롤아웃 오차 구조를 근사합니다.

3. 주요 기여

1. 통합 프레임워크: PRA는 별도로 사전 학습된 토크나이저 없이도 픽셀 공간 AR 생성의 고차원 출력 오류와 학습-추론 불일치를 동시에 해결하는 최초의 확장 가능한 프레임워크입니다.
2. 중간 상태 학습: 본 논문은 VAE의 고정된 잠재 공간에 의존하는 대신, 인과적 AR 모델에 특화되어 최적화된 엔드 투 엔드 학습형 중간 타겟 공간을 도입합니다.
3. 병렬 롤아웃 근사: PRA는 병렬적이고 비순차적인 디코딩 픽셀 입력 구축을 통해 롤아웃 학습(생성된 문맥에 모델을 노출시키는 것)의 이점을 성공적으로 근사합니다.
4. 픽셀 입력, 픽셀 출력 인터페이스: 잠재 공간 AR 모델과 달리, PRA는 원시 픽셀 패치와 직접적인 인터페이스를 유지하여 통합된 이미지 생성 및 이해 응용 분야를 용이하게 합니다.

4. 실험 결과

저자들은 256×256 해상도의 Class-conditional ImageNet-1K 생성 작업에서 PRA를 평가했습니다.

• 최첨단 성능 (State-of-the-Art Performance):
  • PRA-S (135M 파라미터): FID 2.58을 달성하여, 이전의 수십억 파라미터 규모 픽셀 공간 AR 모델(예: FID 3.60을 기록한 FARMER-1.9B)을 능가했습니다.
  • PRA-L (511M 파라미터): FID 1.94를 달성하여 픽셀 공간 AR 모델의 새로운 SOTA를 세웠습니다.
• 절제 연구 (Ablation Studies):
  • 중간 타겟: 문맥 인식형 타겟(접두사 표현 포함)이 로컬 전용 타겟이나 고정된 LDM 인코더보다 우수한 성능을 보였습니다. 중간 차원 $d_z=16$이 최적임을 발견했습니다.
  • 입력 구축: 디코딩된 픽셀 입력을 사용하는 것(PRA)이 깨끗한 정답 픽셀, 정답 중간 상태, 또는 단순 노이즈 주입을 사용하는 것보다 유의미하게 뛰어난 성능을 보였으며, 이는 추론 시의 픽셀 피드백 루프를 일치시키는 것의 필요성을 확인시켜 줍니다.
• 이미지 이해 (Image Understanding):
  • ImageNet 선형 프로빙(linear probing)에 적용했을 때, PRA-L은 68.80%의 Top-1 정확도를 기록하며 잠재 공간 AR 베이스라인(SphereAR-L: 52.19%)과 픽셀 공간 확산 베이스라인(JiT-L: 42.76%)을 모두 앞질렀습니다. 이는 모델이 전이 가능한 시각적 표현을 학습하고 있음을 시사합니다.

5. 의의 및 주장

본 논문은 PRA가 확산 모델 및 잠재 공간 AR 모델과 경쟁하는 데 있어 픽셀 공간 AR의 근본적인 병목 현상을 해결한다고 주장합니다. 저차원 중간 상태를 통해 단일 단계 생성의 난이도를 낮추고, 병렬 롤아웃 근사를 통해 오차 축적을 완화함으로써, PRA는 원시 픽셀로부터 고품질 이미지를 직접 생성할 수 있게 합니다.

저자들은 PRA의 "픽셀 입력, 픽셀 출력" 특성과 강력한 선형 프로빙 성능을 결 통해, PRA가 별도의 토크나이저나 오토인코더가 필요 없는 통합 픽셀 공간 모델(고정밀 이미지 생성과 견고한 시각적 이해가 모두 가능한 모델)을 위한 잠재력을 가지고 있음을 강조합니다.

한계점: 저자들은 PRA가 추가적인 구성 요소(중간 상태 인코더, 픽셀 디코더, 그리고 학습 중 추가적인 병렬 AR 순방향 패스)를 도입한다는 점을 언급했습니다. 이들은 엔드 투 엔드로 학습되지만, 더 넓은 데이터 도메인과 생성 작업에 걸쳐 이 프레임워크를 검증하기 위한 후속 연구가 필요합니다.