Sub-Linear Memory: How to Make Performers SLiM

TL;DR

This paper analyzes Performer-based linear self-attention mechanisms in Transformers, revealing a flexible time-memory tradeoff that enables training with sublinear memory, facilitating deployment on resource-constrained devices.

Contribution

It provides a thorough complexity analysis of linear self-attention, demonstrating that training can be performed with sublinear memory at the cost of increased time, enabling low-memory training and fine-tuning.

Findings

Performers can be trained with sublinear memory usage.

Backward and forward passes are feasible with no approximations.

The approach enables training on low-resource devices.

Abstract

The Transformer architecture has revolutionized deep learning on sequential data, becoming ubiquitous in state-of-the-art solutions for a wide variety of applications. Yet vanilla Transformers are notoriously resource-expensive, requiring $O(L^2)$ in serial time and memory as functions of input length $L$. Recent works proposed various linear self-attention mechanisms, scaling only as $O(L)$ for serial computation. We perform a thorough analysis of recent Transformer mechanisms with linear self-attention, Performers, in terms of overall computational complexity. We observe a remarkable computational flexibility: forward and backward propagation can be performed with no approximations using sublinear memory as a function of $L$ (in addition to negligible storage for the input sequence), at a cost of greater time complexity in the parallel setting. In the extreme case, a Performer consumes only $O(1)$ memory during training, and still requires $O(L)$ time. This discovered time-memory tradeoff can be used for training or, due to complete backward-compatibility, for fine-tuning on a low-memory device, e.g. a smartphone or an earlier-generation GPU, thus contributing towards decentralized and democratized deep learning.