Learning Useful Representations of Recurrent Neural Network Weight Matrices

TL;DR

This paper introduces novel methods for representing RNN weights to improve analysis and downstream task performance, demonstrating the superiority of functionalist approaches in predicting trained tasks.

Contribution

It develops new functionalist approaches for RNN weight representation, provides theoretical insights, and releases datasets for evaluating RNN weight encodings.

Findings

Functionalist approaches outperform mechanistic ones in task prediction.

New datasets enable RNN weight representation learning.

Self-supervised evaluation shows clear advantages of proposed methods.

Abstract

Recurrent Neural Networks (RNNs) are general-purpose parallel-sequential computers. The program of an RNN is its weight matrix. How to learn useful representations of RNN weights that facilitate RNN analysis as well as downstream tasks? While the mechanistic approach directly looks at some RNN's weights to predict its behavior, the functionalist approach analyzes its overall functionality-specifically, its input-output mapping. We consider several mechanistic approaches for RNN weights and adapt the permutation equivariant Deep Weight Space layer for RNNs. Our two novel functionalist approaches extract information from RNN weights by 'interrogating' the RNN through probing inputs. We develop a theoretical framework that demonstrates conditions under which the functionalist approach can generate rich representations that help determine RNN behavior. We release the first two 'model zoo' datasets for RNN weight representation learning. One consists of generative models of a class of formal languages, and the other one of classifiers of sequentially processed MNIST digits. With the help of an emulation-based self-supervised learning technique we compare and evaluate the different RNN weight encoding techniques on multiple downstream applications. On the most challenging one, namely predicting which exact task the RNN was trained on, functionalist approaches show clear superiority.