Large language models store and use goals to guide code actions
Goal Encoding in LLMs
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Large language models (LLMs) can keep track of goals that guide what they do next, like deciding which step to do in a task. The authors found that specific parts inside these models hold information about the next goal to work on. These goal signals can be moved between different programming languages, showing that the model uses flexible plans rather than fixed steps. The model keeps checking the goal by rereading instructions instead of remembering it internally for a long time. The authors tested this idea by changing parts inside the model and seeing if the goals and actions changed accordingly.
What this means in practice
- •For software developers: Improve programming assistants by better tracking and adapting to goals within code generation tasks.
- •For natural language processing engineers: Design more interpretable language models by pinpointing goal representations in attention heads to control outputs.
Authors
Arnab Sen Sharma, Eric Todd, Sheridan Feucht, Andy Arditi, Rohit Gandikota, Natalie Shapira, Gabriele Sarti, Constanza Fierro, Lisa Schut, David Bau
Abstract
We investigate whether and how LLMs encode goals: internal neural representations that precede and guide action, and that are flexible enough to be realized through many different execution paths. We construct a coding setting with multiple sub-tasks that allows us to evaluate goal-seeking behavior and the selection of sub-goals. In this setting, we localize an encoding of the next sub-goal in the outputs of a sparse set of attention heads. The same heads can be used to decode the LM's future actions in natural language (e.g., "I'm checking if a word is a palindrome"). We verify that this encoding is causal for next action and is plastic: transplanting it into a different environment, e.g., another programming language, triggers the same operation realized with a different sequence of tokens. The LM also sustains the encoding throughout the action's execution, which is puzzling because a transformer LM without recurrence cannot carry a state forward for more tokens than it has layers left. We propose that it does not need to. Under a non-recurrent goal-seeking (NRGS) account, the LM frequently re-reads the goal from the task specification, compares it with the progress visible in the growing context, and pursues what remains. We test the predictions of this account with targeted causal interventions.