PaperlayerView exact revision on arXiv

Source-linked AI summary

Agent Attention: On the Integration of Softmax and Linear Attention

Dongchen Han, Tianzhu Ye, Yizeng Han, Zhuofan Xia, Siyuan Pan, Pengfei Wan, Shiji Song, Gao Huang

arXiv:2312.08874v3cs.CV

TL;DR

Global Softmax attention is expressive but computationally expensive for many visual tokens. Agent Attention introduces a small set of agent tokens that aggregate and broadcast information, achieving efficient global modeling and strong results across vision tasks, including improved high-resolution Stable Diffusion generation without extra training.

  • Problem

    Global Softmax attention offers high expressiveness but incurs quadratic computation, limiting its use in visual tasks with many tokens.

  • Method

    Agent Attention uses a small set of agent tokens to aggregate information from keys and values and broadcast it to queries, integrating Softmax and linear attention.

  • Results

    Experiments across image classification, object detection, semantic segmentation, and image generation show substantial improvements, particularly in high-resolution scenarios.

  • Takeaways & Limitations

    Agent Attention provides linear token complexity with global context modeling and can accelerate Stable Diffusion while improving image quality without additional training.

Abstract

from arXiv · show

The attention module is the key component in Transformers. While the global attention mechanism offers high expressiveness, its excessive computational cost restricts its applicability in various scenarios. In this paper, we propose a novel attention paradigm, Agent Attention, to strike a favorable balance between computational efficiency and representation power. Specifically, the Agent Attention, denoted as a quadruple $(Q, A, K, V)$, introduces an additional set of agent tokens $A$ into the conventional attention module. The agent tokens first act as the agent for the query tokens $Q$ to aggregate information from $K$ and $V$, and then broadcast the information back to $Q$. Given the number of agent tokens can be designed to be much smaller than the number of query tokens, the agent attention is significantly more efficient than the widely adopted Softmax attention, while preserving global context modelling capability. Interestingly, we show that the proposed agent attention is equivalent to a generalized form of linear attention. Therefore, agent attention seamlessly integrates the powerful Softmax attention and the highly efficient linear attention. Extensive experiments demonstrate the effectiveness of agent attention with various vision Transformers and across diverse vision tasks, including image classification, object detection, semantic segmentation and image generation. Notably, agent attention has shown remarkable performance in high-resolution scenarios, owning to its linear attention nature. For instance, when applied to Stable Diffusion, our agent attention accelerates generation and substantially enhances image generation quality without any additional training. Code is available at https://github.com/LeapLabTHU/Agent-Attention.

1 Introduction

Agent Attention addresses the quadratic cost of global Softmax attention by routing information through a small set of agent tokens. It preserves global context modeling while improving efficiency and performs across diverse vision tasks, especially high-resolution settings.

  • Transformers have succeeded across vision tasks, but global Softmax attention incurs quadratic computation with respect to token count.
  • Agent Attention introduces agent tokens that aggregate information from keys and values before broadcasting it back to query tokens.This replaces direct query-key interactions with two Softmax attention operations.
  • Agent Attention integrates Softmax and linear attention, combining high expressiveness with low computation complexity.
  • Using far fewer agent tokens than query tokens reduces Softmax attention’s quadratic complexity to linear complexity while retaining global context modeling.
  • Experiments across classification, detection, segmentation, and image generation show substantial improvements, particularly at high resolution.In Stable Diffusion, the method accelerates generation and improves image quality without additional training.

2 Related Works

Related work reduces the cost of visual self-attention through sparse, local, neighborhood, deformable, or dynamically routed attention patterns. These approaches modify where attention is computed to address Softmax attention’s quadratic cost.

  • Softmax attention’s quadratic complexity motivates efficient attention patterns for visual Transformers.
  • PVT reduces computation through sparse global attention by lowering the resolution of keys and values.
  • Swin Transformer confines attention to local windows with shifted windows, while NAT computes attention within feature neighborhoods.
  • DAT, BiFormer, and related methods use deformable or dynamically routed patterns to determine attention regions from the input or query.

3 Preliminaries

Self-attention forms outputs by weighting value tokens according to query-key similarities. Softmax attention has quadratic token complexity, whereas linear attention changes the computation order to achieve linear complexity but faces expressiveness challenges.

  • Self-attention projects input tokens into query, key, and value matrices and computes each output as a similarity-weighted sum of values.
  • The projection matrices W_Q, W_K, and W_V map channel representations into head-specific query, key, and value features.
  • Softmax attention uses exponentiated scaled query-key similarities and requires O(N^2) computation for global pairwise interactions.
  • Linear attention applies mapping functions to queries and keys, then reorders matrix multiplication to reduce token complexity to O(N).
  • Designing an effective linear-attention mapping function is nontrivial, creating a trade-off between computational efficiency and model expressiveness.

4 Agent Transformer

Agent Attention introduces agent tokens that aggregate information from keys and values, then broadcast it to queries through two Softmax operations. This design integrates Softmax and linear attention, providing linear complexity in the number of input features while retaining global context and expressive capability.

  • Agent Attention: Agent Attention is equivalent to a generalized linear attention formulation, integrating Softmax attention with linear attention.The paper expresses the operation as φ_q(Q)φ_k(K)^TV and identifies it as generalized linear attention.
  • Agent Attention: Agent Attention uses agent tokens A to aggregate global information from K and V before broadcasting it back to query tokens Q.The module performs agent aggregation followed by agent broadcast as two Softmax attention operations.
  • Efficiency: With n agent tokens much smaller than N input features, the module has O(Nnd) complexity while preserving global context modeling.Agent tokens can be obtained through pooling, learnable parameters, convolution, deformable points, or token merging; pooling is the default.
  • Module Enhancements: Agent Bias adds spatial information by helping different agent tokens focus on diverse regions.The bias terms B1 and B2 are incorporated into both attention operations and improve results in the reported ablation.
  • Module Enhancements: The full module adds a depth-wise convolution term and has linear complexity in N, with complexity components for projections, agent extraction, agent attention, and DWC.The complexity expression includes 2nNC+2NnC for agent attention and k^2NC for DWC, with k = 3.
  • Experiments: Agent Attention is applied as a plug-in module to DeiT, PVT, Swin, CSwin, and Stable Diffusion, with evaluation including ImageNet-1K classification.Table 1 reports ImageNet-1K results at a default resolution of 224^2, with separate results denoted for 384^2 resolution.

5 Experiments

Experiments across classification, detection, segmentation, diffusion, and high-resolution settings show that Agent Attention improves performance while reducing computational cost. Its benefits are especially pronounced when global receptive fields or increased resolution are used.

  • ImageNet-1K Classification: Agent-PVT-S surpasses PVT-L using 30% of the parameters and 40% of the FLOPs, while inference is 1.4x to 1.7x faster on RTX3090 and A100 GPUs.These results are reported for ImageNet-1K classification.
  • Object Detection: Agent-PVT improves COCO box AP over PVT by +3.9 to +4.7, while Agent-Swin exceeds Swin by up to +1.5 box AP.The comparisons span multiple detection frameworks, schedules, and heads.
  • Semantic Segmentation: Agent-PVT-T and Agent-Swin-T achieve +3.61 and +2.17 mIoU over their counterparts on ADE20K semantic segmentation.The results are reported across SemanticFPN and UperNet backbones.
  • Image Generation: AgentSD accelerates Stable Diffusion generation and improves image quality without extra training, while finetuned agent attention reaches 2.2x acceleration over original Dreambooth.The Stable Diffusion evaluation compares SD, ToMeSD, and AgentSD; the Dreambooth result applies agent attention across diffusion steps.
  • Large Receptive Field and High Resolution: Global receptive fields improve Agent-Swin performance without changing computation complexity, and Agent-DeiT saves 75% of FLOPs for 1024^2-resolution images.Agent-DeiT-S at 448^2 reaches 83.1 accuracy with one quarter of the parameters.
  • Ablation Study: Agent attention improves accuracy by 1.2 in ablation, while fewer agent tokens are suitable for shallow layers but harm deeper-layer performance.The design therefore uses fewer agents in shallow layers and more in deep layers; Agent attention also outperforms other tested linear-attention methods.

6 Conclusion

The conclusion presents Agent Attention as a general vision-Transformer module that combines Softmax attention’s expressiveness with linear attention’s efficiency. It reports consistent gains across vision tasks, especially at high resolution, and improved Stable Diffusion generation without extra training.

  • Conclusion: Agent Attention introduces agent tokens A that collect information from K and V and broadcast it back to query tokens Q.The agents are usually acquired directly from the query space as A = f(Q).
  • Conclusion: Agent Attention integrates Softmax and linear attention, providing high expressive power with low computation complexity across vision Transformer models.The method is presented as a general-purpose alternative to Softmax attention and as a plug-in module.
  • Conclusion: Experiments confirm effectiveness in image classification, semantic segmentation, and object detection, particularly in high-resolution scenarios.The conclusion identifies these tasks as the main evidence for the approach.
  • Conclusion: Agent attention accelerates Stable Diffusion generation and substantially improves image quality without additional training.For Stable Diffusion, the appendix describes replacing Softmax attention with an adjusted agent-attention module.
  • Implementation Details: Agent bias adds spatial information so different agent tokens focus on diverse regions, using parameter-efficient bias components.The agent-bias formulation uses column, row, and block components to construct the full bias.

C.3. Ablation

Ablations show that both Stable Diffusion adjustments improve AgentSD generation quality, while applying agent attention early provides substantial performance gains.

  • C.3. Ablation: Both adjustments made for AgentSD improve generation quality in the ablation experiments.These adjustments are evaluated in Tables 10 and 11.
  • C.3. Ablation: Applying agent attention in the early stages yields substantial performance enhancements.This result is reported in Table 12.

C.4. AgentSD for finetuning

Agent attention can be integrated into DreamBooth finetuning and generation, enabling faster diffusion generation while maintaining or improving image quality under the finetuned setting.

  • C.4. AgentSD for finetuning: Finetuning enables agent attention to be used across all diffusion generation steps without compromising performance.Applying agent attention to all steps only during generation causes a slight image-quality decline.
  • C.4. AgentSD for finetuning: Agent attention applied during both DreamBooth finetuning and generation achieves 2.2x faster generation without sacrificing image quality.Finetuning time and memory costs are also reduced by around 15%.
  • C.4. AgentSD for finetuning: The DreamBooth experiment uses 30 subjects spanning 15 classes with pretrained Stable Diffusion v1.5.Agent attention is applied to all diffusion generation steps, with specified merging and attention-scale settings.
  • C.4. AgentSD for finetuning: Finetuned agent attention matches or surpasses DreamBooth in fidelity and editability, producing clearer details than generation-only agent attention.The visual comparisons include more consistent dog colors, sharper duck-toy depictions, and accurately generated cat eyes.

D. Dataset and Training Setup

The experiments use standard training protocols across classification, detection, segmentation, and linear-attention runtime comparisons, with task-specific datasets and schedules.

  • D. Dataset and Training Setup: ImageNet-1K classification models are trained from scratch for 300 epochs using AdamW, cosine decay, and 20 epochs of linear warm-up.The initial learning rate is 1 × 10^-3 for batch size 1024 and is scaled linearly with batch size.
  • D. Dataset and Training Setup: COCO experiments use 80K training images and 35K validation images, with ImageNet-pretrained backbones and 1x or 3x schedules.The standard 1x schedule trains for 12 epochs, reducing the learning rate at epochs 8 and 11.
  • D. Dataset and Training Setup: Runtime comparisons include other linear-attention methods as additional baselines.The cited setup provides the context for evaluating computational efficiency across methods.
  • D. Dataset and Training Setup: ADE20K semantic-segmentation experiments use 20K training images and 2K validation images with ImageNet-pretrained backbones.UperNet models are trained for 160K iterations with AdamW and linear warm-up.

E. Complete Experimental Results

Across classification, detection, segmentation, and attention-design ablations, agent attention consistently improves baseline performance and remains competitive in runtime comparisons.

  • E. Complete Experimental Results: Visualization comparisons and agent-weight distributions are included alongside the tabulated classification, detection, and segmentation results.The figures show Softmax, linear, and agent attention, while the distribution visualization examines agent-token behavior.
  • E. Complete Experimental Results: Agent attention produces consistent performance improvements across ImageNet-1K classification, COCO detection, and semantic segmentation settings.The experiments cover multiple transformer backbones and detection and segmentation frameworks.
  • E. Complete Experimental Results: Agent attention achieves superior results at comparable speed to other linear-attention methods.Additional runtime comparisons are reported in Fig. 9.
  • E. Complete Experimental Results: Dynamic agent tokens outperform static agents because their input-dependent representations more accurately represent current queries.Pooling is identified as a simple and effective way to obtain dynamic agent tokens.
  • E. Complete Experimental Results: Replacing attention in the first three Swin-T stages improves performance by 1.3, whereas replacing the final stage marginally reduces overall accuracy.The authors attribute the difference to the larger resolutions in the first three stages, which favor global receptive fields.

F. Agent Attention Visualization

The visualizations indicate that agent attention retains Softmax-like attention distributions while using diverse agents that specialize in distinct image regions.

  • F. Agent Attention Visualization: The visualization section compares Softmax attention, linear attention, and agent attention using the feature corresponding to a red-block query.Agent-weight distributions are also visualized for Agent-DeiT-T.
  • F. Agent Attention Visualization: Agent attention produces distributions similar to Softmax attention, whereas linear attention does not generate reasonable distributions.The comparison supports the paper’s claim that agent attention combines Softmax-like expressiveness with linear complexity.
  • F. Agent Attention Visualization: Different agent tokens focus on distinct semantic regions, including heads, sky, bodies, branches, mountains, glasses, masks, and ground.This diversity allows queries to focus on relevant areas during the agent broadcast process.

G. Model Architectures

The paper presents four Agent Attention Transformer architectures: Agent-DeiT, Agent-PVT, Agent-Swin, and Agent-CSwin. These models mainly replace Softmax attention blocks at early stages, leveraging their enlarged receptive field.

  • Four architectures are presented: Agent-DeiT, Agent-PVT, Agent-Swin, and Agent-CSwin.Their designs are documented in Tables 19–23.
  • Softmax attention blocks are mainly replaced with Agent Attention at early Transformer stages.The stated design motivation is the advantage of an enlarged receptive field.
Loading 2312.08874v3…