Tucker bottleneck attention cuts costs for long sequence tasks
Tucker Bottleneck Attention for Multi-Dimensional Sequence Modeling
Machine Learning
Summary
Handling really long sequences with regular self-attention is slow because it requires a lot of calculations. The authors introduce Tucker bottleneck attention (TuBA), which simplifies these calculations by working with compressed versions of the data. This method keeps important connections between parts while speeding up computations and reducing errors. They showed that TuBA performs better and faster than usual attention methods on video and weather data.
What this means in practice
- •For video technology teams: Improve video prediction tasks with faster and more accurate attention mechanisms that handle long sequences efficiently.
- •For meteorological data engineers: Enhance global weather forecasting models by reducing computation and errors using efficient multi-dimensional attention.
Authors
Ryan Solgi, Parsa Madinei, Zheng Zhang
Abstract
The quadratic cost of self-attention limits scalability to long sequences from multidimensional data. We introduce Tucker bottleneck attention (TuBA), which exploits low-rank tensor structure for efficient global token mixing. TuBA projects hidden tensors into compact Tucker cores, performs multi-head self-attention and linear projections on the cores, and writes updates back to the ambient space, enabling subquadratic computation. Its autoregressive extension combines bidirectional interactions within cores with causal attention across cores. On video prediction and global weather forecasting, TuBA achieves favorable accuracy-efficiency trade-offs over standard and efficient attention and task-specific models. Compared to standard self-attention, TuBA reduces error and computation by up to 24.7% and 66.6% for video prediction and 37.1% and 85.1% for autoregressive weather forecasting, with speedups up to 4.27 times. Low-rank Tucker cores and multi-frame generation also outperform full-rank attention and frame-by-frame generation, respectively.