Abstract

기존의 비전 토크나이제이션 방식은 비전 토크나이저의 최적화를 다운스트림 학습과 분리하여 수행하며, 이미지 생성이나 시각적 질의응답과 같은 다양한 태스크 전반에서 시각 토큰이 일반화될 수 있다고 암묵적으로 가정한다. 그러나 저수준 재구성을 목표로 최적화된 비전 토크나이저는, 서로 다른 표현과 의미를 요구하는 다운스트림 태스크에 대해서는 태스크 비의존적이다. 이러한 분리된 패러다임은 중요한 불일치를 초래하는데, 비전 토크나이제이션의 손실 자체가 목표 태스크의 표현 병목이 될 수 있다. 예를 들어, 이미지 내 텍스트를 토크나이즈하는 과정에서 발생한 오류는 해당 텍스트를 인식하거나 생성하는 성능 저하로 이어진다.

이를 해결하기 위해, 우리는 비전 토크나이제이션과 목표 자동회귀 태스크를 공동으로 최적화할 수 있는 엔드투엔드 비전 토크나이저 튜닝 방법인 ETT를 제안한다. 동결된 비전 토크나이저로부터 이산 인덱스만을 사용하는 기존 자동회귀 모델들과 달리, ETT는 토크나이저 코드북의 시각 임베딩을 직접 활용하며, 재구성 목표와 캡션 목표를 함께 사용하여 비전 토크나이저를 엔드투엔드 방식으로 최적화한다. ETT는 기존 대규모 언어 모델의 코드북이나 아키텍처를 수정할 필요 없이 간단하게 구현 및 통합할 수 있다.

광범위한 실험 결과, 제안한 엔드투엔드 비전 토크나이저 튜닝은 동결된 토크나이저를 사용하는 기준선 대비 멀티모달 이해 및 시각 생성 태스크에서 2에서 6퍼센트에 이르는 유의미한 성능 향상을 달성함과 동시에, 기존의 재구성 성능을 유지함을 보여준다. 우리는 이 단순하면서도 강력한 방법이 이미지 생성과 이해를 넘어 멀티모달 파운데이션 모델 전반에 기여할 수 있기를 기대한다.

기존에는 토크나이저 하고 나서 LLM에 넣어서 학습 딸깍 했는데, 이러면 LLM 학습과 토크나이저가 분리되니까, LLM 학습시에도 gradient가 흐르게 하자

1. Introduction

최근 대규모 언어 모델과 멀티모달 사전학습의 급속한 발전으로 인해, 자동회귀 모델링은 자연어 처리 분야를 넘어 비전 및 멀티모달 태스크 전반으로 그 영향력을 확장하고 있다. 자동회귀 모델의 다음 토큰 예측 패러다임 하에서, 멀티모달 학습은 일반적으로 이미지와 텍스트와 같은 서로 다른 모달리티의 데이터를 하나의 통합된 시퀀스 모델링을 위해 압축된 이산 토큰으로 인코딩한다. 예를 들어, 최근 연구인 Emu3는 텍스트, 이미지, 비디오를 모두 이산 토큰으로 토크나이즈한 뒤, 통합된 토큰 공간에서 다음 토큰 예측을 수행한다. 이후의 여러 연구들 또한 이러한 방향을 확장하여 시각 생성 및 인식 성능을 더욱 향상시켰다. 이와 같은 통합된 다음 토큰 예측 패러다임은 대규모 학습과 추론을 모두 지원하는 유연한 멀티모달 학습 프레임워크를 가능하게 한다.

이러한 자동회귀 프레임워크에서 토크나이제이션은 핵심적인 역할을 한다.

이미지나 비디오와 같은 모달리티의 경우, 효율적이면서도 범용적인 토크나이저를 학습하는 것은 매우 도전적인 문제인데, 토크나이제이션은 동시에 압축성과 무손실성을 만족시키기 어렵기 때문이다.
토크나이제이션 과정에서 발생하는 정보 손실은 목표 태스크의 성능 병목이 될 수 있다.
예를 들어, 이미지 내 텍스트를 토크나이즈하는 과정에서의 오류는 해당 텍스트를 인식하거나 생성하는 성능 저하로 직접 이어진다.

그러나 기존의 토크나이제이션 접근법들은 이러한 불일치를 간과한 채, 비전 토크나이저를 다운스트림 학습과 분리하여 별도로 학습한 뒤 그대로 통합한다.

이는 학습된 시각 토큰이 다양한 태스크 전반에서 충분히 일반화될 수 있다는 가정에 기반한다.
특히, 자동인코딩을 목표로 최적화된 비전 토크나이저는 서로 다른 표현과 의미를 요구하는 다운스트림 태스크에 대해 태스크 비의존적이다.
대부분의 토크나이저는 저수준의 픽셀 단위 재구성에 집중하며, 벡터 양자화로 인한 정보 손실로 인해 학습된 표현의 품질은 본질적으로 제한된다.
그 결과, CLIP과 같은 연속적인 고수준 표현을 사용하는 모델들에 비해 시각적 이해 태스크에서 열등한 성능을 보이게 된다.

더 나아가, 기존의 자동회귀 파이프라인들은 일반적으로 비전 토크나이저로부터 얻은 이산 인덱스만을 사용하며, 다운스트림 태스크를 수행하기 위해 대규모 언어 모델 내의 시각 임베딩은 무작위 초기화에 의존한다. 이는 멀티모달 학습에서 핵심적인 요소인 시각 표현 학습과 비전-언어 정렬을 어렵게 만든다.

본 연구에서는 비전 토크나이제이션과 목표 자동회귀 다운스트림 태스크를 공동으로 최적화할 수 있는 엔드투엔드 비전 토크나이저 튜닝 접근법인 ETT를 제안한다.

최근의 비전-언어 모델들이 학습 과정에서 연속적인 비전 인코더를 업데이트함으로써 시각 표현과 비전-언어 정렬을 함께 최적화하는 데서 영감을 받아, 우리는 이산 비전 토크나이저 역시 엔드투엔드 방식으로 튜닝할 수 있음을 제안한다.
이 과정에서 대규모 언어 모델은 비전 토크나이저의 시각적 보조자 역할을 수행한다.
구체적으로, 우리는 이산 인덱스만을 사용하는 대신 비전 토크나이저의 코드북 임베딩을 도입하고, 토큰 단위 캡션 손실을 통합하여 비전 토크나이저의 표현을 직접 최적화한다.

실험 결과, ETT는 멀티모달 이해 및 시각 생성 태스크에서 다운스트림 성능을 크게 향상시키며, 비전 토크나이저가 보다 판별적이고 생성에 유리한 표현을 학습하도록 만든다. 또한 ETT는 비전 토크나이저의 기존 이미지 재구성 성능을 유지할 수 있음을 확인하였다. 더 나아가, 본 방법은 대규모 언어 모델의 기존 텍스트 코드북을 수정하거나, 시각 임베딩 및 비전-언어 정렬을 처음부터 학습하기 위해 임베딩 레이어나 분류 헤드를 확장할 필요 없이 간단하게 구현 및 통합할 수 있다.

2 Related Work

Vision Tokenizer
비전 토크나이저는 이미지나 비디오를 높은 재구성 품질을 유지하면서 이산 토큰으로 양자화하는 것을 목표로 한다. 대표적으로 VQ-VAE는 오토인코더 구조 내에 양자화기를 포함시켜, 연속적인 특징을 이산 표현으로 매핑하도록 학습한다. VQGAN은 지각 손실과 적대적 손실을 결합함으로써 재구성 품질을 향상시켰다. MoVQ는 공간 조건부 정규화를 도입하여 양자화된 벡터를 조절함으로써 고충실도의 이미지 생성을 가능하게 한다. 그 외에도 코드북 활용률을 개선하거나, 보다 진보된 양자화 기법을 도입하는 연구들이 제안되어 왔다. 최근 연구들은 시맨틱 정보를 토크나이저에 통합하여 시각적 표현을 향상시키는 데 초점을 맞추고 있다. 그러나 이러한 접근법들에서도, 다운스트림 태스크에서는 일반적으로 양자화된 이산 인덱스만을 사용하며, 토크나이저 자체는 다운스트림 학습 동안 동결된 상태로 유지된다. 본 연구에서는 IBQ를 강력한 기준선으로 채택하고, 기존의 이산 인덱스 대신 코드북 임베딩을 활용하여 엔드투엔드 튜닝을 수행한다.

Tokenization for Visual Generation and Understanding
이산 표현 기반 시각 생성은 최근 몇 년간 큰 발전을 이루었다. 일부 연구들은 자동회귀 방식을 사용하여 토큰을 순차적으로 예측함으로써 이미지나 비디오를 생성하며, 다른 연구들은 MaskGIT이나 Muse와 같이 마스크드 비전 토큰 모델링을 통해 이미지 생성을 수행한다. 최근에는 동결된 이산 비전 토크나이저를 사용한 상태에서, 하나의 모델 내에서 시각적 이해와 생성을 통합하려는 연구들도 등장하였다. 예를 들어, Emu3는 비디오, 이미지, 텍스트를 하나의 토큰 공간으로 통합하고, 다음 토큰 예측을 통해 다양한 모달리티에 걸친 강력한 이해 및 생성 능력을 달성하였다. Show-o는 자동회귀 모델 내에 마스크 이미지 모델링을 도입하여 이미지 생성을 수행한다. Janus는 이해와 생성 태스크 간의 본질적인 충돌을 완화하기 위해 두 개의 서로 다른 비전 인코더를 각각 사용하는 구조를 제안하였다. 본 연구는 다음 토큰 예측 패러다임에 초점을 맞추어, 비전 토크나이저와 자동회귀 모델을 공동으로 최적화함으로써 시각적 표현을 개선하고, 멀티모달 생성 및 인식 태스크 전반의 성능을 향상시키는 것을 목표로 한다.

3 Methodology

3.1 Vision Tokenizer

Preliminary
VQ 기반 비전 토크나이저는 인코더, 양자화기, 디코더로 구성된다. 인코더는 입력 이미지 I를 H×W×3 차원에서 h×w×D 차원의 특징 맵 f로 투영하며, 여기서 D는 특징 차원이고 h = H/s, w = W/s이며 s는 다운샘플링 비율을 의미한다. 양자화기는 특징 맵 f의 각 특징 벡터를 코드북 B에 포함된 K개의 코드 중 가장 가까운 코드로 매핑하여, h×w×D 차원의 양자화된 임베딩 z를 생성한다. 이후 디코더는 이 양자화된 임베딩을 입력 이미지로 재구성한다. 기존의 VQ 모델들은 코드북 크기와 코드 차원을 동시에 확장하는 데 어려움을 겪어왔다. 선행 연구인 IBQ는 코드 차원을 256으로 확장하면서도 매우 큰 코드북 크기인 262,144를 유지하기 위해, 매 학습 스텝마다 전체 코드북을 동시에 업데이트하는 방식을 최초로 제안하였다.

Vision Tokenizer in ETT
본 연구에서는 이미지 토크나이제이션을 위해 IBQ의 프레임워크를 주로 채택하며, 다운샘플링 비율 s는 16으로 설정한다. 코드북에 포함된 각 이산 토큰의 차원 D는 256이다. 원래의 IBQ 설정을 기반으로 하되, 코드북 크기는 131,072로 조정한다. 토크나이저 학습을 위한 손실 함수 Lvq는 다음과 같이 정의된다.

여기서 Lrec은 픽셀 재구성 손실, Lquant는 인코딩된 특징과 양자화된 임베딩 간의 양자화 손실, Llpips는 LPIPS 기반 지각 손실, LGAN은 PatchGAN에서 유도된 적대적 손실, Lentropy는 엔트로피 손실을 의미한다. λG와 λE는 각각 적대적 손실과 엔트로피 손실의 가중치이다.

3.2 End-to-End Vision Tokenizer Tuning

Discrete Indices to Codebook Embeddings
Emu3와 같은 방법들은 다운스트림 태스크에서 비전 토크나이저의 이산 인덱스만을 사용함으로써, 비전 토크나이저 임베딩이 지니는 풍부한 표현 능력을 활용하지 못한다. 이산 코드북 인덱스에만 의존할 경우, 그래디언트가 전파되지 않아 엔드투엔드 학습이 불가능해진다. 이러한 한계를 극복하기 위해, 우리는 ETT를 제안한다. ETT는 비전 토크나이저의 코드북 임베딩을 대규모 언어 모델과 직접 연결함으로써, 비전 토크나이저에 내재된 풍부한 특징 표현을 효과적으로 활용하는 동시에 엔드투엔드 학습을 가능하게 한다.

LLM Bridges End-to-End Tuning
그림 1에 도시된 바와 같이, 입력 이미지 I가 주어지면 먼저 비전 토크나이저의 코드북을 통해 h×w×D 차원의 양자화된 임베딩 z를 얻는다. 사전학습된 대규모 언어 모델과의 호환성을 보장하기 위해, 우리는 GeLU 활성화 함수를 사용하는 다층 퍼셉트론 기반의 경량 프로젝터를 도입한다. 이 프로젝터는 양자화된 시각 임베딩 z를 대규모 언어 모델의 히든 차원 크기 C에 맞추어 h×w×C 차원의 표현 xI로 매핑한다. 이 구조에서는 사전학습된 대규모 언어 모델과 비전 토크나이저를 포함한 전체 계산 그래프가 미분 가능하게 유지되므로, 그래디언트 기반 최적화를 통해 전체 모델을 엔드투엔드로 학습할 수 있다. 텍스트 입력 T의 경우에는, 사전학습된 대규모 언어 모델의 텍스트 토크나이저와 텍스트 임베딩 레이어를 사용하여 N×C 차원의 텍스트 토큰 임베딩 xT로 변환한다.

Preservation of Reconstructive Capability
엔드투엔드 학습은 비전 토크나이저의 표현을 향상시키지만, 고품질 이미지 생성을 위해서는 기존의 재구성 능력을 유지하는 것이 중요하다. 이를 위해 우리는 전체 학습 목표를 캡션 손실 Lcap과 VQ 손실 Lvq의 결합으로 설정한다. 구체적으로, 이미지 토큰 임베딩 xI와 텍스트 토큰 임베딩 xT를 함께 대규모 언어 모델에 입력한다. 텍스트 토큰에 대해서는 다음과 같은 크로스 엔트로피 손실을 적용한다.

또한, 시각적 재구성을 위해 기존의 Lvq 손실을 그대로 사용한다. 이에 따라, 엔드투엔드 비전 토크나이저 튜닝의 최종 학습 목표는 다음과 같이 정의된다.

여기서 α는 멀티모달 인식 성능과 시각적 재구성 성능 간의 트레이드오프를 조절하는 손실 가중치이다. 대규모 언어 모델과 함께 토크나이저의 인코더와 디코더를 공동으로 학습함으로써, 본 접근법은 재구성 능력을 유지하는 동시에, 멀티모달 이해 및 생성에 효과적인 시맨틱한 시각 토큰을 학습할 수 있도록 한다.

3.3 Training Recipe for Multimodal Generation and Understanding

선행 연구들을 따라, 멀티모달 인식과 생성을 위한 전체 학습 과정은 세 개의 연속적인 학습 단계로 구성된다. 사용되는 학습 데이터는 공개적으로 이용 가능한 이미지 데이터셋과, 이해 및 생성을 위한 다양한 인스트럭션 데이터를 포함하며, 그 세부 내용은 표 1에 정리되어 있다.

Stage 1 Alignment learning
첫 번째 학습 단계의 목적은 비전과 언어 간의 정렬을 효과적으로 구축하는 것이다. 이 단계에서는 사전학습된 대규모 언어 모델과 비전 토크나이저를 모두 동결한 상태에서, 시각 프로젝터 레이어만을 학습하며 이미지-텍스트 캡션 손실 Lcap을 사용한다. 이를 통해 대규모 언어 모델은 토크나이저로부터 직접 시각적 개념과 엔티티를 습득하게 되며, 비전과 언어 모달리티 간의 연결이 형성된다. 구체적으로, 우리는 공개 데이터로부터 수집한 3,200만 개의 이미지-텍스트 쌍으로 구성된 SOL-recap 데이터셋에서 1,200만 장의 이미지 서브셋을 선별하여 사용한다. 이 데이터는 SA-1B, OpenImages, LAION으로부터 수집되었으며, 모든 이미지는 선행 연구를 따라 개선된 캡셔닝 엔진으로 재캡션되었다. 이 단계에서의 고품질 데이터는 학습 안정성과 모달리티 간 정렬 성능을 향상시킨다.

Stage 2 Semantic learning
두 번째 단계는 전체 학습 파이프라인에서 가장 핵심적인 단계로, 엔드투엔드 비전 토크나이저 튜닝이 수행된다. 이 단계에서는 대규모 언어 모델, 시각 프로젝터, 비전 토크나이저의 가중치를 모두 언프리즈하고, 앞서 정의한 캡션 손실 Lcap과 재구성 손실 Lvq를 공동으로 최적화한다. 학습에는 SOL-recap에서 선별한 1,200만 개의 고품질 이미지-텍스트 쌍을 사용하며, 이를 통해 멀티모달 이해와 시각적 재구성 능력을 동시에 학습한다. 이 단계는 비전 토크나이저와 다운스트림 태스크 간의 정렬을 강화하는 동시에, 기존의 재구성 성능을 유지하도록 설계되어 있다.

Stage 3 Post-training
엔드투엔드 튜닝을 통해 향상된 비전 토크나이저를 획득한 이후, 표준적인 포스트 트레이닝 절차를 적용하여 멀티모달 이해와 생성을 수행한다. 이 단계에서는 비전 토크나이저를 동결한 상태에서, 시각 프로젝터와 대규모 언어 모델을 추가로 튜닝하여 인스트럭션 추종 능력을 강화한다. 이해 태스크를 위해서는 SOL-recap, LLaVA-OneVision, Infinity-MM 등 다양한 출처의 고품질 멀티모달 인스트럭션 데이터를 활용한다. 시각 생성 태스크의 경우, Flux 모델을 이용해 생성한 1,400만 개의 AI 생성 샘플과, 이미지 해상도 및 미적 점수를 기준으로 필터링한 1,600만 개의 웹 이미지-텍스트 쌍을 추가로 사용한다.