Unlocking the Future: Exploring Look-Ahead Planning Mechanistic Interpretability in Large Language Models (2024)

Abstract: Planning, as the core module of agents, is crucial in various fields such as embodied agents, web navigation, and tool using. With the development of large language models (LLMs), some researchers treat large language models as intelligent agents to stimulate and evaluate their planning capabilities. However, the planning mechanism is still unclear. In this work, we focus on exploring the look-ahead planning mechanism in large language models from the perspectives of information flow and internal representations. First, we study how planning is done internally by analyzing the multi-layer perception (MLP) and multi-head self-attention (MHSA) components at the last token. We find that the output of MHSA in the middle layers at the last token can directly decode the decision to some extent. Based on this discovery, we further trace the source of MHSA by information flow, and we reveal that MHSA extracts information from spans of the goal states and recent steps. According to information flow, we continue to study what information is encoded within it. Specifically, we explore whether future decisions have been considered in advance in the representation of flow. We demonstrate that the middle and upper layers encode a few short-term future decisions. Overall, our research analyzes the look-ahead planning mechanisms of LLMs, facilitating future research on LLMs performing planning tasks.

## introduction planning is the process of formulating a series of actions to transform a given initial state into a desired goal state @xcite . as the core module of agents @xcite , planning has been widely applied in many fields such as embodied agents @xcite , * corresponding authors. web navigation @xcite and tool using @xcite . with the development of large language models, some researchers treat large language models as intelligent agents to solve complex tasks. this is because large language models may possess some preliminary planning capabilities @xcite . recently, researchers have made efforts to stimulate and evaluate the planning capabilities of large language models. they propose prompt engineering @xcite and instruction fine-tuning @xcite to boost the planning abilities of large language models. additionally, some researchers construct benchmarks such as agentbench @xcite and agentgym @xcite to evaluate the planning capabilities of large models. although they have made some progress, the underlying mechanisms in planning capabilities of large language models remain a largely unexplored frontier. revealing the planning mechanisms of large language models helps to better understand and improve their planning capabilities. therefore, we focus on exploring the underlying mechanisms behind the planning abilities of large language models. in this work, we focus on exploring look-ahead planning mechanisms in large language models. we study the classical planning task blocksworld, which is a fully-observed setting. all entity states are known from the init state and goal state, so exploration is not needed @xcite . as illustrated in figure 1 , given an initial state and a goal state of blocksworld, the model can only pick up or put down one block. the model must generate a sequence of actions to transform the initial state into the goal state, as shown by the green path. however, it is still unclear whether the model, at step t, greedily considers only the action at t + 1 or look-ahead considers the actions at t + 2 and beyond. inspired by psychology, humans engage in look-ahead thinking when making plans @xcite . based on this, we further propose the hypothesis of model look-ahead planning, which is as follows: we design a two-stage paradigm to verify this hypothesis. it can be divided into the finding information flow stage and the probing internal representations stage. the first stage is to analyze the information flow and component functions during planning ( §5). the second stage is examining whether the model stores future information in internal representations ( §6). the specifics are as follows: in summary, our contributions are as follows: ## experimental setup in this paper, we study the blocksworld task in a fully-observed setting where all entity states are known from the init state and goal state, so exploration is not needed @xcite . given a rule r, an initial state s init , a goal state s goal , task description prompts c, the current step t, history a 1 . . . a t , model needs to predict the next action a t+1 in accordance with its generative distribution p(a t+1 | r, s init , s goal , c, a 1 . . . a t ) @xcite . in this paper, all inputs are in text form. all inferences are performed using the teacher-forcing method. previous evaluation works @xcite mainly involved generating a complete plan and then placing it into the environment for assessment. however, since our primary focus is on open-source models, we have reduced the difficulty by using a fill-in-the-blank format for evaluating the models. an example is shown in figure 2 . data previous blocksworld evaluation benchmarks @xcite put the plans generated by models into an environment to verify the correctness. however, existing interpretability methods, such as information flow @xcite , require gold labels. therefore, we synthesize a dataset containing optimal plans, with specific data statistics shown in ## background a transformer-based language model begins by converting an input text into a sequence of n tokens, denoted as s 1 , . . . , s n . each token s i is mapped to a vector x 0 @xmath0. @xmath1, where v is the vocabulary, d is embedding dimension. each layer comprises a multi-head self-attention (mhsa) sublayer followed by a multi-layer perception (mlp) sublayer @xcite . formally, the representation x ℓ i of token i at layer ℓ can be obtained as follows: ## overview of analysis we analyze the look-ahead planning mechanisms of the models from two stages. (1) in the first stage, we explore the internal mechanisms of this process in planning tasks from the perspectives of information flow and component functions. we demonstrate that the middle layer mhsa can directly decode the answers to a certain extent, and we prove that mhsa mainly extracts information from spans of the goal states and recent steps ( §5). ## information flow in planning tasks to trace the source of the correct answer, we begin with the last token. for example, in the first step "pick up", the last token is "up". the model should process the initial state, target state, and history of steps to decide which color to pick up, such as "blue". we analyze this process from two perspectives. (1) first, we study mlp and mhsa functions at the last token by extraction rates @xcite . we find that the output of mhsa in the middle layers can directly decode the correct colors to a certain extent ( §5.1). (2) based on this, we further trace the source of the correct colors by information flow @xcite . from the perspective of early and late planning stages, we prove that mhsa mainly extracts information from spans of the goal states and recent steps ( §5.2). ## internal representations encode planning information based on the previous sections, we discover that mhsa directly extracts answers from the context, but it is still unclear what information is encoded in internal representations. in this section, we demonstrate the look-ahead capability of models from both future decisions existence and history step causality perspectives. (1) for future decisions existence, we use the probing method to probe each layer of the main positions in the context. we find that the accuracy of the current state information gradually decreases as the steps progress. we also find that the middle and upper layers encode future decisions with accuracy decreasing as planning steps increase, proving the look-ahead planning hypothesis ( §6.1). ## conclusion in this paper, we investigate the mechanisms of look-ahead planning in llms through the perspectives of information flow and internal representations. we demonstrate look-ahead planning decisions existence hypothesis. our findings indicate that internal representations of llms encode a few short-term future decisions to some extent when planning is successful. these look-ahead decisions are enriched in the early layers, with their accuracy diminishing as the number of planning steps increases. we demonstrate that mhsa mainly extracts information from the spans of goal states and recent steps. additionally, the output of mhsa in the middle layers at the final token can partially decode the correct decisions.