An All-MLP Sequence Modeling Architecture That Excels at Copying

TL;DR

This paper introduces CausalRN, an all-MLP sequence model capable of copying sequences as effectively as Transformers, by leveraging innovative exponential activation and normalization techniques that enable efficient, scalable in-context retrieval.

Contribution

The paper presents CausalRN, a novel all-MLP architecture that matches Transformer performance on copying tasks through key innovations supporting autoregressive sequence modeling.

Findings

CausalRN can match Transformers in copying tasks.

Exponentially-activated RNs achieve linear time complexity.

Pre-activation normalization enhances memory capacity.

Abstract

Recent work demonstrated Transformers' ability to efficiently copy strings of exponential sizes, distinguishing them from other architectures. We present the Causal Relation Network (CausalRN), an all-MLP sequence modeling architecture that can match Transformers on the copying task. Extending Relation Networks (RNs), we implemented key innovations to support autoregressive sequence modeling while maintaining computational feasibility. We discovered that exponentially-activated RNs are reducible to linear time complexity, and pre-activation normalization induces an infinitely growing memory pool, similar to a KV cache. In ablation study, we found both exponential activation and pre-activation normalization are indispensable for Transformer-level copying. Our findings provide new insights into what actually constitutes strong in-context retrieval.