예측에서 완성까지 자동회귀 이미지 생성의 정교화

초록

TensorAR는 기존 자동회귀(AR) 이미지 생성기가 한 번 예측한 토큰을 수정할 수 없는 한계를 극복한다. 이미지 토큰을 겹치는 윈도우 형태의 텐서(패치 그룹)로 재구성하고, 다음 텐서를 예측하도록 모델을 전환함으로써 이전에 생성된 내용에 반복적인 정교화를 적용한다. 훈련 시에는 코드북 인덱스를 이용한 이산 텐서 노이즈를 주입해 정보 누수를 방지한다. 기존 AR 모델에 최소한의 모듈만 추가하면 바로 적용 가능하며, LlamaGEN·Open‑MAGVIT2·RAR 등에서 FID와 GenEval 점수가 크게 향상된다.

상세 분석

TensorAR는 “다음‑토큰” 예측을 “다음‑텐서” 예측으로 전환함으로써 자동회귀 이미지 생성에 정교화 메커니즘을 도입한다. 핵심 아이디어는 겹치는 윈도우(k‑size)로 이미지 토큰을 묶어 텐서 시퀀스를 만든 뒤, 각 단계에서 하나의 텐서를 출력하도록 모델을 학습시키는 것이다. 겹침(overlap) 덕분에 이후 단계에서 앞서 생성된 토큰을 다시 방문하고 수정할 수 있어, 마치 확산 모델이 여러 단계에 걸쳐 노이즈를 제거하듯 점진적인 품질 향상이 가능하다.

하지만 텐서 기반 예측은 입력 텐서와 출력 텐서 사이에 중복 토큰이 존재하므로, 그대로 학습하면 모델이 단순히 복사해버리는 정보 누수(leakage) 문제가 발생한다. 이를 해결하기 위해 저자들은 이산 텐서 노이즈 스킴을 설계했다. 전통적인 이산 확산 이론을 차용해 각 토큰에 카테고리형 노이즈를 주입하고, 토큰별 노이즈 레벨을 조절해 점진적인 디노이징 과정을 모방한다. 이렇게 하면 모델이 겹치는 토큰을 무작위로 섞인 상태에서 복원하도록 강제되어, 진정한 인과 관계 학습이 이루어진다.

구조적으로는 기존 트랜스포머 기반 AR 모델의 임베딩 레이어와 출력 선형 레이어를 각각 “입력 인코더”(M_enc)와 “출력 디코더”(M_dec)로 래핑한다. 두 모듈은 경량 쿼리 트랜스포머(Q_in, Q_out)와 잔차 연결을 포함해 사전 학습된 파라미터를 그대로 활용하면서도 텐서 입출력을 지원한다. 따라서 모델 아키텍처 자체를 바꾸지 않고도 플러그‑인 형태로 적용할 수 있다(plug‑and‑play).

실험에서는 LlamaGEN, Open‑MAGVIT2, RAR 등 다양한 규모와 조건(클래스‑조건, 텍스트‑조건)에서 TensorAR를 적용했다. 결과는 FID가 평균 510% 개선되고, GenEval 점수가 0.020.04 상승하는 등 전반적인 품질·지연 트레이드오프가 크게 향상되었음을 보여준다. 특히 작은 모델(7B)에서도 큰 모델에 버금가는 성능을 달성해, 정교화 메커니즘이 모델 크기에 크게 의존하지 않음을 시사한다.

한계점으로는 윈도우 크기(k)와 노이즈 스케줄링이 성능에 민감하다는 점, 그리고 겹치는 토큰이 많아질수록 연산량이 증가한다는 점을 들 수 있다. 또한 현재는 토큰 기반 VQ‑AE와 결합된 전제 하에 동작하므로, 토큰화 품질이 전체 성능에 큰 영향을 미친다. 향후 연구에서는 동적 윈도우 조절, 다중 스케일 텐서 정교화, 그리고 비토큰 기반 디코더와의 통합을 탐색할 여지가 있다.