Transformers Can Navigate Mazes With Multi-Step Prediction

要約

言語モデリングで目覚ましい成功を収めたにもかかわらず、シーケンス内の次のトークンを予測するように訓練されたトランスフォーマーは、長期的な計画に苦労しています。
この制限は、迷路のナビゲーションなど、複数のステップを先に計画するための先見の明が必要なタスクで特に顕著です。
ただし、標準的な次の単一トークンの予測目標には、複数のステップを予測したり、これまでにたどったパスを再検討したりするための明示的なメカニズムがありません。
したがって、この研究では、複数のステップを前方（および後方）に明示的に予測することで、トランスフォーマーの迷路ナビゲーションを改善できるかどうかを研究します。
私たちは、パラメーターが一致したトランスフォーマーを同一の設定でゼロからトレーニングし、標準の次トークン予測と MLM-U (複数のステップを前後に明示的に予測する目標) を使用して、さまざまなタイプとサイズの迷路をナビゲートします。
MLM-U は、迷路の種類や複雑さを問わず、標準的な次のトークンの予測と比較して、迷路をナビゲートするトランスフォーマーの能力を大幅に向上させることがわかりました。
また、MLM-U トレーニングは、次のトークン トレーニングと比較して、GPU トレーニング時間の観点からサンプル効率が 4 倍高く、収束が 2 倍高速であることもわかりました。
最後に、より複雑な迷路の場合、より大きな変圧器に拡張することで MLM-U の利点が得られることがわかります。
驚くべきことに、MLM-U でトレーニングされたトランスフォーマーは、A* 検索トレースからの追加の監視を使用して次のトークン予測でトレーニングされたより大きなトランスフォーマーよりも優れていることがわかります。
これらの発見が、変圧器の長期計画能力を向上させる学習目標の約束を強調するものであることを願っています。
コードは https://github.com/facebookresearch/maze_navigation_MLMU にあります。

要約(オリジナル)

Despite their remarkable success in language modeling, transformers trained to predict the next token in a sequence struggle with long-term planning. This limitation is particularly evident in tasks requiring foresight to plan multiple steps ahead such as maze navigation. The standard next single token prediction objective, however, offers no explicit mechanism to predict multiple steps ahead – or revisit the path taken so far. Consequently, in this work we study whether explicitly predicting multiple steps ahead (and backwards) can improve transformers’ maze navigation. We train parameter-matched transformers from scratch, under identical settings, to navigate mazes of varying types and sizes with standard next token prediction and MLM-U, an objective explicitly predicting multiple steps ahead and backwards. We find that MLM-U considerably improves transformers’ ability to navigate mazes compared to standard next token prediction across maze types and complexities. We also find MLM-U training is 4x more sample efficient and converges 2x faster in terms of GPU training hours relative to next token training. Finally, for more complex mazes we find MLM-U benefits from scaling to larger transformers. Remarkably, we find transformers trained with MLM-U outperform larger transformers trained with next token prediction using additional supervision from A* search traces. We hope these findings underscore the promise of learning objectives to advance transformers’ capacity for long-term planning. The code can be found at https://github.com/facebookresearch/maze_navigation_MLMU

arxiv情報

提供元, 利用サービス

arxiv.jp, Google