픽셀 공간 자기회귀 이미지 생성을 위한 병렬 롤아웃 근사

픽셀 공간에서 직접 연속 토큰으로 자기회귀 생성은 토크 차원이 매우 커져 단일 스텝 예측 오류가 커지는 문제와 교사 강제 학습으로 인한 학습-추론 불일치를 동시에 안고 있다. 이 논문은 두 문제를 동시에 해결하는 방법을 도입하여 순차적 롤아웃의 비용 없이 추론 유사 입력을 학습 단계에서 병렬로 구성한다. 그 결과 상대적으로 작은 모델 크기에서도 이전의 대규모 픽셀-스페이스 AR 결과를 능가하는 질적·정량적 성능을 확보했다.

픽셀-스페이스 자기회귀(autoregressive) 이미지 생성은 각 패치를 고차원 연속 벡터로 직접 예측하는 방식이다. 이 때 각 패치가 다음 스텝의 조건(context)이 되므로 단일 스텝에서의 큰 예측 오류가 이후 스텝들로 누적되어 전체 품질을 급격히 저하시킨다. 더구나 학습은 정답 접두사(teacher-forced)를 사용하지만 추론은 모델이 생성한 접두사를 사용하므로 학습-추론 사이에 구조적 불일치가 존재한다.

핵심 해결책은 두 축에서 접근하는 것이다. 출력 측면에서는 고차원 픽셀 패치 대신 저차원 중간 상태를 목표로 삼아 한 스텝의 예측 난이도를 줄인다. 중간 상태는 학습 중에 causal pixel decoder로 역변환되어 픽셀-토큰을 재구성하므로 AR 백본은 여전히 픽셀 입력과 출력을 다루는 외부 인터페이스를 유지한다.

입력 측면에서는 순차적 롤아웃을 대체하는 병렬 복원 경로를 구성하여 추론 시 생성되는 픽셀의 통계적·구조적 특성에 근접한 학습 입력을 만든다. 구체적으로는 각 위치에서 중간 상태를 정규분포 샘플과 보간하여 변형한 뒤 동일한 픽셀 디코더로 복원된 토큰을 AR 입력으로 사용하여 학습-추론 불일치를 줄인다.

이 두 구성은 상호보완적이다. 저차원 중간 상태는 단일 스텝의 error를 낮추고 병렬 복원된 입력은 그 오류가 다음 스텝으로 전이되는 방식을 현실적으로 근사한다. 결과적으로 상대적으로 작은 모델 규모에서도 이전의 대규모 픽셀-스페이스 AR를 능가하는 성능이 관찰되었다.

PRA의 전체 접근은 세 모듈로 구성된다: AR Transformer 백본 f_θ, 목표 인코더 g_ψ로 생성되는 저차원 중간 상태 z, 그리고 causal pixel decoder D_φ이다. 학습 시에는 먼저 teacher-forced forward로 prefix 표현 h_{i-1}=f_θ(x_{<i})를 얻고 이를 g_ψ(x_i, h_{i-1})에 투입해 context-aware한 중간 목표 z_i를 만든다. 이 과정에서 일부 토큰은 p_mask 확률로 마스킹되어 g_ψ가 토큰뿐 아니라 prefix 정보를 활용하도록 유도된다.

중간 목표 z_i는 출력 측면 손실 대상이 되며, 병렬 복원 절차에서는 각 z_i에 대해 z_i^0∼N(0,I), t∼U(t_min,1)를 샘플링하여 z̃_i^t=(1-t) z_i^0 + t z_i를 구성한다. 이 변형된 중간 상태들을 causal pixel decoder D_φ에 투입해 복원된 픽셀 시퀀스 x̄를 얻고, 이 시퀀스를 stop-gradient로 취해 AR 백본의 입력으로 사용한다. 이렇게 하면 모델은 추론 유사한 픽셀 접두사 하에서 중간 상태를 생성하도록 학습된다.

토큰 생성은 token-level diffusion head v_ω로 수행되며 rectified-flow 목표를 사용해 노이즈에서 목표 중간 상태로의 변환을 학습한다. 학습 손실은 L_AR(Flow matching), L_rec(복원 이미지 재구성 손실), 그리고 추가 보조 손실 L_aux의 합으로 구성된다. 전체 학습은 위치별 병렬 연산으로 설계되어 순차 샘플링을 수행하는 기존 롤아웃 대비 계산 비용을 크게 낮춘다.

주요 벤치마크는 class-conditional ImageNet-1K 256×256 생성이며 FID는 50k 샘플로 계산되었다. PRA-S(135M)는 FID 2.58을 기록하여 기존 빌리언-스케일 픽셀-스페이스 AR 결과인 FID 3.60을 능가했고, PRA-L(511M)은 FID 1.94로 픽셀-스페이스 AR 최상위를 기록했다. 비교 실험에서 token dimensionality가 48에서 768으로 커질 때 AR 성능이 악화된 진단 결과를 바탕으로 PRA는 저차원 중간 상태가 단일 스텝 오류를 줄이는 데 결정적임을 보였다.

Ablation 결과는 두 구성 요소가 보완적이라는 점을 보여주었다. 저차원 intermediate target만 적용했을 때와 병렬로 복원된 inference-like inputs만 적용했을 때 각각 성능 향상이 관찰되었고 두 요소를 결합했을 때 가장 큰 이득이 발생했다. 추가로 생성 학습된 AR 백본이 ImageNet 선형 probing에서 AR 및 diffusion 기반 비교군보다 더 높은 정확도를 보였으며 이는 생성을 통한 표현 학습의 추가적 이점을 시사한다.

성능 비교 및 진단 실험은 동일한 모델 스케일과 학습 예산(예: 200 epochs 기준 비교 실험) 하에서 수행되었으며 x x -prediction, input noise injection 등 기존 기법들과의 비교를 통해 PRA의 이득을 정량화했다.

전체 아키텍처는 causal Transformer AR 백본 f_θ, 중간 상태 생성기 g_ψ, causal pixel decoder D_φ, 그리고 token-level flow head v_ω로 구성된다. 입력 이미지는 패치 단위로 분할되어 각 패치가 연속값 토큰 x_i ∈ ℝ^d(예: 16×16×3=768)로 표현되며, 목표 중간 상태는 d_z<d 차원으로 정의된다. 학습 시에는 teacher-forced로 얻은 prefix 표현 h_{i-1}를 g_ψ에 투입하여 context-aware한 z_i를 만든다.

token-level diffusion head은 rectified-flow 목적을 사용하며 학습 시 z_i^{s_i}=(1-s_i) ε_i + s_i z_i 형태의 보간된 상태를 목표로 flow matching 손실을 계산한다. 이 손실은 v_ω(z_i^{s_i}, s_i, h̄_{i-1})가 z_i - ε_i를 근사하도록 L2로 최적화된다. 샘플링은 Euler-Maruyama 유사 적분을 통해 노이즈에서 목표 z로 통합하는 방식으로 수행된다.

병렬 입력 구성에서는 각 위치별로 z̃_i^t를 생성해 D_φ로 다시 디코딩한 복원 토큰 x̄_i를 얻고 이 시퀀스를 stop-gradient로 AR 백본의 입력에 사용한다. 이로 인해 학습은 두 번의 병렬 AR forward를 필요로 하나 순차 롤아웃에 비해 GPU 병렬화 효율이 높아 학습 시간의 폭증을 피할 수 있다.

구현상 세부 파라미터로는 패치 크기(예: 16×16), token dimension(예: 768), 중간 상태 차원 d_z, 마스킹 확률 p_mask, 보간 변수의 범위 t_min 등이 있으며 손실은 L_AR + L_rec + L_aux의 합으로 최종적으로 최적화된다. 추론 단계에서는 순차적 FlowSample 절차로 ẑ_i를 생성하고 D_φ로 디코딩한 픽셀을 누적하여 최종 이미지를 얻는다.