VLM에 내부 월드모델을 붙여 공간 추론을 시각 상태로 푼다

무엇인가

비전언어모델(VLM)은 대부분 언어 공간에서 추론한다. 그런데 물체를 회전시켰을 때 어떻게 보이는지 판단하는 공간 추론은 중간 시각 상태를 머릿속으로 시뮬레이션해야 풀린다. 기존 연구는 잘라낸 이미지나 하이라이트 같은 기존 입력 조작을 중간 단서로 쓰거나, 외부 월드모델을 호출해 상상 이미지를 얻었다. 계획 과제에서는 상태 예측과 행동 계획이 뒤섞여 성능을 분리할 수 없다. 이 논문은 VLM 자체가 중간 시각 상태를 예측하는 내부 월드모델을 갖추고, 그 예측을 실제 추론에 쓰도록 학습될 수 있는지 묻는다.

어떻게 동작하나

구조는 Mixture-of-Transformers(MoT)를 따른다. 이해 브랜치(understanding branch)는 사전학습된 VLM인 Qwen2.5-VL-7B-Instruct로 초기화하고, 생성 브랜치(generation branch) 가중치는 대응하는 이해 브랜치 레이어에서 가져와 초기화한다. 토큰은 모달리티에 따라 라우팅된다. 텍스트와 깨끗한 이미지 토큰은 이해 브랜치로, 노이즈 이미지 토큰은 생성 브랜치로 간다. 생성 브랜치가 노이즈 토큰을 깨끗한 토큰으로 바꾸면 각 노이즈 토큰은 대응하는 깨끗한 토큰으로 교체되고, 어텐션은 두 브랜치의 토큰에 걸쳐 전역으로 계산된다. 경량 MoT는 이해 브랜치를 전체 깊이로 유지하되 생성 브랜치를 얕게 만든다. 기본값은 4개 레이어이며, 이해 브랜치의 연속된 4개 레이어에 정렬해 그 구간에서만 생성 토큰을 처리한다.

무엇과 다른가

추론은 마르코프 과정으로 정식화된다. 시점 j 이전의 추론 상태 H_j는 질문과 지금까지의 텍스트 행동·시각 상태 쌍을 모두 담는다. 정책 모델이 정신적 행동(mental action) T_j를 생성하면, 내부 월드모델이 그 결과 시각 상태를 임베딩 공간에서 예측해 ẑ_j를 내놓는다. 학습 시에는 실제 중간 이미지를 비전 인코더에 통과시킨 E(I_j)가 타깃이 되고, 추론 시에는 월드모델이 임베딩을 직접 생성한다. 이렇게 갱신된 상태가 다음 단계의 텍스트 추론 조건이 되고, n단계 뒤 텍스트 요약과 최종 답을 낸다.

어떻게 쓰나

학습은 두 단계다. 1단계에서는 비전 인코더와 언어 디코더를 포함한 이해 브랜치 전체를 동결하고 생성 브랜치만 학습한다(α=1, β=0). 2단계에서는 비전 인코더만 동결한 채 언어 디코더와 생성 브랜치를 함께 학습한다(α=0, β=1). 생성 브랜치는 rectified flow 목적함수로, 이해 브랜치는 cross-entropy로 최적화하며 전체 손실은 두 항의 가중합이다. 2단계는 온라인 학습으로 진행되어, 이해 브랜치가 정답 시각 토큰이 아니라 생성 브랜치가 만들어낸 토큰을 입력으로 받는다.

전제와 한계

데이터는 프로그램으로 생성한다. 2D 회전 데이터셋 Tetris-2D는 학습 4천 개, 분포 내(ID) 테스트 400개, 분포 외(OOD) 테스트 500개로 구성된다. 3D 회전 데이터셋 Tetris-3D는 학습 1만6천 개에 더해 같은 큐브 모양과 개수가 학습에 나온 쉬운 ID 400개, 개수는 봤지만 모양은 새로운 어려운 ID 400개, 모양이나 개수가 학습에 전혀 없던 OOD 500개를 둔다. 각 샘플은 질문, 질문 이미지, 텍스트 행동과 시각 상태가 교차하는 추론 트레이스, 최종 답을 담고 있으며 중간 상태와 답 모두 검증 가능하다. 상태 예측과 계획이 결합된 ThinkMorph-SpatialNavigation도 학습에 쓰고, VSP 벤치마크의 미로 내비게이션으로 평가한다. 베이스라인은 동일 스텝으로 학습한 Qwen2.5-VL-7B-Instruct SFT, 이산 시각 토큰을 자기회귀 생성하는 LatentUM, 연속 시각 토큰을 쓰는 Mirage, 픽셀을 flow-matching으로 생성하는 ThinkMorph(BAGEL 기반)다. 실험은 H200 8장에서 수행했다.

결과는 Tetris-2D-ID에서 SFT 대비 39.25%p, Tetris-2D-OOD에서 28.8%p 향상이다. 더 어려운 Tetris-3D에서는 SC-ID 19.25%p, C-ID 21.75%p, OOD 16.3%p 앞선다. 학습 곡선에서는 약 2만5천 스텝 부근에 성능 전환이 나타나고, 생성된 시각 상태와 정답 임베딩의 코사인 유사도도 같은 지점에서 빠르게 올라간다. SFT는 그 지점까지는 비슷하게 따라오다가 이후 정체한다. 저자들은 생성 능력이 일정 임계치를 넘어야 월드모델이 실제 이득을 준다고 해석한다.

절제 실험은 이득이 생성된 시각 상태에 실제로 의존함을 보인다. 생성 임베딩으로 정답 임베딩을 top-1 검색했을 때 검색률이 높을수록 정답률이 높았다. 생성 시각 토큰을 전부 0으로 채우거나 무작위 임베딩으로 바꾸면 정확도가 크게 떨어졌고, 토큰 순서를 섞는 것은 영향이 작았다. 생성 브랜치에서 토큰들이 양방향 어텐션을 하고 위치 정보를 이미 갖기 때문이다. 512×512 이미지 기준 기본 시각 토큰 수는 361개(19×19)인데, 토큰을 줄이면 재구성 MSE는 낮아지고 코사인 유사도는 높아지지만 ID 정확도는 토큰이 많을수록 올랐고 OOD 정확도는 4개(2×2)에서 최고였다. 2단계 손실 비율은 α=0.1, β=0.9일 때 전체 성능이 가장 좋았으며, flow 항 가중치를 키우면 토큰 품질은 좋아지지만 다운스트림 성능이 비례해 오르지는 않았다. 두 목적을 처음부터 함께 학습하는 단일 단계 방식은 정확도가 랜덤 수준에 머물렀다. 생성 레이어 배치는 1개층일 때 이해 브랜치 최상위 레이어(23-27)에 정렬한 경우가 가장 좋았고, 4개층으로 늘리면 배치 간 격차가 줄었다. 이해 브랜치와 같은 깊이를 쓰는 원래 MoT 설정은 유의미한 이득이 없었고, 생성 픽셀보다 생성 시각 토큰이 같은 연산 예산에서 더 나은 성능을 냈다.

개발자 관점에서 이 논문이 주는 신호는 명확하다. 중간 시각 상태를 텍스트 추론에 끼워 넣는 것만으로는 부족하고, 생성 품질이 임계치를 넘게 만드는 1단계와 그것을 실제로 쓰게 만드는 2단계를 분리해야 한다. 실무에서 멀티모달 추론 파이프라인을 설계할 때는 세 가지를 확인해야 한다. 중간 시각 상태를 검증 가능한 형태로 만들 수 있는 과제인지, 생성 브랜치를 얕게 붙여도 성능이 유지되는지, 그리고 추론 시 생성 토큰을 제거하거나 교란했을 때 성능이 실제로 떨어지는지다. 마지막 항목이 없으면 모델이 시각 상태 없이 지름길로 답하는 것일 수 있다.

이 발췌 범위에서 저자들은 별도의 한계 절을 두지 않았지만, 본문 곳곳에 전제가 명시돼 있다. 단일 단계 공동 학습의 실패는 "테스트한 설정에서"라는 조건이 붙고, 이득은 생성된 시각 상태에 의존하며 제거·손상 시 성능이 급락한다. 실험은 프로그램으로 합성한 통제된 진단 데이터셋인 Tetris-2D·3D 중심이라 자연 이미지의 공간 추론으로 일반화되는지는 검증되지 않았다. 관련 연구에서 잠재 시각 토큰을 제거해도 영향이 적다는 반론이 제기된 바 있는데, 이 논문은 그와 달리 의존성이 관측됐다고 보고한다. 또한 OOD에서 시각 토큰 수가 4개일 때 최고 성능을 보이는 등 토큰 수와 성능의 관계가 ID와 OOD에서 다르게 나타나 해석에 주의가 필요하다.