From Next Token Prediction to (STRIPS) World Models

Each language version is independently generated for its own context, not a direct translation.

이 논문은 **"인공지능이 단순히 다음 단어를 예측하는 것에서 멈추지 않고, 실제로 세상을 이해하고 계획을 세울 수 있을까?"**라는 질문에 답하기 위해 쓴 연구입니다.

저자들은 인공지능 (특히 '트랜스포머'라는 최신 AI 모델) 이 게임이나 퍼즐 같은 규칙이 명확한 세계를 학습할 때, 단순히 통계적 패턴만 외우는 게 아니라 실제 '세계의 법칙 (World Model)'을 깨우쳐서 새로운 상황에서도 문제를 해결할 수 있는지를 실험했습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

1. 핵심 아이디어: "요리 레시피를 외우는 것 vs 요리의 원리를 이해하는 것"

상상해 보세요. AI 가 '레시피'를 배우는 상황입니다.

기존 AI (단순 예측): "감자를 자르면, 다음에는 양파를 넣어야 해"라고 단순히 다음 단어를 맞추는 것에 집중합니다. 이 AI 는 레시피를 외우기는 하지만, 왜 그 순서인지, 만약 감자가 없으면 어떻게 해야 하는지 같은 원리는 모릅니다.
이 연구의 목표: AI 가 단순히 다음 단어를 맞추는 훈련을 하더라도, 결국 **"감자, 양파, 소금"이라는 재료와 "자르다, 볶다"라는 행동 사이의 논리적 관계 (STRIPS 모델)**를 스스로 깨닫게 해서, 보지 못한 새로운 상황에서도 요리를 할 수 있게 만드는 것입니다.

2. 두 가지 새로운 AI 설계도 (아키텍처)

저자들은 이 목표를 달성하기 위해 두 가지 다른 방식의 AI 를 만들었습니다.

① STRIPS 트랜스포머 (규칙을 딱딱 맞게 설계한 AI)

비유: "수학 공식이 적힌 공책"
이 AI 는 처음부터 논리 규칙을 몸속에 심어두었습니다. 마치 "A 가 있으면 B 가 사라진다"는 식의 수학적 법칙을 하드코딩해 둔 것과 같습니다.
장점: 이론적으로 완벽하게 작동할 수 있습니다.
단점: 너무 딱딱해서 배우기 어렵습니다. 마치 수학 공식을 암기하듯 많은 데이터와 노력이 필요하며, 실수하기 쉽습니다.

② 스틱 브레이킹 (Stick-Breaking) 트랜스포머 (유연하게 배우는 AI)

비유: "유연한 명상가"
이 AI 는 특별한 규칙을 미리 심어두지 않았습니다. 대신, **주의를 기울이는 방식 (어텐션)**을 아주 특별한 방법으로 바꿨습니다.
핵심 기술 (스틱 브레이킹): 과거의 정보를 볼 때, **"가장 최근의 중요한 정보"**에 집중하고 나머지는 잘라버리는 (Stick-breaking) 방식을 사용합니다. 마치 긴 이야기에서 "가장 마지막에 일어난 일"이 현재 상황을 결정한다는 직관을 가진 것과 같습니다.
결과: 놀랍게도 이 AI 가 더 잘 배우고 더 잘 일반화했습니다. 규칙을 강제로 주입하지 않아도 스스로 논리를 찾아냈기 때문입니다.

3. 실험 결과: "보지 못한 세상에서도 길을 찾을 수 있을까?"

연구진은 5 가지 다른 세계 (블록 쌓기, 페리 운송, 퍼즐, 미로, 물류 등) 에서 AI 를 훈련시켰습니다.

기존 AI (일반 트랜스포머): 짧은 길이의 훈련 데이터에서는 잘했지만, 길이가 긴 새로운 상황에서는 길을 잃었습니다. (단순히 외운 것이라서 긴 이야기의 끝을 예측하지 못함)
이 연구의 AI (특히 스틱 브레이킹): 훈련 데이터보다 훨씬 길고 복잡한 상황에서도 완벽하게 작동했습니다.
가장 놀라운 점: AI 가 학습한 내용을 다시 **기존의 고전적인 계획 알고리즘 (STRIPS 플랜너)**에 넣었을 때, 보지 못한 수백만 가지의 새로운 상황에서도 완벽한 해결책을 찾아냈습니다.

4. 결론: 왜 이 연구가 중요한가?

이 논문은 **"AI 가 다음 단어를 예측하는 훈련만으로도, 실제로 세상을 이해하고 논리적으로 사고할 수 있는 '세계 모델'을 만들 수 있다"**는 것을 증명했습니다.

기존의 생각: AI 는 그냥 통계적 확률만 계산하는 '말 잘하는 원숭이'일 뿐이다.
이 연구의 발견: AI 는 훈련을 통해 **실제 세계의 인과관계 (원인과 결과)**를 학습할 수 있으며, 이를 통해 새로운 문제를 스스로 해결할 수 있다.

한 줄 요약:

"AI 에게 단순히 다음 단어를 맞추게 하는 훈련을 시켰더니, AI 가 스스로 '세상의 법칙'을 깨우쳐서 보지 못한 새로운 퍼즐도 척척 해결하는 마법을 부렸습니다. 특히, 복잡한 규칙을 강요하지 않고 유연하게 학습하게 한 방식이 가장 효과적이었습니다."

이 연구는 인공지능이 단순한 챗봇을 넘어, 실제로 복잡한 문제를 해결하고 계획을 세울 수 있는 진정한 '지성'으로 발전할 수 있는 가능성을 보여줍니다.

Each language version is independently generated for its own context, not a direct translation.

이 논문은 **다음 토큰 예측 (Next Token Prediction)**을 통해 **STRIPS 기반의 세계 모델 (World Models)**을 학습할 수 있는지, 그리고 이러한 모델이 실제 **계획 (Planning)**을 지원하는지 검증하는 연구입니다. 저자들은 제어된 기호적 (symbolic) 환경에서 행동 궤적 (action traces) 만으로부터 STRIPS 행동 모델을 학습하고, 그 정확성을 정밀하게 평가하는 두 가지 아키텍처를 제안합니다.

주요 내용은 다음과 같습니다.

1. 연구 문제 및 배경

문제 정의: 최근 대규모 언어 모델 (LLM) 이 다음 토큰을 예측하는 과정에서 세계 모델을 학습하는지 여부에 대한 논의가 활발합니다. 그러나 기존 연구들은 주로 체스나 오델로 같은 게임 환경에서 잠재 상태 표현을 학습하는 데 그쳤으며, 이 표현들이 실제 **계획 (Planning)**을 수행하기에 충분한지 여부는 불명확했습니다.
목표: 명제 논리 기반의 STRIPS 세계 모델을 학습하여, 학습된 모델이 기존 STRIPS 플래너 (off-the-shelf planners) 와 호환되도록 하고, 훈련 시 보지 못한 초기 상태와 목표에 대해 계획이 가능한지 확인하는 것입니다.
학습 태스크: 상태 관측치 없이 **행동 (Action) 과 그 적용 가능성 (Applicability)**에 대한 정보만 포함된 긍정적 (적용 가능) 및 부정적 (적용 불가) 행동 궤적 (traces) 을 입력받아, 해당 궤적이 모델과 일치하는지 분류하는 이진 분류 문제를 다룹니다.

2. 제안된 방법론 (Methodology)

저자들은 두 가지 아키텍처를 제안하여 비교 분석했습니다.

A. STRIPS Transformer (기호적 정렬 모델)

개념: STRIPS 구조와 이론적으로 연결된 하드 어텐션 (Hard Attention) 트랜스포머를 기반으로 합니다. B-RASP (Boolean RASP) 형식 언어와 STRIPS 도메인의 관계에 기반합니다.
작동 원리:
- 각 도메인 원자 (Atom) 에 대해 별도의 어텐션 헤드를 할당합니다.
- 각 원자의 현재 진리값을 결정하는 **가장 최근의 영향력 있는 이전 행동 (Most recent affecting action)**을 식별하는 데 초점을 맞춥니다.
- Stick-breaking Attention을 사용하여 마스크된 하드 어텐션을 구현하며, 이는 특정 원자에 대해 가장 최근의 수정 행동을 선택하는 논리를 기호적으로 구현합니다.
- 모델 파라미터가 STRIPS의 전제조건 (preconditions) 과 효과 (effects) 와 직접적으로 매핑되도록 설계되어 기호적 인덕티브 바이어스 (Inductive Bias) 가 강합니다.

B. Stick-Breaking (SB) Transformer (표준 트랜스포머 변형)

개념: 명시적인 기호 구조를 내장하지 않은 표준 디코더 스타일 트랜스포머입니다.
주요 변경점:
- 기존 Softmax 어텐션과 위치 인코딩 (Positional Encoding) 을 제거합니다.
- Stick-breaking Attention을 도입하여, Softmax 대신 가중치를 계산합니다. 이는 최근성 편향 (recency bias) 을 유지하면서 미분 가능한 방식으로 하드 어텐션을 근사합니다.
- STRIPS 모델을 직접 인코딩하지는 않지만, 학습된 표현을 통해 STRIPS 모델을 추출할 수 있습니다.

C. 학습 및 계획 추출 (Learning & Extraction)

학습: 긍정/부정 궤적에 대해 Focal Loss 를 사용하여 행동의 적용 가능성을 예측하도록 학습합니다.
상태 정보 인코딩: 궤적에 상태 정보가 없으므로, Setup Actions를 도입합니다.
- init-p: 초기 상태를 인코딩 (원자 추가).
- test-p: 최종 상태의 원자 진리값을 테스트 (원자 전제조건 확인).
모델 추출: 학습된 트랜스포머 파라미터에서 STRIPS 모델 ( $M' = \langle F', A' \rangle$ $M^{'} = ⟨ F^{'}, A^{'} ⟩$ ) 을 추출합니다.
- STRIPS Transformer 는 파라미터를 이진화하여 직접 추출.
- SB Transformer 는 test-p 행동의 예측을 통해 상태 전이를 재구성하고, 다수결 원칙으로 STRIPS 모델을 추출합니다.

3. 주요 실험 결과 (Results)

다섯 가지 고전적 계획 도메인 (Blocksworld, Ferry, Npuzzle, Maze, Logistics) 에서 실험을 수행했습니다.

일반화 능력 (Generalization):
- SB Transformer는 훈련 데이터보다 훨씬 긴 궤적 (Long traces) 에서도 거의 완벽한 정확도를 보이며, **기존 Softmax 트랜스포머 (Sinusoidal, RoPE)**는 긴 시퀀스로 일반화하는 데 실패했습니다.
- Stick-breaking Attention이 긴 시퀀스 일반화에 필수적임이 입증되었습니다.
계획 성능 (Planning Performance):
- 두 모델 모두 추출된 STRIPS 모델을 통해 **기존 플래너 (Mimir, FF)**와 연동하여, 훈련 시 보지 않은 지수적으로 많은 초기 상태와 목표에 대해 성공적으로 계획을 수립했습니다.
- SB Transformer가 STRIPS Transformer 보다 학습이 더 쉽고, 더 높은 일반화 성능을 보였습니다. STRIPS Transformer 는 강한 기호적 바이어스를 가졌음에도 불구하고 최적화가 어렵고 더 큰 데이터셋이 필요했습니다.
조합적 추론 (Compositional Reasoning):
- 두 모델 모두 학습된 기호 모델을 통해 새로운 상태와 목표에 대한 조합적 추론이 가능함을 보였습니다.
- 특히 SB Transformer 는 test-p 행동을 통해 학습된 표현이 실제 도메인 원자와 정렬 (Alignment) 되어, 추출된 모델이 원래 도메인과 동일한 원자 집합을 사용하도록 했습니다.

4. 주요 기여 및 의의

다음 토큰 예측과 세계 모델의 연결: 단순한 통계적 패턴 매칭을 넘어, 트랜스포머가 다음 토큰 예측을 통해 실제 계획이 가능한 기호적 세계 모델을 학습할 수 있음을 증명했습니다.
Stick-breaking Attention 의 유효성: 표준 Softmax 어텐션은 긴 시퀀스 일반화에 실패하지만, Stick-breaking Attention 을 사용하면 기호적 세계 모델을 학습하고 긴 시퀀스에서도 강력한 일반화 성능을 발휘함을 보였습니다.
효율적인 모델 추출: 학습된 신경망 모델에서 정확한 STRIPS 모델을 추출하여 기존 고전적 플래너와 호환되게 함으로써, 신경망의 학습 능력과 기호적 계획의 강점을 결합했습니다.
데이터 효율성 및 확장성: SB Transformer 는 STRIPS Transformer 보다 적은 데이터로도 우수한 성능을 보이며, 훈련 시 보지 않은 지수적인 수의 문제实例에 대해 일반화되는 것을 확인했습니다.

5. 결론

이 연구는 트랜스포머 아키텍처가 다음 토큰 예측을 통해 STRIPS 형태의 세계 모델을 학습할 수 있음을 보여주었습니다. 특히 Stick-breaking Attention을 활용한 표준 트랜스포머가 명시적인 기호 구조를 가진 모델보다 더 효율적이고 강력하게 일반화하며, 이를 통해 추출된 모델은 실제 계획 작업에서 높은 성능을 발휘합니다. 이는 LLM 이 단순한 언어 생성기를 넘어, 복잡한 의사결정 및 계획 작업을 위한 내부 세계 모델을 학습할 수 있는 가능성을 제시하는 중요한 결과입니다.