Towards Modular Algorithm Induction
This work addresses the challenge of algorithm induction for AI systems, representing an incremental advancement in neural network-based algorithmic learning.
The paper tackles the problem of learning algorithms from input-output examples by introducing a modular neural network architecture called Main, which achieves perfect generalization to inputs of much longer lengths than those used in training across five algorithmic tasks.
We present a modular neural network architecture Main that learns algorithms given a set of input-output examples. Main consists of a neural controller that interacts with a variable-length input tape and learns to compose modules together with their corresponding argument choices. Unlike previous approaches, Main uses a general domain-agnostic mechanism for selection of modules and their arguments. It uses a general input tape layout together with a parallel history tape to indicate most recently used locations. Finally, it uses a memoryless controller with a length-invariant self-attention based input tape encoding to allow for random access to tape locations. The Main architecture is trained end-to-end using reinforcement learning from a set of input-output examples. We evaluate Main on five algorithmic tasks and show that it can learn policies that generalizes perfectly to inputs of much longer lengths than the ones used for training.