잠재 제어 역학 학습을 통한 마스크 이미지 생성 가속화

마스크 기반 이미지 생성 모델(MIGM)은 뛰어난 성능에도 불구하고 반복적인 Transformer 연산으로 인해 추론 속도가 매우 느렸다. 이 논문은 특징 벡터의 변화량을 예측하는 경량 모델을 도입하여, 품질 저하 없이 생성 속도를 4배 이상 높이는 실질적인 최적화 방안을 제시했다.

마스크 이미지 생성 모델(MIGM)은 빈 캔버스에서 시작해 토큰을 하나씩 채워가는 과정을 거친다. 이때 각 단계에서 생성되는 내부 특징(feature)들은 서로 매우 유사하지만, 기존 방식은 매번 무거운 Transformer 전체를 다시 실행하여 이 특징을 계산한다. 이는 마치 아주 조금씩 변하는 그림을 그리면서 매번 처음부터 모든 붓질을 다시 하는 것과 같은 비효율을 초래한다.

이 논문의 핵심 아이디어는 '특징이 어떻게 변할지'만 예측하는 가벼운 지름길 모델을 만드는 것이다. 하지만 단순히 이전 특징만 보고 다음을 예측하면 생성의 무작위성을 반영하지 못해 결과가 뭉개지게 된다. 연구진은 매 단계 어떤 토큰이 선택되었는지(샘플링 결과)를 지름길 모델의 입력으로 넣어줌으로써, 특징이 나아갈 정확한 방향을 가이드하도록 설계했다.

결과적으로 무거운 원본 모델은 가끔씩만 호출하여 오차를 교정하고, 대부분의 단계는 가벼운 지름길 모델이 처리하게 함으로써 전체 연산량을 획기적으로 줄였다. 이는 고해상도 이미지 생성 시 발생하는 막대한 컴퓨팅 자원 소모를 해결할 수 있는 중요한 전환점이 된다.

전체 생성 과정을 상태 공간 모델(State-Space Model)로 정의한다. 현재 단계 $t_i$의 특징 $f_{t_i}$와 샘플링된 토큰 $x_{t_i}$가 주어질 때, 다음 단계의 특징 $f_{t_{i+1}}$은 $f_{t_i} + S_ heta(f_{t_i}, x_{t_i}, t_i)$로 계산된다. 여기서 $S_ heta$는 학습 가능한 경량 숏컷 모델이다.

숏컷 모델은 Cross-Attention 레이어와 Self-Attention 레이어로 구성된다. [이전 특징 벡터와 샘플링된 토큰 임베딩을 입력으로] → [Cross-Attention을 통해 토큰 정보를 특징 벡터에 주입하고] → [Self-Attention과 시간 임베딩을 통해 특징의 진화 방향을 결정하여] → [최종적으로 특징 변화량 $\Delta f$를 출력한다]. 이 값은 잔차 연결(Residual Connection)을 통해 이전 특징에 더해진다.

학습 시에는 원본 모델(Base Model)을 동결하고 숏컷 모델만 MSE Loss를 통해 학습시킨다. 추론 시에는 오차 누적을 방지하기 위해 $N$단계 중 $B$번만 원본 모델을 호출하는 'Full Step'을 수행하고, 나머지 단계는 숏컷 모델을 사용하는 'Shortcut Step'으로 구성하는 스케줄링 전략을 사용한다.

Lumina-DiMOO 모델을 사용한 텍스트-이미지 생성 실험에서, 64단계 생성 기준 4.01배의 가속을 달성하면서도 ImageReward 0.90(원본 0.91)을 기록하여 품질 저하가 거의 없음을 확인했다. 가속 배율을 5.79배까지 높여도 원본 대비 우수한 효율성을 유지했다.

MaskGIT 모델의 경우, 15단계 생성 시 1.94배 가속하며 FID 8.90을 기록했다. 특히 32단계 설정에서는 FID 6.84를 달성하여 원본 모델의 최적 설정(15단계, FID 7.60)보다 더 뛰어난 품질과 속도 균형을 보여주었다.

기존의 훈련이 필요 없는 가속 기법인 ReCAP, TaylorSeer 등과 비교했을 때, 동일한 속도에서 훨씬 높은 이미지 품질(CLIPScore, UniPercept-IQA 등)을 제공하며 파레토 최적(Pareto Frontier)을 갱신했다.

MIGM의 마지막 레이어 특징이 연속적인 단계 사이에서 매우 높은 코사인 유사도(0.95 이상)를 보인다는 점에 착안했다. 또한, 특징 궤적의 변화량이 입력 변화량에 비례하는 국소적 립시츠 연속성(Local Lipschitz behavior)을 가짐을 실험적으로 증명하여 경량 모델로도 충분히 모사가 가능함을 이론적으로 뒷받침했다.

숏컷 모델은 원본 모델 대비 파라미터 수를 1/37(Lumina-DiMOO 기준)로 줄였으며, 연산 지연 시간(Latency)은 약 1/30 수준으로 단축했다. 연산 효율을 극대화하기 위해 입력 특징을 낮은 차원으로 투영하는 Bottleneck 구조(비율 R=2)를 채택했다.

학습 시 단순 MSE Loss만으로도 KL Divergence 등 복잡한 손실 함수보다 우수한 성능을 보였으며, 이는 잠재 역학 자체가 학습하기에 충분히 매끄러운 구조를 가지고 있음을 시사한다. 추론 시에는 $1 + \lfloor jN/B \rfloor$ 단계마다 원본 모델을 호출하여 분포 편향(Distribution Shift) 문제를 해결했다.

숏컷 모델을 연속적으로 과도하게 사용할 경우 오차가 누적되어 이미지가 흐릿해지거나 구조가 무너질 수 있으므로, 주기적으로 원본 모델을 호출하는 Full Step 예산(Budget B)을 적절히 설정해야 한다.