(1D) 순서화된 토큰을 통한 효율적인 테스트 시간 검색 가능성 연구

기존의 2D 그리드 방식 토큰 구조는 생성 중간 단계에서 전체적인 의미를 파악하기 어려워 추론 시점의 최적화가 힘들었다. 이 논문은 1D 순서화된 토큰 구조가 생성 과정에서 '거시적 의미에서 세부 묘사'로 이어지는 계층적 정보를 제공하여, 별도의 추가 학습 없이도 검색 알고리즘을 통해 생성 품질을 획기적으로 높일 수 있음을 증명했다.

기존의 이미지 생성 모델은 이미지를 바둑판 모양의 2D 그리드로 나누고 왼쪽 위부터 오른쪽 아래로 순차적으로 토큰을 예측한다. 이 방식은 생성 중간 단계에서 이미지의 일부분(예: 왼쪽 상단 벽면)만 보여주기 때문에, 전체 이미지가 프롬프트와 일치하는지 판단하는 검증기(Verifier)가 제대로 작동하기 어렵다. 이는 마치 퍼즐의 한 구석 조각만 보고 전체 그림을 맞추려는 것과 같다.

반면 1D 순서화된 토큰(FlexTok 등)은 이미지를 계층적으로 압축한다. 첫 번째 토큰이 이미지 전체의 대략적인 구도와 핵심 피사체 정보를 담고, 뒤로 갈수록 세부적인 질감과 묘사를 더하는 방식이다. 따라서 생성 초기 단계의 토큰만으로도 '이 이미지가 고양이인지 강아지인지'를 판단할 수 있는 전역적 의미(Global Semantics)가 형성된다.

이러한 구조적 특징 덕분에 추론 시점에 여러 후보 토큰 중 검증기 점수가 높은 경로를 선택하는 '검색'이 가능해진다. 초기 토큰에서 이미 목표와 일치하는 방향을 잡을 수 있기 때문에, 잘못된 생성을 조기에 차단하고 유망한 경로에 연산 자원을 집중함으로써 최종 출력물의 품질을 비약적으로 높일 수 있다.

본 연구는 1D 순서화된 토큰 구조의 이점을 극대화하기 위해 Search-over-Tokens(SoTo) 프레임워크를 설계했다. 핵심은 Autoregressive(AR) 모델을 확률적 가이드(Prior)로 사용하고, 외부 검증기(Verifier)를 통해 각 단계의 생성물을 평가하며 검색을 수행하는 구조이다.

검색 알고리즘으로는 Beam Search를 주력으로 사용한다. 각 단계 t에서 k개의 유망한 부분 시퀀스(Beams)를 유지하며, 모델의 확률 분포 p(xt | x<t)에서 M개의 후보 토큰을 샘플링한다. 이후 각 후보를 이미지로 복원(Detokenize)하여 검증기 g(x)의 점수를 계산하고 상위 k개를 선택하는 과정을 반복한다. [부분 토큰 시퀀스 입력 → Detokenizer를 통한 중간 이미지 복원 → 검증기 점수 계산 → 상위 k개 유지 → 최적 경로 탐색]

이 과정에서 1D 토큰은 Nested Dropout 기법으로 학습되어, 시퀀스의 어느 지점에서 끊더라도 유효한 이미지를 복원할 수 있다. 이는 검증기가 생성 중간 단계마다 실시간으로 피드백을 줄 수 있게 하여 검색의 효율성을 극대화한다. 반면 2D 그리드 방식은 중간 단계에서 나머지 부분을 0으로 채워야 하므로 검증 신호가 부정확해지는 한계가 있다.

실험 결과, 1D 순서화된 토큰 기반 모델(FlexTok)은 동일한 연산량 조건에서 2D 그리드 기반 모델(Janus 등)보다 훨씬 강력한 테스트 시간 확장성(TTS)을 보였다. 특히 Beam Search를 적용했을 때 CLIPScore와 ImageReward 지표에서 2D 모델 대비 압도적인 성능 향상을 기록했다.

제로샷 제어 실험에서는 추가 학습 없이 DreamSim 검증기만을 사용하여 참조 이미지의 개념을 유지하며 새로운 이미지를 생성하는 데 성공했다. FlexTok은 DINO-I 지표에서 +18.4점의 향상을 보인 반면, 2D 모델인 Janus는 Lookahead Search를 동원하고도 +5.9점 향상에 그쳐 토큰 구조의 중요성을 입증했다. 또한 530M 파라미터의 작은 모델이 충분한 검색 연산을 사용할 경우, 검색을 하지 않은 3.4B 파라미터의 대형 모델보다 더 높은 성능을 낼 수 있음을 확인하여 모델 크기와 추론 연산 간의 파레토 최적 관계를 도출했다.

본 연구는 1D 순서화된 토큰이 검색 격차(Search Gap)를 줄이는 원리를 수학적으로 분석했다. 검색 격차 Δ는 휴리스틱 오류(Heuristic Error) Bt에 의해 바운드되는데, 1D 토큰은 중간 복원 이미지와 최종 완결 이미지 사이의 재구성 오차를 최소화하도록 설계되어 Bt를 낮게 유지한다. [중간 토큰 복원 이미지와 최종 이미지의 L2 거리 → Lipschitz 연속성 가정 하의 검증기 점수 차이 → 검색 격차의 상한선 결정]

구현 측면에서는 FlexTok 아키텍처를 기본으로 하며, 212M부터 3.4B까지 다양한 크기의 AR 모델을 평가했다. 검증기로는 CLIPScore, ImageReward, HPSv2, Grounded-SAM 등을 개별 또는 앙상블 형태로 사용했다. 특히 2D 그리드 토큰의 한계를 보완하기 위해 미래 토큰을 미리 예측해보는 Lookahead Search를 비교군으로 설정하여, 1D 구조가 본질적으로 검색에 더 효율적인 표현 방식임을 기술적으로 증명했다.

테스트 시간 검색은 외부 검증기의 견고성에 크게 의존하므로, 검증기의 약점을 악용하여 시각적으로는 부자연스럽지만 점수만 높은 이미지가 생성되는 '검증기 해킹(Verifier Hacking)' 현상이 발생할 수 있다. 또한, AR 모델이 학습 과정에서 전혀 보지 못한 정보(Prior Bottleneck)는 검색을 통해서도 완전히 복구하기 어렵다는 한계가 있다.