Better World Models Can Lead to Better Post-Training Performance
Created by
- Haebom
Category
Empty
저자
Prakhar Gupta, Henry Conklin, Sarah-Jane Leslie, Andrew Lee
개요
본 연구는 Transformer 모델의 내부 표현과 다운스트림 능력에 명시적인 세계 모델링 목표가 미치는 영향을 다양한 훈련 단계에서 연구합니다. 2x2x2 루빅스 큐브를 사용하여 (1) 명시적인 사전 훈련된 세계 모델이 모델의 잠재 표현에 어떻게 영향을 미치는지, (2) 세계 모델 품질이 강화 학습 사후 훈련 후 모델의 성능에 어떻게 영향을 미치는지 질문합니다. 표준 다음 토큰 예측과 두 가지 명시적인 세계 모델링 전략(i) 상태 예측 사전 훈련 및 (ii) 결합된 상태 예측 + 다음 토큰 목표)을 비교하고, 그룹 상대 정책 최적화(GRPO)가 사후 훈련으로 적용된 후 작업 성능을 평가합니다. 선형 프로브와 인과적 개입을 통해 표현 품질을 평가합니다. 명시적인 세계 모델링은 더 선형적으로 디코딩 가능하고 인과적으로 조작 가능한 상태 표현을 생성합니다. 특히, 개선된 상태 표현은 GRPO에 더 높은 이점을 가져다주며, 특히 더 어려운 큐브 상태에서 그렇습니다. 본 연구 결과는 상태 표현을 개선하는 것이 시퀀스 계획 작업에 대한 사후 훈련의 효과를 향상시킬 수 있음을 나타냅니다.
시사점, 한계점
•
시사점:
◦
명시적인 세계 모델링은 Transformer 모델의 잠재 표현을 개선하여 선형 디코딩 가능성과 인과적 조작성을 높입니다.
◦
개선된 상태 표현은 강화 학습 기반의 사후 훈련(GRPO)의 성능을 향상시킵니다.
◦
특히 어려운 큐브 상태에서 GRPO의 성능 향상이 두드러집니다.
◦
상태 표현의 개선은 시퀀스 계획 작업의 효율성을 높일 수 있습니다.
•
한계점:
◦
연구는 2x2x2 루빅스 큐브라는 특정 도메인에 초점을 맞추고 있습니다.
◦
다른 복잡한 문제로 일반화 가능성에 대한 추가 연구가 필요합니다.
◦
사전 훈련 및 사후 훈련 과정에 대한 자세한 분석이 필요합니다.
