SPARTAN: Sparse Hierarchical Memory for Parameter-Efficient Transformers

TL;DR

SPARTAN introduces a hierarchical sparse memory architecture for transformers that enables efficient fine-tuning on edge devices by only updating memory components, significantly reducing storage and increasing inference speed.

Contribution

It proposes a novel hierarchical sparse memory design that allows parameter-efficient fine-tuning of pre-trained language models on edge devices, outperforming existing methods in speed and comparable accuracy.

Findings

Over 90% inference speedup on Raspberry Pi 4

Outperforms PE baselines by 0.1 points on GLUE

Trains 34% faster in few-shot settings

Abstract

Fine-tuning pre-trained language models (PLMs) achieves impressive performance on a range of downstream tasks, and their sizes have consequently been getting bigger. Since a different copy of the model is required for each task, this paradigm is infeasible for storage-constrained edge devices like mobile phones. In this paper, we propose SPARTAN, a parameter efficient (PE) and computationally fast architecture for edge devices that adds hierarchically organized sparse memory after each Transformer layer. SPARTAN freezes the PLM parameters and fine-tunes only its memory, thus significantly reducing storage costs by re-using the PLM backbone for different tasks. SPARTAN contains two levels of memory, with only a sparse subset of parents being chosen in the first level for each input, and children cells corresponding to those parents being used to compute an output representation. This sparsity combined with other architecture optimizations improves SPARTAN's throughput by over 90% during inference on a Raspberry Pi 4 when compared to PE baselines (adapters) while also outperforming the latter by 0.1 points on the GLUE benchmark. Further, it can be trained 34% faster in a few-shot setting, while performing within 0.9 points of adapters. Qualitative analysis shows that different parent cells in SPARTAN specialize in different topics, thus dividing responsibility efficiently.