"Attamba: Attending To Multi-Token States"

Attamba, proposed in this paper, compresses transformer memory by chunking tokens into efficient state representations

So it transforms long sequences into bite-sized pieces while keeping the meaning intact

Attamba introduces a novel way to make transformer models more efficient by using State Space Models (SSMs) to compress chunks of tokens before attention computation. This reduces memory usage and computational costs while maintaining model quality, addressing the quadratic scaling problem in transformers without sacrificing performance.

-----

https://arxiv.org/abs/2411.17685

🤖 Original Problem:

Transformer models struggle with long sequences due to quadratic computational complexity in attention mechanisms. This makes processing extended contexts expensive in terms of both memory and computation.

-----

🔧 Solution in this Paper:

→ Attamba replaces traditional key-value projection matrices with SSM blocks that compress multiple tokens into single states.

→ The architecture uses cyclic token chunking to reduce bias from fixed boundaries and maintains leading tokens for recent context.

→ SSMs process chunks of tokens in linear time, creating compressed representations that preserve essential contextual information.

→ The model can flexibly adjust chunk sizes and leading token counts to balance between efficiency and performance.

-----

💡 Key Insights:

→ SSMs can effectively compress token chunks while preserving important contextual information

→ Cyclic chunking across layers improves model quality by 5% compared to fixed boundaries

→ The architecture is robust to even randomized chunk boundaries

→ Leading tokens preservation significantly improves model quality (8.5% with L=P)

-----

📊 Results:

→ 24% improved perplexity compared to similar-sized transformers

→ 4x smaller KV-Cache and Attention FLOPs with only 5% perplexity trade-off

→ 8x KV-compression with 10% perplexity trade-off when using chunk size 8