Why investigate neural networks that plan?
Cross-posted to Substack.
Over the last two years at FAR AI, I and some others[1] have worked on reverse-engineering a neural network that learned, purely through reinforcement learning, to play Sokoban (see papers 3, 2, 1). I wanted to understand whether neural networks choose actions explicitly based on their future consequences, and how they represent the goals they use to decide between these consequences.
We succeeded at this, for the recurrent neural network we used. We now understand most of the planning-relevant parts of the Sokoban-playing recurrent neural network, and found that it explicitly does consider more than one plan for reaching the multiple goals, and chooses between these plans based on how many steps they take. The plans are constructed using particular ways which encode the transition rules of the game of Sokoban, forming a model of the world.
In the next few days, I will write what we learned about how this network works, how do we know it is true, and what it means for planning capabilities in other neural networks such as transformers.
What is learned search?
A process which minimizes a function is an optimizer. For example, to train a neural network, we iteratively minimize some loss function. The algorithms to do this are hand-coded, and we know which objective they’re optimizing: it’s the loss that we’ve coded in. Another classical optimizer in AI is the tree search. In a game such as chess, tree search algorithms try a large combination of possible future moves, evaluate what the board looks like after that, and pick the move that results in the best evaluation.
A learned optimizer is an optimizer which is learned by a NN through some training process, instead of hand-coded. Because machine learning algorithms are typically themselves optimizers, Hubinger et al. (2019) named an optimizer that is learned a mesa-optimizer. For example, if we’re successfully training a NN to play chess, it might become an optimizer for winning at chess. (As opposed to becoming a compressed lookup table which learns to associate positions with likely good moves, without any internal structure.) If a NN learns to do search in the sense of looking for sequences of moves that will win, and then executing those, then it has learned to do search, so we say it implements learned search.[2]
Largely with my close collaborators Taufeeque and Thomas Bush. ↩︎
In 2019, the central problem of AI alignment was ‘choosing the correct reward function’, so the authors focused on the fact that we would not know which function a neural network which learns end-to-end optimizes. They called the problem of ensuring the reward from the mesa-optimizer matches the training reward ‘inner alignment’, to separate it from choosing the right training reward, ‘outer alignment’. In reality, hand-coded NN training isn’t actually minimizing the loss, it’s just written that way on the page. Purely ‘minimizing loss’ doesn’t predict that NNs would generalize well, so it cannot be the right explanation. Instead, the pressure of reward optimization somehow selects higher-performing models. Therefore decomposing alignment into ‘choose the perfect reward function’ and ‘ensure the resulting NN optimizes perfectly for the training reward function’ is very counterproductive. ↩︎