Using Multi-task and Transfer Learning to Solve Working Memory Tasks
This work addresses the challenge of improving generalization in working memory tasks for AI systems, though it appears incremental as it builds on existing memory-augmented neural network concepts.
The authors tackled the problem of solving complex working memory tasks by proposing the Memory-Augmented Encoder-Solver (MAES) architecture, which achieved task-size generalization by handling inputs 50 times longer than in training and outperformed models like LSTM, NTM, and DNC across all tasks.
We propose a new architecture called Memory-Augmented Encoder-Solver (MAES) that enables transfer learning to solve complex working memory tasks adapted from cognitive psychology. It uses dual recurrent neural network controllers, inside the encoder and solver, respectively, that interface with a shared memory module and is completely differentiable. We study different types of encoders in a systematic manner and demonstrate a unique advantage of multi-task learning in obtaining the best possible encoder. We show by extensive experimentation that the trained MAES models achieve task-size generalization, i.e., they are capable of handling sequential inputs 50 times longer than seen during training, with appropriately large memory modules. We demonstrate that the performance achieved by MAES far outperforms existing and well-known models such as the LSTM, NTM and DNC on the entire suite of tasks.