Skip to yearly menu bar Skip to main content


Poster

Towards an Understanding of Stepwise Inference in Transformers: A Synthetic Graph Navigation Model

Mikail Khona · Maya Okawa · Jan Hula · Rahul Ramesh · Kento Nishi · Robert Dick · Ekdeep Singh Lubana · Hidenori Tanaka


Abstract:

Stepwise inference, such as scratchpads and chain-of-thought (CoT), is an important capability of large language models, where a model decomposes a complex task into a sequence of manageable subproblems. However, despite the significant gain in performance, the underlying mechanisms of stepwise inference have remained elusive. To address this gap, we propose to study auto-regressive Transformer models solving a graph navigation problem, where a model is tasked with traversing a path from a start to a goal node on a synthetically generated graph. Through this simple, controllable, and interpretable framework of graph navigation, we empirically reproduce and analyze several phenomena observed at scale: (i) the stepwise inference reasoning gap, the cause of which we find in the structure of the training data; (ii) a diversity-accuracy tradeoff as sampling temperature varies; (iii) a simplicity bias in the model's output; and (iv) combinatorial generalization, failure on length generalization and a primacy bias with in-context exemplars.Overall, this work introduces a grounded synthetic framework for studying stepwise inference and offers mechanistic hypotheses that lay the foundation for a deeper understanding of this phenomenon.

Live content is unavailable. Log in and register to view live content