Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping (SearchFormer)

[arXiv](2024/02/21 version v1)

 

Abstract

Transformer를 통해 maze, sokoban puzzle과 같은 복잡한 의사 결정 작업 해결

 

 

Problem Setup

문제: 미로 탐색과 sokoban puzzle

직접 해보기. 이거 생각보다 재밌음;; 시간 순삭 조심;;

 

Generating execution traces of A∗ search

A*(A-star로 읽음) 알고리즘을 통해 2가지 토큰 시퀀스를 생성할 수 있다.

• Search-augmented sequence: <prompt><trace><plan> 형태로 execution trace, optimal plan을 포함한다.
• Solution-only sequence: <prompt><plan> 형태.

 

Training a Transformer model

T5, RoPE

 

인코더는 <prompt>를 처리하고 디코더는 <trace><plan> 또는 <plan>를 처리한다.

 

Cross-entropy로 손실을 계산하며

모델의 이전 단계 출력이 아닌 ground truth를 입력으로 사용하는 teacher forcing 방법을 사용한다.

 

또한 손실을 응답 시퀀스의 길이인 m_d로 나누어 최대한 짧은 응답을 생성하도록 권장한다.

 

Moving past algorithm imitation via search dynamics bootstrapping

비용이 같은 노드 간의 결정과 하위 노드가 확장되는 순서를 무작위로 한 비결정론적 A* 알고리즘에 대해 Search-augmented model(<trace><plan>)을 훈련한다.

 

이 모델로 시퀀스를 생성하고, 생성된 시퀀스 중 optimal plan으로 끝나면서 trace가 원본 시퀀스보다 짧은 시퀀스를 fine-tuning dataset에 포함시킨다.

 

이 데이터로 fine-tuning 된 모델을 SearchFormer라고 한다. 또한 이 과정을 반복적으로 수행하여 모델의 탐색 성능을 더욱 향상시킬 수 있다.