Compact Recurrent Transformer with Persistent Memory

TL;DR

The paper addresses the challenge of scaling Transformers to long sequences by introducing the Compact Recurrent Transformer (CRT), which processes short local segments with a shallow Transformer while compressing long-range information into a single persistent memory vector via an RNN. This design enables comparable or superior performance to full-length Transformers on language tasks (Word PTB, WikiText-103) with reduced segment lengths and substantially lower FLOPs, and achieves state-of-the-art results on the Toyota Smarthome video dataset. The approach balances local parallelism with efficient global information flow, validated through complexity analyses, gradient stability considerations, and cross-domain experiments, highlighting strong potential for edge computing and low-SWaP deployments. Overall, CRT demonstrates that explicit, vectorized persistent memory can capture long-range context with markedly improved efficiency, broadening the practical applicability of Transformer-based models.

Abstract

The Transformer architecture has shown significant success in many language processing and visual tasks. However, the method faces challenges in efficiently scaling to long sequences because the self-attention computation is quadratic with respect to the input length. To overcome this limitation, several approaches scale to longer sequences by breaking long sequences into a series of segments, restricting self-attention to local dependencies between tokens within each segment and using a memory mechanism to manage information flow between segments. However, these approached generally introduce additional compute overhead that restricts them from being used for applications where limited compute memory and power are of great concern (such as edge computing). We propose a novel and efficient Compact Recurrent Transformer (CRT), which combines shallow Transformer models that process short local segments with recurrent neural networks to compress and manage a single persistent memory vector that summarizes long-range global information between segments. We evaluate CRT on WordPTB and WikiText-103 for next-token-prediction tasks, as well as on the Toyota Smarthome video dataset for classification. CRT achieves comparable or superior prediction results to full-length Transformers in the language datasets while using significantly shorter segments (half or quarter size) and substantially reduced FLOPs. Our approach also demonstrates state-of-the-art performance on the Toyota Smarthome video dataset.

Figures (6)

• Figure 1: Compact Recurrent Transformer Architecture
• Figure 2: Transformer with Memory Token
• Figure 3: Recurrent Memory Architecture
• Figure 4: Compact Recurrent ViT Architecture
• Figure 5: Comparing the CRT model to the baseline Transformer and Transformer-XL for a 3-layer model in terms of PPL and FLOPs on different segment sizes for the Word PDB dataset. Given the same segment length, the CRT outperforms the baseline Transformer and Transformer-XL models. An interesting observation is that The CRT model using the smaller segment size of 17 tokens performs about as well as the Transformer and Transformer-XL models with segment sizes of 70 tokens.

Theorems & Definitions (4)

• Lemma 1
• Lemma 2
• Theorem 3
• proof