Using Multi-task and Transfer Learning to Solve Working Memory Tasks

TL;DR

This paper introduces MAES, a neural architecture that leverages multi-task and transfer learning to effectively solve complex working memory tasks, demonstrating superior generalization and performance over existing models.

Contribution

The paper presents MAES, a novel differentiable architecture with dual controllers and shared memory, showing that multi-task learning enhances encoder quality and enables large-scale task generalization.

Findings

MAES achieves 50x longer sequence handling than training data.

MAES outperforms LSTM, NTM, and DNC on all tested tasks.

Multi-task learning improves encoder effectiveness.

Abstract

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.