- Transformer와 LLM은 대화·이미지 이해·코드 완성에서는 강하지만, 다단계 계획과 고차 추론에서는 안정적인 성능을 내기 어려움
- 이 논문은 계획 과제와 최적 해법을 토큰 시퀀스로 만들고, A*가 문제를 푸는 실행 추적까지 학습 데이터에 넣음
- Searchformer는 A*의 탐색 과정을 먼저 모방한 뒤, 최적 계획을 유지하면서 더 짧은 탐색 시퀀스를 만들도록 미세조정됨
- Sokoban 실험에서 Searchformer 계열 모델은 테스트 과제의 93.7% 를 해결했고, A* 기준 구현보다 평균 26.8% 적은 탐색 단계를 사용함
- 실행 추적은 생성 시퀀스를 10×~100× 늘리는 부담이 있지만, 더 큰 solution-only 모델보다 적은 학습 시퀀스로도 미지 과제에서 최적 계획을 더 자주 생성함
Transformer가 잘하는 일과 약한 계획 과제
- Transformer 기반 아키텍처는 여러 작업에서 높은 성능을 보임
- 사람 수준 대화
- 고품질 이미지 이해
- 비디오 생성
- 멀티모달 생성
- 코드 완성
- 인터넷 규모 데이터로 학습된 LLM 같은 모델은 실제 사용 사례에서 잘 일반화할 수 있음
- 하지만 계획 및 추론 과제에서는 여전히 한계가 남아 있음
- LLM은 다단계 계획 과제에서 부족한 모습을 보임
- 고차 추론 수행에서도 어려움이 나타남
단계적 사고 프롬프트의 한계
- 최근 접근들은 Transformer가 중간 “생각”을 생성한 뒤 답하도록 유도해 성능을 높이려 함
- Chain-of-Thought(CoT) 프롬프팅과 Tree-of-thoughts(ToT) 는 모델이 단계적으로 “생각”하도록 장려함
- 이런 기법은 자주 효과적이지만, self-enforcing 같은 이유로 오히려 성능을 낮출 수도 있음
- 한 데이터셋에서 잘 통하던 방식이 다른 데이터셋에서는 실패할 수 있음
- 공간 추론과 수학 추론처럼 필요한 추론 유형이 달라지는 경우가 예시임
- Transformer와 LLM이 계획, 다단계 의사결정, 추론을 안정적으로 수행하게 하는 방법은 여전히 활발한 연구 주제임
A* 탐색 동역학을 학습 데이터로 넣는 방식
- 이 접근은 Transformer가 복잡한 계획 과제를 더 강건하게 풀도록 학습시키는 데 초점을 둠
- 모델은 LLM처럼 단어 시퀀스가 주어졌을 때 다음 단어를 예측하도록 학습됨
- 실험은 합성 언어와 합성 어휘를 사용하는 합성 생성 데이터셋에서 수행됨
- 계획 과제와 최적 해법 계획은 토큰이라고 부르는 단어 시퀀스로 표현됨
- A*가 수행한 계산 과정은 실행 추적 토큰 시퀀스로 기록됨
- 실행 추적은 A*의 탐색 동역학을 담은 시퀀스 데이터셋을 구성함
- Transformer는 탐색이 보강된 시퀀스를 통해 A*의 탐색 동역학과 최적 계획을 함께 인코딩한 토큰 시퀀스를 생성하도록 학습됨
Searchformer의 학습 절차
- 최종 모델인 Searchformer는 두 단계로 만들어짐
- 먼저 Transformer가 A*의 탐색 과정을 모방하도록 학습됨
- 이후 최적 계획을 출력하면서도 더 적은 탐색 단계 안에서 계획을 찾도록 미세조정됨
- 이 과정을 search dynamics bootstrapping이라고 부름
- 목표는 A* 기준 구현보다 더 적은 탐색 단계로 복잡한 계획 과제를 해결하는 Transformer를 얻는 것임
Sokoban 실험과 일반화 성능
- Sokoban 퍼즐에서 Searchformer 계열 모델은 전체 테스트 과제의 93.7% 를 해결함
- 평균 탐색 단계는 A* 기준 구현보다 26.8% 더 적음
- 과제 복잡도, 데이터셋 크기, 모델 크기를 통제한 실험에서 실행 추적 포함의 효과가 확인됨
- 실행 추적을 학습 데이터에 넣으면 생성 시퀀스 길이가 10×~100× 늘어남
- 그럼에도 독립 테스트 과제 세트에서는 성능이 증가함
- search-augmented 모델은 더 큰 solution-only 모델보다 학습 시퀀스가 10배 적어도 미지 과제에서 최적 계획을 더 자주 생성함
- search-augmented 모델은 과제 설명, 해법, 실행 추적을 포함한 데이터로 학습됨
- solution-only 모델은 과제 설명과 과제 해법만 포함한 시퀀스로 학습됨
- 이 결과는 A*의 탐색 동역학을 Transformer 학습 과정에 포함하면 계획 과제 성능을 높일 수 있음을 보여줌