A Close Look At World Model Recovery In Supervised Fine-Tuned LLM Planners

기술 요약: 지도 미세 조정된 LLM 플래너에서의 월드 모델 회복에 대한 심층 분석

문제 정의
지도 미세 조정(Supervised Fine-Tuning, SFT)이 대규모 언어 모델(LLM)의 엔드 투 엔드(end-to-end) 계획 능력을 향상시켰으나, 이러한 모델들이 해결하는 계획 문제에 대해 진정으로 표현하고 추론할 수 있는 내부 표현을 학습하는지는 여전히 불분 clear합니다. 구체적으로, 미세 조정된 LLM이 "월드 모델(world model)"—즉, 상태 전이와 행동 유효성에 대한 내부적 이해—을 실제로 회복하는 것인지, 아니면 단순히 유효한 행동 시퀀스를 모방하는 법을 배우는 것인지에 대한 의문이 남아 있습니다. 기존 연구들은 LLM이 제로샷(zero-shot) 설정에서 계획 수립에 어려움을 겪으며, 견고한 추론 메커니즘보다는 지름길 솔루션을 학습할 가능성이 있다고 시사합니다. 본 연구는 유효한 계획 예시에 대한 SFT가 고전적 계획 도메인에서 월드 모델 회복을 유도하는지 조사합니다.

방법론
저자들은 두 가지 상호 보완적인 정의를 사용하여 월드 모델 회복을 평가하는 프레임워크를 도입합니다:

1. 내부 표현(Internal Representations): LLM이 상태 술어(state predicates)의 참/거짓 값과 행동의 유효성을 인코딩하는 선형 표현을 학습하는지 여부.
2. 생성 능력(Generative Capabilities): LLM이 유효한 행동 시퀀스에는 높은 확률을 할당하고, 유효하지 않은 시퀀스에는 낮은 확률을 할당하는지 여부.

본 연구는 gemma2-9b-instruct 모델을 활용하여, 두 가지 고전적 계획 환경인 Blocksworld와 Logistics에서 엔드 투 엔드 계획을 생성하도록 미세 조정된 모델(이른바 "Plansformers")을 사용합니다. 저자들은 회복에 미치는 영향을 결정하기 위해 세 가지 뚜렷한 학습 데이터 분포를 실험합니다:

• 최적 계획(Optimal Plans): 휴리스틱 솔버(Fast Downward)에 의해 생성됨.
• 랜덤 워크 계획(Random Walk Plans): 최대 스텝 수에 도달할 때까지 무작위로 유효한 행동을 샘플링하여 생성됨.
• 랜덤 워크++ 계획(Random Walk++ Plans): 사이클을 방지하고 상태 공간 커버리지를 개선하기 위해, 이전에 방문했던 상태로 이어지는 행동을 무효화하는 수정된 랜덤 워크 방식.

또한, 저자들은 중간 상태를 행동 시퀀스 사이에 삽입하여 근거(groundedness)를 강화하는 **상태 체인 오브 소트(State Chain-of-Thought, State-CoT)**를 탐구합니다.

회복을 측정하기 위해 저자들은 다음을 채택합니다:

• 선형 프로브(Linear Probes): 행동 유효성과 상태 술어의 참 값을 예측하기 위해 은닉층 활성화값에 대해 훈련된 로지스틱 회귀 모델.
• 생성적 분류(Generative Classification): 부분적인 계획이 주어졌을 때 모델의 토큰 확률이 유효한 행동을 유효하지 않은 행동보다 높게 순위 매기는지 평가하는 지표.
• 분포 외(Out-of-Distribution, OOD) 테스트: 훈련 중에 본 것보다 복잡도가 증가한(예: 블록 수가 더 많은) 계획 인스턴스에 대한 평가.

주요 결과

1. 유효성의 선형 인코딩: 유효한 행동 시퀀스에 대한 지도 미세 조정은 LLM이 내부 표현에 행동 유효성을 선형적으로 인코딩할 수 있게 합니다. 프로브 F1 점수는 모든 학습 분포에 대해 중간 및 상위 레이어에서 완벽에 가까운 수준(~1.0)에 도달합니다.
2. 부분적 상태 표현: 모델은 일부 상태 술어(예: arm-empty, at ?pkg ?loc)의 참 값을 선형적으로 인코딩하는 법을 배우지만, 성능은 술어의 복잡도에 따라 크게 달라집니다. 인자가 적은 단순한 술어가 더 복잡한 술어보다 더 정확하게 인코딩됩니다.
3. 내부 표현과 생성 지표 간의 괴리: 내부 표현과 생성 출력 사이의 불일치는 중요한 발견입니다. 최적 계획으로 미세 조정된 모델은 행동 유효성에 대한 고품질의 내부 표현(높은 프로브 점수)을 보유하고 있지만, 생성 중에 유효한 행동과 유효하지 않은 행동을 구분하기 위해 토큰 확률을 사용하는 데 실패합니다. 반대로, 랜덤 워크 데이터로 미세 조정된 모델은 우수한 생성적 분류 성능(유효한 행동을 유효하지 않은 행동보다 높게 순위 매김)을 보이지만, OOD 일반화에는 어려움을 겪습니다.
4. 상태 공간 커버리지의 영향: 사이클을 제거하여 더 넓은 상태 공간 커버리지를 보장하는 Random Walk++ 데이터는 가장 강력한 월드 모델 회복을 보여줍니다. 이는 표준 랜덤 워크나 최적 계획에 비해 인-도메인(in-distribution) 생성적 분류와 OOD 계획 작업 모두에서 우수한 성능을 달성합니다.
5. OOD 일반화: 행동 유효성에 대한 내부 표현은 분포 변화(OOD 데이터)에 대해 견고하게 유지되지만, 표준 랜덤 워크로 훈련된 모델의 생성 성능은 붕괴됩니다. Random Walk++ 및 최적 계획으로 훈련된 모델은 OOD 변화에 대해 더 나은 회복력을 보이지만, 문제 복잡도가 증가함에 따라 성능은 여전히 저하됩니다.

의의 및 주장
본 논문은 해석 가능성 기법을 계획 LLM에 적용하기 위한 "레시피"를 제공한다고 주장하며, 내부 표현 분석과 생성적 평가를 결합한 총체적인 관점을 제시합니다. 저자들의 결론은 다음과 같습니다:

• 유효한 계획에 대한 SFT는 LLM이 행동 유효성과 특정 상태 술어를 선형적으로 인코딩하는 내부 표현을 학습하는 데 충분하며, 이는 일종의 월드 모델 회복을 시사합니다.
• 그러나 내부 표현의 존재가 모델이 생성 확률에서 이를 효과적으로 사용할 것임을 보장하지는 않으며, 훈련 분포(특히 상태 공간 커버리지)가 이러한 정렬에 큰 영향을 미칩니다.
• 연구 결과는 LLM이 계획 월드 모델을 부분적으로 회복할 수 있음을 입증함으로써, 계획 능력을 향나시키기 위한 미세 조정의 생존 가능성에 대해 "낙관적"입니다. 하지만 저자들은 이러한 표현과 모델의 계획 과정 사이의 인과 관계가 아직 증명되지 않았으며, 더 큰 파라미터 수나 자연어 계획으로의 일반화는 여전히 미해결 과제라는 점을 언급하며 한계점 또한 명시하고 있습니다.

저자들은 OOD 계획 문제를 해결했거나 인과적 추론을 증명했다고 주장하는 것이 아니라, 미세 조정된 LLM이 생성 능력과는 구별되면서도 그와 관련된 구조화된 내부 표현을 개발한다는 것에 대한 경험적 증거를 제공하고 있다고 명시적으로 밝히고 있습니다.