"Tensor Product Attention Is All You Need"

Factorizing attention: The key to unlocking longer contexts in LLMs.

Breaking attention into Lego blocks

TPA (Tensor Product Attention) reduces memory usage in LLMs by factorizing attention components into smaller parts while improving model performance compared to standard attention mechanisms.

-----

https://arxiv.org/abs/2501.06425

🤔 Original Problem:

LLMs struggle with processing long sequences due to massive memory requirements from key-value (KV) caches during inference, limiting their practical applications.

-----

🔧 Solution in this Paper:

→ TPA factorizes queries, keys, and values into low-rank tensor products that depend on input tokens

→ Instead of storing full attention matrices, it keeps only compact factorized components

→ The method seamlessly integrates with RoPE positional embeddings by pre-rotating token factors

→ TPA introduces T6, a new transformer architecture that implements this factorized attention mechanism

-----

💡 Key Insights:

→ Attention matrices can be effectively compressed using contextual tensor decomposition

→ Memory efficiency doesn't have to compromise model quality

→ Existing attention methods (MHA, MQA, GQA) are special cases of TPA

→ RoPE compatibility enables easy adoption in modern architectures

-----

📊 Results:

→ 10x reduction in KV cache size during inference

→ Outperforms baseline models on perplexity metrics

→ Improved downstream task performance across multiple benchmarks

→ Enables processing of significantly longer sequences under fixed memory constraints