🤖 AI

Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL

이 논문은 Masked Diffusion Language Models (MDLM)를 에이전트 강화 학습을 위한 우수하고 조종 가능한 텍스트 기반 월드 모델로 소개하며, 이는 대규모 데이터셋과 새로운 GRPO 훈련 프레임워크에 의해 뒷받침되어 더 큰 규모의 LLM들과 비교했을 때 자기회귀적 편향을 극복하고 현저히 높은 롤아웃 다양성과 제로샷 전이 성능을 달성합니다.

원저자: Darshan Deshpande

게시일 2026-07-21

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Darshan Deshpande

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 로봇에게 복잡한 비디오 게임(예를 들어, 거대하고 변화무쌍한 도시를 탐험하거나 바쁜 커피숍을 운영하는 것)을 가르치려 한다고 상상해 보세요. 학습을 위해 로봇은 수백만 번의 연습이 필요합니다. 현실 세계라면 매 연습 실행마다 물리적인 커피숍을 새로 지어야 할 텐데, 이는 불가능한 일입니다. 그래서 과학자들은 로봇이 무언가를 망가뜨리지 않고도 시도하고, 실패하고, 배울 수 있는 디지털 세계인 '시뮬레이터'를 사용합니다. 이것이 바로 **강화 학습(Reinforcement Learning, RL)**의 핵심입니다. 즉, 샌드박스 안에서 시행착오를 통해 배우는 AI입니다.

하지만 여기 문제가 있습니다. 이러한 샌드박스를 구축하는 것은 매우 어렵습니다. 보통 인간이 모든 규칙과 시나리오를 직접 설계해야 하는데, 이는 속도가 느릴 뿐만 아니라 로봇이 접할 수 있는 상황의 다양성을 제한합니다. 만약 로봇이 특정 시나리오에 너무 익숙해지면, 실제로 사고하는 법을 배우는 대신 정확한 단계를 암기하여 '부정행위'를 하기 시작합니다. 이는 마치 연습 시험의 답을 통째로 외웠지만, 실제 시험에서 질문이 조금만 달라져도 낙제하는 학생과 같습니다. 이를 해결하기 위해 연구자들은 스스로 디지털 세계를 상상하고 시뮬레이션할 수 있는 '월드 모델(World Models)'을 구축하려고 노력하고 있습니다. 즉, AI가 스스로 다양한 연습 시나리오를 끊임없이 만들어내도록 하는 것입니다. 여기서 핵심적인 질문은, AI 시뮬레이터가 복잡하고 규칙이 많은 환경에서 다음에 무슨 일이 일어날지를 혼란 없이 예측할 만큼 충분히 똑똑해질 수 있는가 하는 점입니다.

이 논문은 새로운 종류의 AI 시뮬레이터인 **마스크 확산 언어 모델(Masked Diffusion Language Model, MDLM)**을 소개하며, 이 모델이 우리가 흔히 사용하는 표준 AI 모델보다 훨씬 더 뛰어난 '게임 마스터'라고 주장합니다.

문제점: 일방통행로 vs 양방향 통행로

이것이 왜 중요한지 이해하려면, 표준 AI(이를 자기회귀(Autoregressive, AR) 모델이라고 부릅니다)가 어떻게 생각하는지 그려보십시오. 그것은 마치 사람이 왼쪽에서 오른쪽으로 한 단어씩 이야기를 써 내려가는 것과 같습니다. 일단 단어를 쓰면 다시 돌아가서 수정할 수 없습니다. 만약 "은행이 카드를 거절했습니다"라고 썼다면, 그 문장에 갇히게 됩니다. 나중에 "은행이 카드를 승인했습니다"라고 썼어야 했다는 것을 깨닫더라도 이미 늦었습니다. 이야기는 이미 망가졌습니다. 복잡한 시뮬레이션에서 이는 **전역적 비일관성(global incoherence)**이라는 문제를 일으킵니다. AI는 완벽한 도입부를 작성할 수는 있지만, 결말에 다다랐을 때 세부 사항이 세상의 규칙과 맞지 않을 수 있습니다. 이는 마치 범인의 이름은 잊어버린 채 사건을 해결한 탐정이나, 요리를 다 해놓고 가스불 끄는 것을 잊어버린 요리사와 같습니다.

저자들은 이러한 표준적인 '왼쪽에서 오른쪽으로' 가는 방식이 병목 현상을 일으킨다고 제안합니다. 대신 그들은 **마스크 확산(Masked Diffusion)**이라는 다른 접근 방식을 제안합니다. 글을 처음부터 끝까지 쓰는 대신, 페이지의 일부 단어가 검은 상자(마스크) 뒤에 숨겨져 있다고 상상해 보십시오. 당신의 임무는 그 숨겨진 단어들을 추측하는 것입니다. 하지만 여기에는 비결이 있습니다. 당신은 상자 앞의 단어와 상자 뒤의 단어를 모두 보고 추측할 수 있습니다. 전체 그림을 가이드 삼아 중간, 끝, 혹은 시작 부분을 동시에 채워 넣을 수 있습니다. 이를 통해 AI는 시작 부분이 끝 부분과 일치하는지 확인할 수 있으며, 전체 이야기가 서로 조화를 이루도록 보장할 수 있습니다.

실험: 더 나은 샌드박스 구축하기

연구진은 단순히 말로만 설명하지 않고 직접 구현했습니다. 그들은 코딩 도구부터 고객 서비스 시나리오에 이르기까지 9가지 서로 다른 환경에 걸친 239,403개의 방대한 '연습 실행(궤적)' 데이터셋을 구축했습니다. 그들은 이 데이터로 새로운 마스크 확산 모델을 훈련시킨 후, 현존하는 최고의 '왼쪽에서 오른쪽으로' 방식의 표준 모델들과 대결시켰습니다.

결과는 놀라웠습니다. 마스크 확산 모델은 훨씬 더 작았음에도 불구하고(일부는 단 80억 개의 파라미터였던 반면, 이들이 이긴 표준 모델은 300억 개가 넘었습니다), 훨씬 더 일관적이고 현실적인 시뮬레이션을 만들어냈습니다.

이것을 비유하자면, 표준적인 거대 모델은 한 번에 한 글자씩 타이핑하는 매우 빠른 타자수와 같습니다. 초반에 오타를 내면 문장 전체가 망가집니다. 반면, 새로운 마스크 확산 모델은 찰흙 덩어리 전체를 한꺼번에 바라보며, 최종 조각상이 모든 각도에서 제대로 보이도록 여러 곳을 동시에 깎아 나가는 조각가와 같습니다.

결과: 더 똑똑한 연습, 더 뛰어난 플레이어

진정한 테스트는 이 시뮬레이터를 사용하여 실제 AI 에이전트(학습하려는 '로봇')를 훈련시킬 때 이루어졌습니다. 그들은 마스크 확산 모델이 만든 세상에서 에이전트들이 연습하게 한 다음, 완전히 새로운 환경(ScienceWorld, ALFWorld, AppWorld 등)에서 테스트했습니다.

결과는 매우 인상적이었습니다. 마스크 확산 시뮬레이터로 훈련된 에이전트들은 기존 표준 모델로 훈련된 에이전트들에 비해 성공률이 최대 **47%**까지 향상되었습니다. 이는 연구자들이 해당 새로운 환경에 대해 구체적으로 가르치지 않았음에도 불구하고 달성한 성과였습니다. 에이전트들은 연습 환경이 더 다양하고 논리적이었기 때문에 더 나은 일반적 기술을 자연스럽게 습득한 것입니다.

예를 들어, 복잡한 과업에서 수행 능력이 떨어졌던(성공률 단 5.7%) 12억 개의 파라미터를 가진 작은 AI 에이전트가, 새 시뮬레이터로 훈련한 후에는 **53.6%**의 성공률로 급등했습니다. 이는 엄청난 도약이며, 시뮬레이터가 에이전트가 적응하는 법을 배울 수 있도록 훨씬 더 풍부하고 다양한 도전 과제를 제공했음을 시사합니다.

이것이 중요한 이유

이 논문은 AI 시뮬레이터를 구축하는 방식이 바뀌어야 한다고 제안합니다. 기존의 '왼쪽에서 오른쪽으로 쓰는' 방식은 한계에 부딪혀, AI가 반복적인 루프에 빠지거나 시뮬레이션을 망가뜨리는 논리적 오류를 범하게 만듭니다. '전체 그림을 한꺼번에 보는 방식(양방향 디노이징)'으로 전환함으로써, 우리는 더 똑똑할 뿐만 아니라 더 효율적인 시뮬레이터를 만들 수 있습니다.

저자들은 또한 전문가들에게 시뮬레이션을 평가하게 하는 '인간 검증'을 실시했습니다. 마스크 확산 모델은 현실성(5점 만점에 4.75점)과 정확성(4.25점)에서 매우 높은 점수를 받았으며, 이는 인간이 시뮬레이션된 세계가 믿을 만하고 유용하다고 느꼈음을 의미합니다.

하지만 논문은 이것이 아직 완벽하게 완성된 제품은 아니라는 점을 주의 깊게 언급합니다. 새로운 모델은 여전히 올바른 API 키를 생성하거나, 아주 긴 아이템 목록을 처리할 때 혼란을 겪는 등 특정 세부 사항에서 어려움을 겪기도 합니다. 그러나 핵심적인 발견은 명확합니다. AI가 미래를 생각하는 방식—일방통행로에서 양방향 통행로로—을 바꿈으로써, 우리는 더 나은 훈련장을 만들 수 있습니다. 이는 소프트웨어 버그를 수정하는 것부터 고객 서비스를 관리하는 것까지, 복잡한 업무를 수행할 수 있는 차세대 AI 에이전트가 더 빠르게, 더 똑똑하게 발전할 수 있음을 의미하며, 이는 그들이 더 논리적인 세상에서 연습했기 때문입니다.

기술 요약: 조종 가능한 텍스트 기반 월드 모델로서의 마스크드 디퓨전 언어 모델

문제 정의

에이전트 시스템을 위한 강화 학습(RL)에는 다양하고 적응적인 훈련 환경이 필요합니다. 그러나 현재의 환경은 고정된 작업 및 보상 난이도를 가진 수동 제작 방식에 의존하는 경우가 많아, 모델이 개선됨에 따라 효과적인 훈련 신호가 감소하는 문제를 겪습니다. 에이전트가 롤아웃 과정에서 더 다양한 상황에 직면할수록, 특정 워크플로우나 도구 구조에 과적합되는 모드 붕괴(mode collapse) 현상이 발생합니다. 월드 모델(미래 상태를 예측하는 생성 모델)은 훈련 데이터를 다양화하기 위한 확장 가능한 해결책을 제시하지만, 월드 모델로 사용되는 기존의 자기회귀(autoregressive, AR) 언어 모델은 근본적인 좌에서 우로 향하는 편향(left-to-right bias) 문제를 가집니다. 이러한 인과적 제약은 전역적으로 상호 의존적인 상태 앵커(예: 도구 스키마, 이전 턴, 예상 결과)에 동시에 조건화하는 것을 방해하며, 결과적으로 접두사(prefix)는 일관되지만 전체적으로는 부조화스러운 출력(예: 성공적인 환불 액션과 에러 상태가 함께 나타나는 경우)을 생성하게 만듭니다.

방법론

1. 텍스트 기반 월드 모델링의 정식화

저자들은 월드 모델링을 조종 가능한 전이 역학(steerable transition-dynamics) 문제로 정식화합니다. 이들은 환경 상태 $s_t$ 를 다섯 가지의 근거 있는 구성 요소로 분해합니다:

$e_0$ : 초기 환경 상태 (파일, 데이터베이스, UI).
$c_t$ : 작업 컨텍스트 (목표, 히스토리).
$\mathcal{T}$ : 도구 스키마 (시그니처 및 제약 조건).
$\mathcal{R}$ : 도메인 규칙 (유효한 전이).
$D$ : 조종 지시문 (안전 제약, 상위 수준 지침).

목표는 $\mathcal{T}, \mathcal{R}, D$ 가 궤적 전체에 걸쳐 고정된 상태에서 분포 $p_\theta(c_{t+1} | e_0, c_{\le t}, a_{\le t}; \mathcal{T}, \mathcal{R}, D)$ 를 학습하는 것입니다.

2. 마스크드 디퓨전 언어 모델 (MDLMs)

본 논문은 AR 생성 대신 **마스크드 디퓨전 언어 모델(MDLMs)**을 활용합니다.

메커니즘: MDLMs는 손상된 토큰 시퀀스를 반복적으로 디노이징(denoising)하는 법을 학습합니다. 훈련 중에 깨끗한 다음 상태 시퀀스는 확률 $\tau$ 로 마스킹되며, 모델은 히스토리 및 마스킹되지 않은 컨텍스트(미래 앵커 포함)에 조건화되어 마스킹된 토큰을 예측합니다.
AR 대비 장점: AR 모델이 좌에서 우로 향하는 순서에 의존해야 하는 것과 달리, MDLMs는 양방향 디노이징을 수행합니다. 이들은 상태 내의 이전 및 이후 컨텍스트 모두에 조건화될 수 있습니다 (예: 알려진 status: "error" 앵커를 사용하여 error_code와 resolution 생성을 유도). 이를 통해 병렬 토큰 생성이 가능해지며, 접두사로 인한 오류 누적을 방-지합니다.
다양성: 디노이징 과정은 토큰 선택과 생성 순서 모두에 확률성을 도입하여, AR 디코더에서 흔히 발생하는 "접두사 모드 붕괴(prefix mode collapse)"를 방지합니다.

3. 데이터셋 및 훈련 프레임워크

데이터셋: 9개의 오픈 소스 환경(예: SWE-SMITH, TAUBENCH, TOOLATHLON)과 12개의 프런티어 모델 제품군(Qwen, GPT, Claude, Gemini 포함)을 아우르는 239,403개의 근거 있는 상태-액션 궤적으로 구성된 큐레이션된 데이터셋을 사용합니다.
훈련: 모델은 코사인 스케줄을 가진 AdamW을 사용하여 미세 조정되었습니다. 저자들은 결정론적 상태 체크를 갖춘 플러그 앤 플레이 GRPO (Group Relative Policy Optimization) 훈련 프레임워크를 도입하여, MDLM이 생성한 롤아웃 상에서 RL 에이전트를 훈련합니다.

주요 기여

정식화 및 데이터셋: 텍스트 기반 월드 모델링을 다섯 가지 특정 구성 요소를 가진 조종 가능한 전이 문제로 체계적으로 정식화한 최초의 연구이며, 239k 궤적을 포함한 대규모 데이터셋을 제공합니다.
비교 연구: AR LM과 MDLMs 간의 포괄적인 비교를 통해, MDLMs가 양방향 앵커 인식 디노이징을 통해 우수한 일관성과 근거성을 달성함을 입증했습니다.
성능 향상: MDLMs(특히 SDAR-8B와 같은 8B 파라미터 모델)가 월드 모델링 충실도 측면에서 4배 더 많은 파라미터를 가진 AR 모델(예: Qwen-3.5-35B)을 능가함을 보여줍니다.
제로샷 전이(Zero-Shot Transfer): MDLM이 생성한 롤아웃을 사용하여 환경별 미세 조정 없이 분포 외(OOD) 환경(ScienceWorld, ALFWorld, AppWorld)에서 RL 에이전트를 성공적으로 훈련했습니다.
인간 평가: 현실성, 정확성, 유용성에 대한 엄격한 인간 평가를 통해 MDLM 롤아웃의 신뢰성을 확인했습니다.

결과

생성 충실도 (RQ1)

지표: BLEU-1, ROUGE-L, MAUVE(분포 정렬)를 사용하여 평가했습니다.
결과: 제로샷 설정에서 MDLMs는 분포 정렬(MAUVE) 측면에서 AR 대응 모델보다 우수했습니다. 미세 조정 후에는 MDLMs가 모든 분할에서 압도했습니다.
- SDAR-8B(8B 파라미터)는 인도메인 MAUVE(0.982 vs 0.932) 및 OCCUBENCH MAUVE(0.979 vs 0.960)에서 Qwen-3.5-35B-A3B(35B 파라미터)를 능가했습니다.
- 이는 모델 용량이 아닌 **인과적 귀납 편향(causal inductive bias)**이 LLM 기반 월드 모델링의 주요 병목 구간임을 시사합니다.

다운스트림 RL 성능 (RQ2)

설정: Qwen 기반(AR) 또는 SDAR 기반(MDLM) 월드 모델의 롤아웃을 사용하여 에이전트(LFM2.5-1.2B, Qwen3-4B, Mistral-7B)를 GRPO로 훈련했습니다.
결과: SDAR-WM 롤아웃으로 훈련했을 때 베이스라인 및 Qwen-WM 롤아웃 대비 상당한 절대적 이득을 얻었습니다:
- ALFWorld: LFM2.5-1.2B는 SDAR-WM을 사용하여 5.7%(베이스)에서 **53.6%(+47.9% 절대적 증가)**로 향상되었습니다.
- ScienceWorld: Mistral-7B는 3.3%에서 **48.4%(+45.1% 절대적 증가)**로 향상되었습니다.
- 이러한 개선은 테스트 환경에 대한 제로샷 설정에서 발생했습니다.

인간 평가 (RQ3)

지표: 현실성, 결과 정확성, 훈련 유용성 (1–5 리커트 척도).
결과: SDAR는 높은 점수를 기록했습니다: 4.75(현실성), 4.25(정확성), 4.50(유용성).
일치도: 모든 차원에서 Krippendorff's $\alpha$ 가 0.89를 초과하여, 강력한 평가자 간 일치도를 나타냈습니다.
실패 사례: 평가자들은 숫자 키 타입 강제 변환, 손상된 API 키, 그리고 초기 실패 후 반복적인 에러로 붕괴되는 경향 등을 지적했습니다.

의의 및 주장

본 논문은 마스크드 디퓨전 언어 모델이 자기회귀 모델에 비해 텍스트 기반 월드 모델링을 위한 우수한 아키텍처임을 주장합니다. 양방향 조건화를 활용함으로써, MDLMs는 도구 스키마 및 도메인 규칙과 전역적 일관성을 유지하며 다양하고 근거 있는 환경 시뮬레이션을 생성할 수 있습니다.

저자들은 이 접근 방식이 에이전트 RL을 위한 확장 가능하고, 다양하며, 조종 가능한 훈련 신호를 가능하게 한다고 단언합니다. 결정적으로, 높은 충실도의 월드 모델링을 위해 거대한 파라미터 수가 반드시 필요한 것은 아니며, 오히려 아키텍처의 선택(디퓨전 vs. 자기회귀)이 핵심적인 요소임을 입증했습니다. 이 연구는 "상상 속의 세계에서의 RL"을 향한 경로를 제공하며, 에이전트가 수동으로 제작된 정적인 훈련장이 아닌, 시뮬레이션된 다양한 환경에서 훈련할 수 있게 합니다.

한계점 인정: 논문은 MDLMs가 유망함에도 불구하고, 여전히 API 키와 같은 구조화된 필드에서 어려움을 겪으며 복잡한 탐색 작업에서 "페이지네이션 드리프트(pagination drift)"를 보인다고 언급했습니다. 이러한 문제는 베이스 모델의 안전 정렬(safety alignment) 및 디퓨전의 블록 크기 제약에 기인한 것으로 보이며, 향후 최적화가 필요한 영역으로 제시되었습니다.