This repository contains notebooks for replicating the results of the paper "Structured World Representations in Maze-Solving Transformers". The paper can currently be found on ArXiv as 2312.02566. Our accompanying blog post can be found at unsearch.org/research/01_maze_transformer_world_representations (also see the twitter thread).
This repository contains only finalized notebooks for the ArXiv and UniReps 2023 papers, which depend on our base libraries to work:
- maze-dataset for dataset generation, manipulation, and visualization
- maze-transformer for model training, evaluation, and analysis
This repository will not be updated with future results, but the base libraries may change. For latest updates, see unsearch.org.
To install this project with dependencies, we recommend using poetry, but pip is also supported. Simply clone and install:
git clone https://github.com/understanding-search/structured-representations-maze-transformers
cd structured-representations-maze-transformers
# install via poetry
poetry install
# or with pip
pip install -e .Notebooks and required model files are located in the notebooks directory. If using poetry, select the initialized virtual environment and run the notebook.
Notebooks may generate a variety of files during runtime, such as test datasets and figures.
Please feel free to submit an issue if you have any questions, comments, or trouble running the notebooks -- we are happy to help!
Please cite this work as:
@misc{ivanitskiy2023swrmt,
title={Structured World Representations in Maze-Solving Transformers},
author={Michael Igorevich Ivanitskiy and Alex F. Spies and Tilman Räuker and Guillaume Corlouer and Chris Mathwin and Lucia Quirke and Can Rager and Rusheb Shah and Dan Valentine and Cecilia Diniz Behn and Katsumi Inoue and Samy Wu Fung},
year={2023},
eprint={2312.02566},
url={https://arxiv.org/abs/2312.02566},
archivePrefix={arXiv},
primaryClass={cs.LG}
}