Folding Attention: Memory and Power Optimization for On-Device Transformer-based Streaming Speech Recognition

TL;DR

The paper identifies memory and power as the principal bottlenecks in on-device streaming Transformer ASR, shifting focus from attention-score optimization to the linear projection layers in self-attention and feedforward networks. It introduces folding attention, which splits tokens into sub-tokens via a folding factor $N$, reducing linear-layer parameter counts by up to a factor of $N$ and memory by up to a factor of $N^2$, at the cost of a modest increase in attention-score computation. Empirical results on LibriSpeech and an in-house dataset show folding attention achieves 12–24% reductions in model size and 11–23% reductions in power on LibriSpeech, and 14–23% size and 13–21% power reductions on the in-house data, all while preserving accuracy and maintaining streaming-friendly RT-F. These findings enhance the feasibility of deploying high-quality, low-power, on-device streaming ASR systems by addressing the memory-read/write bottlenecks inherent to Transformer attention components.

Abstract

Transformer-based models excel in speech recognition. Existing efforts to optimize Transformer inference, typically for long-context applications, center on simplifying attention score calculations. However, streaming speech recognition models usually process a limited number of tokens each time, making attention score calculation less of a bottleneck. Instead, the bottleneck lies in the linear projection layers of multi-head attention and feedforward networks, constituting a substantial portion of the model size and contributing significantly to computation, memory, and power usage. To address this bottleneck, we propose folding attention, a technique targeting these linear layers, significantly reducing model size and improving memory and power efficiency. Experiments on on-device Transformer-based streaming speech recognition models show that folding attention reduces model size (and corresponding memory consumption) by up to 24% and power consumption by up to 23%, all without compromising model accuracy or computation overhead.

Figures (6)

• Figure 1: Standard self-attention (single-head as an example).
• Figure 2: Folding attention (folding factor of 2 and single-head as an example). The subsequent feedforward networks (not shown in the diagram) can be repositioned prior to the unfolding operator to achieve a size reduction (fourfold in this example).
• Figure 3: Model size vs. word error rate on LibriSpeech.
• Figure 4: Power vs. word error rate on LibriSpeech.
• Figure 5: Model size vs. word error rate on the in-house dataset.