Revisiting Bi-Linear State Transitions in Recurrent Neural Networks

TL;DR

This paper explores bilinear operations in recurrent neural networks, showing they serve as a natural inductive bias for modeling active state transitions, especially in state tracking tasks of increasing complexity.

Contribution

It provides theoretical and empirical evidence that bilinear state updates are a natural and effective inductive bias for representing hidden state evolution in RNNs.

Findings

Bilinear operations improve state tracking performance.

Hierarchical structure of bilinear state updates aligns with task complexity.

Linear RNNs like Mamba occupy the simplest level of this hierarchy.

Abstract

The role of hidden units in recurrent neural networks is typically seen as modeling memory, with research focusing on enhancing information retention through gating mechanisms. A less explored perspective views hidden units as active participants in the computation performed by the network, rather than passive memory stores. In this work, we revisit bilinear operations, which involve multiplicative interactions between hidden units and input embeddings. We demonstrate theoretically and empirically that they constitute a natural inductive bias for representing the evolution of hidden states in state tracking tasks. These are the simplest type of tasks that require hidden units to actively contribute to the behavior of the network. We also show that bilinear state updates form a natural hierarchy corresponding to state tracking tasks of increasing complexity, with popular linear recurrent networks such as Mamba residing at the lowest-complexity center of that hierarchy.