마스크드 확산 언어모델, Agent RL을 위한 텍스트 월드 모델로 부상

도입

Agent 강화학습에서 점점 더 큰 병목은 모델이 아니라 환경이다. 손으로 설계한 환경은 다양성과 난이도에 한계가 있고, 모델 성능이 올라갈수록 고정된 과제와 보상은 충분한 학습 신호를 주지 못한다. 장기 과제에서 보상이 희소하면 Agent가 특정 워크플로우나 도구 사용 구조에만 갇히는 문제도 생긴다. 이런 배경에서 환경 상태를 시뮬레이션하는 월드 모델은 훈련 다양성을 확장하는 유력한 방법으로 다시 주목받고 있다.

이 논문의 핵심 질문은 텍스트 기반 월드 모델이 반드시 자기회귀 언어모델이어야 하느냐는 것이다. 자기회귀 모델은 왼쪽에서 오른쪽으로 토큰을 생성하기 때문에 일반 텍스트 생성에는 자연스럽다. 그러나 Agent 환경에서 다음 상태는 단순한 문장이 아니다. 도구 스키마, 이전 대화, 도메인 규칙, 초기 조건, 기대되는 결과를 동시에 만족해야 하는 제약된 상태다. 저자들은 Masked Diffusion Language Models, 즉 MDLM이 양방향 denoising을 통해 이런 전역 앵커를 더 잘 반영할 수 있다고 본다.

핵심 요점

• 문제의 재정의: 논문은 텍스트 월드 모델링을 조향 가능한 상태 전이 동역학 문제로 정식화한다. 초기 상태, 과제 문맥, 도구 스키마, 도메인 규칙, steering directives로 구성 요소를 나눈다.
• 데이터 구성: 9개의 오픈소스 환경과 12개의 frontier model family에 걸쳐 239,403개의 grounded state-action trajectory를 구축했다.
• 모델 비교: MDLM은 앵커를 고려한 양방향 denoising을 통해 coherence, groundedness, 검증된 rollout diversity에서 자기보다 4배 이상 큰 파라미터 규모의 자기회귀 LLM보다 더 나은 결과를 냈다고 보고된다. 추론 지연은 비슷한 수준으로 제시됐다.
• 훈련 프레임워크: deterministic state checks를 포함한 플러그앤플레이 GRPO 훈련 프레임워크를 도입해, 시뮬레이션 상태가 어긋나는 문제를 줄이려 했다.
• OOD 전이: ScienceWorld, ALFWorld, AppWorld 세 가지 분포 밖 환경에서 LFM2.5, Qwen3, Mistral 등 1.2B~7B Agent 백본을 대상으로 zero-shot ablation을 수행했다. 환경별 추가 미세조정 없이 baseline 대비 최대 47%의 절대 향상이 보고됐다.

의미와 영향

이 연구의 의미는 확산 언어모델이 또 하나의 벤치마크에서 좋았다는 데 그치지 않는다. 더 큰 메시지는 Agent용 월드 모델에서는 생성 방식 자체가 중요하다는 점이다. 환경 상태는 여러 제약을 동시에 만족해야 하며, 순차 생성만으로는 긴 rollout에서 모순이 누적되기 쉽다.

MDLM의 양방향 denoising은 전역 제약을 먼저 유지하면서 세부 상태를 채워 넣는 방식에 가깝다. 따라서 도구 호출, 이전 상태, 결과의 타당성이 중요한 Agent RL에서는 자기회귀 모델보다 더 조향하기 쉬운 시뮬레이터가 될 가능성이 있다.

물론 한계도 있다. 결과는 논문에서 구성한 특정 환경과 궤적 데이터에 기반하며, 실제 도구 생태계나 장시간 실행 과제, 고위험 도메인에서도 같은 장점이 유지되는지는 추가 검증이 필요하다. 그럼에도 이 논문은 텍스트 월드 모델의 기본 선택지를 다시 생각하게 만드는 중요한 시도다.

출처: Hugging Face Daily Papers