Flowformer: Linearizing Transformers with Conservation Flows

TL;DR

Flowformer introduces a novel linear attention mechanism based on flow network theory, enabling efficient processing of long sequences across various domains without relying on specific inductive biases.

Contribution

This paper presents Flow-Attention, a flow conservation-based linear attention mechanism that enhances transformer scalability and generality.

Findings

Achieves linear complexity in attention computation.

Performs well across long sequences, vision, NLP, and reinforcement learning.

Does not depend on inductive biases like locality.

Abstract

Transformers based on the attention mechanism have achieved impressive success in various areas. However, the attention mechanism has a quadratic complexity, significantly impeding Transformers from dealing with numerous tokens and scaling up to bigger models. Previous methods mainly utilize the similarity decomposition and the associativity of matrix multiplication to devise linear-time attention mechanisms. They avoid degeneration of attention to a trivial distribution by reintroducing inductive biases such as the locality, thereby at the expense of model generality and expressiveness. In this paper, we linearize Transformers free from specific inductive biases based on the flow network theory. We cast attention as the information flow aggregated from the sources (values) to the sinks (results) through the learned flow capacities (attentions). Within this framework, we apply the property of flow conservation into attention and propose the Flow-Attention mechanism of linear complexity. By respectively conserving the incoming flow of sinks for source competition and the outgoing flow of sources for sink allocation, Flow-Attention inherently generates informative attentions without using specific inductive biases. Empowered by the Flow-Attention, Flowformer yields strong performance in linear time for wide areas, including long sequence, time series, vision, natural language, and reinforcement learning. The code and settings are available at this repository: https://github.com/thuml/Flowformer.