Sub-Linear Memory: How to Make Performers SLiM
This work addresses the high memory consumption of Transformers, making them more accessible for training or fine-tuning on resource-constrained devices, which is beneficial for democratizing deep learning.
The Transformer architecture, while powerful, is resource-intensive, scaling quadratically with input length L. This paper analyzes Performers, a linear self-attention mechanism, and finds that they can achieve sub-linear memory usage during training, down to O(1) memory, by trading off with increased time complexity.
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.