TinyFormer: Efficient Transformer Design and Deployment on Tiny Devices

TL;DR

TinyFormer tackles the challenge of running transformers on ultra-resource-constrained MCUs by integrating hardware-aware neural architecture search (SuperNAS) with hardware-aware sparsification (SparseNAS) and a deployment engine (SparseEngine). The framework automatically designs sparse transformer hybrids that fit tight memory/storage budgets and auto-generates MCU-optimized code, delivering up to 12.2x speedups over CMSIS-NN while maintaining strong accuracy (e.g., 96.1% on CIFAR-10 with 1MB storage/320KB RAM). Key innovations include blockwise pruning with mixed block sizes, Scaled-LayerNorm for integer-only inference, blockwise run-length sparse coding, and a code-generation pipeline from ONNX to C++ for STM32 MCUs. The results demonstrate practical viability of accurate, efficient transformers in TinyML, enabling new edge applications with guaranteed latency and substantial resource savings.

Abstract

Developing deep learning models on tiny devices (e.g. Microcontroller units, MCUs) has attracted much attention in various embedded IoT applications. However, it is challenging to efficiently design and deploy recent advanced models (e.g. transformers) on tiny devices due to their severe hardware resource constraints. In this work, we propose TinyFormer, a framework specifically designed to develop and deploy resource-efficient transformer models on MCUs. TinyFormer consists of SuperNAS, SparseNAS, and SparseEngine. Separately, SuperNAS aims to search for an appropriate supernet from a vast search space. SparseNAS evaluates the best sparse single-path transformer model from the identified supernet. Finally, SparseEngine efficiently deploys the searched sparse models onto MCUs. To the best of our knowledge, SparseEngine is the first deployment framework capable of performing inference of sparse transformer models on MCUs. Evaluation results on the CIFAR-10 dataset demonstrate that TinyFormer can design efficient transformers with an accuracy of 96.1% while adhering to hardware constraints of 1MB storage and 320KB memory. Additionally, TinyFormer achieves significant speedups in sparse inference, up to 12.2x comparing to the CMSIS-NN library. TinyFormer is believed to bring powerful transformers into TinyML scenarios and to greatly expand the scope of deep learning applications

Figures (17)

• Figure 1: Accuracy and resource usage comparison among different compact models when evaluated on CIFAR-10. Tiny devices only have MB-level storage and KB-level memory, which is a huge difference between the resources available on edge or cloud platforms.
• Figure 2: TinyFormer is a hardware-aware framework. SuperNAS is co-designed with SparseNAS to produce sparse models with transformers under resource limits. SparseEngine enables sparse inference on MCUs.
• Figure 3: Our supernet architecture design. We design four types of choice modules: Downsample Module, MobileNetV2 Module, Transformer Module and Pooling Module. Each choice module contains $2$ or $3$ architecture candidates inside. The single-path model is sampled from the supernet with only one architecture candidate invoked per choice module.
• Figure 5: Blockwise pruning method. We abbreviate depthwise convolution in 2d as DW-Conv2d. $W_{in}$ and $W_{out}$ denote the dimensions of linear layer, $C_{in}$ means the number of input channels and $W_{width}$ is the kernel width of depthwise convolution.
• Figure 6: The workflow overview of SparseEngine. Sparse models obtained from SparseNAS are first analyzed for sparse coding and memory allocation. Then, different kinds of layers' computations are transformed via automatic code generation for the following compilation and execution on targeted MCUs.