PaperlayerView exact revision on arXiv

Source-linked AI summary

Flowformer: Linearizing Transformers with Conservation Flows

Haixu Wu, Jialong Wu, Jiehui Xu, Jianmin Wang, Mingsheng Long

arXiv:2202.06258v2cs.LGcs.AI

TL;DR

Quadratic attention limits Transformers on long sequences and larger models, while linear approaches can risk trivial attention or rely on specific inductive biases. Flowformer uses conserved source-to-sink flows in Flow-Attention and reports strong performance in linear time across five benchmark areas.

  • Problem

    Quadratic attention limits Transformers on long sequences and scaling to bigger models, while linear methods face the challenge of avoiding trivial attention without sacrificing generality.

    Source: §1 · p. 1

  • Method

    Flow-Attention treats values as sources and results as sinks, using flow conservation for source competition and sink allocation.

    Source: §3 · p. 3, §3.2 · p. 5

  • Results

    Flowformer reports strong performance in linear time across five benchmarks spanning long sequences, language, vision, time series, and reinforcement learning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Flowformer achieves linear complexity and avoids degenerated attention without specific inductive biases across the reported areas.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Flowformer addresses the challenge of making attention linear without losing informative token competition or relying on specific inductive biases. It applies flow conservation to attention and reports strong performance in linear time across five benchmark areas.

  • Quadratic pairwise-relation modeling makes standard Transformers computationally prohibitive for long sequences and scaling to bigger models.Source: §1 · p. 1
  • The paper models attention as information flowing from values as sources to results as sinks through learned attention capacities.Source: §1 · p. 2
  • Flow-Attention conserves incoming sink flow for source competition and outgoing source flow for sink allocation, avoiding trivial attention without specific inductive biases.Source: §1 · p. 2
  • Flowformer achieves strong performance in linear time across five benchmarks covering long sequences, language, vision, time series, and reinforcement learning.Source: §1 · p. 2

2. Preliminaries

Canonical attention aggregates values according to query-key similarities, with softmax competition helping prevent trivial attention but pairwise computation scaling quadratically. Efficient linear approaches reduce this cost through similarity decomposition or sparsification, while Flowformer combines decomposition with flow conservation.

2.1. General View of Attention Mechanism

  • Canonical attention uses query-key similarity to aggregate values, with softmax competition preventing trivial attention but QK^T computation costing O(nmd).Source: §2.1 · p. 2

2.2. Efficient and Linear Transformers

  • Similarity-decomposition methods use feature-map projections and matrix associativity to replace direct query-key calculation with key-value multiplication and linear sequence-length complexity.Source: §2.2 · pp. 2–3
  • Linear Transformer, RFA, and Performer avoid softmax through feature or kernel methods, while approximation error or specific kernel requirements remain concerns for some methods.Source: §2.2 · p. 3
  • Attention-sparsification methods reduce modeled relations, but the cited approaches sacrifice information utilization and face an efficiency-performance dilemma.Source: §2.2 · p. 3
  • Flowformer combines similarity decomposition with flow conservation to achieve token competition without trivial attention or specific inductive biases.Source: §2.2 · p. 3

3. Method

Flowformer reframes attention as information flowing from values to results and applies flow conservation to create source competition and sink allocation. This produces informative attention with linear complexity without specific inductive biases.

3.2. Flow-Attention Mechanism

  • Flow-Attention applies flow conservation to create source competition and sink allocation, avoiding trivial attention without specific inductive biases.The method fixes flow resources to induce competition among sources and allocates incoming flow capacity across sinks.Source: §3 · p. 3, §3.2 · pp. 4–5

3.1. Attention Mechanism: A Flow Network View

  • The flow-network view treats values as sources, results as sinks, and query-key similarities as positive flow capacities.A non-negative element-wise projection φ defines the similarity function φ(Q)φ(K)^T for linear complexity.Source: §3.1 · p. 3

3.2. Flow-Attention Mechanism

  • Incoming flow conservation makes source capacities indicate source importance, while outgoing flow conservation determines each sink’s allocation capacity.The method sets information entering or leaving other layers to a default value of 1 to fix the resource.Source: §3.2 · pp. 4–5
  • The mechanism reweights sources, aggregates their information through associative matrix multiplication, then filters each sink’s result using its allocation capacity.Both Softmax and Sigmoid operations can be computed in linear time.Source: §3.2 · p. 5
  • Replacing Transformer attention with Flow-Attention gives Flowformer linear complexity while retaining generality without specific inductive biases.The flow-capacity calculations and associative aggregation are linear in sequence length.Source: §3.2 · p. 5

4. Experiments

Across five benchmarks spanning long-sequence modeling, language, vision, time series, and reinforcement learning, the experiments evaluate Flowformer’s performance and efficiency. Results report strong or competitive performance across these tasks, with linear complexity in the vision comparison.

4. Experiments

  • Flowformer shows strong performance in linear time across five benchmarks covering long sequences, language, vision, time series, and reinforcement learning.Language modeling and reinforcement learning also test the causal version of Flow-Attention.Source: §4 · p. 5

4.1. Long Sequence Modeling

  • Flowformer achieves state-of-the-art results on Long-Range Arena’s ListOps and Retrieval tasks and competitive results on the others.Its overall score improves from 55.23 to 56.48, and it surpasses the vanilla Transformer on all five tasks. Competition and allocation yield average promotions of 1.23 and 0.9, respectively; it also outperforms Performer by a large margin (51.41→56.48) with competitive inference speed and better training efficiency.Source: §4.1 · pp. 5–6, Table 2 · p. 6

4.2. Language Modeling

  • Flowformer achieves the best language-modeling performance and outperforms the vanilla Transformer on WikiText-103.The autoregressive experiment also tests causal Flow-Attention; ablations report perplexity reductions of 1.4 for competition and 0.4 for allocation.Source: §4.2 · pp. 6–7, Table 4 · p. 6

4.3. Image Recognition

  • On ImageNet-1K, Flowformer is the best linear Transformer and exceeds the vanilla Transformer in Top-1 accuracy, 80.4 versus 78.7.The comparison reports strong Top-1 and Top-5 accuracy with linear complexity. Flow-Attention adds no model parameters; DeiT+Flowformer changes Top-1 from 79.8 to 80.0 and FLOPs from 4.6 to 4.2.Source: §4.3 · pp. 7–8, Table 5 · p. 7

4.4. Time Series Classification

  • Flowformer achieves the best performance on the time-series classification benchmark and is the only method reported to beat classical methods on average.It surpasses Rocket, identified as the state-of-the-art classical method in the setup.Source: §4.4 · p. 8, Table 6 · p. 8

4.5. Reinforcement Learning

  • In offline reinforcement learning, Flowformer remains competitive with Decision Transformer, with averaged rewards of 72.2 versus 73.5.The reported comparison contrasts this result with lower averaged rewards for Reformer and Performer, both 63.9 and 63.8, respectively.Source: §4.5 · p. 9, Table 7 · p. 9

5. Conclusions

Flowformer applies flow conservation to linearize attention, producing source competition and sink allocation without specific inductive biases. The conclusion reports strong performance across several areas and identifies scaling up for general-purpose pretraining as future work.

5. Conclusions

  • Flowformer uses flow conservation to give sources competition and sinks allocation while filtering aggregated source information.The paper says this yields linear complexity and avoids degenerated attention without specific inductive biases.Source: §5 · p. 9
  • Flowformer is reported to perform strongly across vision, language, long sequence, time series, and reinforcement learning.Source: §5 · p. 9
  • Future work will scale up Flowformer for general-purpose pre-trained models and a wider range of upstream and downstream tasks.Source: §5 · p. 9

A. Implementation Details

The appendix presents pseudocode for normal and causal Flow-Attention. It specifies Causal-Dot-Product for information aggregation in the causal version.

A.1. Pseudo-code for Flow-Attention

  • The causal version adopts Causal-Dot-Product for information aggregation.The cited method is attributed to Katharopoulos et al. (2020).Source: §A.1 · p. 13
  • Algorithm 1 presents the multi-head normal version of Flow-Attention.Source: §A.1 · p. 13
  • In both pseudocode excerpts, Q, K, and V are split; one annotates Q ∈ Rn×h×d, while the other annotates Q, K, and V ∈ Rn×h×d.Source: §A.1 · p. 13
  • Algorithm 2 presents the multi-head causal version of Flow-Attention.Source: §A.1 · p. 13

A.2. Flowformer Architecture

The architecture section frames Flowformer as a stack of L layers operating on a length-n input. It identifies the embedded input, layer outputs, and hidden representations.

A.2. Flowformer Architecture

  • Flowformer's architecture is specified for L layers and a length-n input X with d_input dimensions.Source: §A.2 · p. 13
  • The initial representation X0 is the embedding of the raw input, with shape n × d.Source: §A.2 · p. 13
  • Each layer output X_l has d channels, and Z_l denotes that layer's hidden representation.Source: §A.2 · p. 13

A.3. Experiment Configuration

The appendix describes configurations for ImageNet-1K, UEA, and D4RL, while Long-Range Arena and WikiText-103 follow their official public protocols. The tasks span hierarchical image classification, varied time-series subsets, and offline reinforcement learning environments.

A.3. Experiment Configuration

  • Long-Range Arena and WikiText-103 experiments follow their official public protocols.Source: §A.3 · p. 14
  • ImageNet-1K uses 224 × 224 images split into patches, hierarchical Transformer stages with convolutional downsampling, and global average pooling for classification.The architecture omits a classification token and is summarized in Table 8.Source: §A.3 · p. 14, Table 8 · p. 14
  • UEA includes time-series subsets varying in dimension, length, data sufficiency, and number of classes.The section presents the benchmark as enabling a comprehensive comparison.Source: §A.3 · p. 14, Table 9 · p. 14
  • D4RL experiments train on datasets collected by different policies for HalfCheetah, Hopper, and Walker, following Decision Transformer configurations.D4RL is described as offline reinforcement learning from pre-collected data, with actions then performed in an online environment.Source: §A.3 · p. 14, Figure 6 · p. 14

B.1. Ablation on Activate Function for Non-negative Operation

The activation-function ablation compares non-negative choices on Long-Range Arena and reports the sigmoid as the strongest average performer. The non-negative requirement follows from the flow-network formulation.

B.1. Ablation on Activate Function for Non-negative Operation

  • The sigmoid achieves the best averaged performance among the tested activation functions.The comparison covers elu(·) + 1.0, Relu(·), and sigmoid on Long-Range Arena.Source: §B.1 · p. 15, Table 10 · p. 15
  • All tested activation functions achieve state-of-the-art performance on Long-Range Arena.The tested choices are elu(·) + 1.0, Relu(·), and sigmoid.Source: §B.1 · p. 15, Table 10 · p. 15
  • The flow-network non-negativity requirement motivates the activation-function choices.The paper sets φ to sigmoid to satisfy that requirement.Source: §B.1 · p. 15

B.2. Ablation on Activate Functions for Competition and Allocation

The section selects Softmax for competition and Sigmoid for allocation, then relates these flow-conservation choices to attention visualizations. The visualizations describe how the mechanism focuses on informative content in images.

B.2. Ablation on Activate Functions for Competition and Allocation

  • Softmax-Sigmoid is the best of the four tested competition-allocation combinations.The paper concludes Softmax is better suited to competition and Sigmoid to allocation.Source: §B.2 · p. 15, Table 11 · p. 15

C. Preliminaries of Flow Network

  • Normalizing incoming sink flow and outgoing source flow to 1 matches the assumed exchange with adjacent layers.Values arrive from the previous layer and results pass to the next; the design follows flow conservation.Source: C. Preliminaries of Flow Network · p. 15

D. More Attention Visualizations

D.1. Image Recognition

  • Flowformer’s image attention captures essential image regions, while allocation weights cover informative parts for classification.Figure 9 compares learned attention; Figure 7 shows learned allocation weights.Source: §D.1 · pp. 15–16, Figure 7 · p. 16

D.2. Time Series Classification

On SpokenArabicDigits, Flowformer and cosFormer tie at 98.8% accuracy, while the visualizations show Flowformer attending to class-distinguishing time-series phases. The comparisons also describe distractions or missed phases in other attention models.

D.2. Time Series Classification

  • 98.8% accuracy is reported for both Flowformer and cosFormer on SpokenArabicDigits.Transformer is reported at 98.4% and Linear Transformer at 98.0%.Source: §D.2 · p. 16
  • The visual comparison reports distraction by unimportant fluctuations for vanilla Transformer and a missed second “zero” phase for Linear Transformer and cosFormer.These behaviors are described for Figure 8.Source: §D.2 · p. 16, Figure 8 · p. 16

Related papers

Loading 2202.06258v2…