Theory, Analysis, and Best Practices for Sigmoid Self-Attention

2024/09/06 by Jason Ramapuram, Ramapuram, Jason, Federico Danieli +20 · 1 voice · 22 citations
Psychology · #Artificial intelligence #Computer science #Identity, Memory, and Therapy #Psychology #Sigmoid function

paper · pdf · doi:10.48550/arxiv.2409.04431

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2024/09/06 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Attention is a key part of the transformer architecture. It is a sequence-to-sequence mapping that transforms each sequence element into a weighted sum of values. The weights are typically obtained as the softmax of dot products between keys and queries. Recent work has explored alternatives to softmax attention in transformers, such as ReLU and sigmoid activations. In this work, we revisit sigmoid attention and conduct an in-depth theoretical and empirical analysis. Theoretically, we prove that transformers with sigmoid attention are universal function approximators and benefit from improved regularity compared to softmax attention. Through detailed empirical analysis, we identify stabilization of large initial attention norms during the early stages of training as a crucial factor for the successful training of models with sigmoid attention, outperforming prior attempts. We also introduce FLASHSIGMOID, a hardware-aware and memory-efficient implementation of sigmoid attention yielding a 17% inference kernel speed-up over FLASHATTENTION2 on H100 GPUs. Experiments across language, vision, and speech show that properly normalized sigmoid attention matches the strong performance of softmax attention on a wide range of domains and scales, which previous attempts at sigmoid attention were unable to fully achieve. Our work unifies prior art and establishes best practices for sigmoid attention as a drop-in softmax replacement in transformers.

Cited by

• The Devil is in the Spectrum: Mitigating Representation Collapse in LLMs via Topologically Regularized Side-Path
• A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training
• STAR: Skeletal Token Alignment and Rearrangement for Interaction Recognition
• Variational-Ising-Attention (VIA):TailoredAttentionMattersfor Science
• All-in-One ASR: Unifying Encoder-Decoder Models of CTC, Attention, and Transducer in Dual-Mode ASR
• Sliced ReLU attention: Quasi-linear contextual expressivity via sorting
• Weaver: Kronecker Product Approximations of Spatiotemporal Attention for Traffic Network Forecasting
• Dimension-Free Minimax Rates for Learning Pairwise Interactions in Attention-Style Models
• Forget Attention: Importance-Aware Attention Is All You Need
• Value-State Gated Attention for Mitigating Extreme-Token Phenomena in Transformers
• RACE Attention: A Strictly Linear-Time Attention Layer for Training on Outrageously Large Contexts
• ENSI: Efficient Non-Interactive Secure Inference for Large Language Models
• Your Attention Matters: to Improve Model Robustness to Noise and Spurious Correlations
• Intra-DP: A High Performance Collaborative Inference System for Mobile Edge Computing
• Scaling Context Requires Rethinking Attention
• Dynamic Injection of Entity Knowledge into Dense Retrievers
• Unpacking Softmax: How Temperature Drives Representation Collapse, Compression, and Generalization
• Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
• Always Skip Attention
• The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks
• Screening Is Enough
• Attention Needs to Focus: A Unified Perspective on Attention Allocation

Discussions

• Excited to see Sigmoid Attention accepted at ICLR 2025 !! Make attention ~18% faster with a drop-in replacement 🚀 Code: github.com/apple/ml-sig... Paper arxiv.org/abs/2409.04431 [bsky, 28 points, 1 comments]

Related

• Sigmoid Self-Attention has Lower Sample Complexity than Softmax Self-Attention: A Mixture-of-Experts Perspective
• Armour: Generalizable Compact Self-Attention for Vision Transformers
• LASER: Attention with Exponential Transformation
• Transformer Based Linear Attention with Optimized GPU Kernel Implementation
• SimA: Simple Softmax-free Attention for Vision Transformers