Looped World Models

H.A. Lu, Z.L.V. Wei, Q. Zhang, J. Zeng, et al., "Looped World Models," arXiv:2606.18208, 2026.

저자

Hongyuan Lu와 Z.L. Victor Wei가 리딩 컨트리뷰터를 맡았습니다. Lu는 홍콩중문대(CUHK) 박사과정 학생이자 FaceMind Research Asia의 공동창업자로, 대형 언어 모델의 텍스트 빈도 법칙(Adam's Law)과 멀티링궐 추론을 연구해온 인물입니다. 시니어 저자 Wai Lam은 CUHK 시스템공학·공학경영학과 교수이자 CUHK Text Mining Group 책임자로, NLP·텍스트 마이닝 분야에서 오랜 이력을 가집니다. 30명 이상의 저자 목록은 FaceMind가 산학 협력 형태의 대규모 연구 조직임을 보여줍니다.

배경

세계 모델(World Model)은 에이전트가 보지 않은 미래 상태를 예측하고 시뮬레이션하는 핵심 구성 요소입니다. PlaNet과 Dreamer 계열이 순환 상태 공간 모델(RSSM)로 픽셀 기반 환경 예측을 개척했고, IRIS와 TransDreamer는 트랜스포머를 도입해 장거리 의존성을 처리했습니다.

그러나 공통적인 딜레마가 있습니다. 긴 지평(long-horizon) 롤아웃에서 충실한 예측을 생성하려면 깊은 연산이 필요하지만, 모델이 깊어질수록 파라미터 수가 폭발하고 배포 비용이 올라갑니다. 기존 접근은 모델 크기와 데이터 양이라는 두 축으로만 스케일링을 이해했습니다.

이 논문은 세 번째 축을 제안합니다: 잠재 깊이(iterative latent depth), 즉 추론 시 몇 번 반복하느냐입니다. 파라미터 공유 루프 아키텍처는 NLP에서 Universal Transformer, ALBERT, Ouro 등으로 파라미터 효율성을 입증했지만 세계 모델에 적용된 사례가 없었습니다. LoopWM은 이 공백을 채웁니다.

아키텍처

전체 구조

LoopWM은 4개 모듈로 구성됩니다.

• \(E_\varphi(o_k)\): 관측 인코더. 환경 관측 \(o_k\)를 잠재 임베딩 \(e_k \in \mathbb{R}^d\)로 변환합니다.
• \(A_\psi(a_k)\): 액션 임베더. 행동 \(a_k\)를 잠재 벡터 \(u_k \in \mathbb{R}^d\)로 변환합니다.
• \(L_\theta(h_{k-1}, e_k, u_k)\): 루프드 다이나믹스 코어. 이전 상태 \(h_{k-1}\)와 현재 관측·행동을 받아 \(T\)번 반복 정제 후 \(h_k\)를 출력합니다.
• \(D_\xi(h_k)\): 예측 헤드. 잠재 상태 \(h_k\)에서 다음 관측 \(\hat{o}_{k+1}\), 보상 \(\hat{r}_k\), 지속 신호 \(\hat{c}_k\)를 생성합니다.

전체 순전파는 다음과 같습니다.

\[e_k = E_\varphi(o_k),\quad u_k = A_\psi(a_k),\quad h_k = L_\theta(h_{k-1}, e_k, u_k),\quad (\hat{o}_{k+1}, \hat{r}_k, \hat{c}_k) = D_\xi(h_k)\]

루프드 다이나믹스 코어

다이나믹스 코어 \(L_\theta\)는 세 블록으로 나뉩니다.

프렐류드 \(\mathcal{P}\): \(L_P\)개 트랜스포머 레이어로 이전 잠재 상태와 현재 관측·행동 임베딩을 합쳐 컨디셔닝 신호를 만듭니다.

\[e = \text{LN}(\mathcal{P}([h_{k-1}; e_k; u_k])) \in \mathbb{R}^d\]

순환 블록 \(\mathcal{R}\): \(L_R\)개 파라미터 공유 트랜스포머 레이어를 \(T\)번 반복 적용합니다. 각 반복 \(t\)에서의 갱신 규칙은 다음과 같습니다.

\[h^{(t+1)} = \bar{A} h^{(t)} + \bar{B} e + \bar{R}(h^{(t)}, e)\]

\(\bar{A}\)는 이전 잠재 상태가 얼마나 보존되는지를, \(\bar{B}\)는 컨디셔닝 신호의 영향을 제어합니다. \(\bar{R}\)은 멀티헤드 어텐션과 피드포워드 레이어를 포함하는 비선형 트랜스포머 연산입니다. 결정적으로, \(\mathcal{R}\)의 파라미터는 모든 \(T\)번 반복에서 공유됩니다. 계산 깊이가 파라미터 수와 독립적인 이유입니다.

코다 \(\mathcal{C}\): \(L_C\)개 레이어(비공유)로 최종 잠재 상태 \(h^{(T)}\)를 처리해 \(h_k\)를 출력합니다.

스펙트럼 안정성 제약

루프 횟수 \(T\)가 늘어날수록 잠재 상태가 발산하지 않도록 \(\bar{A}\)의 스펙트럼 놈을 1 미만으로 제한합니다. Parcae(Prairie et al., 2026)를 따라 연속시간 음의 대각 행렬을 이산화합니다.

\[A := \text{diag}(-\exp(a)), \quad a \in \mathbb{R}^d \text{ (학습 가능)}\]

\[\bar{A} = \exp(\Delta \cdot A), \quad \Delta \in \mathbb{R}^d_{>0} \text{ (학습 가능)}\]

\(A\)의 대각 원소가 모두 음수이므로 \(\Delta \cdot A\)의 원소도 모두 음수이고, \(\exp\)를 취하면 \((0, 1)\) 구간에 들어옵니다. 따라서 \(\bar{A}\)는 대각 원소가 모두 \((0, 1)\)인 대각 행렬이 되어 \(\rho(\bar{A}) < 1\)이 구성 자체로 보장됩니다. 그레디언트 클리핑이나 후처리 정규화 없이도 롤아웃이 발산하지 않습니다.

변수 깊이 학습

학습 시 루프 횟수 \(T\)를 학습 가능한 평균 \(\mu_\text{rec}\)을 갖는 포아송 분포에서 샘플링합니다.

\[T \sim \text{Poisson}(\mu_\text{rec})\]

기존 연구들이 미니배치 단위로 \(T\)를 샘플링했던 것과 달리, 여기서는 마이크로배치 내 각 시퀀스마다 독립적으로 샘플링합니다. 이 방식이 학습 손실의 분산을 줄이고 손실 스파이크를 제거합니다. 역전파는 메모리 비용 절감을 위해 \(\mu_\text{bwd} = \lceil\mu_\text{rec}/2\rceil\)번만 수행합니다.

세계 모델 손실은 관측·보상·지속 예측을 결합합니다.

\[\mathcal{L}_\text{wm} = \mathbb{E}_{T \sim \text{Poisson}(\mu_\text{rec})}\left[\sum_{k=1}^{K} \mathcal{L}_\text{obs}(o_{k+1}, \hat{o}_{k+1}) + \lambda_r \mathcal{L}_\text{rew}(r_k, \hat{r}_k) + \lambda_c \mathcal{L}_\text{cont}(c_k, \hat{c}_k)\right]\]

조기 종료

추론 시 각 루프 반복 \(t\)마다 경량 MLP 게이트로 종료 여부를 판단합니다.

\[g^{(t)} = \sigma\!\left(w_g^\top h^{(t)} + b_g\right)\]

\(g^{(t)}\)가 임계값 \(\tau\)를 초과하면 루프를 조기 종료하고 \(h^{(t)}\)를 최종 잠재 상태로 사용합니다. 단순한 전이는 루프 1~2회만 돌고, 복잡한 전이는 더 많은 반복을 할당받습니다. 100레이어 고정 깊이 기준선과 비교했을 때, 단순 전이 구간에서는 최대 \(25\times\) 추론 FLOP 절감이 가능합니다.

학습 시에는 게이트가 극단적 해(항상 첫 반복에서 종료하거나 절대 종료하지 않음)로 수렴하지 않도록 엔트로피 정규화를 추가합니다.

\[\mathcal{L}_\text{ent} = -\alpha \, \mathbb{E}\!\left[\sum_{t=1}^{T} H\!\left(g^{(t)}\right)\right]\]

LoopWM-DD: 지연 디코딩

표준 세계 모델은 매 환경 스텝 \(k\)마다 디코더 \(D_\xi\)를 호출합니다. 이 방식은 두 가지 비효율을 만듭니다: (i) 중간 스텝에서도 픽셀 수준 재구성을 위한 표현 용량을 소비하고, (ii) 다단계 액션 시퀀스에서 잠재 추론이 중단됩니다.

LoopWM-DD는 \(K\)단계 롤아웃 전체를 잠재 공간에서만 수행하고 디코더를 마지막 스텝에서만 한 번 호출합니다.

\[u_k = A_\psi(a_k), \quad h_{k+1} = L_\theta^\text{core}(h_k, u_k) \quad (k = 0, \ldots, K-1)\]

\[(\hat{o}_K, \hat{r}_K, \hat{c}_K) = D_\xi(h_K)\]

전체 유효 깊이는 \(K \times T\) 파라미터 공유 트랜스포머 적용이지만, 디코더는 단 한 번만 통과합니다.

중간 디코딩 지도 없이 잠재 상태가 의미 있는 정보를 유지하도록 두 가지 제약을 추가합니다.

잠재 일관성 손실: 중간 관측 \(o_k\)를 동결 인코더로 처리해 얻은 참조 임베딩 \(e_k^* = \text{sg}(E_\varphi(o_k))\)에 잠재 궤적을 정렬합니다.

\[\mathcal{L}_\text{consist} = \frac{1}{K-1}\sum_{k=1}^{K-1}\left\|g_\omega(h_k) - e_k^*\right\|_2^2\]

스펙트럼 수축 예산: 잠재 상태가 \(K\)단계에 걸쳐 발산하지 않도록 누적 변화량에 상한 패널티를 부과합니다.

학습은 \(K = 1\)에서 시작해 단계적으로 지평을 늘리는 커리큘럼을 따릅니다. \(K \times T\)번의 파라미터 공유 적용을 거치는 역전파가 처음부터 긴 지평에서는 불안정하기 때문입니다.

결과

평가는 텍스트 기반 세계 모델링 벤치마크 두 개에서 이루어졌습니다. ScienceWorld는 과학 개념이 포함된 텍스트 환경에서 5개 연속 행동을 입력받아 최종 상태를 예측하는 과제이고, AlfWorld는 가정 환경에서의 텍스트 탐색·조작 과제입니다. LoopWM은 약 1B 파라미터 규모입니다.

ScienceWorld 주요 결과 (14개 태스크 평균)

LoopWM은 ScienceWorld 전체에서 EM 기준 claude-opus-4-6-max 대비 +21.2%p를 기록합니다. Lifespan 태스크에서는 claude-opus의 0%를 100%로 끌어올렸고, Boil·Conductivity에서도 격차가 뚜렷합니다. 단, LifeStages 태스크는 EM 0%로 가장 취약한 지점으로 남습니다.

AlfWorld 주요 결과

AlfWorld에서는 양상이 다릅니다. EM 기준으로는 claude-opus-4-6-max(53.0%)가 LoopWM(51.6%)을 소폭 앞서고, Entity 점수는 gemini가 90.2%로 가장 높습니다. LoopWM은 Token F1과 BLEU-4에서는 1위이지만 Entity 예측에서는 분명한 약점을 보입니다. 저자들은 이 엔티티 점수 개선이 향후 최적화의 핵심 방향이라고 밝혔습니다.

지연 디코딩의 누적 효과

지연 디코딩의 이점은 롤아웃이 쌓일수록 커집니다. ScienceWorld에서 LoopWM-DD와 gemini-3-flash-preview-thinking의 EM 격차(Table 5)는 Step 1(+73.2%), Step 3(+103.6%), Step 5(+113.8%)로 스텝이 길어질수록 확대됩니다. 단순히 1스텝 예측을 잘하는 게 아니라, 중간 디코딩 없이 잠재 공간에서 연속 추론을 수행했을 때 더 정확한 표현을 유지한다는 증거입니다.

회고

저자들이 논문에서 직접 인정한 한계 사항들입니다.

첫째, 제시 범위의 한계입니다. 저자들은 더 폭넓은 실험 결과를 이미 확보했지만 이 버전에서는 핵심 아키텍처 논제 정립에 집중하며 일부만 공개했다고 밝혔습니다. 연속 시각 환경에서의 훈련 가능성은 내부적으로 확인했지만 본문에는 포함되지 않았습니다.

둘째, 스케일링 법칙 미완성입니다. Step 1~5 실험이 잠재 깊이가 의미 있는 스케일링 차원임을 보여주지만, 더 넓은 태스크·컴퓨팅 범위에 걸친 완전한 스케일링 법칙 분석은 미완입니다.

셋째, 포지셔닝 분석 부재입니다. RSSM 계열, 자기회귀 비디오 토큰 세계 모델, 확산 기반 세계 모델과의 더 명시적인 비교 분석이 필요합니다. 논문 자체도 이 부분의 보강 가치를 인정합니다.

넷째, Entity 점수 약점입니다. AlfWorld에서 Entity 예측이 다른 메트릭 대비 뒤처지는 패턴이 일관되게 나타나며, 저자들은 이를 미래 최적화 방향으로 지목했습니다.

이 외에 외부 관점에서 주목할 점: 1B 파라미터 모델이 100B+ 규모의 범용 LLM을 특화 세계 모델링 과제에서 앞서는 것은 전문화 학습의 효과이지, 아키텍처 자체의 우위만으로 설명되지는 않습니다. "100× 파라미터 효율"이라는 주장은 이 맥락에서 읽어야 합니다.

정리

• LoopWM은 파라미터를 공유하는 트랜스포머 블록을 반복 적용해 잠재 환경 상태를 정제하는 세계 모델입니다. 스펙트럼 제약(\(\rho(\bar{A}) < 1\))이 구성 자체로 롤아웃 안정성을 보장하고, 조기 종료 게이트가 예측 복잡도에 따라 연산량을 자동 조절합니다.
• LoopWM-DD는 \(K\)단계 롤아웃 전체를 잠재 공간에서만 수행하고 마지막에 한 번만 디코딩합니다. 이 방식의 이점은 롤아웃 길이가 늘어날수록 커집니다.
• 이 논문은 잠재 깊이를 모델 크기·데이터 양과 독립적인 세 번째 스케일링 축으로 제안합니다. 구체적인 스케일링 법칙은 후속 연구 과제로 남습니다.